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

pymc_extras/statespace/models/VARMAX.py

@@ -1,15 +1,19 @@
-from collections.abc import Sequence
-from typing import Any
-
 import numpy as np
 import pytensor
 import pytensor.tensor as pt
 
 from pytensor.compile.mode import Mode
-from pytensor.tensor.slinalg import solve_discrete_lyapunov
-
+from pytensor.tensor.linalg import solve_discrete_lyapunov
+
+from pymc_extras.statespace.core.properties import (
+    Coord,
+    Data,
+    Parameter,
+    Shock,
+    State,
+)
 from pymc_extras.statespace.core.statespace import PyMCStateSpace
-from pymc_extras.statespace.models.utilities import make_default_coords, validate_names
+from pymc_extras.statespace.models.utilities import validate_names
 from pymc_extras.statespace.utils.constants import (
     ALL_STATE_AUX_DIM,
     ALL_STATE_DIM,
@@ -212,110 +216,96 @@ class BayesianVARMAX(PyMCStateSpace):
             "sigma_obs": k_endog * self.measurement_error,
         }
 
-    @property
-    def param_names(self):
-        names = ["x0", "P0", "ar_params", "ma_params", "state_cov", "sigma_obs"]
-        if self.stationary_initialization:
-            names.remove("P0")
-            names.remove("x0")
-        if not self.measurement_error:
-            names.remove("sigma_obs")
-        if self.p == 0:
-            names.remove("ar_params")
-        if self.q == 0:
-            names.remove("ma_params")
-
-        # Add exogenous regression coefficents rather than remove, since we might have to handle
-        # several (if self.exog_state_names is a dict)
-        if isinstance(self.exog_state_names, list):
-            names.append("beta_exog")
-        elif isinstance(self.exog_state_names, dict):
-            names.extend([f"beta_{name}" for name in self.exog_state_names.keys()])
+    def set_parameters(self) -> Parameter | tuple[Parameter, ...] | None:
+        k_endog = self.k_endog
+        k_states = self.k_states
+        k_posdef = self.k_posdef
 
-        return names
-
-    @property
-    def param_info(self) -> dict[str, dict[str, Any]]:
-        info = {
-            "x0": {
-                "shape": (self.k_states,),
-                "constraints": None,
-            },
-            "P0": {
-                "shape": (self.k_states, self.k_states),
-                "constraints": "Positive Semi-definite",
-            },
-            "sigma_obs": {
-                "shape": (self.k_endog, self.k_endog),
-                "constraints": "Positive Semi-definite",
-            },
-            "state_cov": {
-                "shape": (self.k_posdef, self.k_posdef),
-                "constraints": "Positive Semi-definite",
-            },
-            "ar_params": {
-                "shape": (self.k_endog, self.p, self.k_endog),
-                "constraints": "None",
-            },
-            "ma_params": {
-                "shape": (self.k_endog, self.q, self.k_endog),
-                "constraints": "None",
-            },
-        }
+        parameters = []
 
-        if isinstance(self.exog_state_names, list):
-            k_exog = len(self.exog_state_names)
-            info["beta_exog"] = {
-                "shape": (self.k_endog, k_exog),
-                "constraints": "None",
-            }
+        if not self.stationary_initialization:
+            parameters.append(
+                Parameter(
+                    name="x0",
+                    shape=(k_states,),
+                    dims=(ALL_STATE_DIM,),
+                    constraints=None,
+                )
+            )
+            parameters.append(
+                Parameter(
+                    name="P0",
+                    shape=(k_states, k_states),
+                    dims=(ALL_STATE_DIM, ALL_STATE_AUX_DIM),
+                    constraints="Positive Semi-definite",
+                )
+            )
 
-        elif isinstance(self.exog_state_names, dict):
-            for name, exog_names in self.exog_state_names.items():
-                k_exog = len(exog_names)
-                info[f"beta_{name}"] = {
-                    "shape": (k_exog,),
-                    "constraints": "None",
-                }
+        if self.p > 0:
+            parameters.append(
+                Parameter(
+                    name="ar_params",
+                    shape=(k_endog, self.p, k_endog),
+                    dims=(OBS_STATE_DIM, AR_PARAM_DIM, OBS_STATE_AUX_DIM),
+                    constraints=None,
+                )
+            )
 
-        for name in self.param_names:
-            info[name]["dims"] = self.param_dims[name]
+        if self.q > 0:
+            parameters.append(
+                Parameter(
+                    name="ma_params",
+                    shape=(k_endog, self.q, k_endog),
+                    dims=(OBS_STATE_DIM, MA_PARAM_DIM, OBS_STATE_AUX_DIM),
+                    constraints=None,
+                )
+            )
 
