pymc-extras 0.7.0__py3-none-any.whl → 0.8.0__py3-none-any.whl

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

@@ -1,13 +1,21 @@
+import warnings
+
 import numpy as np
 
 from pytensor import tensor as pt
-from pytensor.tensor.slinalg import block_diag
-
+from pytensor.tensor.linalg import block_diag
+
+from pymc_extras.statespace.core.properties import (
+    Coord,
+    Parameter,
+    Shock,
+    State,
+)
 from pymc_extras.statespace.models.structural.core import Component
 from pymc_extras.statespace.models.structural.utils import _frequency_transition_block
 
 
-class CycleComponent(Component):
+class Cycle(Component):
     r"""
     A component for modeling longer-term cyclical effects
 
@@ -64,7 +72,7 @@ class CycleComponent(Component):
     effects, such as business cycles, and that the seasonal component be used for shorter term effects, such as
     weekly or monthly seasonality.
 
-    Unlike a FrequencySeasonality component, the length of a CycleComponent can be estimated.
+    Unlike a FrequencySeasonality component, the length of a Cycle can be estimated.
 
     **Multivariate Support:**
     For multivariate time series with k endogenous variables, the component creates:
@@ -89,8 +97,8 @@ class CycleComponent(Component):
         data = np.random.normal(size=(100, 1))
 
         # Build the structural model
-        grw = st.LevelTrendComponent(order=1, innovations_order=1)
-        cycle = st.CycleComponent(
+        grw = st.LevelTrend(order=1, innovations_order=1)
+        cycle = st.Cycle(
             "business_cycle", cycle_length=12, estimate_cycle_length=False, innovations=True, dampen=True
         )
         ss_mod = (grw + cycle).build()
@@ -117,7 +125,7 @@ class CycleComponent(Component):
     .. code:: python
 
         # Multivariate cycle component
-        cycle = st.CycleComponent(
+        cycle = st.Cycle(
             name='business_cycle',
             cycle_length=12,
             estimate_cycle_length=False,
@@ -188,18 +196,111 @@ class CycleComponent(Component):
         obs_state_idx = np.zeros(k_states)
         obs_state_idx[slice(0, k_states, 2)] = 1
 
+        state_names = [f"{f}_{name}" for f in ["Cos", "Sin"]]
+
         super().__init__(
             name=name,
             k_endog=k_endog,
             k_states=k_states,
             k_posdef=k_posdef,
+            base_state_names=state_names,
             measurement_error=False,
             combine_hidden_states=True,
             obs_state_idxs=obs_state_idx,
-            observed_state_names=observed_state_names,
+            base_observed_state_names=observed_state_names,
             share_states=share_states,
         )
 
+    def set_states(self) -> State | tuple[State, ...] | None:
+        k_endog_effective = 1 if self.share_states else self.k_endog
+
+        base_names = self.base_state_names
+        observed_state_names = self.base_observed_state_names
+
+        if self.share_states:
+            state_names = [f"{name}[shared]" for name in base_names]
+        else:
+            state_names = [
+                f"{name}[{var_name}]" if k_endog_effective > 1 else name
+                for var_name in observed_state_names
+                for name in base_names
+            ]
+        hidden_states = [State(name=name, observed=False, shared=True) for name in state_names]
+        observed_states = [
+            State(name=name, observed=True, shared=False) for name in observed_state_names
+        ]
+        return *hidden_states, *observed_states
+
+    def set_parameters(self) -> Parameter | tuple[Parameter, ...] | None:
+        k_endog = self.k_endog
+        k_endog_effective = 1 if self.share_states else k_endog
+
+        cycle_param = Parameter(
+            name=f"params_{self.name}",
+            shape=(2,) if k_endog_effective == 1 else (k_endog_effective, 2),
+            dims=(f"state_{self.name}",)
+            if k_endog_effective == 1
+            else (f"endog_{self.name}", f"state_{self.name}"),
+            constraints=None,
+        )
+
+        params_container = [cycle_param]
+
+        if self.estimate_cycle_length:
+            length_param = Parameter(
+                name=f"length_{self.name}",
+                shape=() if k_endog_effective == 1 else (k_endog_effective,),
+                dims=None if k_endog_effective == 1 else (f"endog_{self.name}",),
+                constraints="Positive, non-zero",
+            )
+            params_container.append(length_param)
+
+        if self.dampen:
+            dampen_param = Parameter(
+                name=f"dampening_factor_{self.name}",
+                shape=() if k_endog_effective == 1 else (k_endog_effective,),
+                dims=None if k_endog_effective == 1 else (f"endog_{self.name}",),
+                constraints="0 < x ≤ 1",
+            )
+            params_container.append(dampen_param)
+
+        if self.innovations:
+            sigma_param = Parameter(
+                name=f"sigma_{self.name}",
+                shape=() if k_endog_effective == 1 else (k_endog_effective,),
+                dims=None if k_endog_effective == 1 else (f"endog_{self.name}",),
+                constraints="Positive",
+            )
+            params_container.append(sigma_param)
+
+        return tuple(params_container)
+
+    def set_shocks(self) -> Shock | tuple[Shock, ...] | None:
+        if self.innovations:
+            return tuple(Shock(name=name) for name in self.state_names)
+        return None
+
+    def set_coords(self) -> Coord | tuple[Coord, ...] | None:
+        k_endog = self.k_endog
+        k_endog_effective = 1 if self.share_states else k_endog
+        base_names = tuple(f"{f}_{self.name}" for f in ["Cos", "Sin"])
+        observed_state_names = self.observed_state_names
+
+        state_coords = Coord(
+            dimension=f"state_{self.name}",
+            labels=base_names
+            if k_endog_effective == 1
+            else (f"Cos_{self.name}", f"Sin_{self.name}"),
+        )
+
+        coord_container = [state_coords]
+
+        if k_endog_effective != 1:
+            endog_coords = Coord(dimension=f"endog_{self.name}", labels=observed_state_names)
+            coord_container.append(endog_coords)
+
+        return tuple(coord_container)
+
     def make_symbolic_graph(self) -> None:
         k_endog = self.k_endog
         k_endog_effective = 1 if self.share_states else k_endog
@@ -250,76 +351,14 @@ class CycleComponent(Component):
             # explicitly set state cov to 0 when no innovations
             self.ssm["state_cov", :, :] = pt.zeros((self.k_posdef, self.k_posdef))
 
-    def populate_component_properties(self):
-        k_endog = self.k_endog
-        k_endog_effective = 1 if self.share_states else k_endog
-
-        base_names = [f"{f}_{self.name}" for f in ["Cos", "Sin"]]
 
-        if self.share_states:
-            self.state_names = [f"{name}[shared]" for name in base_names]
-        else:
-            self.state_names = [
-                f"{name}[{var_name}]" if k_endog_effective > 1 else name
-                for var_name in self.observed_state_names
-                for name in base_names
-            ]
-
-        self.param_names = [f"params_{self.name}"]
-
-        if k_endog_effective == 1:
-            self.param_dims = {f"params_{self.name}": (f"state_{self.name}",)}
-            self.coords = {f"state_{self.name}": base_names}
-            self.param_info = {
-                f"params_{self.name}": {
-                    "shape": (2,),
-                    "constraints": None,
-                    "dims": (f"state_{self.name}",),
-                }
-            }
-        else:
-            self.param_dims = {f"params_{self.name}": (f"endog_{self.name}", f"state_{self.name}")}
-            self.coords = {
-                f"state_{self.name}": [f"Cos_{self.name}", f"Sin_{self.name}"],
-                f"endog_{self.name}": self.observed_state_names,
-            }
-            self.param_info = {
-                f"params_{self.name}": {
-                    "shape": (k_endog_effective, 2),
-                    "constraints": None,
-                    "dims": (f"endog_{self.name}", f"state_{self.name}"),
-                }
-            }
-
-        if self.estimate_cycle_length:
-            self.param_names += [f"length_{self.name}"]
-            self.param_info[f"length_{self.name}"] = {
-                "shape": () if k_endog_effective == 1 else (k_endog_effective,),
-                "constraints": "Positive, non-zero",
-                "dims": None if k_endog_effective == 1 else (f"endog_{self.name}",),
-            }
-
-        if self.dampen:
-            self.param_names += [f"dampening_factor_{self.name}"]
-            self.param_info[f"dampening_factor_{self.name}"] = {
-                "shape": () if k_endog_effective == 1 else (k_endog_effective,),
-                "constraints": "0 < x ≤ 1",
-                "dims": None if k_endog_effective == 1 else (f"endog_{self.name}",),
-            }
-
-        if self.innovations:
-            self.param_names += [f"sigma_{self.name}"]
-            if k_endog_effective == 1:
-                self.param_info[f"sigma_{self.name}"] = {
-                    "shape": (),
-                    "constraints": "Positive",
-                    "dims": None,
-                }
-            else:
-                self.param_dims[f"sigma_{self.name}"] = (f"endog_{self.name}",)
-                self.param_info[f"sigma_{self.name}"] = {
-                    "shape": (k_endog_effective,),
-                    "constraints": "Positive",
-                    "dims": (f"endog_{self.name}",),
-                }
-            self.shock_names = self.state_names.copy()
+def __getattr__(name: str):
+    if name == "CycleComponent":
+        warnings.warn(
+            "CycleComponent is deprecated and will be removed in a future release. "
+            "Use Cycle instead.",
+            FutureWarning,
+            stacklevel=2,
+        )
+        return Cycle
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

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

@@ -1,12 +1,20 @@
+import warnings
+
 import numpy as np
 import pytensor.tensor as pt
 
+from pymc_extras.statespace.core.properties import (
+    Coord,
+    Parameter,
+    Shock,
+    State,
+)
 from pymc_extras.statespace.models.structural.core import Component
 from pymc_extras.statespace.models.structural.utils import order_to_mask
 from pymc_extras.statespace.utils.constants import POSITION_DERIVATIVE_NAMES
 
 
-class LevelTrendComponent(Component):
+class LevelTrend(Component):
     r"""
     Level and trend component of a structural time series model
 
