Allow mutlivariate ETS models

Author: jessegrabowski

pymc_experimental/statespace/models/ETS.py

@@ -10,6 +10,7 @@
     ALL_STATE_AUX_DIM,
     ALL_STATE_DIM,
     ETS_SEASONAL_DIM,
+    OBS_STATE_AUX_DIM,
     OBS_STATE_DIM,
 )
 
@@ -176,12 +177,15 @@ class BayesianETS(PyMCStateSpace):
     def __init__(
         self,
         order: tuple[str, str, str] | None = None,
+        endog_names: str | list[str] | None = None,
+        k_endog: int = 1,
         trend: bool = True,
         damped_trend: bool = False,
         seasonal: bool = False,
         seasonal_periods: int | None = None,
         measurement_error: bool = False,
         use_transformed_parameterization: bool = False,
+        dense_innovation_covariance: bool = False,
         filter_type: str = "standard",
         verbose: bool = True,
     ):
@@ -214,13 +218,26 @@ def __init__(
         if self.seasonal and self.seasonal_periods is None:
             raise ValueError("If seasonal is True, seasonal_periods must be provided.")
 
+        if endog_names is not None:
+            endog_names = list(endog_names)
+            k_endog = len(endog_names)
+        else:
+            endog_names = [f"data_{i}" for i in range(k_endog)] if k_endog > 1 else ["data"]
+
+        self.endog_names = endog_names
+
+        if dense_innovation_covariance and k_endog == 1:
+            dense_innovation_covariance = False
+
+        self.dense_innovation_covariance = dense_innovation_covariance
+
         k_states = (
             2
             + int(trend)
             + int(seasonal) * (seasonal_periods if seasonal_periods is not None else 0)
-        )
-        k_posdef = 1
-        k_endog = 1
+        ) * k_endog
+
+        k_posdef = k_endog
 
         super().__init__(
             k_endog,
@@ -243,6 +260,7 @@ def param_names(self):
             "gamma",
             "phi",
             "sigma_state",
+            "state_cov",
             "sigma_obs",
         ]
         if not self.trend:
@@ -256,6 +274,11 @@ def param_names(self):
         if not self.measurement_error:
             names.remove("sigma_obs")
 
+        if self.dense_innovation_covariance:
+            names.remove("sigma_state")
+        else:
+            names.remove("state_cov")
+
         return names
 
     @property
@@ -283,27 +306,34 @@ def param_info(self) -> dict[str, dict[str, Any]]:
                 "constraints": "Positive",
             },
             "alpha": {
-                "shape": None,
+                "shape": None if self.k_endog == 1 else (self.k_endog,),
                 "constraints": "0 < alpha < 1",
             },
             "beta": {
-                "shape": None,
+                "shape": None if self.k_endog == 1 else (self.k_endog,),
                 "constraints": "0 < beta < 1"
                 if not self.use_transformed_parameterization
                 else "0 < beta < alpha",
             },
             "gamma": {
-                "shape": None,
+                "shape": None if self.k_endog == 1 else (self.k_endog,),
                 "constraints": "0 < gamma< 1"
                 if not self.use_transformed_parameterization
                 else "0 < gamma < (1 - alpha)",
             },
             "phi": {
-                "shape": None,
+                "shape": None if self.k_endog == 1 else (self.k_endog,),
                 "constraints": "0 < phi < 1",
             },
         }
 
+        if self.dense_innovation_covariance:
+            del info["sigma_state"]
+            info["state_cov"] = {
+                "shape": (self.k_posdef, self.k_posdef),
+                "constraints": "Positive Semi-definite",
+            }
+
         for name in self.param_names:
             info[name]["dims"] = self.param_dims.get(name, None)
 
@@ -317,15 +347,22 @@ def state_names(self):
         if self.seasonal:
             states += [f"L{i}.season" for i in range(1, self.seasonal_periods + 1)]
 
+        if self.k_endog > 1:
+            states = [f"{name}_{state}" for name in self.endog_names for state in states]
+
         return states
 
     @property
     def observed_states(self):
-        return ["data"]
+        return self.endog_names
 
     @property
     def shock_names(self):
-        return ["innovation"]
+        return (
+            ["innovation"]
+            if self.k_endog == 1
+            else [f"{name}_innovation" for name in self.endog_names]
+        )
 
     @property
     def param_dims(self):
