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

I think what happens is because our data classes are immutable the __post_init__ runs right after our merge/add because we always return new objects of the same dataclass and it see that there are duplicate keys even though the merge/add method had allowed them via allow_duplicates.

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

I'm leaning more towards just using observed_state_names instead. If you are okay with that I am going to remove observed_states

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?

So here what is happening for shared states _get_combined_shapes counts the unique names between components. It needs to combine the state names without raising but it needs the unique names and not duplicates. It is subtle but if you look closely the StateInfo merge is an override that allows unique merge even when allow_duplicates is False compared to the base class that will raise an error. Obviously, this may not be the best way to handle this and maybe we need to look at _get_combined_shapes and see if it makes more sense to alter how that works.

Can we adjust the base class to handle the behavior you actually want?

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

My bad on the indentation. Will fix in next commit.

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

My bad on the indentation. Will fix in next commit.

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?

I agree with you. I will change it to _set_data_info

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?

Yes, you are correct it's either going to be a dataclass or None. I will update the logic.

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?

I think that would be nice. Right now, there is this intermediate conversion step that needs to take place.

Ok let's leave it for the next PR then and make the priority here just getting each model to be represented in the new way

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

Agreed. I will update the properties to derive from the <foo>_info sources of truth.

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

Yes, correct. I will fix this error.

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?

I don't think we still need _combine_property. Everything can be handled as you say with left.merge(right). I think that would also be cleaner and we already handle improper merges with non-identical objects in the base class merge method.

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?

No, I am sorry about that. I just wasn't sure where I could place it. Should I stick it in /statespace/utils/data_tools or /statespace/utils/coord_tools?

Yeah anywhere is fine. It could also just go in the file where it's used.

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!

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!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jessegrabowski commented Nov 7, 2025 •

edited

Loading