@@ -145,6 +153,7 @@ class LevelTrendComponent(Component):
                 f"lower order model, explicitly omit the zeros."
             )
         k_states = max_state
+        state_names = POSITION_DERIVATIVE_NAMES[:max_state]
 
         if isinstance(innovations_order, int):
             n = innovations_order
@@ -164,7 +173,8 @@ class LevelTrendComponent(Component):
             k_endog=k_endog,
             k_states=k_states * k_endog if not share_states else k_states,
             k_posdef=k_posdef * k_endog if not share_states else k_posdef,
-            observed_state_names=observed_state_names,
+            base_state_names=state_names,
+            base_observed_state_names=observed_state_names,
             measurement_error=False,
             combine_hidden_states=False,
             obs_state_idxs=np.tile(
@@ -173,71 +183,102 @@ class LevelTrendComponent(Component):
             share_states=share_states,
         )
 
-    def populate_component_properties(self):
-        k_endog = self.k_endog
-        k_endog_effective = 1 if self.share_states else k_endog
-
-        k_states = self.k_states // k_endog_effective
-        k_posdef = self.k_posdef // k_endog_effective
-
-        name_slice = POSITION_DERIVATIVE_NAMES[:k_states]
-        self.param_names = [f"initial_{self.name}"]
-        base_names = [name for name, mask in zip(name_slice, self._order_mask) if mask]
+    def set_states(self) -> State | tuple[State, ...] | None:
+        observed_state_names = self.base_observed_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_state_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 obs_name in observed_state_names
+                for name in self.base_state_names
             ]
 
