Update autoregressive component and tests

jessegrabowski · jessegrabowski · commit a6327b77b78d · 2025-07-06T18:16:30.000+08:00
diff --git a/pymc_extras/statespace/models/structural/components/autoregressive.py b/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 = self.observed_state_names.copy()
+        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:
diff --git a/tests/statespace/models/structural/components/test_autoregressive.py b/tests/statespace/models/structural/components/test_autoregressive.py
@@ -1,8 +1,10 @@
 import numpy as np
+import pytensor
 import pytest
 
 from numpy.testing import assert_allclose
 from pytensor import config
+from pytensor.graph.basic import explicit_graph_inputs
 
 from pymc_extras.statespace.models import structural as st
 from tests.statespace.models.structural.conftest import _assert_basic_coords_correct
@@ -11,37 +13,126 @@
 
 @pytest.mark.parametrize("order", [1, 2, [1, 0, 1]], ids=["AR1", "AR2", "AR(1,0,1)"])
 def test_autoregressive_model(order, rng):
-    ar = st.AutoregressiveComponent(order=order)
+    k = sum(order) if isinstance(order, list) else order
+    ar = st.AutoregressiveComponent(order=order).build(verbose=False)
     params = {
-        "ar_params": np.full((sum(ar.order),), 0.5, dtype=config.floatX),
-        "sigma_ar": 0.0,
+        "auto_regressive_params": np.full((k,), 0.5, dtype=config.floatX),
+        "auto_regressive_sigma": 0.1,
+        "initial_state_cov": np.eye(k),
     }
 
-    x, y = simulate_from_numpy_model(ar, rng, params, steps=100)
-
     # Check coords
-    ar.build(verbose=False)
     _assert_basic_coords_correct(ar)
+
     lags = np.arange(len(order) if isinstance(order, list) else order, dtype="int") + 1
     if isinstance(order, list):
         lags = lags[np.flatnonzero(order)]
-    assert_allclose(ar.coords["ar_lag"], lags)
+    assert_allclose(ar.coords["auto_regressive_lag"], lags)
 
 
-def test_autoregressive_multiple_observed(rng):
+def test_autoregressive_multiple_observed_build(rng):
     ar = st.AutoregressiveComponent(order=3, observed_state_names=["data_1", "data_2"])
     mod = ar.build(verbose=False)
 
