1. added add and merge methods to base class

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

Author: Dekermanjian

pymc_extras/statespace/core/properties.py

@@ -1,9 +1,11 @@
+from __future__ import annotations
+
 import warnings
 
 from collections.abc import Iterator
 from copy import deepcopy
 from dataclasses import dataclass, fields
-from typing import Generic, Self, TypeVar
+from typing import TYPE_CHECKING, Generic, Self, TypeVar
 
 from pymc_extras.statespace.core import PyMCStateSpace
 from pymc_extras.statespace.utils.constants import (
@@ -15,6 +17,9 @@
     SHOCK_DIM,
 )
 
+if TYPE_CHECKING:
+    from pymc_extras.statespace.models.structural.core import Component
+
 
 @dataclass(frozen=True)
 class Property:
@@ -62,7 +67,7 @@ def __getitem__(self, key: str) -> T:
     def __contains__(self, key: object) -> bool:
         return key in self._index
 
-    def __iter__(self) -> Iterator[str]:
+    def __iter__(self) -> Iterator[T]:
         return iter(self.items)
 
     def __len__(self) -> int:
@@ -71,11 +76,24 @@ def __len__(self) -> int:
     def __str__(self) -> str:
         return f"{self.key_field}s: {list(self._index.keys())}"
 
+    def add(self, new_item: T):
+        return type(self)([*self.items, new_item])
+
+    def merge(self, other: Self, allow_duplicates: bool = False) -> Self:
+        if not isinstance(other, type(self)):
+            raise TypeError(f"Cannot merge {type(other).__name__} with {type(self).__name__}")
+
+        overlapping = set(self.names) & set(other.names)
+        if overlapping and not allow_duplicates:
+            raise ValueError(f"Duplicate names found: {overlapping}")
+
+        return type(self)(list(self.items) + list(other.items))
+
     @property
     def names(self) -> tuple[str, ...]:
         return tuple(self._index.keys())
 
-    def copy(self) -> "Info[T]":
+    def copy(self) -> Info[T]:
         return deepcopy(self)
 
 
@@ -92,21 +110,6 @@ class ParameterInfo(Info[Parameter]):
     def __init__(self, parameters: list[Parameter]):
         super().__init__(items=tuple(parameters), key_field="name")
 
-    def add(self, parameter: Parameter) -> "ParameterInfo":
-        # return a new ParameterInfo with parameter appended
-        return ParameterInfo(parameters=[*list(self.items), parameter])
-
-    def merge(self, other: "ParameterInfo", allow_duplicates: bool = False) -> "ParameterInfo":
-        """Combine parameters from two ParameterInfo objects."""
-        if not isinstance(other, ParameterInfo):
-            raise TypeError(f"Cannot merge {type(other).__name__} with ParameterInfo")
-
-        overlapping = set(self.names) & set(other.names)
-        if overlapping and not allow_duplicates:
-            raise ValueError(f"Duplicate parameter names found: {overlapping}")
-
-        return ParameterInfo(parameters=list(self.items) + list(other.items))
-
 
 @dataclass(frozen=True)
 class Data(Property):
@@ -132,21 +135,6 @@ def exogenous_names(self) -> tuple[str, ...]:
     def __str__(self) -> str:
         return f"data: {[d.name for d in self.items]}\nneeds exogenous data: {self.needs_exogenous_data}"
 
-    def add(self, data: Data) -> "DataInfo":
-        # return a new DataInfo with data appended
-        return DataInfo(data=[*list(self.items), data])
-
-    def merge(self, other: "DataInfo", allow_duplicates: bool = False) -> "DataInfo":
-        """Combine data from two DataInfo objects."""
-        if not isinstance(other, DataInfo):
-            raise TypeError(f"Cannot merge {type(other).__name__} with DataInfo")
-
-        overlapping = set(self.names) & set(other.names)
-        if overlapping and not allow_duplicates:
-            raise ValueError(f"Duplicate data names found: {overlapping}")
-
-        return DataInfo(data=list(self.items) + list(other.items))
-
 
 @dataclass(frozen=True)
 class Coord(Property):
@@ -169,7 +157,7 @@ def __str__(self) -> str:
 
     @classmethod
     def default_coords_from_model(
-        cls, model: PyMCStateSpace
+        cls, model: PyMCStateSpace | Component
     ) -> (
         Self
     ):  # TODO: Need to figure out how to include Component type was causing circular import issues
@@ -192,21 +180,6 @@ def default_coords_from_model(
     def to_dict(self):
         return {coord.dimension: coord.labels for coord in self.items if len(coord.labels) > 0}
 
-    def add(self, coord: Coord) -> "CoordInfo":
-        # return a new CoordInfo with data appended
-        return CoordInfo(coords=[*list(self.items), coord])
-
-    def merge(self, other: "CoordInfo", allow_duplicates: bool = False) -> "CoordInfo":
-        """Combine data from two CoordInfo objects."""
-        if not isinstance(other, CoordInfo):
-            raise TypeError(f"Cannot merge {type(other).__name__} with CoordInfo")
-
-        overlapping = set(self.names) & set(other.names)
-        if overlapping and not allow_duplicates:
-            raise ValueError(f"Duplicate coord names found: {overlapping}")
-
-        return CoordInfo(coords=list(self.items) + list(other.items))
-
 
 @dataclass(frozen=True)
 class State(Property):
@@ -237,11 +210,7 @@ def observed_state_names(self) -> tuple[State, ...]:
     def unobserved_state_names(self) -> tuple[State, ...]:
         return tuple(s.name for s in self.items if not s.observed)
 
-    def add(self, state: State) -> "StateInfo":
-        # return a new StateInfo with state appended
-        return StateInfo(states=[*list(self.items), state])
-
-    def merge(self, other: "StateInfo", allow_duplicates: bool = False) -> "StateInfo":
+    def merge(self, other: StateInfo, allow_duplicates: bool = False) -> StateInfo:
         """Combine states from two StateInfo objects."""
         if not isinstance(other, StateInfo):
             raise TypeError(f"Cannot merge {type(other).__name__} with StateInfo")
@@ -270,18 +239,3 @@ class Shock(Property):
 class ShockInfo(Info[Shock]):
     def __init__(self, shocks: list[Shock]):
         super().__init__(items=tuple(shocks), key_field="name")
-
-    def add(self, shock: Shock) -> "ShockInfo":
-        # return a new ShockInfo with shock appended
-        return ShockInfo(shocks=[*list(self.items), shock])
-
-    def merge(self, other: "ShockInfo", allow_duplicates: bool = False) -> "ShockInfo":
-        """Combine shocks from two ShockInfo objects."""
-        if not isinstance(other, ShockInfo):
-            raise TypeError(f"Cannot merge {type(other).__name__} with ShockInfo")
-
-        overlapping = set(self.names) & set(other.names)
-        if overlapping and not allow_duplicates:
-            raise ValueError(f"Duplicate shock names found: {overlapping}")
-
-        return ShockInfo(shocks=list(self.items) + list(other.items))

pymc_extras/statespace/models/structural/core.py

@@ -34,6 +34,7 @@
     ALL_STATE_DIM,
     LONG_MATRIX_NAMES,
 )
+from pymc_extras.statespace.utils.message_tools import snake_to_pascal
 
 _log = logging.getLogger(__name__)
 floatX = config.floatX
@@ -815,8 +816,10 @@ def _combine_property(self, other, name, allow_duplicates=True):
             )
 
         if not is_dataclass(self_prop):
+            # TODO: This works right now because we are only passing <foo>_info info names into _combine_property
+            # If we don't follow that schema moving forward this will break.
             raise TypeError(
-                f"All component properties are expected to be dataclasses, but found {type(self_prop)}"
+                f"Component properties are expected to be {snake_to_pascal(name)}, but found {type(self_prop)}"
                 f"for property {name} of {self} and {type(other_prop)} for {other}'"
             )
 

pymc_extras/statespace/utils/message_tools.py

@@ -0,0 +1,5 @@
+import re
+
+
+def snake_to_pascal(s: str) -> str:
+    return re.sub(r"(?:^|_)([a-z])", lambda m: m.group(1).upper(), s)

tests/statespace/core/test_properties.py

@@ -68,7 +68,7 @@ def test_data_info_needs_exogenous_and_str():
 def test_coord_info_make_defaults_from_component_and_types():
     class DummyComponent:
         state_names = ["x1", "x2"]
-        observed_state_names = ["x2"]
+        observed_states = ["x2"]
         shock_names = ["eps1"]
 
     ci = CoordInfo.default_coords_from_model(DummyComponent())
@@ -117,3 +117,26 @@ def test_info_is_iterable_and_unpackable():
 
     a, b = info.items
     assert a.name == "p1" and b.name == "p2"
+
+
+def test_info_add_method():
+    a_param = Parameter(name="a", shape=(1,), dims=("dim",))
+    param_info = ParameterInfo(parameters=[a_param])
+
+    b_param = Parameter(name="b", shape=(1,), dims=("dim",))
+
+    new_param_info = param_info.add(new_item=b_param)
+
+    assert new_param_info.names == ("a", "b")
+
+
+def test_info_merge_method():
+    a_param = Parameter(name="a", shape=(1,), dims=("dim",))
+    a_param_info = ParameterInfo(parameters=[a_param])
+
+    b_param = Parameter(name="b", shape=(1,), dims=("dim",))
+    b_param_info = ParameterInfo(parameters=[b_param])
+
+    new_param_info = a_param_info.merge(b_param_info)
+
+    assert new_param_info.names == ("a", "b")