-        return {name: info[name] for name in self.param_names}
+        parameters.append(
+            Parameter(
+                name="state_cov",
+                shape=(k_posdef, k_posdef),
+                dims=(SHOCK_DIM, SHOCK_AUX_DIM),
+                constraints="Positive Semi-definite",
+            )
+        )
 
-    @property
-    def data_info(self) -> dict[str, dict[str, Any]]:
-        info = {}
+        if self.measurement_error:
+            parameters.append(
+                Parameter(
+                    name="sigma_obs",
+                    shape=(k_endog,),
+                    dims=(OBS_STATE_DIM,),
+                    constraints="Positive",
+                )
+            )
 
+        # Handle exogenous parameters
         if isinstance(self.exog_state_names, list):
-            info = {
-                "exogenous_data": {
-                    "dims": (TIME_DIM, EXOG_STATE_DIM),
-                    "shape": (None, self.k_exog),
-                }
-            }
-
+            k_exog = len(self.exog_state_names)
+            parameters.append(
+                Parameter(
+                    name="beta_exog",
+                    shape=(k_endog, k_exog),
+                    dims=(OBS_STATE_DIM, EXOG_STATE_DIM),
+                    constraints=None,
+                )
+            )
         elif isinstance(self.exog_state_names, dict):
-            info = {
-                f"{endog_state}_exogenous_data": {
-                    "dims": (TIME_DIM, f"{EXOG_STATE_DIM}_{endog_state}"),
-                    "shape": (None, len(exog_names)),
-                }
-                for endog_state, exog_names in self.exog_state_names.items()
-            }
+            for name, exog_names in self.exog_state_names.items():
+                k_exog = len(exog_names)
+                parameters.append(
+                    Parameter(
+                        name=f"beta_{name}",
+                        shape=(k_exog,),
+                        dims=(f"{EXOG_STATE_DIM}_{name}",),
+                        constraints=None,
+                    )
+                )
 
-        return info
+        return tuple(parameters)
 
-    @property
-    def data_names(self) -> list[str]:
-        if isinstance(self.exog_state_names, list):
-            return ["exogenous_data"]
-        elif isinstance(self.exog_state_names, dict):
-            return [f"{endog_state}_exogenous_data" for endog_state in self.exog_state_names.keys()]
-        return []
-
-    @property
-    def state_names(self):
+    def set_states(self) -> State | tuple[State, ...] | None:
         state_names = self.endog_names.copy()
         state_names += [
             f"L{i + 1}_{state}" for i in range(self.p - 1) for state in self.endog_names
@@ -324,66 +314,65 @@ class BayesianVARMAX(PyMCStateSpace):
             f"L{i + 1}_{state}_innov" for i in range(self.q) for state in self.endog_names
         ]
 
-        return state_names
+        hidden_states = [State(name=name, observed=False) for name in state_names]
 
-    @property
-    def observed_states(self):
-        return self.endog_names
+        # The first k_endog states are observed
+        observed_states = [State(name=name, observed=True) for name in self.endog_names]
 
-    @property
-    def shock_names(self):
-        return self.endog_names
+        return *hidden_states, *observed_states
 
-    @property
-    def default_priors(self):
-        raise NotImplementedError
+    def set_shocks(self) -> Shock | tuple[Shock, ...] | None:
+        return tuple(Shock(name=name) for name in self.endog_names)
 
-    @property
-    def coords(self) -> dict[str, Sequence]:
-        coords = make_default_coords(self)
-        if self.p > 0:
-            coords.update({AR_PARAM_DIM: list(range(1, self.p + 1))})
-        if self.q > 0:
-            coords.update({MA_PARAM_DIM: list(range(1, self.q + 1))})
+    def set_data_info(self) -> tuple[Data, ...] | None:
+        data = []
 
         if isinstance(self.exog_state_names, list):
-            coords[EXOG_STATE_DIM] = self.exog_state_names
+            k_exog = len(self.exog_state_names)
+            data.append(
+                Data(
+                    name="exogenous_data",
+                    shape=(None, k_exog),
+                    dims=(TIME_DIM, EXOG_STATE_DIM),
+                    is_exogenous=True,
+                )
+            )
         elif isinstance(self.exog_state_names, dict):
-            for name, exog_names in self.exog_state_names.items():
-                coords[f"{EXOG_STATE_DIM}_{name}"] = exog_names
+            for endog_state, exog_names in self.exog_state_names.items():
+                k_exog = len(exog_names)
+                data.append(
+                    Data(
+                        name=f"{endog_state}_exogenous_data",
+                        shape=(None, k_exog),
+                        dims=(TIME_DIM, f"{EXOG_STATE_DIM}_{endog_state}"),
+                        is_exogenous=True,
+                    )
+                )
 
-        return coords
+        return tuple(data)
 
-    @property
-    def param_dims(self):
-        coord_map = {
-            "x0": (ALL_STATE_DIM,),
-            "P0": (ALL_STATE_DIM, ALL_STATE_AUX_DIM),
-            "sigma_obs": (OBS_STATE_DIM,),
-            "state_cov": (SHOCK_DIM, SHOCK_AUX_DIM),
-            "ar_params": (OBS_STATE_DIM, AR_PARAM_DIM, OBS_STATE_AUX_DIM),
-            "ma_params": (OBS_STATE_DIM, MA_PARAM_DIM, OBS_STATE_AUX_DIM),
-        }
+    def set_coords(self) -> Coord | tuple[Coord, ...] | None:
+        coords = list(self.default_coords())
 
-        if not self.measurement_error:
-            del coord_map["sigma_obs"]
-        if self.p == 0:
-            del coord_map["ar_params"]
-        if self.q == 0:
-            del coord_map["ma_params"]
-        if self.stationary_initialization:
-            del coord_map["P0"]
-            del coord_map["x0"]
+        # AR/MA param coords
+        if self.p > 0:
+            coords.append(Coord(dimension=AR_PARAM_DIM, labels=tuple(range(1, self.p + 1))))
+
+        if self.q > 0:
+            coords.append(Coord(dimension=MA_PARAM_DIM, labels=tuple(range(1, self.q + 1))))
 
+        # Exogenous coords
         if isinstance(self.exog_state_names, list):
-            coord_map["beta_exog"] = (OBS_STATE_DIM, EXOG_STATE_DIM)
+            coords.append(Coord(dimension=EXOG_STATE_DIM, labels=tuple(self.exog_state_names)))
         elif isinstance(self.exog_state_names, dict):
-            # If each state has its own exogenous variables, each parameter needs it own dim, since we expect the
-            # dim labels to all be different (otherwise we'd be in the list case).
-            for name in self.exog_state_names.keys():
-                coord_map[f"beta_{name}"] = (f"{EXOG_STATE_DIM}_{name}",)
+            for name, exog_names in self.exog_state_names.items():
+                coords.append(Coord(dimension=f"{EXOG_STATE_DIM}_{name}", labels=tuple(exog_names)))
 
-        return coord_map
+        return tuple(coords)
+
+    @property
+    def default_priors(self):
+        raise NotImplementedError
 
     def add_default_priors(self):
         raise NotImplementedError

pymc_extras/statespace/models/__init__.py

@@ -1,6 +1,13 @@
 from pymc_extras.statespace.models import structural
+from pymc_extras.statespace.models.DFM import BayesianDynamicFactor
 from pymc_extras.statespace.models.ETS import BayesianETS
 from pymc_extras.statespace.models.SARIMAX import BayesianSARIMAX
 from pymc_extras.statespace.models.VARMAX import BayesianVARMAX
 
-__all__ = ["BayesianSARIMAX", "BayesianVARMAX", "BayesianETS", "structural"]
+__all__ = [
+    "BayesianSARIMAX",
+    "BayesianVARMAX",
+    "BayesianETS",
+    "BayesianDynamicFactor",
+    "structural",
+]

pymc_extras/statespace/models/structural/__init__.py

@@ -1,21 +1,43 @@
+import warnings
+
 from pymc_extras.statespace.models.structural.components.autoregressive import (
-    AutoregressiveComponent,
+    Autoregressive,
 )
-from pymc_extras.statespace.models.structural.components.cycle import CycleComponent
-from pymc_extras.statespace.models.structural.components.level_trend import LevelTrendComponent
+from pymc_extras.statespace.models.structural.components.cycle import Cycle
+from pymc_extras.statespace.models.structural.components.level_trend import LevelTrend
 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 import Regression
 from pymc_extras.statespace.models.structural.components.seasonality import (
     FrequencySeasonality,
     TimeSeasonality,
 )
 
+_DEPRECATED_NAMES = {
+    "LevelTrendComponent": LevelTrend,
+    "CycleComponent": Cycle,
+    "RegressionComponent": Regression,
+    "AutoregressiveComponent": Autoregressive,
+}
+
+
+def __getattr__(name: str):
+    if name in _DEPRECATED_NAMES:
+        warnings.warn(
+            f"{name} is deprecated and will be removed in a future release. "
+            f"Use {_DEPRECATED_NAMES[name].__name__} instead.",
+            FutureWarning,
+            stacklevel=2,
+        )
+        return _DEPRECATED_NAMES[name]
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
 __all__ = [
-    "AutoregressiveComponent",
-    "CycleComponent",
+    "Autoregressive",
+    "Cycle",
     "FrequencySeasonality",
-    "LevelTrendComponent",
+    "LevelTrend",
     "MeasurementError",
-    "RegressionComponent",
+    "Regression",
     "TimeSeasonality",
 ]

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

