Align structural naming registry with prefixed state space names

Author: opherdonchin

pymc_extras/statespace/models/structural/core.py

@@ -117,7 +117,7 @@ class StructuralTimeSeries(PyMCStateSpace):
     def __init__(
         self,
         ssm: PytensorRepresentation,
-        name: str,
+        name: str | None,
         coords_info: CoordInfo,
         param_info: ParameterInfo,
         data_info: DataInfo,
@@ -184,10 +184,31 @@ def __init__(
             verbose=verbose,
             measurement_error=measurement_error,
             mode=mode,
+            name=name,
         )
 
-        self._tensor_variable_info = tensor_variable_info
-        self._tensor_data_info = tensor_data_info
+        if name is None:
+            self._tensor_variable_info = tensor_variable_info
+            self._tensor_data_info = tensor_data_info
+        else:
+            self._tensor_variable_info = SymbolicVariableInfo(
+                symbolic_variables=tuple(
+                    SymbolicVariable(
+                        name=self.prefixed_name(symbolic_variable.name),
+                        symbolic_variable=symbolic_variable.symbolic_variable,
+                    )
+                    for symbolic_variable in tensor_variable_info
+                )
+            )
+            self._tensor_data_info = SymbolicDataInfo(
+                symbolic_data=tuple(
+                    SymbolicData(
+                        name=self.prefixed_name(symbolic_data.name),
+                        symbolic_data=symbolic_data.symbolic_data,
+                    )
+                    for symbolic_data in tensor_data_info
+                )
+            )
         self._component_info = component_info.copy()
         self._exog_names = data_info.exogenous_names
         self._needs_exog_data = data_info.needs_exogenous_data

tests/statespace/models/structural/test_core.py

@@ -18,6 +18,42 @@
 RTOL = 0 if floatX.endswith("64") else 1e-6
 
 
+def _build_named_structural_model(name: str):
+    return (
+        st.LevelTrend(order=1, innovations_order=1)
+        + st.Regression(name="reg", state_names=["x"])
+        + st.MeasurementError(name="obs")
+    ).build(name=name, verbose=False)
+
+
+def test_structural_name_propagates_to_base_and_scopes_p0():
+    ss_mod = _build_named_structural_model(name="m1")
+
+    assert ss_mod.name == "m1"
+    assert "P0" in ss_mod.param_names
+    assert ss_mod.prefixed_name("P0") in ss_mod._name_to_variable
+    assert "P0" not in ss_mod._name_to_variable
+
+
+def test_named_structural_models_do_not_collide_in_placeholder_registries():
+    with pm.Model():
+        m1 = _build_named_structural_model(name="m1")
+        m2 = _build_named_structural_model(name="m2")
+
+    var_keys_1 = set(m1._name_to_variable)
+    var_keys_2 = set(m2._name_to_variable)
+    data_keys_1 = set(m1._name_to_data)
+    data_keys_2 = set(m2._name_to_data)
+
+    assert var_keys_1.isdisjoint(var_keys_2)
+    assert data_keys_1.isdisjoint(data_keys_2)
+
+    assert var_keys_1 == {m1.prefixed_name(name) for name in m1.param_names}
+    assert var_keys_2 == {m2.prefixed_name(name) for name in m2.param_names}
+    assert data_keys_1 == {m1.prefixed_name(name) for name in m1.data_names}
+    assert data_keys_2 == {m2.prefixed_name(name) for name in m2.data_names}
+
+
 def test_add_components():
     ll = st.LevelTrend(order=2)
     se = st.TimeSeasonality(name="seasonal", season_length=12)