@@ -339,11 +376,23 @@ def param_dims(self):
             "seasonal_param": (ETS_SEASONAL_DIM,),
         }
 
+        if self.dense_innovation_covariance:
+            del coord_map["sigma_state"]
+            coord_map["state_cov"] = (OBS_STATE_DIM, OBS_STATE_AUX_DIM)
+
         if self.k_endog == 1:
             coord_map["sigma_state"] = None
             coord_map["sigma_obs"] = None
             coord_map["initial_level"] = None
             coord_map["initial_trend"] = None
+        else:
+            coord_map["alpha"] = (OBS_STATE_DIM,)
+            coord_map["beta"] = (OBS_STATE_DIM,)
+            coord_map["gamma"] = (OBS_STATE_DIM,)
+            coord_map["phi"] = (OBS_STATE_DIM,)
+            coord_map["initial_seasonal"] = (OBS_STATE_DIM, ETS_SEASONAL_DIM)
+            coord_map["seasonal_param"] = (OBS_STATE_DIM, ETS_SEASONAL_DIM)
+
         if not self.measurement_error:
             del coord_map["sigma_obs"]
         if not self.seasonal:
@@ -360,6 +409,8 @@ def coords(self) -> dict[str, Sequence]:
         return coords
 
     def make_symbolic_graph(self) -> None:
+        k_states_each = self.k_states // self.k_endog
+
         P0 = self.make_and_register_variable(
             "P0", shape=(self.k_states, self.k_states), dtype=floatX
         )
@@ -368,21 +419,37 @@ def make_symbolic_graph(self) -> None:
         initial_level = self.make_and_register_variable(
             "initial_level", shape=(self.k_endog,) if self.k_endog > 1 else (), dtype=floatX
         )
-        self.ssm["initial_state", 1] = initial_level
+
+        initial_states = [pt.zeros(k_states_each) for _ in range(self.k_endog)]
+        if self.k_endog == 1:
+            initial_states = [pt.set_subtensor(initial_states[0][1], initial_level)]
+        else:
+            initial_states = [
+                pt.set_subtensor(initial_state[1], initial_level[i])
+                for i, initial_state in enumerate(initial_states)
+            ]
 
         # The shape of R can be pre-allocated, then filled with the required parameters