@@ -1,12 +1,19 @@
+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 AR_PARAM_DIM
 
 
-class AutoregressiveComponent(Component):
+class Autoregressive(Component):
     r"""
     Autoregressive timeseries component
 
@@ -58,8 +65,8 @@ class AutoregressiveComponent(Component):
         import pymc as pm
         import pytensor.tensor as pt
 
-        trend = st.LevelTrendComponent(order=1, innovations_order=0)
-        ar = st.AutoregressiveComponent(2)
+        trend = st.LevelTrend(order=1, innovations_order=0)
+        ar = st.Autoregressive(2)
         ss_mod = (trend + ar).build()
 
         with pm.Model(coords=ss_mod.coords) as model:
@@ -94,6 +101,8 @@ class AutoregressiveComponent(Component):
         self.order = order
         self.ar_lags = ar_lags
 
+        state_names = [f"L{i + 1}_{name}" for i in range(k_states)]
+
         super().__init__(
             name=name,
             k_endog=k_endog,
@@ -101,61 +110,82 @@ class AutoregressiveComponent(Component):
             k_posdef=k_posdef,
             measurement_error=True,
             combine_hidden_states=True,
-            observed_state_names=observed_state_names,
+            base_observed_state_names=observed_state_names,
+            base_state_names=state_names,
             obs_state_idxs=np.tile(np.r_[[1.0], np.zeros(k_states - 1)], k_endog_effective),
             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  # this is also the number of AR lags
-        base_names = [f"L{i + 1}_{self.name}" for i in range(k_states)]
+    def set_states(self) -> State | tuple[State, ...] | None:
+        base_names = self.base_state_names
+        observed_state_names = self.base_observed_state_names
 
         if self.share_states:
-            self.state_names = [f"{name}[shared]" for name in base_names]
-            self.shock_names = [f"{self.name}[shared]"]
+            state_names = [f"{name}[shared]" for name in base_names]
         else:
-            self.state_names = [
+            state_names = [
                 f"{name}[{state_name}]"
-                for state_name in self.observed_state_names
+                for state_name in observed_state_names
                 for name in base_names
             ]
-            self.shock_names = [
-                f"{self.name}[{obs_name}]" for obs_name in self.observed_state_names
-            ]
 
-        self.param_names = [f"params_{self.name}", f"sigma_{self.name}"]
-        self.param_dims = {f"params_{self.name}": (f"lag_{self.name}",)}
-        self.coords = {f"lag_{self.name}": self.ar_lags.tolist()}
+        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_effective > 1:
-            self.param_dims[f"params_{self.name}"] = (
+    def set_parameters(self) -> Parameter | tuple[Parameter, ...] | 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  # this is also the number of AR lags
+
+        ar_param = Parameter(
+            name=f"params_{self.name}",
+            shape=(k_endog_effective, k_states) if k_endog_effective > 1 else (k_states,),
+            dims=(f"lag_{self.name}",)
+            if k_endog_effective == 1
+            else (
                 f"endog_{self.name}",
                 f"lag_{self.name}",
-            )
-            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"params_{self.name}": {
-                "shape": (k_endog_effective, k_states) if k_endog_effective > 1 else (k_states,),
-                "constraints": None,
-                "dims": (AR_PARAM_DIM,)
-                if k_endog_effective == 1
-                else (
-                    f"endog_{self.name}",
-                    f"lag_{self.name}",
-                ),
-            },
-            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,
-            },
-        }
+            ),
+            constraints=None,
+        )
+
+        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 ar_param, sigma_param
+
+    def set_shocks(self) -> Shock | tuple[Shock, ...] | None:
+        observed_state_names = self.observed_state_names
+
+        if self.share_states:
+            shock_names = [f"{self.name}[shared]"]
+        else:
+            shock_names = [f"{self.name}[{obs_name}]" for obs_name in observed_state_names]
+
+        return tuple(Shock(name=name) for name in shock_names)
+
+    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.observed_state_names
+
+        lag_coord = Coord(dimension=f"lag_{self.name}", labels=self.ar_lags.tolist())
+
+        coord_container = [lag_coord]
+
+        if k_endog_effective > 1:
+            endog_coord = Coord(dimension=f"endog_{self.name}", labels=observed_state_names)
+            coord_container.append(endog_coord)
+
+        return tuple(coord_container)
 
     def make_symbolic_graph(self) -> None:
         k_endog = self.k_endog
@@ -211,3 +241,15 @@ class AutoregressiveComponent(Component):
 
         cov_idx = ("state_cov", *np.diag_indices(k_posdef))
         self.ssm[cov_idx] = sigma_ar**2
+
+
+def __getattr__(name: str):
+    if name == "AutoregressiveComponent":
+        warnings.warn(
+            "AutoregressiveComponent is deprecated and will be removed in a future release. "
+            "Use Autoregressive instead.",
+            FutureWarning,
+            stacklevel=2,
+        )
+        return Autoregressive
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")