+    assert mod.k_endog == 2
+    assert mod.k_states == 6
+    assert mod.k_posdef == 2
+
+    assert mod.state_names == [
+        "L1[data_1]",
+        "L2[data_1]",
+        "L3[data_1]",
+        "L1[data_2]",
+        "L2[data_2]",
+        "L3[data_2]",
+    ]
+
+    assert mod.shock_names == ["data_1", "data_2"]
+
     params = {
-        "ar_params": np.full(
+        "auto_regressive_params": np.full(
             (
                 2,
                 sum(ar.order),
             ),
             0.5,
             dtype=config.floatX,
         ),
-        "sigma_ar": np.ones((2,)) * 1e-3,
+        "auto_regressive_sigma": np.array([0.05, 0.12]),
     }
+    _, _, _, _, T, Z, R, _, Q = mod._unpack_statespace_with_placeholders()
+    input_vars = explicit_graph_inputs([T, Z, R, Q])
+    fn = pytensor.function(
+        inputs=list(input_vars),
+        outputs=[T, Z, R, Q],
+        mode="FAST_COMPILE",
+    )
+
+    T, Z, R, Q = fn(**params)
+
+    np.testing.assert_allclose(
+        T,
+        np.array(
+            [
+                [0.5, 0.5, 0.5, 0.0, 0.0, 0.0],
+                [1.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+                [0.0, 1.0, 0.0, 0.0, 0.0, 0.0],
+                [0.0, 0.0, 0.0, 0.5, 0.5, 0.5],
+                [0.0, 0.0, 0.0, 1.0, 0.0, 0.0],
+                [0.0, 0.0, 0.0, 0.0, 1.0, 0.0],
+            ]
+        ),
+    )
+
+    np.testing.assert_allclose(
+        Z, np.array([[1.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 1.0, 0.0, 0.0]])
+    )
+
+    np.testing.assert_allclose(
+        R, np.array([[1.0, 0.0], [0.0, 0.0], [0.0, 0.0], [0.0, 1.0], [0.0, 0.0], [0.0, 0.0]])
+    )
+
+    np.testing.assert_allclose(Q, np.diag([0.05**2, 0.12**2]))
+
+
+def test_autoregressive_multiple_observed_data(rng):
+    ar = st.AutoregressiveComponent(order=1, observed_state_names=["data_1", "data_2", "data_3"])
+    mod = ar.build(verbose=False)
+
+    params = {
+        "auto_regressive_params": np.array([0.9, 0.8, 0.5]).reshape((3, 1)),
+        "auto_regressive_sigma": np.array([0.05, 0.12, 0.22]),
+        "initial_state_cov": np.eye(3),
+    }
+
+    # Recover the AR(1) coefficients from the simulated data via OLS
+    x, y = simulate_from_numpy_model(mod, rng, params, steps=2000)
+    for i in range(3):
+        ols_coefs = np.polyfit(y[:-1, i], y[1:, i], 1)
+        np.testing.assert_allclose(ols_coefs[0], params["auto_regressive_params"][i, 0], atol=1e-1)
+
+
+def test_add_autoregressive_different_observed():
+    mod_1 = st.AutoregressiveComponent(order=1, name="ar1", observed_state_names=["data_1"])
+    mod_2 = st.AutoregressiveComponent(name="ar6", order=6, observed_state_names=["data_2"])
+
+    mod = (mod_1 + mod_2).build(verbose=False)
+
+    print(mod.coords)
+
+    assert mod.k_endog == 2
+    assert mod.k_states == 7
+    assert mod.k_posdef == 2
+    assert mod.state_names == [
+        "L1[data_1]",
+        "L1[data_2]",
+        "L2[data_2]",
+        "L3[data_2]",
+        "L4[data_2]",
+        "L5[data_2]",
+        "L6[data_2]",
+    ]
 
-    x, y = simulate_from_numpy_model(ar, rng, params, steps=100)
+    assert mod.shock_names == ["data_1", "data_2"]
+    assert mod.coords["ar1_lag"] == [1]
+    assert mod.coords["ar6_lag"] == [1, 2, 3, 4, 5, 6]
diff --git a/tests/statespace/models/structural/conftest.py b/tests/statespace/models/structural/conftest.py
@@ -24,5 +24,6 @@ def _assert_basic_coords_correct(mod):
     assert mod.coords[SHOCK_DIM] == mod.shock_names
     assert mod.coords[SHOCK_AUX_DIM] == mod.shock_names
     expected_obs = mod.observed_state_names if hasattr(mod, "observed_state_names") else ["data"]
+
     assert mod.coords[OBS_STATE_DIM] == expected_obs
     assert mod.coords[OBS_STATE_AUX_DIM] == expected_obs
diff --git a/tests/statespace/models/structural/test_against_statsmodels.py b/tests/statespace/models/structural/test_against_statsmodels.py
@@ -404,20 +404,20 @@ def create_structural_model_and_equivalent_statsmodel(
         components.append(comp)
 
     if autoregressive is not None:
-        ar_names = [f"L{i+1}.data" for i in range(autoregressive)]
+        ar_names = [f"L{i+1}" for i in range(autoregressive)]
         ar_params = rng.normal(size=(autoregressive,)).astype(floatX)
         if autoregressive == 1:
             ar_params = ar_params.item()
         sigma2 = np.abs(rng.normal()).astype(floatX)
 
         params["ar_params"] = ar_params
-        params["sigma_ar"] = np.sqrt(sigma2)
+        params["ar_sigma"] = np.sqrt(sigma2)
         expected_param_dims["ar_params"] += (AR_PARAM_DIM,)
         expected_coords[AR_PARAM_DIM] += tuple(list(range(1, autoregressive + 1)))
         expected_coords[ALL_STATE_DIM] += ar_names
         expected_coords[ALL_STATE_AUX_DIM] += ar_names
-        expected_coords[SHOCK_DIM] += ["data_ar_innovation"]
-        expected_coords[SHOCK_AUX_DIM] += ["data_ar_innovation"]
+        expected_coords[SHOCK_DIM] += ["data"]
+        expected_coords[SHOCK_AUX_DIM] += ["data"]
 
         sm_params["sigma2.ar"] = sigma2
         for i, rho in enumerate(ar_params):
