pymc-extras 0.4.0__py3-none-any.whl → 0.4.1__py3-none-any.whl

pymc_extras/statespace/models/structural/components/cycle.py

@@ -43,6 +43,11 @@ class CycleComponent(Component):
         Names of the observed state variables. For univariate time series, defaults to ``["data"]``.
         For multivariate time series, specify a list of names for each endogenous variable.
 
+    share_states: bool, default False
+        Whether latent states are shared across the observed states. If True, there will be only one set of latent
+        states, which are observed by all observed states. If False, each observed state has its own set of
+        latent states. This argument has no effect if `k_endog` is 1.
+
     Notes
     -----
     The cycle component is very similar in implementation to the frequency domain seasonal component, expect that it
@@ -155,6 +160,7 @@ class CycleComponent(Component):
         dampen: bool = False,
         innovations: bool = True,
         observed_state_names: list[str] | None = None,
+        share_states: bool = False,
     ):
         if observed_state_names is None:
             observed_state_names = ["data"]
@@ -167,6 +173,7 @@ class CycleComponent(Component):
             cycle = int(cycle_length) if cycle_length is not None else "Estimate"
             name = f"Cycle[s={cycle}, dampen={dampen}, innovations={innovations}]"
 
+        self.share_states = share_states
         self.estimate_cycle_length = estimate_cycle_length
         self.cycle_length = cycle_length
         self.innovations = innovations
@@ -175,8 +182,8 @@ class CycleComponent(Component):
 
         k_endog = len(observed_state_names)
 
-        k_states = 2 * k_endog
-        k_posdef = 2 * k_endog
+        k_states = 2 if share_states else 2 * k_endog
+        k_posdef = 2 if share_states else 2 * k_endog
 
         obs_state_idx = np.zeros(k_states)
         obs_state_idx[slice(0, k_states, 2)] = 1
@@ -190,21 +197,26 @@ class CycleComponent(Component):
             combine_hidden_states=True,
             obs_state_idxs=obs_state_idx,
             observed_state_names=observed_state_names,
+            share_states=share_states,
         )
 
     def make_symbolic_graph(self) -> None:
+        k_endog = self.k_endog
+        k_endog_effective = 1 if self.share_states else k_endog
+
         Z = np.array([1.0, 0.0]).reshape((1, -1))
-        design_matrix = block_diag(*[Z for _ in range(self.k_endog)])
+        design_matrix = block_diag(*[Z for _ in range(k_endog_effective)])
         self.ssm["design", :, :] = pt.as_tensor_variable(design_matrix)
 
         # selection matrix R defines structure of innovations (always identity for cycle components)
         # when innovations=False, state cov Q=0, hence R @ Q @ R.T = 0
         R = np.eye(2)  # 2x2 identity for each cycle component
-        selection_matrix = block_diag(*[R for _ in range(self.k_endog)])
+        selection_matrix = block_diag(*[R for _ in range(k_endog_effective)])
         self.ssm["selection", :, :] = pt.as_tensor_variable(selection_matrix)
 
         init_state = self.make_and_register_variable(
-            f"{self.name}", shape=(self.k_endog, 2) if self.k_endog > 1 else (self.k_states,)
+            f"params_{self.name}",
+            shape=(k_endog_effective, 2) if k_endog_effective > 1 else (self.k_states,),
         )
         self.ssm["initial_state", :] = init_state.ravel()
 
@@ -219,19 +231,19 @@ class CycleComponent(Component):
             rho = 1
 
         T = rho * _frequency_transition_block(lamb, j=1)
-        transition = block_diag(*[T for _ in range(self.k_endog)])
+        transition = block_diag(*[T for _ in range(k_endog_effective)])
         self.ssm["transition"] = pt.specify_shape(transition, (self.k_states, self.k_states))
 
         if self.innovations:
-            if self.k_endog == 1:
+            if k_endog_effective == 1:
                 sigma_cycle = self.make_and_register_variable(f"sigma_{self.name}", shape=())
                 self.ssm["state_cov", :, :] = pt.eye(self.k_posdef) * sigma_cycle**2
             else:
                 sigma_cycle = self.make_and_register_variable(
-                    f"sigma_{self.name}", shape=(self.k_endog,)
+                    f"sigma_{self.name}", shape=(k_endog_effective,)
                 )
                 state_cov = block_diag(
-                    *[pt.eye(2) * sigma_cycle[i] ** 2 for i in range(self.k_endog)]
+                    *[pt.eye(2) * sigma_cycle[i] ** 2 for i in range(k_endog_effective)]
                 )
                 self.ssm["state_cov"] = pt.specify_shape(state_cov, (self.k_states, self.k_states))
         else:
