LevelTrend state/shock names depend on component name

Author: jessegrabowski

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
@@ -114,7 +113,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,35 +165,46 @@ def populate_component_properties(self):
         k_posdef = self.k_posdef // k_endog
 
         name_slice = POSITION_DERIVATIVE_NAMES[:k_states]
-        self.param_names = ["initial_trend"]
+        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
+            f"{name}[{obs_name}]" if k_endog > 1 else name
+            for obs_name in self.observed_state_names
+            for name in base_names
         ]
-        self.param_dims = {"initial_trend": ("trend_state",)}
-        self.coords = {"trend_state": base_names}
+        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["trend_state"] = (
-                "trend_endog",
-                "trend_state",
+            self.param_dims[f"{self.name}_state"] = (
+                f"{self.name}_endog",
+                f"{self.name}_state",
             )
-            self.param_dims = {"initial_trend": ("trend_endog", "trend_state")}
-            self.coords["trend_endog"] = self.observed_state_names
+            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 = {"initial_trend": {"shape": shape, "constraints": None}}
+        self.param_info = {f"{self.name}_initial": {"shape": shape, "constraints": None}}
 
         if self.k_posdef > 0:
-            self.param_names += ["sigma_trend"]
-            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["sigma_trend"] = (
-                ("trend_shock",) if k_endog == 1 else ("trend_endog", "trend_shock")
+            self.shock_names = [
+                f"{name}[{obs_name}]" if k_endog > 1 else 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["trend_shock"] = self.shock_names
-            self.param_info["sigma_trend"] = {
+            self.coords[f"{self.name}_shock"] = self.shock_names
+            self.param_info[f"{self.name}_sigma"] = {
                 "shape": (k_posdef,) if k_endog == 1 else (k_endog, k_posdef),
                 "constraints": "Positive",
             }
@@ -208,28 +218,34 @@ def make_symbolic_graph(self) -> None:
         k_posdef = self.k_posdef // k_endog
 
         initial_trend = self.make_and_register_variable(
-            "initial_trend",
+            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(
-                "sigma_trend",
+                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)

tests/statespace/models/structural/components/test_level_trend.py

@@ -1,4 +1,5 @@
 import numpy as np
+import pytensor
 
 from numpy.testing import assert_allclose
 from pytensor import config
@@ -13,15 +14,15 @@
 
 def test_level_trend_model(rng):
     mod = st.LevelTrendComponent(order=2, innovations_order=0)
-    params = {"initial_trend": [0.0, 1.0]}
+    params = {"level_trend_initial": [0.0, 1.0]}
     x, y = simulate_from_numpy_model(mod, rng, params)
 
     assert_allclose(np.diff(y), 1, atol=ATOL, rtol=RTOL)
 
     # Check coords
     mod = mod.build(verbose=False)
     _assert_basic_coords_correct(mod)
-    assert mod.coords["trend_state"] == ["level", "trend"]
+    assert mod.coords["level_trend_state"] == ["level", "trend"]
 
 
 def test_level_trend_multiple_observed_construction():
@@ -33,12 +34,22 @@ def test_level_trend_multiple_observed_construction():
     assert mod.k_states == 6
     assert mod.k_posdef == 3
 
-    assert mod.coords["trend_state"] == ["level", "trend"]
-    assert mod.coords["trend_endog"] == ["data_1", "data_2", "data_3"]
+    assert mod.coords["level_trend_state"] == ["level", "trend"]
+    assert mod.coords["level_trend_endog"] == ["data_1", "data_2", "data_3"]
 
-    Z = mod.ssm["design"].eval()
-    T = mod.ssm["transition"].eval()
-    R = mod.ssm["selection"].eval()
+    assert mod.state_names == [
+        "level[data_1]",
+        "trend[data_1]",
+        "level[data_2]",
+        "trend[data_2]",
+        "level[data_3]",
+        "trend[data_3]",
+    ]
+    assert mod.shock_names == ["level_shock[data_1]", "level_shock[data_2]", "level_shock[data_3]"]
+
+    Z, T, R = pytensor.function(
+        [], [mod.ssm["design"], mod.ssm["transition"], mod.ssm["selection"]], mode="FAST_COMPILE"
+    )()
 
     np.testing.assert_allclose(
         Z,
@@ -84,8 +95,64 @@ def test_level_trend_multiple_observed(rng):
     mod = st.LevelTrendComponent(
         order=2, innovations_order=0, observed_state_names=["data_1", "data_2", "data_3"]
     )
-    params = {"initial_trend": np.array([[0.0, 1.0], [0.0, 2.0], [0.0, 3.0]])}
+    params = {"level_trend_initial": np.array([[0.0, 1.0], [0.0, 2.0], [0.0, 3.0]])}
 
     x, y = simulate_from_numpy_model(mod, rng, params)
     assert (np.diff(y, axis=0) == np.array([[1.0, 2.0, 3.0]])).all().all()
     assert (np.diff(x, axis=0) == np.array([[1.0, 0.0, 2.0, 0.0, 3.0, 0.0]])).all().all()
+
+
+def test_add_level_trend_with_different_observed():
+    mod_1 = st.LevelTrendComponent(
+        name="ll", order=2, innovations_order=[0, 1], observed_state_names=["data_1"]
+    )
+    mod_2 = st.LevelTrendComponent(
+        name="grw", order=1, innovations_order=[1], observed_state_names=["data_2"]
+    )
+
+    mod = (mod_1 + mod_2).build(verbose=False)
+    assert mod.k_endog == 2
+    assert mod.k_states == 3
+    assert mod.k_posdef == 2
+
+    assert mod.coords["ll_state"] == ["level", "trend"]
+    assert mod.coords["grw_state"] == ["level"]
+
+    assert mod.state_names == ["level[data_1]", "trend[data_1]", "level[data_2]"]
+    assert mod.shock_names == ["trend_shock[data_1]", "level_shock[data_2]"]
+
+    Z, T, R = pytensor.function(
+        [], [mod.ssm["design"], mod.ssm["transition"], mod.ssm["selection"]], mode="FAST_COMPILE"
+    )()
+
+    np.testing.assert_allclose(
+        Z,
+        np.array(
+            [
+                [1.0, 0.0, 0.0],
+                [0.0, 0.0, 1.0],
+            ]
+        ),
+    )
+
+    np.testing.assert_allclose(
+        T,
+        np.array(
+            [
+                [1.0, 1.0, 0.0],
+                [0.0, 1.0, 0.0],
+                [0.0, 0.0, 1.0],
+            ]
+        ),
+    )
+
+    np.testing.assert_allclose(
+        R,
+        np.array(
+            [
+                [0.0, 0.0],
+                [1.0, 0.0],
+                [0.0, 1.0],
+            ]
+        ),
+    )

tests/statespace/models/structural/test_against_statsmodels.py

@@ -220,7 +220,7 @@ def create_structural_model_and_equivalent_statsmodel(
 
     if level:
         level_trend_order[0] = 1
-        expected_coords["trend_state"] += [
+        expected_coords["level_state"] += [
             "level",
         ]
         expected_coords[ALL_STATE_DIM] += [
@@ -231,7 +231,7 @@ def create_structural_model_and_equivalent_statsmodel(
         ]
         if stochastic_level:
             level_trend_innov_order[0] = 1
-            expected_coords["trend_shock"] += ["level"]
+            expected_coords["level_shock"] += ["level"]
             expected_coords[SHOCK_DIM] += [
                 "level",
             ]
@@ -241,7 +241,7 @@ def create_structural_model_and_equivalent_statsmodel(
 
     if trend:
         level_trend_order[1] = 1
-        expected_coords["trend_state"] += [
+        expected_coords["level_state"] += [
             "trend",
         ]
         expected_coords[ALL_STATE_DIM] += [
@@ -253,12 +253,12 @@ def create_structural_model_and_equivalent_statsmodel(
 
         if stochastic_trend:
             level_trend_innov_order[1] = 1
-            expected_coords["trend_shock"] += ["trend"]
+            expected_coords["level_shock"] += ["trend"]
             expected_coords[SHOCK_DIM] += ["trend"]
             expected_coords[SHOCK_AUX_DIM] += ["trend"]
 
     if level or trend:
-        expected_param_dims["initial_trend"] += ("trend_state",)
+        expected_param_dims["level_initial"] += ("level_state",)
         level_value = np.where(
             level_trend_order,
             rng.normal(
@@ -272,13 +272,13 @@ def create_structural_model_and_equivalent_statsmodel(
         max_order = np.flatnonzero(level_value)[-1].item() + 1
         level_trend_order = level_trend_order[:max_order]
 
-        params["initial_trend"] = level_value[:max_order]
+        params["level_initial"] = level_value[:max_order]
         sm_init["level"] = level_value[0]
         sm_init["trend"] = level_value[1]
 
         if sum(level_trend_innov_order) > 0:
-            expected_param_dims["sigma_trend"] += ("trend_shock",)
-            params["sigma_trend"] = np.sqrt(sigma_level_value2)
+            expected_param_dims["level_sigma"] += ("level_shock",)
+            params["level_sigma"] = np.sqrt(sigma_level_value2)
 
         sigma_level_value = sigma_level_value2.tolist()
         if stochastic_level: