made necessary changes to get the regression component test to pass using the new dataclasses API

Author: Dekermanjian

pymc_extras/statespace/core/properties.py

@@ -3,7 +3,6 @@
 from typing import Generic, Self, TypeVar
 
 from pymc_extras.statespace.core import PyMCStateSpace
-from pymc_extras.statespace.models.structural.core import Component
 from pymc_extras.statespace.utils.constants import (
     ALL_STATE_AUX_DIM,
     ALL_STATE_DIM,
@@ -87,6 +86,21 @@ 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") -> "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:
+            raise ValueError(f"Duplicate parameter names found: {overlapping}")
+
+        return ParameterInfo(parameters=list(self.items) + list(other.items))
+
 
 @dataclass(frozen=True)
 class Data(Property):
@@ -108,6 +122,21 @@ def needs_exogenous_data(self) -> bool:
     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") -> "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:
+            raise ValueError(f"Duplicate data names found: {overlapping}")
+
+        return DataInfo(data=list(self.items) + list(other.items))
+
 
 @dataclass(frozen=True)
 class Coord(Property):
@@ -129,7 +158,11 @@ def __str__(self) -> str:
         return base
 
     @classmethod
-    def default_coords_from_model(cls, model: Component | PyMCStateSpace) -> Self:
+    def default_coords_from_model(
+        cls, model: PyMCStateSpace
+    ) -> (
+        Self
+    ):  # TODO: Need to figure out how to include Component type was causing circular import issues
         states = tuple(model.state_names)
         obs_states = tuple(model.observed_state_names)
         shocks = tuple(model.shock_names)
@@ -149,6 +182,21 @@ def default_coords_from_model(cls, model: Component | PyMCStateSpace) -> Self:
     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") -> "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:
+            raise ValueError(f"Duplicate coord names found: {overlapping}")
+
+        return CoordInfo(coords=list(self.items) + list(other.items))
+
 
 @dataclass(frozen=True)
 class State(Property):
@@ -171,6 +219,21 @@ def __str__(self) -> str:
     def observed_states(self) -> tuple[State, ...]:
         return tuple(s for s in self.items if 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") -> "StateInfo":
+        """Combine states from two StateInfo objects."""
+        if not isinstance(other, StateInfo):
+            raise TypeError(f"Cannot merge {type(other).__name__} with StateInfo")
+
+        overlapping = set(self.names) & set(other.names)
+        if overlapping:
+            raise ValueError(f"Duplicate state names found: {overlapping}")
+
+        return StateInfo(states=list(self.items) + list(other.items))
+
 
 @dataclass(frozen=True)
 class Shock(Property):
@@ -181,3 +244,18 @@ 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") -> "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:
+            raise ValueError(f"Duplicate shock names found: {overlapping}")
+
+        return ShockInfo(shocks=list(self.items) + list(other.items))

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

@@ -2,6 +2,18 @@
 
 from pytensor import tensor as pt
 
+from pymc_extras.statespace.core.properties import (
+    Coord,
+    CoordInfo,
+    Data,
+    DataInfo,
+    Parameter,
+    ParameterInfo,
+    Shock,
+    ShockInfo,
+    State,
+    StateInfo,
+)
 from pymc_extras.statespace.models.structural.core import Component
 from pymc_extras.statespace.utils.constants import TIME_DIM
 
@@ -194,64 +206,110 @@ def make_symbolic_graph(self) -> None:
             row_idx, col_idx = np.diag_indices(self.k_states)
             self.ssm["state_cov", row_idx, col_idx] = sigma_beta.ravel() ** 2
 
-    def populate_component_properties(self) -> None:
+    def _set_parameters(self) -> None:
         k_endog = self.k_endog
         k_endog_effective = 1 if self.share_states else k_endog
+        k_states = self.k_states // k_endog_effective
+
+        beta_param_name = f"beta_{self.name}"
+        beta_param_shape = (k_endog_effective, k_states) if k_endog_effective > 1 else (k_states,)
+        beta_param_dims = (
+            (f"endog_{self.name}", f"state_{self.name}")
+            if k_endog_effective > 1
+            else (f"state_{self.name}",)
+        )
+
+        beta_param_constraints = None
+        beta_parameter = Parameter(
+            name=beta_param_name,
+            shape=beta_param_shape,
+            dims=beta_param_dims,
+            constraints=beta_param_constraints,
+        )
 
+        if self.innovations:
+            sigma_param_name = f"sigma_beta_{self.name}"
+            sigma_param_dims = (f"state_{self.name}",)
+            sigma_param_shape = (k_states,)
+            sigma_param_constraints = "Positive"
+
+            sigma_parameter = Parameter(
+                name=sigma_param_name,
+                shape=sigma_param_shape,
+                dims=sigma_param_dims,
+                constraints=sigma_param_constraints,
+            )
+
+            self.param_info = ParameterInfo(parameters=[beta_parameter, sigma_parameter])
+            self.param_names = self.param_info.names
+        else:
+            self.param_info = ParameterInfo(parameters=[beta_parameter])
+            self.param_names = self.param_info.names
+
+    def _set_data(self) -> None:
+        k_endog = self.k_endog
+        k_endog_effective = 1 if self.share_states else k_endog
         k_states = self.k_states // k_endog_effective
 
+        data_name = f"data_{self.name}"
+        data_shape = (None, k_states)
+        data_dims = (TIME_DIM, f"state_{self.name}")
+
+        data_prop = Data(name=data_name, shape=data_shape, dims=data_dims, is_exogenous=True)
+        self.data_info = DataInfo(data=[data_prop])
+        self.data_names = self.data_info.names
+
+    def _set_shocks(self) -> None:
         if self.share_states:
-            self.shock_names = [f"{state_name}_shared" for state_name in self.state_names]
+            shock_names = [f"{state_name}_shared" for state_name in self.state_names]
         else:
-            self.shock_names = self.state_names
+            shock_names = self.state_names
 
-        self.param_names = [f"beta_{self.name}"]
-        self.data_names = [f"data_{self.name}"]
-        self.param_dims = {
-            f"beta_{self.name}": (f"endog_{self.name}", f"state_{self.name}")
-            if k_endog_effective > 1
-            else (f"state_{self.name}",)
-        }
+        self.shock_info = ShockInfo(shocks=[Shock(name=name) for name in shock_names])
+        self.shock_names = self.shock_info.names
 
-        base_names = self.state_names
+    def _set_states(self) -> None:
+        self.base_names = self.state_names
 
         if self.share_states:
-            self.state_names = [f"{name}[{self.name}_shared]" for name in base_names]
+            state_names = [f"{name}[{self.name}_shared]" for name in self.base_names]
+            self.state_info = StateInfo(
+                states=[State(name=name, observed=True, shared=True) for name in state_names]
+            )
+            self.state_names = self.state_info.names
         else:
-            self.state_names = [
+            state_names = [
                 f"{name}[{obs_name}]"
                 for obs_name in self.observed_state_names
-                for name in base_names
+                for name in self.base_names
             ]
+            self.state_info = StateInfo(
+                states=[State(name=name, observed=True, shared=False) for name in state_names]
+            )
+            self.state_names = self.state_info.names
 
-        self.param_info = {
-            f"beta_{self.name}": {
-                "shape": (k_endog_effective, k_states) if k_endog_effective > 1 else (k_states,),
-                "constraints": None,
-                "dims": (f"endog_{self.name}", f"state_{self.name}")
-                if k_endog_effective > 1
-                else (f"state_{self.name}",),
-            },
-        }
-
-        self.data_info = {
-            f"data_{self.name}": {
-                "shape": (None, k_states),
-                "dims": (TIME_DIM, f"state_{self.name}"),
-            },
-        }
-        self.coords = {
-            f"state_{self.name}": base_names,
-            f"endog_{self.name}": self.observed_state_names,
-        }
+    def _set_coords(self) -> None:
+        regression_state_coord = Coord(
+            dimension=f"state_{self.name}", labels=[state for state in self.base_names]
+        )
+        endogenous_state_coord = Coord(
+            dimension=f"endog_{self.name}", labels=[state for state in self.observed_state_names]
+        )
 
-        if self.innovations:
-            self.param_names += [f"sigma_beta_{self.name}"]
-            self.param_dims[f"sigma_beta_{self.name}"] = (f"state_{self.name}",)
-            self.param_info[f"sigma_beta_{self.name}"] = {
-                "shape": (k_states,),
-                "constraints": "Positive",
-                "dims": (f"state_{self.name}",)
-                if k_endog_effective == 1
-                else (f"endog_{self.name}", f"state_{self.name}"),
-            }
+        self.coords = CoordInfo(coords=[regression_state_coord, endogenous_state_coord])
+
+    def populate_component_properties(self) -> None:
+        # Set parameter info
+        self._set_parameters()
+
+        # Set data info
+        self._set_data()
+
+        # Set shock info
+        self._set_shocks()
+
+        # Set states info
+        self._set_states()
+
+        # Set coordinates info
+        self._set_coords()