Skip to content

Conversation

@opherdonchin
Copy link

Summary

As discussed in [#602](#602), this PR adds an optional name parameter to PyMCStateSpace. When provided, it is used to prefix all variables and registered data created by that state-space model. This makes it possible to use multiple state-space models inside one PyMC model without name collisions.

What this changes

  • Add name: str | None = None to PyMCStateSpace.__init__.

  • Add _prefix_name(self, var_name: str) -> str, which:

    • returns var_name unchanged if self.name is None
    • otherwise returns "{name}_{var_name}".
  • Apply _prefix_name when:

    • creating pm.Deterministic variables,
    • registering Kalman filter and likelihood outputs,
    • constructing names in SequenceMvNormal,
    • naming H_jittered, x0_slice, P0_slice, forecast_*, irf, initial_shock,
    • building names for posterior predictive variables.
  • Update data helpers (add_data_to_active_model, register_data_with_pymc) so that state-space models can optionally pass prefixed data names.

Backwards compatibility

  • If name is not provided, everything behaves exactly as before.
  • Existing code that uses a single PyMCStateSpace should work without any changes.

Example

ssm1 = PyMCStateSpace(..., name="vision")
ssm2 = PyMCStateSpace(..., name="motor")

This results in variables such as "vision_obs", "motor_obs", "vision_data", etc., allowing both models to coexist cleanly in the same PyMC model.

Testing

  • Added tests/statespace/test_register_variable.py to check the basic naming and variable-registration behavior.
  • Attempted to run pytest tests/statespace in a clean pip install -e ".[dev]" environment.
  • Test collection currently stops with FutureWarning errors from pytensor.graph.basic.* imports (these have moved to pytensor.graph.traversal.*). This appears unrelated to this PR and originates from existing imports in the codebase and test suite.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants