merged resolved conflicts

Author: Dekermanjian

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

@@ -65,7 +65,7 @@ class AutoregressiveComponent(Component):
     def __init__(
         self,
         order: int = 1,
-        name: str = "AutoRegressive",
+        name: str = "auto_regressive",
         observed_state_names: list[str] | None = None,
     ):
         if observed_state_names is None:
@@ -92,27 +92,30 @@ def __init__(
         )
 
     def populate_component_properties(self):
+        k_states = self.k_states // self.k_endog
+
         self.state_names = [
-            f"L{i + 1}.{state_name}"
-            for i in range(self.k_states)
+            f"L{i + 1}[{state_name}]"
             for state_name in self.observed_state_names
+            for i in range(k_states)
         ]
-        self.shock_names = [f"{name}_{self.name}_innovation" for name in self.observed_state_names]
-        self.param_names = ["ar_params", "sigma_ar"]
-        self.param_dims = {"ar_params": (AR_PARAM_DIM,)}
-        self.coords = {AR_PARAM_DIM: self.ar_lags.tolist()}
+
+        self.shock_names = [f"{self.name}[{obs_name}]" for obs_name in self.observed_state_names]
+        self.param_names = [f"{self.name}_params", f"{self.name}_sigma"]
+        self.param_dims = {f"{self.name}_params": (f"{self.name}_lag",)}
+        self.coords = {f"{self.name}_lag": self.ar_lags.tolist()}
 
         if self.k_endog > 1:
