Represent statespace metadata with dataclasses #607

Dekermanjian · 2025-11-02T17:50:36Z

This is a draft proposal for #598

The idea is to handle each component separately using _set_{component} methods and all information are stored using data classes for easy mapping.

I believe this will simplify our tests of these components and will reduce redundancies where we have the same information spread across multiple sub-components like data_names and data_info.

@jessegrabowski let me know what you think I put a little notebook together to showcase the changes.

review-notebook-app · 2025-11-02T17:50:42Z

Check out this pull request on

See visual diffs & provide feedback on Jupyter Notebooks.

Powered by ReviewNB

jessegrabowski

This is a great first pass, much cleaner than what we have now.

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

jessegrabowski · 2025-11-06T20:06:12Z

We can also keep all of the existing properties like state_names, shock_names, state_dims, etc, but move them to the base class and just extract the requested info from the relevant Info objects.

jessegrabowski · 2025-11-07T14:10:12Z

Reflecting on it, I am convinced this is the way to go. It's 1000x more ergonomic. I made some changes to your initial code to make the API more "dictionary like", and to reduce code duplication. I moved everything to statespace/core/properties.py, because this is ultimately going to replace what we have in both the core models and the components.

Dekermanjian · 2025-11-07T14:37:48Z

@jessegrabowski, this is looking really cool! What can I do to help push this forward?

jessegrabowski · 2025-11-07T15:04:06Z

Delete the new regression_dataclass.py and simply refactor regression.py to use the new stuff.

We should keep your notebook with the plan to add it as a new example for the docs. Or it can be merged into the custom statespace notebook. So that should also be updated to import from the new properties.py file

Dekermanjian · 2025-11-07T15:15:27Z

Perfect! I'll work on that today!! It is really looking cool!

…sing the new dataclasses API

pymc_extras/statespace/core/properties.py

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

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

pymc_extras/statespace/models/structural/core.py

pymc_extras/statespace/models/utilities.py

tests/statespace/core/test_properties.py

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

tests/statespace/models/structural/conftest.py

Dekermanjian · 2025-11-11T12:20:08Z

@jessegrabowski, I agree with all of your comments above. I am going to start making those changes.

…uplicate with warning 2. removed unnecessary imports from __init__ after deleting regression_dataclass 3. updated components and structural classes to only utilize dataclasses and pull other objects from <foo>_info dataclasses 4. updated tests to conform to dataclass api

pymc_extras/statespace/core/properties.py

2. created tests for add and merge methods 3. added utility to convert from snake to pascal and integrated it in error messaging

… and placed default shoch and state setters

jessegrabowski

Incomplete review, I'll continue tomorrow AM

jessegrabowski · 2025-12-08T01:24:41Z

pymc_extras/statespace/core/properties.py

+            # if key in index:
+            #     raise ValueError(f"Duplicate {self.key_field} '{key}' detected.") # This needs to be possible for shared states


That shouldn't happen here though, it should come up in merge or add right? And we handle it there with the allow_duplicates flag

jessegrabowski · 2025-12-08T01:25:41Z

pymc_extras/statespace/core/properties.py

+            raise AttributeError(f"Items missing attribute '{self.key_field}': {missing_attr}")
+        object.__setattr__(self, "_index", index)
+
+    def _key(self, item: T) -> str:


Is this used?

pymc_extras/statespace/core/properties.py

jessegrabowski · 2025-12-08T01:31:11Z

pymc_extras/statespace/core/properties.py

+    @property
+    def observed_states(self) -> tuple[State, ...]:  # Is this needed??
+        return tuple(s for s in self.items if s.observed)


If you want to keep it (as an alias for observed_state_names), then pick one to be the "canonical" name and just return that one from the other ones, rather than re-writing the loop in several places

jessegrabowski · 2025-12-08T01:32:05Z

pymc_extras/statespace/core/properties.py

+    def unobserved_state_names(self) -> tuple[State, ...]:
+        return tuple(s.name for s in self.items if not s.observed)
+
+    def merge(self, other: StateInfo, allow_duplicates: bool = False) -> StateInfo:


Why doesn't the base class version work?

jessegrabowski · 2025-12-08T01:36:54Z

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

+            )
+
+            self.param_info = ParameterInfo(parameters=[beta_parameter, sigma_parameter])
+            self.param_names = self.param_info.names


Suggested change

self.param_names = self.param_info.names

jessegrabowski · 2025-12-08T01:37:01Z

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

+            self.param_names = self.param_info.names
+        else:
+            self.param_info = ParameterInfo(parameters=[beta_parameter])
+            self.param_names = self.param_info.names


Suggested change

self.param_names = self.param_info.names

self.param_names = self.param_info.names

jessegrabowski · 2025-12-12T23:52:37Z

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

+            self.param_info = ParameterInfo(parameters=[beta_parameter])
+            self.param_names = self.param_info.names

+    def _set_data(self) -> None:


nit: the name _set_data is too close to pm.set_data, which changes the actual model data. _set_data_info here?

jessegrabowski · 2025-12-13T00:04:18Z

pymc_extras/statespace/models/structural/core.py

