Delay dropping data names from states/coords until .build

Author: jessegrabowski

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

@@ -168,9 +168,7 @@ def populate_component_properties(self):
         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}]" if k_endog > 1 else name
-            for obs_name in self.observed_state_names
-            for name in base_names
+            f"{name}[{obs_name}]" for obs_name in self.observed_state_names for name in base_names
         ]
         self.param_dims = {f"{self.name}_initial": (f"{self.name}_state",)}
         self.coords = {f"{self.name}_state": base_names}
@@ -193,7 +191,7 @@ def populate_component_properties(self):
                 name for name, mask in zip(name_slice, self.innovations_order) if mask
             ]
             self.shock_names = [
-                f"{name}[{obs_name}]" if k_endog > 1 else name
+                f"{name}[{obs_name}]"
                 for obs_name in self.observed_state_names
                 for name in shock_base_names
             ]

pymc_extras/statespace/models/structural/core.py

@@ -76,16 +76,19 @@ def __init__(
         param_names, param_dims, param_info = self._add_inital_state_cov_to_properties(
             param_names, param_dims, param_info, k_states
         )
-        self._state_names = state_names.copy()
-        self._data_names = data_names.copy()
-        self._shock_names = shock_names.copy()
-        self._param_names = param_names.copy()
-        self._param_dims = param_dims.copy()
+
+        self._state_names = self._strip_data_names_if_unambiguous(state_names, k_endog)
+        self._data_names = self._strip_data_names_if_unambiguous(data_names, k_endog)
+        self._shock_names = self._strip_data_names_if_unambiguous(shock_names, k_endog)
+        self._param_names = self._strip_data_names_if_unambiguous(param_names, k_endog)
+        self._param_dims = param_dims
 
         default_coords = make_default_coords(self)
         coords.update(default_coords)
 
-        self._coords = coords
+        self._coords = {
+            k: self._strip_data_names_if_unambiguous(v, k_endog) for k, v in coords.items()
+        }
         self._param_info = param_info.copy()
         self._data_info = data_info.copy()
         self.measurement_error = measurement_error
@@ -122,6 +125,25 @@ def __init__(
         P0 = self.make_and_register_variable("P0", shape=(self.k_states, self.k_states))
         self.ssm["initial_state_cov"] = P0
 
+    def _strip_data_names_if_unambiguous(self, names: list[str], k_endog: int):
+        """
+        State names from components should always be of the form name[data_name], in the case that the component is
+        associated with multiple observed states. Not doing so leads to ambiguity -- we might have two level states,
+        but which goes to which observed component? So we set `level[data_1]` and `level[data_2]`.
+
+        In cases where there is only one observed state (when k_endog == 1), we can strip the data part and just use
+        the state name. This is a bit cleaner.
+        """
+        if k_endog == 1:
+            [data_name] = self.observed_states
+            return [
+                name.replace(f"[{data_name}]", "") if isinstance(name, str) else name
+                for name in names
+            ]
+
+        else:
+            return names
+
     @staticmethod
     def _add_inital_state_cov_to_properties(param_names, param_dims, param_info, k_states):
         param_names += ["P0"]

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

@@ -119,7 +119,7 @@ def test_add_level_trend_with_different_observed():
     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]"]
+    assert mod.shock_names == ["trend[data_1]", "level[data_2]"]
 
     Z, T, R = pytensor.function(
         [], [mod.ssm["design"], mod.ssm["transition"], mod.ssm["selection"]], mode="FAST_COMPILE"