Add shared_state argument to MeasurementError

Author: jessegrabowski

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
@@ -113,25 +122,32 @@ def __init__(
         )
 
     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

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

@@ -1,4 +1,7 @@
 import numpy as np
+import pytensor
+
+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
@@ -19,6 +22,45 @@ def test_measurement_error_multiple_observed():
     assert mod.param_dims["sigma_obs"] == ("endog_obs",)
 
 
+def test_measurement_error_share_states():
+    mod = st.MeasurementError("obs", observed_state_names=["data_1", "data_2"], share_states=True)
+    mod.build(verbose=False)
+
+    assert mod.k_endog == 2
+    assert mod.param_names == ["sigma_obs", "P0"]
+    assert "endog_obs" not in mod.coords
+
+    # Check that the parameter is shared across the observed states
+    assert mod.param_info["sigma_obs"]["shape"] == ()
+
+    outputs = mod.ssm["obs_cov"]
+
+    H = pytensor.function(list(explicit_graph_inputs([outputs])), outputs)(sigma_obs=np.array(0.5))
+    np.testing.assert_allclose(H, np.diag([0.5, 0.5]) ** 2)
+
+
+def test_measurement_error_shared_and_not_shared():
+    shared = st.MeasurementError(
+        "error_shared", observed_state_names=["data_1", "data_2"], share_states=True
+    )
+    individual = st.MeasurementError("error_individual", observed_state_names=["data_1", "data_2"])
+    mod = (shared + individual).build(verbose=False)
+
+    assert mod.k_endog == 2
+    assert mod.param_names == ["sigma_error_shared", "sigma_error_individual", "P0"]
+    assert mod.coords["endog_error_individual"] == ["data_1", "data_2"]
+
+    assert mod.param_info["sigma_error_shared"]["shape"] == ()
+    assert mod.param_info["sigma_error_individual"]["shape"] == (2,)
+
+    outputs = mod.ssm["obs_cov"]
+
+    H = pytensor.function(list(explicit_graph_inputs([outputs])), outputs)(
+        sigma_error_shared=np.array(0.5), sigma_error_individual=np.array([0.1, 0.9])
+    )
+    np.testing.assert_allclose(H, np.diag([0.5, 0.5]) ** 2 + np.diag([0.1, 0.9]) ** 2)
+
+
 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"])