-        R = pt.zeros((self.k_states, self.k_posdef))
+        R = pt.zeros((self.k_states // self.k_endog, 1))
+
+        alpha = self.make_and_register_variable(
+            "alpha", shape=() if self.k_endog == 1 else (self.k_endog,), dtype=floatX
+        )
 
-        alpha = self.make_and_register_variable("alpha", shape=(), dtype=floatX)
-        R = pt.set_subtensor(R[1, 0], alpha)  # and l_t = ... + alpha * e_t
+        if self.k_endog == 1:
+            # The R[0, 0] entry needs to be adjusted for a shift in the time indices. Consider the (A, N, N) model:
+            # y_t = l_{t-1} + e_t
+            # l_t = l_{t-1} + alpha * e_t
+            R_list = [pt.set_subtensor(R[1, 0], alpha)]  # and l_t = ... + alpha * e_t
 
-        # The R[0, 0] entry needs to be adjusted for a shift in the time indices. Consider the (A, N, N) model:
-        # y_t = l_{t-1} + e_t
-        # l_t = l_{t-1} + alpha * e_t
-        # We want the first equation to be in terms of time t on the RHS, because our observation equation is always
-        # y_t = Z @ x_t. Re-arranging equation 2, we get l_{t-1} = l_t - alpha * e_t --> y_t = l_t + e_t - alpha * e_t
-        # --> y_t = l_t + (1 - alpha) * e_t
-        R = pt.set_subtensor(R[0, :], (1 - alpha))
+            # We want the first equation to be in terms of time t on the RHS, because our observation equation is always
+            # y_t = Z @ x_t. Re-arranging equation 2, we get l_{t-1} = l_t - alpha * e_t --> y_t = l_t + e_t - alpha * e_t
+            # --> y_t = l_t + (1 - alpha) * e_t
+            R_list = [pt.set_subtensor(R[0, :], (1 - alpha)) for R in R_list]
+        else:
+            # If there are multiple endog, clone the basic R matrix and modify the appropriate entries
+            R_list = [pt.set_subtensor(R[1, 0], alpha[i]) for i in range(self.k_endog)]
+            R_list = [pt.set_subtensor(R[0, :], (1 - alpha[i])) for i, R in enumerate(R_list)]
 
         # Shock and level component always exists, the base case is e_t = e_t and l_t = l_{t-1}
         T_base = pt.as_tensor_variable(np.array([[0.0, 0.0], [0.0, 1.0]]))
@@ -391,77 +458,134 @@ def make_symbolic_graph(self) -> None:
             initial_trend = self.make_and_register_variable(
                 "initial_trend", shape=(self.k_endog,) if self.k_endog > 1 else (), dtype=floatX
             )
-            self.ssm["initial_state", 2] = initial_trend
 
-            beta = self.make_and_register_variable("beta", shape=(), dtype=floatX)
+            if self.k_endog == 1:
+                initial_states = [pt.set_subtensor(initial_states[0][2], initial_trend)]
+            else:
+                initial_states = [
+                    pt.set_subtensor(initial_state[2], initial_trend[i])
+                    for i, initial_state in enumerate(initial_states)
+                ]
+            beta = self.make_and_register_variable(
+                "beta", shape=() if self.k_endog == 1 else (self.k_endog,), dtype=floatX
+            )
             if self.use_transformed_parameterization:
-                R = pt.set_subtensor(R[2, 0], beta)
+                param = beta
+            else:
+                param = alpha * beta
+            if self.k_endog == 1:
+                R_list = [pt.set_subtensor(R[2, 0], param) for R in R_list]
             else:
-                R = pt.set_subtensor(R[2, 0], alpha * beta)
+                R_list = [pt.set_subtensor(R[2, 0], param[i]) for i, R in enumerate(R_list)]
 
             # If a trend is requested, we have the following transition equations (omitting the shocks):
             # l_t = l_{t-1} + b_{t-1}
             # b_t = b_{t-1}
             T_base = pt.as_tensor_variable(([0.0, 0.0, 0.0], [0.0, 1.0, 1.0], [0.0, 0.0, 1.0]))
 
         if self.damped_trend:
-            phi = self.make_and_register_variable("phi", shape=(), dtype=floatX)
+            phi = self.make_and_register_variable(
+                "phi", shape=() if self.k_endog == 1 else (self.k_endog,), dtype=floatX
+            )
             # We are always in the case where we have a trend, so we can add the dampening parameter to T_base defined
             # in that branch. Transition equations become:
             # l_t = l_{t-1} + phi * b_{t-1}
             # b_t = phi * b_{t-1}
-            T_base = pt.set_subtensor(T_base[1:, 2], phi)
+            if self.k_endog > 1:
+                T_base = [pt.set_subtensor(T_base[1:, 2], phi[i]) for i in range(self.k_endog)]
+            else:
+                T_base = pt.set_subtensor(T_base[1:, 2], phi)
 
-        T_components = [T_base]
+        T_components = (
+            [T_base for _ in range(self.k_endog)] if not isinstance(T_base, list) else T_base
+        )
 
         if self.seasonal:
             initial_seasonal = self.make_and_register_variable(
-                "initial_seasonal", shape=(self.seasonal_periods,), dtype=floatX
+                "initial_seasonal",
+                shape=(self.seasonal_periods,)
+                if self.k_endog == 1
+                else (self.k_endog, self.seasonal_periods),
+                dtype=floatX,
             )
-
-            self.ssm["initial_state", 2 + int(self.trend) :] = initial_seasonal
-
-            gamma = self.make_and_register_variable("gamma", shape=(), dtype=floatX)
-
-            if self.use_transformed_parameterization:
-                param = gamma
+            if self.k_endog == 1:
+                initial_states = [
+                    pt.set_subtensor(initial_states[0][2 + int(self.trend) :], initial_seasonal)
+                ]
             else:
-                param = (1 - alpha) * gamma
+                initial_states = [
+                    pt.set_subtensor(initial_state[2 + int(self.trend) :], initial_seasonal[i])
+                    for i, initial_state in enumerate(initial_states)
+                ]
 
-            R = pt.set_subtensor(R[2 + int(self.trend), 0], param)
+            gamma = self.make_and_register_variable(
+                "gamma", shape=() if self.k_endog == 1 else (self.k_endog,), dtype=floatX
+            )
 
+            param = gamma if self.use_transformed_parameterization else (1 - alpha) * gamma
             # Additional adjustment to the R[0, 0] position is required. Start from:
             # y_t = l_{t-1} + s_{t-m} + e_t
             # l_t = l_{t-1} + alpha * e_t
             # s_t = s_{t-m} + gamma * e_t
             # Solve for l_{t-1} and s_{t-m} in terms of l_t and s_t, then substitute into the observation equation:
             # y_t = l_t + s_t - alpha * e_t - gamma * e_t + e_t --> y_t = l_t + s_t + (1 - alpha - gamma) * e_t
-            R = pt.set_subtensor(R[0, 0], R[0, 0] - param)
+
+            if self.k_endog == 1:
+                R_list = [pt.set_subtensor(R[2 + int(self.trend), 0], param) for R in R_list]
+                R_list = [pt.set_subtensor(R[0, 0], R[0, 0] - param) for R in R_list]
+
+            else:
+                R_list = [
+                    pt.set_subtensor(R[2 + int(self.trend), 0], param[i])
+                    for i, R in enumerate(R_list)
+                ]
+                R_list = [
+                    pt.set_subtensor(R[0, 0], R[0, 0] - param[i]) for i, R in enumerate(R_list)
+                ]
 
             # The seasonal component is always going to look like a TimeFrequency structural component, see that
             # docstring for more details
-            T_seasonal = pt.eye(self.seasonal_periods, k=-1)
-            T_seasonal = pt.set_subtensor(T_seasonal[0, -1], 1.0)
-            T_components += [T_seasonal]
+            T_seasonals = [pt.eye(self.seasonal_periods, k=-1) for _ in range(self.k_endog)]
+            T_seasonals = [pt.set_subtensor(T_seasonal[0, -1], 1.0) for T_seasonal in T_seasonals]
+
+            # Organize the components so it goes T1, T_seasonal_1, T2, T_seasonal_2, etc.
+            T_components = [
+                matrix[i] for i in range(self.k_endog) for matrix in [T_components, T_seasonals]
+            ]
 
-        self.ssm["selection"] = R
+        x0 = pt.concatenate(initial_states, axis=0)
+        R = pt.linalg.block_diag(*R_list)
+
+        self.ssm["initial_state"] = x0
+        self.ssm["selection"] = pt.specify_shape(R, shape=(self.k_states, self.k_posdef))
 
         T = pt.linalg.block_diag(*T_components)
         self.ssm["transition"] = pt.specify_shape(T, (self.k_states, self.k_states))
 
-        Z = np.zeros((self.k_endog, self.k_states))
-        Z[0, 0] = 1.0  # innovation
-        Z[0, 1] = 1.0  # level
-        if self.seasonal:
-            Z[0, 2 + int(self.trend)] = 1.0
+        Zs = [np.zeros((self.k_endog, self.k_states // self.k_endog)) for _ in range(self.k_endog)]
+        for i, Z in enumerate(Zs):
+            Z[i, 0] = 1.0  # innovation
+            Z[i, 1] = 1.0  # level
+            if self.seasonal:
+                Z[i, 2 + int(self.trend)] = 1.0
+
+        Z = pt.concatenate(Zs, axis=1)
+
         self.ssm["design"] = Z
 
         # Set up the state covariance matrix
-        state_cov_idx = ("state_cov", *np.diag_indices(self.k_posdef))
-        state_cov = self.make_and_register_variable(
-            "sigma_state", shape=() if self.k_posdef == 1 else (self.k_posdef,), dtype=floatX
-        )
-        self.ssm[state_cov_idx] = state_cov**2
+        if self.dense_innovation_covariance:
+            state_cov = self.make_and_register_variable(
+                "state_cov", shape=(self.k_posdef, self.k_posdef), dtype=floatX
+            )
+            self.ssm["state_cov"] = state_cov
+
+        else:
+            state_cov_idx = ("state_cov", *np.diag_indices(self.k_posdef))
+            state_cov = self.make_and_register_variable(
+                "sigma_state", shape=() if self.k_posdef == 1 else (self.k_posdef,), dtype=floatX
+            )
+            self.ssm[state_cov_idx] = state_cov**2
 
         if self.measurement_error:
             obs_cov_idx = ("obs_cov", *np.diag_indices(self.k_endog))