Add name parameter to PyMCStateSpace and update related methods for unique graph naming

Create a test to ensure multiple state space models can coexist with distinct names.

Author: opherdonchin

pymc_extras/statespace/core/statespace.py

@@ -261,6 +261,7 @@ def __init__(
         verbose: bool = True,
         measurement_error: bool = False,
         mode: str | None = None,
+        name: str | None = None,
     ):
         self._fit_coords: dict[str, Sequence[str]] | None = None
         self._fit_dims: dict[str, Sequence[str]] | None = None
@@ -274,6 +275,7 @@ def __init__(
         self.k_endog = k_endog
         self.k_states = k_states
         self.k_posdef = k_posdef
+        self.name = name
         self.measurement_error = measurement_error
         self.mode = mode
 
@@ -305,6 +307,12 @@ def __init__(
             console = Console()
             console.print(self.requirement_table)
 
+    def graph_name(self, base: str) -> str:
+        if not self.name:
+            return base
+        prefix = f"{self.name}_"
+        return base if base.startswith(prefix) else f"{self.name}_{base}"
+
     def _populate_properties(self) -> None:
         self._set_parameters()
         self._set_states()
@@ -649,14 +657,16 @@ def make_and_register_variable(
                 f"parameters."
             )
 
-        if name in self._tensor_variable_info:
+        gname = self.graph_name(name)
+
+        if gname in self._tensor_variable_info:
             raise ValueError(
-                f"{name} is already a registered placeholder variable with shape "
-                f"{self._tensor_variable_info[name].type.shape}"
+                f"{gname} is already a registered placeholder variable with shape "
+                f"{self._tensor_variable_info[gname].type.shape}"
             )
 
-        placeholder = pt.tensor(name, shape=shape, dtype=dtype)
-        tensor_var = SymbolicVariable(name=name, symbolic_variable=placeholder)
+        placeholder = pt.tensor(gname, shape=shape, dtype=dtype)
+        tensor_var = SymbolicVariable(name=gname, symbolic_variable=placeholder)
         self._tensor_variable_info = self._tensor_variable_info.add(tensor_var)
         return placeholder
 
@@ -685,18 +695,20 @@ def make_and_register_data(
         """
         if name not in self.data_names:
             raise ValueError(
-                f"{name} is not a model parameter. All placeholder variables should correspond to model "
-                f"parameters."
+                f"{name} is not a model data-variable. All placeholder variables should correspond to model "
+                f"data-variables."
             )
 
-        if name in self._tensor_data_info:
+        gname = self.graph_name(name)
+
+        if gname in self._tensor_data_info:
             raise ValueError(
-                f"{name} is already a registered placeholder variable with shape "
-                f"{self._tensor_data_info[name].type.shape}"
+                f"{gname} is already a registered placeholder variable with shape "
+                f"{self._tensor_data_info[gname].type.shape}"
             )
 
-        placeholder = pt.tensor(name, shape=shape, dtype=dtype)
-        tensor_data = SymbolicData(name=name, symbolic_data=placeholder)
+        placeholder = pt.tensor(gname, shape=shape, dtype=dtype)
+        tensor_data = SymbolicData(name=gname, symbolic_data=placeholder)
         self._tensor_data_info = self._tensor_data_info.add(tensor_data)
         return placeholder
 
@@ -800,11 +812,12 @@ def _save_exogenous_data_info(self):
         """
         pymc_mod = modelcontext(None)
         for data_name in self.data_names:
-            data = pymc_mod[data_name]
+            gname = self.graph_name(data_name)
+            data = pymc_mod[gname]
             self._fit_exog_data[data_name] = {
-                "name": data_name,
+                "name": gname,
                 "value": data.get_value(),
-                "dims": pymc_mod.named_vars_to_dims.get(data_name, None),
+                "dims": pymc_mod.named_vars_to_dims.get(gname, None),
             }
 
     def _insert_random_variables(self):
@@ -843,9 +856,10 @@ def _insert_random_variables(self):
         found_params = []
         with pymc_model:
             for param_name in self.param_names:
-                param = getattr(pymc_model, param_name, None)
+                gname = self.graph_name(param_name)
+                param = getattr(pymc_model, gname, None)
                 if param is not None:
-                    found_params.append(param.name)
+                    found_params.append(param_name)
 
         missing_params = list(set(self.param_names) - set(found_params))
         if len(missing_params) > 0:
@@ -880,9 +894,10 @@ def _insert_data_variables(self):
         found_data = []
         with pymc_model:
             for data_name in data_names:
-                data = getattr(pymc_model, data_name, None)
+                gname = self.graph_name(data_name)
+                data = getattr(pymc_model, gname, None)
                 if data is not None:
-                    found_data.append(data.name)
+                    found_data.append(data_name)
 
         missing_data = list(set(data_names) - set(found_data))
         if len(missing_data) > 0:
@@ -1046,6 +1061,7 @@ def build_statespace_graph(
             obs_coords=obs_coords,
             register_data=register_data,
             missing_fill_value=missing_fill_value,
+            data_name=self.graph_name("data"),
         )
 
         filter_outputs = self.kalman_filter.build_graph(
@@ -1145,11 +1161,12 @@ def _build_dummy_graph(self) -> None:
         """
 
         def infer_variable_shape(name):
-            shape = self._name_to_variable[name].type.shape
+            gname = self.graph_name(name)
+            shape = self._name_to_variable[gname].type.shape
             if not any(dim is None for dim in shape):
                 return shape
 
-            dim_names = self._fit_dims.get(name, None)
+            dim_names = self._fit_dims.get(gname, None)
             if dim_names is None:
                 raise ValueError(
                     f"Could not infer shape for {name}, because it was not given coords during model"
@@ -1166,9 +1183,9 @@ def infer_variable_shape(name):
 
         for name in self.param_names:
             pm.Flat(
-                name,
+                self.graph_name(name),
                 shape=infer_variable_shape(name),
-                dims=self._fit_dims.get(name, None),
+                dims=self._fit_dims.get(self.graph_name(name), None),
             )
 
     def _kalman_filter_outputs_from_dummy_graph(
@@ -1208,14 +1225,14 @@ def _kalman_filter_outputs_from_dummy_graph(
         self._insert_random_variables()
 
         for name in self.data_names:
-            if name not in pm_mod:
+            if self.graph_name(name) not in pm_mod:
                 pm.Data(**self._fit_exog_data[name])
 
         self._insert_data_variables()
 
         for name in self.data_names:
             if name in scenario.keys():
-                pm.set_data({name: scenario[name]})
+                pm.set_data({self.graph_name(name): scenario[name]})
 
         x0, P0, c, d, T, Z, R, H, Q = self.unpack_statespace()
 
@@ -1230,6 +1247,7 @@ def _kalman_filter_outputs_from_dummy_graph(
             obs_coords=obs_coords,
             data_dims=data_dims,
             register_data=True,
+            data_name=self.graph_name("data"),
         )
 
         filter_outputs = self.kalman_filter.build_graph(
@@ -1786,7 +1804,7 @@ def sample_statespace_matrices(
             self._insert_random_variables()
 
             for name in self.data_names:
-                pm.Data(**self.data_info[name])
+                pm.Data(name=self.graph_name(name), **self.data_info[name])
 
             self._insert_data_variables()
             matrices = self.unpack_statespace()
@@ -1852,6 +1870,7 @@ def sample_filter_outputs(
                 n_obs=self.ssm.k_endog,
                 obs_coords=obs_coords,
                 register_data=True,
+                data_name=self.graph_name("data"),
             )
 
             filter_outputs = self.kalman_filter.build_graph(
@@ -2283,12 +2302,15 @@ def _build_forecast_model(
             mu, cov = grouped_outputs[group_idx]
 
             sub_dict = {
-                data_var: pt.as_tensor_variable(data_var.get_value(), name="data")
+                data_var: pt.as_tensor_variable(
+                    data_var.get_value(), name=self.graph_name("data")
+                )
                 for data_var in forecast_model.data_vars
             }
 
             missing_data_vars = np.setdiff1d(
-                ar1=[*self.data_names, "data"], ar2=[k.name for k, _ in sub_dict.items()]
+                ar1=[*[self.graph_name(name) for name in self.data_names], self.graph_name("data")],
+                ar2=[k.name for k, _ in sub_dict.items()],
             )
             if missing_data_vars.size > 0:
                 raise ValueError(f"{missing_data_vars} data used for fitting not found!")
@@ -2466,8 +2488,9 @@ def forecast(
         with forecast_model:
             if scenario is not None:
                 dummy_obs_data = np.zeros((len(forecast_index), self.k_endog))
+                scoped_scenario = {self.graph_name(name): value for name, value in scenario.items()}
                 pm.set_data(
-                    scenario | {"data": dummy_obs_data},
+                    scoped_scenario | {self.graph_name("data"): dummy_obs_data},
                     coords={"data_time": np.arange(len(forecast_index))},
                 )
 

pymc_extras/statespace/utils/data_tools.py

@@ -121,7 +121,7 @@ def preprocess_pandas_data(data, n_obs, obs_coords=None, check_column_names=Fals
         return preprocess_numpy_data(data.values, n_obs, obs_coords)
 
 
-def add_data_to_active_model(values, index, data_dims=None):
+def add_data_to_active_model(values, index, data_dims=None, data_name="data"):
     pymc_mod = modelcontext(None)
     if data_dims is None:
         data_dims = [TIME_DIM, OBS_STATE_DIM]
@@ -146,7 +146,7 @@ def add_data_to_active_model(values, index, data_dims=None):
     else:
         data_shape = (None, values.shape[-1])
 
-    data = pm.Data("data", values, dims=data_dims, shape=data_shape)
+    data = pm.Data(data_name, values, dims=data_dims, shape=data_shape)
 
     return data
 
@@ -178,7 +178,13 @@ def mask_missing_values_in_data(values, missing_fill_value=None):
 
 
 def register_data_with_pymc(
-    data, n_obs, obs_coords, register_data=True, missing_fill_value=None, data_dims=None
+    data,
+    n_obs,
+    obs_coords,
+    register_data=True,
+    missing_fill_value=None,
+    data_dims=None,
+    data_name="data",
 ):
     if isinstance(data, pt.TensorVariable | TensorSharedVariable):
         values, index = preprocess_tensor_data(data, n_obs, obs_coords)
@@ -192,7 +198,7 @@ def register_data_with_pymc(
     data, nan_mask = mask_missing_values_in_data(values, missing_fill_value)
 
     if register_data:
-        data = add_data_to_active_model(data, index, data_dims)
+        data = add_data_to_active_model(data, index, data_dims, data_name=data_name)
     else:
-        data = pytensor.shared(data, name="data")
+        data = pytensor.shared(data, name=data_name)
     return data, nan_mask

tests/statespace/test_namespace.py

@@ -0,0 +1,13 @@
+import pymc as pm
+
+from pymc_extras.statespace.core.statespace import PyMCStateSpace
+
+
+def test_two_statespace_models_can_coexist_with_names(monkeypatch):
+    monkeypatch.setattr(PyMCStateSpace, "make_symbolic_graph", lambda self: None)
+
+    with pm.Model():
+        ssm_a = PyMCStateSpace(k_endog=1, k_states=1, k_posdef=1, name="a")
+        ssm_b = PyMCStateSpace(k_endog=1, k_states=1, k_posdef=1, name="b")
+
+        assert ssm_a.graph_name("data") != ssm_b.graph_name("data")