+        # TODO: discuss if copying is still needed since these are now immutable
+        self._coord_info = coords_info.copy()


Indeed, it shouldn't be necessary

jessegrabowski · 2025-12-13T00:06:27Z

pymc_extras/statespace/models/structural/core.py

+        if is_dataclass(param_info):
+            param_info = param_info.add(initial_state_cov_param)


If param_info isn't a dataclass at this point what else can it be? None? Can we just check for that instead?

jessegrabowski · 2025-12-13T00:09:16Z

pymc_extras/statespace/models/structural/core.py

-        self.state_names = list(state_names) if state_names is not None else []
-        self.observed_state_names = (
-            list(observed_state_names) if observed_state_names is not None else []
+        self.param_info = ParameterInfo(


Should we change the component signature to just take the Info objects directly?

jessegrabowski · 2025-12-13T00:22:31Z

pymc_extras/statespace/models/structural/core.py

+        self.state_names = self.state_info.unobserved_state_names
+        self.observed_state_names = self.state_info.observed_state_names
+        self.param_names = self.param_info.names
+        self.data_names = [d.name for d in self.data_info if not d.is_exogenous]
+        self.exog_names = self.data_info.exogenous_names
+        self.shock_names = self.shock_info.names

-        self.param_info = {}
-        self.data_info = {}
+        self.coords = self.coord_info.to_dict()
+        self.param_dims = [p.dims for p in self.param_info]

-        self.param_counts = {}
+        self.needs_exog_data = self.data_info.needs_exogenous_data


These should all be properties

jessegrabowski · 2025-12-13T00:23:26Z

pymc_extras/statespace/models/structural/core.py

+    def _set_parameters(self) -> None:
+        raise NotImplementedError
+
+    def _set_shocks(self) -> None:
+        return ShockInfo(shocks=[Shock(name=f"shock_{n}") for n in range(self.k_posdef or 0)])
+
+    def _set_states(self) -> None:


These don't return None

jessegrabowski · 2025-12-13T00:24:25Z

pymc_extras/statespace/models/structural/core.py


-        if not isinstance(self_prop, list | dict):
+        if not is_dataclass(self_prop):
+            # TODO: This works right now because we are only passing <foo>_info info names into _combine_property


Do we still need _combine_property? Can we just call left.merge(right) since everything is an Info?

jessegrabowski · 2025-12-13T00:25:30Z

pymc_extras/statespace/utils/message_tools.py

@@ -0,0 +1,5 @@
+import re


Do we need a whole new file for this?

proposal for updating propogate_component_properties using data classes

a4dacd8

Dekermanjian marked this pull request as draft November 2, 2025 17:50

Dekermanjian added enhancements New feature or request request discussion statespace labels Nov 2, 2025

jessegrabowski reviewed Nov 6, 2025

View reviewed changes

Iterate on proposal

7f32a48

jessegrabowski changed the title ~~proposal for updating propogate_component_properties using data classes~~ Represent statespace metadata with dataclasses Nov 7, 2025

Jesse Grabowski and others added 3 commits November 7, 2025 11:45

Fix iterator, add to_dict method to CoordsInfo

d65fc0a

Add observed_states helper to StateInfo

c6a48fc

made necessary changes to get the regression component test to pass u…

92e333f

…sing the new dataclasses API

jessegrabowski requested changes Nov 11, 2025

View reviewed changes

jessegrabowski reviewed Nov 15, 2025

View reviewed changes

pymc_extras/statespace/core/properties.py Show resolved Hide resolved

Dekermanjian added 2 commits November 16, 2025 09:44

1. added add and merge methods to base class

228acff

2. created tests for add and merge methods 3. added utility to convert from snake to pascal and integrated it in error messaging

removed data & coords setters in _set<foo> medthod in Component class…

1ae433f

… and placed default shoch and state setters

jessegrabowski requested changes Dec 8, 2025

View reviewed changes

jessegrabowski requested changes Dec 13, 2025

View reviewed changes

		# if key in index:
		# raise ValueError(f"Duplicate {self.key_field} '{key}' detected.") # This needs to be possible for shared states

	self.param_names = self.param_info.names
	self.param_names = self.param_info.names

		# TODO: discuss if copying is still needed since these are now immutable
		self._coord_info = coords_info.copy()

		if is_dataclass(param_info):
		param_info = param_info.add(initial_state_cov_param)

Represent statespace metadata with dataclasses #607

Are you sure you want to change the base?

Represent statespace metadata with dataclasses #607

Uh oh!

Conversation

Dekermanjian commented Nov 2, 2025

Uh oh!

review-notebook-app bot commented Nov 2, 2025

Uh oh!

jessegrabowski left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

jessegrabowski commented Nov 6, 2025

Uh oh!

jessegrabowski commented Nov 7, 2025

Uh oh!

Dekermanjian commented Nov 7, 2025

Uh oh!

jessegrabowski commented Nov 7, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Dekermanjian commented Nov 7, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Dekermanjian commented Nov 11, 2025

Uh oh!

Uh oh!

jessegrabowski left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

jessegrabowski commented Nov 7, 2025 •

edited

Loading