Allow multiple observed states in measurement error component

Author: jessegrabowski

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

@@ -64,16 +64,22 @@ def __init__(
     def populate_component_properties(self):
         self.param_names = [f"sigma_{self.name}"]
         self.param_dims = {}
+        self.coords = {}
+
+        if self.k_endog > 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": (),
+                "shape": (self.k_endog,) if self.k_endog > 1 else (),
                 "constraints": "Positive",
-                "dims": None,
+                "dims": (f"endog_{self.name}",) if self.k_endog > 1 else None,
             }
         }
 
     def make_symbolic_graph(self) -> None:
-        sigma_shape = ()
+        sigma_shape = () if self.k_endog == 1 else (self.k_endog,)
         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]]

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

@@ -1,3 +1,5 @@
+import numpy as np
+
 from pymc_extras.statespace.models import structural as st
 from tests.statespace.models.structural.conftest import _assert_basic_coords_correct
 
@@ -8,3 +10,23 @@ def test_measurement_error(rng):
 
     _assert_basic_coords_correct(mod)
     assert "sigma_obs" in mod.param_names
+
+
+def test_measurement_error_multiple_observed():
+    mod = st.MeasurementError("obs", observed_state_names=["data_1", "data_2"])
+    assert mod.k_endog == 2
+    assert mod.coords["endog_obs"] == ["data_1", "data_2"]
+    assert mod.param_dims["sigma_obs"] == ("endog_obs",)
+
+
+def test_build_with_measurement_error_subset():
+    ll = st.LevelTrendComponent(order=2, observed_state_names=["data_1", "data_2", "data_3"])
+    me = st.MeasurementError("obs", observed_state_names=["data_1", "data_3"])
+    mod = (ll + me).build()
+
+    H = mod.ssm["obs_cov"]
+    assert H.type.shape == (3, 3)
+    np.testing.assert_allclose(
+        H.eval({"sigma_obs": [1.0, 3.0]}),
+        np.array([[1.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 9.0]]),
+    )

tests/statespace/models/structural/test_against_statsmodels.py

@@ -416,8 +416,8 @@ def create_structural_model_and_equivalent_statsmodel(
         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] += ["ar_innovation"]
-        expected_coords[SHOCK_AUX_DIM] += ["ar_innovation"]
+        expected_coords[SHOCK_DIM] += ["data_ar_innovation"]
+        expected_coords[SHOCK_AUX_DIM] += ["data_ar_innovation"]
 
         sm_params["sigma2.ar"] = sigma2
         for i, rho in enumerate(ar_params):

tests/statespace/models/structural/test_core.py

@@ -64,6 +64,16 @@ def test_add_components():
         assert_allclose(all_mat, np.concatenate([ll_mat, se_mat], axis=axis), atol=ATOL, rtol=RTOL)
 
 
+def test_add_components_multiple_observed():
+    ll = st.LevelTrendComponent(order=2, observed_state_names=["data_1", "data_2"])
+    me = st.MeasurementError(name="obs", observed_state_names=["data_1", "data_2"])
+
+    mod = (ll + me).build()
+
+    for property in ["param_names", "shock_names", "param_info", "coords", "param_dims"]:
+        assert [x in getattr(mod, property) for x in getattr(ll, property)]
+
+
 @pytest.mark.skipif(floatX.endswith("32"), reason="Prior covariance not PSD at half-precision")
 def test_extract_components_from_idata(rng):
     time_idx = pd.date_range(start="2000-01-01", freq="D", periods=100)