@@ -239,33 +251,41 @@ class CycleComponent(Component):
             self.ssm["state_cov", :, :] = pt.zeros((self.k_posdef, self.k_posdef))
 
     def populate_component_properties(self):
-        self.state_names = [
-            f"{f}_{self.name}[{var_name}]" if self.k_endog > 1 else f"{f}_{self.name}"
-            for var_name in self.observed_state_names
-            for f in ["Cos", "Sin"]
-        ]
+        k_endog = self.k_endog
+        k_endog_effective = 1 if self.share_states else k_endog
+
+        base_names = [f"{f}_{self.name}" for f in ["Cos", "Sin"]]
+
+        if self.share_states:
+            self.state_names = [f"{name}[shared]" for name in base_names]
+        else:
+            self.state_names = [
+                f"{name}[{var_name}]" if k_endog_effective > 1 else name
+                for var_name in self.observed_state_names
+                for name in base_names
+            ]
 
-        self.param_names = [f"{self.name}"]
+        self.param_names = [f"params_{self.name}"]
 
-        if self.k_endog == 1:
-            self.param_dims = {self.name: (f"state_{self.name}",)}
-            self.coords = {f"state_{self.name}": self.state_names}
+        if k_endog_effective == 1:
+            self.param_dims = {f"params_{self.name}": (f"state_{self.name}",)}
+            self.coords = {f"state_{self.name}": base_names}
             self.param_info = {
-                f"{self.name}": {
+                f"params_{self.name}": {
                     "shape": (2,),
                     "constraints": None,
                     "dims": (f"state_{self.name}",),
                 }
             }
         else:
-            self.param_dims = {self.name: (f"endog_{self.name}", f"state_{self.name}")}
+            self.param_dims = {f"params_{self.name}": (f"endog_{self.name}", f"state_{self.name}")}
             self.coords = {
                 f"state_{self.name}": [f"Cos_{self.name}", f"Sin_{self.name}"],
                 f"endog_{self.name}": self.observed_state_names,
             }
             self.param_info = {
-                f"{self.name}": {
-                    "shape": (self.k_endog, 2),
+                f"params_{self.name}": {
+                    "shape": (k_endog_effective, 2),
                     "constraints": None,
                     "dims": (f"endog_{self.name}", f"state_{self.name}"),
                 }
@@ -274,22 +294,22 @@ class CycleComponent(Component):
         if self.estimate_cycle_length:
             self.param_names += [f"length_{self.name}"]
             self.param_info[f"length_{self.name}"] = {
-                "shape": () if self.k_endog == 1 else (self.k_endog,),
+                "shape": () if k_endog_effective == 1 else (k_endog_effective,),
                 "constraints": "Positive, non-zero",
-                "dims": None if self.k_endog == 1 else f"endog_{self.name}",
+                "dims": None if k_endog_effective == 1 else (f"endog_{self.name}",),
             }
 
         if self.dampen:
             self.param_names += [f"dampening_factor_{self.name}"]
             self.param_info[f"dampening_factor_{self.name}"] = {
-                "shape": () if self.k_endog == 1 else (self.k_endog,),
+                "shape": () if k_endog_effective == 1 else (k_endog_effective,),
                 "constraints": "0 < x ≤ 1",
-                "dims": None if self.k_endog == 1 else (f"endog_{self.name}",),
+                "dims": None if k_endog_effective == 1 else (f"endog_{self.name}",),
             }
 
         if self.innovations:
             self.param_names += [f"sigma_{self.name}"]
-            if self.k_endog == 1:
+            if k_endog_effective == 1:
                 self.param_info[f"sigma_{self.name}"] = {
                     "shape": (),
                     "constraints": "Positive",
@@ -298,7 +318,7 @@ class CycleComponent(Component):
             else:
                 self.param_dims[f"sigma_{self.name}"] = (f"endog_{self.name}",)
                 self.param_info[f"sigma_{self.name}"] = {
-                    "shape": (self.k_endog,),
+                    "shape": (k_endog_effective,),
                     "constraints": "Positive",
                     "dims": (f"endog_{self.name}",),
                 }

pymc_extras/statespace/models/structural/components/level_trend.py

@@ -13,13 +13,11 @@ class LevelTrendComponent(Component):
     Parameters
     ----------
     order : int
-
         Number of time derivatives of the trend to include in the model. For example, when order=3, the trend will
         be of the form ``y = a + b * t + c * t ** 2``, where the coefficients ``a, b, c`` come from the initial
         state values.
 
     innovations_order : int or sequence of int, optional
-
         The number of stochastic innovations to include in the model. By default, ``innovations_order = order``
 
     name : str, default "level_trend"
@@ -28,6 +26,11 @@ class LevelTrendComponent(Component):
     observed_state_names : list[str] | None, default None
         List of strings for observed state labels. If None, defaults to ["data"].
 
+    share_states: bool, default False
+        Whether latent states are shared across the observed states. If True, there will be only one set of latent
+        states, which are observed by all observed states. If False, each observed state has its own set of
+        latent states. This argument has no effect if `k_endog` is 1.
+
     Notes
     -----
     This class implements the level and trend components of the general structural time series model. In the most
@@ -120,7 +123,10 @@ class LevelTrendComponent(Component):
         innovations_order: int | list[int] | None = None,
         name: str = "level_trend",
         observed_state_names: list[str] | None = None,
+        share_states: bool = False,
     ):
+        self.share_states = share_states
+
         if innovations_order is None:
             innovations_order = order
 
@@ -156,37 +162,51 @@ class LevelTrendComponent(Component):
         super().__init__(
             name,
             k_endog=k_endog,
-            k_states=k_states * k_endog,
-            k_posdef=k_posdef * k_endog,
+            k_states=k_states * k_endog if not share_states else k_states,
+            k_posdef=k_posdef * k_endog if not share_states else k_posdef,
             observed_state_names=observed_state_names,
             measurement_error=False,
             combine_hidden_states=False,
-            obs_state_idxs=np.tile(np.array([1.0] + [0.0] * (k_states - 1)), k_endog),
+            obs_state_idxs=np.tile(
+                np.array([1.0] + [0.0] * (k_states - 1)), k_endog if not share_states else 1
+            ),
+            share_states=share_states,
         )
 
     def populate_component_properties(self):
         k_endog = self.k_endog
-        k_states = self.k_states // k_endog
-        k_posdef = self.k_posdef // k_endog
+        k_endog_effective = 1 if self.share_states else k_endog
+
+        k_states = self.k_states // k_endog_effective
+        k_posdef = self.k_posdef // k_endog_effective
 
         name_slice = POSITION_DERIVATIVE_NAMES[:k_states]
         self.param_names = [f"initial_{self.name}"]
         base_names = [name for name, mask in zip(name_slice, self._order_mask) if mask]
-        self.state_names = [
-            f"{name}[{obs_name}]" for obs_name in self.observed_state_names for name in base_names
-        ]
+
+        if self.share_states:
+            self.state_names = [f"{name}[{self.name}_shared]" for name in base_names]
+        else:
+            self.state_names = [
+                f"{name}[{obs_name}]"
+                for obs_name in self.observed_state_names
+                for name in base_names
+            ]
+
         self.param_dims = {f"initial_{self.name}": (f"state_{self.name}",)}
         self.coords = {f"state_{self.name}": base_names}
 
         if k_endog > 1:
+            self.coords[f"endog_{self.name}"] = self.observed_state_names
+
+        if k_endog_effective > 1:
             self.param_dims[f"state_{self.name}"] = (
                 f"endog_{self.name}",
                 f"state_{self.name}",
             )
             self.param_dims = {f"initial_{self.name}": (f"endog_{self.name}", f"state_{self.name}")}
-            self.coords[f"endog_{self.name}"] = self.observed_state_names
 
-        shape = (k_endog, k_states) if k_endog > 1 else (k_states,)
+        shape = (k_endog_effective, k_states) if k_endog_effective > 1 else (k_states,)
         self.param_info = {f"initial_{self.name}": {"shape": shape, "constraints": None}}
 
         if self.k_posdef > 0:
@@ -196,20 +216,23 @@ class LevelTrendComponent(Component):
                 name for name, mask in zip(name_slice, self.innovations_order) if mask
             ]
 
-            self.shock_names = [
-                f"{name}[{obs_name}]"
-                for obs_name in self.observed_state_names
-                for name in base_shock_names
-            ]
+            if self.share_states:
+                self.shock_names = [f"{name}[{self.name}_shared]" for name in base_shock_names]
+            else:
+                self.shock_names = [
+                    f"{name}[{obs_name}]"
+                    for obs_name in self.observed_state_names
+                    for name in base_shock_names
+                ]
 
             self.param_dims[f"sigma_{self.name}"] = (
                 (f"shock_{self.name}",)
-                if k_endog == 1
+                if k_endog_effective == 1
                 else (f"endog_{self.name}", f"shock_{self.name}")
             )
             self.coords[f"shock_{self.name}"] = base_shock_names
             self.param_info[f"sigma_{self.name}"] = {
-                "shape": (k_posdef,) if k_endog == 1 else (k_endog, k_posdef),
+                "shape": (k_posdef,) if k_endog_effective == 1 else (k_endog_effective, k_posdef),
                 "constraints": "Positive",
             }
 
@@ -218,12 +241,14 @@ class LevelTrendComponent(Component):
 
     def make_symbolic_graph(self) -> None:
         k_endog = self.k_endog
-        k_states = self.k_states // k_endog
-        k_posdef = self.k_posdef // k_endog
+        k_endog_effective = 1 if self.share_states else k_endog
+
+        k_states = self.k_states // k_endog_effective
+        k_posdef = self.k_posdef // k_endog_effective
 
         initial_trend = self.make_and_register_variable(
             f"initial_{self.name}",
-            shape=(k_states,) if k_endog == 1 else (k_endog, k_states),
+            shape=(k_states,) if k_endog_effective == 1 else (k_endog, k_states),
         )
         self.ssm["initial_state", :] = initial_trend.ravel()
 
@@ -231,27 +256,34 @@ class LevelTrendComponent(Component):
         T = pt.zeros((k_states, k_states))[triu_idx[0], triu_idx[1]].set(1)
 
         self.ssm["transition", :, :] = pt.specify_shape(
-            pt.linalg.block_diag(*[T for _ in range(k_endog)]), (self.k_states, self.k_states)
+            pt.linalg.block_diag(*[T for _ in range(k_endog_effective)]),
+            (self.k_states, self.k_states),
         )
 
         R = np.eye(k_states)
         R = R[:, self.innovations_order]
 
         self.ssm["selection", :, :] = pt.specify_shape(
-            pt.linalg.block_diag(*[R for _ in range(k_endog)]), (self.k_states, self.k_posdef)
+            pt.linalg.block_diag(*[R for _ in range(k_endog_effective)]),
+            (self.k_states, self.k_posdef),
         )
 
         Z = np.array([1.0] + [0.0] * (k_states - 1)).reshape((1, -1))
 
-        self.ssm["design", :, :] = pt.specify_shape(
-            pt.linalg.block_diag(*[Z for _ in range(k_endog)]), (self.k_endog, self.k_states)
-        )
+        if self.share_states:
+            self.ssm["design", :, :] = pt.specify_shape(
+                pt.join(0, *[Z for _ in range(k_endog)]), (self.k_endog, self.k_states)
+            )
+        else:
+            self.ssm["design", :, :] = pt.specify_shape(
+                pt.linalg.block_diag(*[Z for _ in range(k_endog)]), (self.k_endog, self.k_states)
+            )
 
         if k_posdef > 0:
             sigma_trend = self.make_and_register_variable(
                 f"sigma_{self.name}",
-                shape=(k_posdef,) if k_endog == 1 else (k_endog, k_posdef),
+                shape=(k_posdef,) if k_endog_effective == 1 else (k_endog, k_posdef),
             )
-            diag_idx = np.diag_indices(k_posdef * k_endog)
+            diag_idx = np.diag_indices(k_posdef * k_endog_effective)
             idx = np.s_["state_cov", diag_idx[0], diag_idx[1]]
             self.ssm[idx] = (sigma_trend**2).ravel()

pymc_extras/statespace/models/structural/components/measurement_error.py

@@ -17,6 +17,10 @@ class MeasurementError(Component):
         Name of the measurement error component. Default is "MeasurementError".
     observed_state_names : list[str] | None, optional
         Names of the observed variables. If None, defaults to ["data"].
+    share_states: bool, default False
+        Whether latent states are shared across the observed states. If True, there will be only one set of latent
+        states, which are observed by all observed states. If False, each observed state has its own set of
+        latent states. This argument has no effect if `k_endog` is 1.
 
     Notes
     -----
@@ -93,11 +97,16 @@ class MeasurementError(Component):
     """
 
     def __init__(
-        self, name: str = "MeasurementError", observed_state_names: list[str] | None = None
+        self,
+        name: str = "MeasurementError",
+        observed_state_names: list[str] | None = None,
+        share_states: bool = False,
     ):
         if observed_state_names is None:
             observed_state_names = ["data"]
 
+        self.share_states = share_states
+
         k_endog = len(observed_state_names)
         k_states = 0
         k_posdef = 0
@@ -110,28 +119,36 @@ class MeasurementError(Component):
             measurement_error=True,
             combine_hidden_states=False,
             observed_state_names=observed_state_names,
+            share_states=share_states,
         )
 
     def populate_component_properties(self):
+        k_endog = self.k_endog
+        k_endog_effective = 1 if self.share_states else k_endog
+
         self.param_names = [f"sigma_{self.name}"]
         self.param_dims = {}
         self.coords = {}
 
-        if self.k_endog > 1:
+        if k_endog_effective > 1:
             self.param_dims[f"sigma_{self.name}"] = (f"endog_{self.name}",)
             self.coords[f"endog_{self.name}"] = self.observed_state_names
 
         self.param_info = {
             f"sigma_{self.name}": {
-                "shape": (self.k_endog,) if self.k_endog > 1 else (),
+                "shape": (k_endog_effective,) if k_endog_effective > 1 else (),
                 "constraints": "Positive",
-                "dims": (f"endog_{self.name}",) if self.k_endog > 1 else None,
+                "dims": (f"endog_{self.name}",) if k_endog_effective > 1 else None,
             }
         }
 
     def make_symbolic_graph(self) -> None:
-        sigma_shape = () if self.k_endog == 1 else (self.k_endog,)
+        k_endog = self.k_endog
+        k_endog_effective = 1 if self.share_states else k_endog
+
+        sigma_shape = () if k_endog_effective == 1 else (k_endog_effective,)
         error_sigma = self.make_and_register_variable(f"sigma_{self.name}", shape=sigma_shape)
+
         diag_idx = np.diag_indices(self.k_endog)
         idx = np.s_["obs_cov", diag_idx[0], diag_idx[1]]
         self.ssm[idx] = error_sigma**2

pymc_extras/statespace/models/structural/components/regression.py

@@ -31,6 +31,11 @@ class RegressionComponent(Component):
         Whether to include stochastic innovations in the regression coefficients,
         allowing them to vary over time. If True, coefficients follow a random walk.
 
+    share_states: bool, default False
+        Whether latent states are shared across the observed states. If True, there will be only one set of latent
+        states, which are observed by all observed states. If False, each observed state has its own set of
+        latent states.
+
     Notes
     -----
     This component implements regression with exogenous variables in a structural time series
@@ -107,7 +112,10 @@ class RegressionComponent(Component):
         state_names: list[str] | None = None,
         observed_state_names: list[str] | None = None,
         innovations=False,
+        share_states: bool = False,
     ):
+        self.share_states = share_states
+
         if observed_state_names is None:
             observed_state_names = ["data"]
 
@@ -121,9 +129,10 @@ class RegressionComponent(Component):
         super().__init__(
             name=name,
             k_endog=k_endog,
-            k_states=k_states * k_endog,
-            k_posdef=k_posdef * k_endog,
+            k_states=k_states * k_endog if not share_states else k_states,
+            k_posdef=k_posdef * k_endog if not share_states else k_posdef,
             state_names=self.state_names,
+            share_states=share_states,
             observed_state_names=observed_state_names,
             measurement_error=False,
             combine_hidden_states=False,
@@ -153,10 +162,12 @@ class RegressionComponent(Component):
 
     def make_symbolic_graph(self) -> None:
         k_endog = self.k_endog
-        k_states = self.k_states // k_endog
+        k_endog_effective = 1 if self.share_states else k_endog
+
+        k_states = self.k_states // k_endog_effective
 
         betas = self.make_and_register_variable(
-            f"beta_{self.name}", shape=(k_endog, k_states) if k_endog > 1 else (k_states,)
+            f"beta_{self.name}", shape=(k_endog, k_states) if k_endog_effective > 1 else (k_states,)
         )
         regression_data = self.make_and_register_data(f"data_{self.name}", shape=(None, k_states))
 
@@ -164,43 +175,61 @@ class RegressionComponent(Component):
         self.ssm["transition", :, :] = pt.eye(self.k_states)
         self.ssm["selection", :, :] = pt.eye(self.k_states)
 
-        Z = pt.linalg.block_diag(*[pt.expand_dims(regression_data, 1) for _ in range(k_endog)])
-        self.ssm["design"] = pt.specify_shape(
-            Z, (None, k_endog, regression_data.type.shape[1] * k_endog)
-        )
+        if self.share_states:
+            self.ssm["design"] = pt.specify_shape(
+                pt.join(1, *[pt.expand_dims(regression_data, 1) for _ in range(k_endog)]),
+                (None, k_endog, self.k_states),
+            )
+        else:
+            Z = pt.linalg.block_diag(*[pt.expand_dims(regression_data, 1) for _ in range(k_endog)])
+            self.ssm["design"] = pt.specify_shape(
+                Z, (None, k_endog, regression_data.type.shape[1] * k_endog)
+            )
 
         if self.innovations:
             sigma_beta = self.make_and_register_variable(
-                f"sigma_beta_{self.name}", (k_states,) if k_endog == 1 else (k_endog, k_states)
+                f"sigma_beta_{self.name}",
+                (k_states,) if k_endog_effective == 1 else (k_endog, k_states),
             )
             row_idx, col_idx = np.diag_indices(self.k_states)
             self.ssm["state_cov", row_idx, col_idx] = sigma_beta.ravel() ** 2
 
     def populate_component_properties(self) -> None:
         k_endog = self.k_endog
-        k_states = self.k_states // k_endog
+        k_endog_effective = 1 if self.share_states else k_endog
+
+        k_states = self.k_states // k_endog_effective
 
-        self.shock_names = self.state_names
+        if self.share_states:
+            self.shock_names = [f"{state_name}_shared" for state_name in self.state_names]
+        else:
+            self.shock_names = self.state_names
 
         self.param_names = [f"beta_{self.name}"]
         self.data_names = [f"data_{self.name}"]
         self.param_dims = {
             f"beta_{self.name}": (f"endog_{self.name}", f"state_{self.name}")
-            if k_endog > 1
+            if k_endog_effective > 1
             else (f"state_{self.name}",)
         }
 
         base_names = self.state_names
-        self.state_names = [
-            f"{name}[{obs_name}]" for obs_name in self.observed_state_names for name in base_names
-        ]
+
+        if self.share_states:
+            self.state_names = [f"{name}[{self.name}_shared]" for name in base_names]
+        else:
+            self.state_names = [
+                f"{name}[{obs_name}]"
+                for obs_name in self.observed_state_names
+                for name in base_names
+            ]
 
         self.param_info = {
             f"beta_{self.name}": {
-                "shape": (k_endog, k_states) if k_endog > 1 else (k_states,),
+                "shape": (k_endog_effective, k_states) if k_endog_effective > 1 else (k_states,),
                 "constraints": None,
                 "dims": (f"endog_{self.name}", f"state_{self.name}")
-                if k_endog > 1
+                if k_endog_effective > 1
                 else (f"state_{self.name}",),
             },
         }
@@ -223,6 +252,6 @@ class RegressionComponent(Component):
                 "shape": (k_states,),
                 "constraints": "Positive",
                 "dims": (f"state_{self.name}",)
-                if k_endog == 1
+                if k_endog_effective == 1
                 else (f"endog_{self.name}", f"state_{self.name}"),
             }