1. Updated dataclasses to include copy method and replaced raise on duplicate 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

Author: Dekermanjian

pymc_extras/statespace/core/properties.py

@@ -1,4 +1,7 @@
+import warnings
+
 from collections.abc import Iterator
+from copy import deepcopy
 from dataclasses import dataclass, fields
 from typing import Generic, Self, TypeVar
 
@@ -36,8 +39,8 @@ def __post_init__(self):
                 missing_attr.append(item)
                 continue
             key = getattr(item, self.key_field)
-            if key in index:
-                raise ValueError(f"Duplicate {self.key_field} '{key}' detected.")
+            # if key in index:
+            #     raise ValueError(f"Duplicate {self.key_field} '{key}' detected.") # This needs to be possible for shared states
             index[key] = item
         if missing_attr:
             raise AttributeError(f"Items missing attribute '{self.key_field}': {missing_attr}")
@@ -72,6 +75,9 @@ def __str__(self) -> str:
     def names(self) -> tuple[str, ...]:
         return tuple(self._index.keys())
 
+    def copy(self) -> "Info[T]":
+        return deepcopy(self)
+
 
 @dataclass(frozen=True)
 class Parameter(Property):
@@ -90,13 +96,13 @@ 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":
+    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:
+        if overlapping and not allow_duplicates:
             raise ValueError(f"Duplicate parameter names found: {overlapping}")
 
         return ParameterInfo(parameters=list(self.items) + list(other.items))
@@ -119,20 +125,24 @@ def __init__(self, data: list[Data]):
     def needs_exogenous_data(self) -> bool:
         return any(d.is_exogenous for d in self.items)
 
+    @property
+    def exogenous_names(self) -> tuple[str, ...]:
+        return tuple(d.name for d in self.items if d.is_exogenous)
+
     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":
+    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:
+        if overlapping and not allow_duplicates:
             raise ValueError(f"Duplicate data names found: {overlapping}")
 
         return DataInfo(data=list(self.items) + list(other.items))
@@ -164,7 +174,7 @@ def default_coords_from_model(
         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)
+        obs_states = tuple(model.observed_states)
         shocks = tuple(model.shock_names)
 
         dim_to_labels = (
@@ -186,13 +196,13 @@ 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":
+    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:
+        if overlapping and not allow_duplicates:
             raise ValueError(f"Duplicate coord names found: {overlapping}")
 
         return CoordInfo(coords=list(self.items) + list(other.items))
@@ -216,21 +226,37 @@ def __str__(self) -> str:
         )
 
     @property
-    def observed_states(self) -> tuple[State, ...]:
+    def observed_states(self) -> tuple[State, ...]:  # Is this needed??
         return tuple(s for s in self.items if s.observed)
 
+    @property
+    def observed_state_names(self) -> tuple[State, ...]:
+        return tuple(s.name for s in self.items if s.observed)
+
+    @property
+    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") -> "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")
 
         overlapping = set(self.names) & set(other.names)
-        if overlapping:
-            raise ValueError(f"Duplicate state names found: {overlapping}")
+        if overlapping and not allow_duplicates:
+            # This is necessary for shared states
+            warnings.warn(
+                f"Duplicate state names found: {overlapping}. Merge will ONLY retain unique states",
+                UserWarning,
+            )
+            return StateInfo(
+                states=list(self.items)
+                + [item for item in other.items if item.name not in overlapping]
+            )
 
         return StateInfo(states=list(self.items) + list(other.items))
 
@@ -249,13 +275,13 @@ 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":
+    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:
+        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/__init__.py

@@ -5,9 +5,6 @@
 from pymc_extras.statespace.models.structural.components.level_trend import LevelTrendComponent
 from pymc_extras.statespace.models.structural.components.measurement_error import MeasurementError
 from pymc_extras.statespace.models.structural.components.regression import RegressionComponent
-from pymc_extras.statespace.models.structural.components.regression_dataclass import (
-    RegressionComponent as RegressionComponentDataClass,
-)
 from pymc_extras.statespace.models.structural.components.seasonality import (
     FrequencySeasonality,
     TimeSeasonality,
@@ -20,6 +17,5 @@
     "LevelTrendComponent",
     "MeasurementError",
     "RegressionComponent",
-    "RegressionComponentDataClass",
     "TimeSeasonality",
 ]

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

@@ -211,33 +211,23 @@ def _set_parameters(self) -> None:
         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,
+            name=f"beta_{self.name}",
+            shape=(k_endog_effective, k_states) if k_endog_effective > 1 else (k_states,),
+            dims=(
+                (f"endog_{self.name}", f"state_{self.name}")
+                if k_endog_effective > 1
+                else (f"state_{self.name}",)
+            ),
+            constraints=None,
         )
 
         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,
+                name=f"sigma_beta_{self.name}",
+                shape=(k_states,),
+                dims=(f"state_{self.name}",),
+                constraints="Positive",
             )
 
             self.param_info = ParameterInfo(parameters=[beta_parameter, sigma_parameter])
@@ -251,11 +241,12 @@ def _set_data(self) -> None:
         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)
+        data_prop = Data(
+            name=f"data_{self.name}",
+            shape=(None, k_states),
+            dims=(TIME_DIM, f"state_{self.name}"),
+            is_exogenous=True,
+        )
         self.data_info = DataInfo(data=[data_prop])
         self.data_names = self.data_info.names
 
@@ -274,19 +265,35 @@ def _set_states(self) -> None:
         if self.share_states:
             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]
+                states=[State(name=name, observed=False, shared=True) for name in state_names]
+            )
+            self.state_info = self.state_info.merge(
+                StateInfo(
+                    states=[
+                        State(name=name, observed=True, shared=False)
+                        for name in self.observed_state_names
+                    ]
+                )
             )
-            self.state_names = self.state_info.names
+            self.state_names = self.state_info.unobserved_state_names
         else:
             state_names = [
                 f"{name}[{obs_name}]"
                 for obs_name in self.observed_state_names
                 for name in self.base_names
             ]
             self.state_info = StateInfo(
-                states=[State(name=name, observed=True, shared=False) for name in state_names]
+                states=[State(name=name, observed=False, shared=False) for name in state_names]
             )
-            self.state_names = self.state_info.names
+            self.state_info = self.state_info.merge(
+                StateInfo(
+                    states=[
+                        State(name=name, observed=True, shared=False)
+                        for name in self.observed_state_names
+                    ]
+                )
+            )
+            self.state_names = self.state_info.unobserved_state_names
 
     def _set_coords(self) -> None:
         regression_state_coord = Coord(
@@ -296,20 +303,11 @@ def _set_coords(self) -> None:
             dimension=f"endog_{self.name}", labels=[state for state in self.observed_state_names]
         )
 
-        self.coords = CoordInfo(coords=[regression_state_coord, endogenous_state_coord])
+        self.coords_info = 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()