-        self.param_dims = {f"initial_{self.name}": (f"state_{self.name}",)}
-        self.coords = {f"state_{self.name}": base_names}
+        hidden_states = [State(name=name, observed=False, shared=True) for name in state_names]
+        observed_states = [
+            State(name=name, observed=True, shared=False) for name in observed_state_names
+        ]
+        return *hidden_states, *observed_states
 
-        if k_endog > 1:
-            self.coords[f"endog_{self.name}"] = self.observed_state_names
+    def set_parameters(self) -> Parameter | tuple[Parameter, ...] | None:
+        k_endog = self.k_endog
+        k_endog_effective = 1 if self.share_states else k_endog
 
-        if k_endog_effective > 1:
-            self.param_dims[f"state_{self.name}"] = (
-                f"endog_{self.name}",
-                f"state_{self.name}",
-            )
-            self.param_dims = {f"initial_{self.name}": (f"endog_{self.name}", f"state_{self.name}")}
+        k_states = self.k_states // k_endog_effective
+        k_posdef = self.k_posdef // k_endog_effective
 
-        shape = (k_endog_effective, k_states) if k_endog_effective > 1 else (k_states,)
-        self.param_info = {f"initial_{self.name}": {"shape": shape, "constraints": None}}
+        initial_param = Parameter(
+            name=f"initial_{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.k_posdef > 0:
-            self.param_names += [f"sigma_{self.name}"]
+            sigma_param = Parameter(
+                name=f"sigma_{self.name}",
+                shape=(k_posdef,) if k_endog_effective == 1 else (k_endog_effective, k_posdef),
+                dims=(f"shock_{self.name}",)
+                if k_endog_effective == 1
+                else (f"endog_{self.name}", f"shock_{self.name}"),
+                constraints="Positive",
+            )
+            return initial_param, sigma_param
+        else:
+            return (initial_param,)
 
+    def set_shocks(self) -> Shock | tuple[Shock, ...] | 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
+        name_slice = POSITION_DERIVATIVE_NAMES[:k_states]
+
+        if self.k_posdef > 0:
             base_shock_names = [
                 name for name, mask in zip(name_slice, self.innovations_order) if mask
             ]
 
             if self.share_states:
-                self.shock_names = [f"{name}[{self.name}_shared]" for name in base_shock_names]
+                shock_names = [f"{name}[{self.name}_shared]" for name in base_shock_names]
             else:
-                self.shock_names = [
+                shock_names = [
                     f"{name}[{obs_name}]"
                     for obs_name in self.observed_state_names
                     for name in base_shock_names
                 ]
 
-            self.param_dims[f"sigma_{self.name}"] = (
-                (f"shock_{self.name}",)
-                if k_endog_effective == 1
-                else (f"endog_{self.name}", f"shock_{self.name}")
+            return tuple(Shock(name=name) for name in shock_names)
+        return None
+
+    def set_coords(self) -> Coord | tuple[Coord, ...] | 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
+        name_slice = POSITION_DERIVATIVE_NAMES[:k_states]
+
+        base_names = [name for name, mask in zip(name_slice, self._order_mask) if mask]
+
+        base_shock_names = [name for name, mask in zip(name_slice, self.innovations_order) if mask]
+
+        state_coord = Coord(dimension=f"state_{self.name}", labels=tuple(base_names))
+
+        coords_container = [state_coord]
+
+        if k_endog > 1:
+            endog_coord = Coord(
+                dimension=f"endog_{self.name}",
+                labels=self.observed_state_names,
             )
-            self.coords[f"shock_{self.name}"] = base_shock_names
-            self.param_info[f"sigma_{self.name}"] = {
-                "shape": (k_posdef,) if k_endog_effective == 1 else (k_endog_effective, k_posdef),
-                "constraints": "Positive",
-            }
+            coords_container.append(endog_coord)
 
-        for name in self.param_names:
-            self.param_info[name]["dims"] = self.param_dims[name]
+        if self.k_posdef > 0:
+            shock_coord = Coord(dimension=f"shock_{self.name}", labels=tuple(base_shock_names))
+            coords_container.append(shock_coord)
+
+        return tuple(coords_container)
 
     def make_symbolic_graph(self) -> None:
         k_endog = self.k_endog
@@ -287,3 +328,15 @@ class LevelTrendComponent(Component):
             diag_idx = np.diag_indices(k_posdef * k_endog_effective)
             idx = np.s_["state_cov", diag_idx[0], diag_idx[1]]
             self.ssm[idx] = (sigma_trend**2).ravel()
+
+
+def __getattr__(name: str):
+    if name == "LevelTrendComponent":
+        warnings.warn(
+            "LevelTrendComponent is deprecated and will be removed in a future release. "
+            "Use LevelTrend instead.",
+            FutureWarning,
+            stacklevel=2,
+        )
+        return LevelTrend
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

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

@@ -1,5 +1,9 @@
 import numpy as np
 
+from pymc_extras.statespace.core.properties import (
+    Coord,
+    Parameter,
+)
 from pymc_extras.statespace.models.structural.core import Component
 
 
@@ -49,7 +53,7 @@ class MeasurementError(Component):
         import pymc as pm
         import pytensor.tensor as pt
 
-        trend = st.LevelTrendComponent(order=2, innovations_order=1)
+        trend = st.LevelTrend(order=2, innovations_order=1)
         error = st.MeasurementError()
 
         ss_mod = (trend + error).build()
@@ -79,7 +83,7 @@ class MeasurementError(Component):
 
     .. code:: python
 
-        trend = st.LevelTrendComponent(order=2, innovations_order=1)
+        trend = st.LevelTrend(order=2, innovations_order=1)
         seasonal = st.TimeSeasonality(season_length=12, innovations=True)
         error = st.MeasurementError()
 
@@ -118,29 +122,35 @@ class MeasurementError(Component):
             k_posdef,
             measurement_error=True,
             combine_hidden_states=False,
-            observed_state_names=observed_state_names,
+            base_observed_state_names=observed_state_names,
             share_states=share_states,
         )
 
-    def populate_component_properties(self):
+    def set_parameters(self) -> Parameter | tuple[Parameter, ...] | None:
         k_endog = self.k_endog
         k_endog_effective = 1 if self.share_states else k_endog
 
-        self.param_names = [f"sigma_{self.name}"]
-        self.param_dims = {}
-        self.coords = {}
+        sigma_param = Parameter(
+            name=f"sigma_{self.name}",
+            shape=(k_endog_effective,) if k_endog_effective > 1 else (),
+            dims=(f"endog_{self.name}",) if k_endog_effective > 1 else None,
+            constraints="Positive",
+        )
+
+        return (sigma_param,)
+
+    def set_coords(self) -> Coord | tuple[Coord, ...] | None:
+        k_endog = self.k_endog
+        k_endog_effective = 1 if self.share_states else k_endog
+        observed_state_names = self.base_observed_state_names
+
+        coords_container = []
 
         if k_endog_effective > 1:
-            self.param_dims[f"sigma_{self.name}"] = (f"endog_{self.name}",)
-            self.coords[f"endog_{self.name}"] = self.observed_state_names
-
-        self.param_info = {
-            f"sigma_{self.name}": {
-                "shape": (k_endog_effective,) if k_endog_effective > 1 else (),
-                "constraints": "Positive",
-                "dims": (f"endog_{self.name}",) if k_endog_effective > 1 else None,
-            }
-        }
+            endog_coord = Coord(dimension=f"endog_{self.name}", labels=observed_state_names)
+            coords_container.append(endog_coord)
+
+        return tuple(coords_container) if coords_container else None
 
     def make_symbolic_graph(self) -> None:
         k_endog = self.k_endog