-            self.param_dims["ar_params"] = (
+            self.param_dims[f"{self.name}_params"] = (
                 f"{self.name}_endog",
                 AR_PARAM_DIM,
             )
-            self.param_dims["sigma_ar"] = (f"{self.name}_endog",)
+            self.param_dims[f"{self.name}_sigma"] = (f"{self.name}_endog",)
 
             self.coords[f"{self.name}_endog"] = self.observed_state_names
 
         self.param_info = {
-            "ar_params": {
+            f"{self.name}_params": {
                 "shape": (self.k_states,) if self.k_endog == 1 else (self.k_endog, self.k_states),
                 "constraints": None,
                 "dims": (AR_PARAM_DIM,)
@@ -122,7 +125,7 @@ def populate_component_properties(self):
                     AR_PARAM_DIM,
                 ),
             },
-            "sigma_ar": {
+            f"{self.name}_sigma": {
                 "shape": () if self.k_endog == 1 else (self.k_endog,),
                 "constraints": "Positive",
                 "dims": None if self.k_endog == 1 else (f"{self.name}_endog",),
@@ -136,10 +139,10 @@ def make_symbolic_graph(self) -> None:
 
         k_nonzero = int(sum(self.order))
         ar_params = self.make_and_register_variable(
-            "ar_params", shape=(k_nonzero,) if k_endog == 1 else (k_endog, k_nonzero)
+            f"{self.name}_params", shape=(k_nonzero,) if k_endog == 1 else (k_endog, k_nonzero)
         )
         sigma_ar = self.make_and_register_variable(
-            "sigma_ar", shape=() if k_endog == 1 else (k_endog,)
+            f"{self.name}_sigma", shape=() if k_endog == 1 else (k_endog,)
         )
 
         if k_endog == 1:

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

@@ -1,6 +1,7 @@
 import numpy as np
 
 from pytensor import tensor as pt
+from pytensor.tensor.slinalg import block_diag
 from scipy import linalg
 
 from pymc_extras.statespace.models.structural.core import Component
@@ -96,7 +97,6 @@ class CycleComponent(Component):
 
             cycle_strength = pm.Normal("business_cycle", dims=ss_mod.param_dims["business_cycle"])
             cycle_length = pm.Uniform('business_cycle_length', lower=6, upper=12)
-
             sigma_cycle = pm.HalfNormal('sigma_business_cycle', sigma=1)
 
             ss_mod.build_statespace_graph(data)
@@ -124,13 +124,15 @@ class CycleComponent(Component):
         with pm.Model(coords=ss_mod.coords) as model:
             P0 = pm.Deterministic("P0", pt.eye(ss_mod.k_states), dims=ss_mod.param_dims["P0"])
             # Initial states: shape (3, 2) for 3 variables, 2 states each
-            cycle_init = pm.Normal('business_cycle', dims=('business_cycle_endog', 'business_cycle_state'))
+            cycle_init = pm.Normal('business_cycle', dims=ss_mod.param_dims["business_cycle"])
 
             # Dampening factor: scalar (shared across variables)
-            dampening = pm.Uniform('business_cycle_dampening_factor', lower=0.8, upper=1.0)
+            dampening = pm.Beta("business_cycle_dampening_factor", 2, 2)
 
             # Innovation variances: shape (3,) for variable-specific variances
-            sigma_cycle = pm.HalfNormal('sigma_business_cycle', dims=('business_cycle_endog',))
+            sigma_cycle = pm.HalfNormal(
+                "sigma_business_cycle", dims=ss_mod.param_dims["sigma_business_cycle"]
+            )
 
             ss_mod.build_statespace_graph(data)
             idata = pm.sample()
@@ -220,12 +222,8 @@ def make_symbolic_graph(self) -> None:
         if self.k_endog == 1:
             self.ssm["transition", :, :] = T
         else:
-            # can't make the linalg.block_diag logic work here
-            # doing it manually for now
-            for i in range(self.k_endog):
-                start_idx = i * 2
-                end_idx = (i + 1) * 2
-                self.ssm["transition", start_idx:end_idx, start_idx:end_idx] = T
+            transition = block_diag(*[T for _ in range(self.k_endog)])
+            self.ssm["transition"] = pt.specify_shape(transition, (self.k_states, self.k_states))
 
         if self.innovations:
             if self.k_endog == 1:
@@ -235,16 +233,20 @@ def make_symbolic_graph(self) -> None:
                 sigma_cycle = self.make_and_register_variable(
                     f"sigma_{self.name}", shape=(self.k_endog,)
                 )
-                # can't make the linalg.block_diag logic work here
-                # doing it manually for now
-                for i in range(self.k_endog):
-                    start_idx = i * 2
-                    end_idx = (i + 1) * 2
-                    Q_block = pt.eye(2) * sigma_cycle[i] ** 2
-                    self.ssm["state_cov", start_idx:end_idx, start_idx:end_idx] = Q_block
+                state_cov = block_diag(
+                    *[pt.eye(2) * sigma_cycle[i] ** 2 for i in range(self.k_endog)]
+                )
+                self.ssm["state_cov"] = pt.specify_shape(state_cov, (self.k_states, self.k_states))
 
     def populate_component_properties(self):
-        self.state_names = [f"{self.name}_{f}" for f in ["Cos", "Sin"]]
+        if self.k_endog == 1:
+            self.state_names = [f"{self.name}_{f}" for f in ["Cos", "Sin"]]
+        else:
+            # For multivariate cycles, create state names for each observed state
+            self.state_names = []
+            for var_name in self.observed_state_names:
+                self.state_names.extend([f"{self.name}_{var_name}_{f}" for f in ["Cos", "Sin"]])
+
         self.param_names = [f"{self.name}"]
 
         if self.k_endog == 1:
@@ -260,7 +262,7 @@ def populate_component_properties(self):
         else:
             self.param_dims = {self.name: (f"{self.name}_endog", f"{self.name}_state")}
             self.coords = {
-                f"{self.name}_state": self.state_names,
+                f"{self.name}_state": [f"{self.name}_Cos", f"{self.name}_Sin"],
                 f"{self.name}_endog": self.observed_state_names,
             }
             self.param_info = {

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

@@ -1,6 +1,5 @@
 import numpy as np
-
-from scipy import linalg
+import pytensor.tensor as pt
 
 from pymc_extras.statespace.models.structural.core import Component
 from pymc_extras.statespace.models.structural.utils import order_to_mask
@@ -13,7 +12,6 @@ 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
@@ -114,7 +112,7 @@ def __init__(
         self,
         order: int | list[int] = 2,
         innovations_order: int | list[int] | None = None,
-        name: str = "LevelTrend",
+        name: str = "level_trend",
         observed_state_names: list[str] | None = None,
     ):
         if innovations_order is None:
@@ -166,37 +164,44 @@ def populate_component_properties(self):
         k_posdef = self.k_posdef // k_endog
 
         name_slice = POSITION_DERIVATIVE_NAMES[:k_states]
-        self.param_names = [f"initial_{self.name}"]
+        self.param_names = [f"{self.name}_initial"]
         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
         ]
-        self.param_dims = {f"initial_{self.name}": (f"{self.name}_state",)}
+        self.param_dims = {f"{self.name}_initial": (f"{self.name}_state",)}
         self.coords = {f"{self.name}_state": base_names}
 
         if k_endog > 1:
             self.param_dims[f"{self.name}_state"] = (
                 f"{self.name}_endog",
                 f"{self.name}_state",
             )
-            self.param_dims = {f"initial_{self.name}": (f"{self.name}_endog", f"{self.name}_state")}
+            self.param_dims = {f"{self.name}_initial": (f"{self.name}_endog", f"{self.name}_state")}
             self.coords[f"{self.name}_endog"] = self.observed_state_names
 
         shape = (k_endog, k_states) if k_endog > 1 else (k_states,)
-        self.param_info = {f"initial_{self.name}": {"shape": shape, "constraints": None}}
+        self.param_info = {f"{self.name}_initial": {"shape": shape, "constraints": None}}
 
         if self.k_posdef > 0:
-            self.param_names += [f"sigma_{self.name}"]
-            self.shock_names = [
+            self.param_names += [f"{self.name}_sigma"]
+
+            shock_base_names = [
                 name for name, mask in zip(name_slice, self.innovations_order) if mask
             ]
-            self.param_dims[f"sigma_{self.name}"] = (
+            self.shock_names = [
+                f"{name}[{obs_name}]"
+                for obs_name in self.observed_state_names
+                for name in shock_base_names
+            ]
+
+            self.param_dims[f"{self.name}_sigma"] = (
                 (f"{self.name}_shock",)
                 if k_endog == 1
                 else (f"{self.name}_endog", f"{self.name}_shock")
             )
             self.coords[f"{self.name}_shock"] = self.shock_names
-            self.param_info[f"sigma_{self.name}"] = {
+            self.param_info[f"{self.name}_sigma"] = {
                 "shape": (k_posdef,) if k_endog == 1 else (k_endog, k_posdef),
                 "constraints": "Positive",
             }
@@ -210,28 +215,34 @@ def make_symbolic_graph(self) -> None:
         k_posdef = self.k_posdef // k_endog
 
         initial_trend = self.make_and_register_variable(
-            f"initial_{self.name}",
+            f"{self.name}_initial",
             shape=(k_states,) if k_endog == 1 else (k_endog, k_states),
         )
         self.ssm["initial_state", :] = initial_trend.ravel()
 
-        triu_idx = np.triu_indices(k_states)
-        T = np.zeros((k_states, k_states))
-        T[triu_idx[0], triu_idx[1]] = 1
+        triu_idx = pt.triu_indices(k_states)
+        T = pt.zeros((k_states, k_states))[triu_idx[0], triu_idx[1]].set(1)
 
-        self.ssm["transition"] = linalg.block_diag(*[T for _ in range(k_endog)])
+        self.ssm["transition", :, :] = pt.specify_shape(
+            pt.linalg.block_diag(*[T for _ in range(k_endog)]), (self.k_states, self.k_states)
+        )
 
         R = np.eye(k_states)
         R = R[:, self.innovations_order]
 
-        self.ssm["selection", :, :] = linalg.block_diag(*[R for _ in range(k_endog)])
+        self.ssm["selection", :, :] = pt.specify_shape(
+            pt.linalg.block_diag(*[R for _ in range(k_endog)]), (self.k_states, self.k_posdef)
+        )
 
         Z = np.array([1.0] + [0.0] * (k_states - 1)).reshape((1, -1))
-        self.ssm["design"] = linalg.block_diag(*[Z for _ in range(k_endog)])
+
+        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}",
+                f"{self.name}_sigma",
                 shape=(k_posdef,) if k_endog == 1 else (k_endog, k_posdef),
             )
             diag_idx = np.diag_indices(k_posdef * k_endog)