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

pymc_extras/statespace/core/statespace.py

@@ -19,6 +19,22 @@ from rich.box import SIMPLE_HEAD
 from rich.console import Console
 from rich.table import Table
 
+from pymc_extras.statespace.core.properties import (
+    Coord,
+    CoordInfo,
+    Data,
+    DataInfo,
+    Parameter,
+    ParameterInfo,
+    Shock,
+    ShockInfo,
+    State,
+    StateInfo,
+    SymbolicData,
+    SymbolicDataInfo,
+    SymbolicVariable,
+    SymbolicVariableInfo,
+)
 from pymc_extras.statespace.core.representation import PytensorRepresentation
 from pymc_extras.statespace.filters import (
     KalmanSmoother,
@@ -69,6 +85,22 @@ def _verify_group(group):
         raise ValueError(f'Argument "group" must be one of "prior" or "posterior", found {group}')
 
 
+def _validate_property(props, property_name, expected_type):
+    if isinstance(props, expected_type) or props is None:
+        return
+    elif not isinstance(props, tuple | list):
+        raise TypeError(
+            f"The {property_name} property must be a {expected_type.__name__} or a "
+            f"list/tuple of {expected_type.__name__} instances."
+        )
+
+    if not all(isinstance(prop, expected_type) for prop in props):
+        raise TypeError(
+            f"All elements of the {property_name} property must be instances of "
+            f"{expected_type.__name__}."
+        )
+
+
 class PyMCStateSpace:
     r"""
     Base class for Linear Gaussian Statespace models in PyMC.
@@ -236,8 +268,8 @@ class PyMCStateSpace:
         self._fit_exog_data: dict[str, dict] = {}
 
         self._needs_exog_data = None
-        self._name_to_variable = {}
-        self._name_to_data = {}
+        self._tensor_variable_info = SymbolicVariableInfo()
+        self._tensor_data_info = SymbolicDataInfo()
 
         self.k_endog = k_endog
         self.k_states = k_states
@@ -245,6 +277,8 @@ class PyMCStateSpace:
         self.measurement_error = measurement_error
         self.mode = mode
 
+        self._populate_properties()
+
         # All models contain a state space representation and a Kalman filter
         self.ssm = PytensorRepresentation(k_endog, k_states, k_posdef)
 
@@ -271,17 +305,107 @@ class PyMCStateSpace:
             console = Console()
             console.print(self.requirement_table)
 
+    def _populate_properties(self) -> None:
+        self._set_parameters()
+        self._set_states()
+        self._set_shocks()
+        self._set_coords()
+        self._set_data_info()
+
+    def set_parameters(self) -> Parameter | tuple[Parameter, ...] | None:
+        """
+        Provides parameter metadata to the model.
+
+        Optional. Default implementation sets an empty ParameterInfo.
+        Child classes can override to define model parameters.
+        """
+        return
+
+    def _set_parameters(self) -> None:
+        param_list = self.set_parameters()
+        self._param_info = ParameterInfo(parameters=param_list)
+
+    def set_states(self) -> State | tuple[State, ...] | None:
+        """
+        Provide state metadata to the model.
+
+        Optional. Default implementation creates generic state names (state_0, state_1, ...).
+        Child classes can override to provide meaningful state names.
+        """
+        hidden_states = [
+            State(name=f"hidden_state_{name}", observed=False) for name in range(self.k_states or 0)
+        ]
+        observed_states = [
+            State(name=f"observed_state_{name}", observed=True) for name in range(self.k_endog or 0)
+        ]
+        return *hidden_states, *observed_states
+
+    def _set_states(self) -> None:
+        states = self.set_states()
+        _validate_property(states, "states", State)
+
+        if isinstance(states, State):
+            states = (states,)
+
+        self._state_info = StateInfo(states=states)
+
+    def set_shocks(self) -> Shock | tuple[Shock, ...] | None:
+        """
+        Provide shock metadata to the model.
+
+        Optional. Default implementation creates generic shock names (shock_0, shock_1, ...).
+        Child classes can override to provide meaningful shock names.
+        """
+        return tuple(Shock(name=f"shock_{i}") for i in range(self.k_posdef))
+
+    def _set_shocks(self) -> None:
+        shocks = self.set_shocks()
+        _validate_property(shocks, "shocks", Shock)
+        if isinstance(shocks, Shock):
+            shocks = (shocks,)
+        self._shock_info = ShockInfo(shocks=shocks)
+
+    def default_coords(self) -> tuple[Coord, ...]:
+        return CoordInfo.default_coords_from_model(self).items
+
+    def set_coords(self) -> Coord | tuple[Coord, ...] | None:
+        """
+        Provide coordinates to the model.
+
+        Optional. Default implementation sets an empty CoordInfo.
+        Child classes can override to provide model-specific coordinates.
+        """
+        return self.default_coords()
+
+    def _set_coords(self) -> None:
+        coords = self.set_coords()
+        _validate_property(coords, "coords", Coord)
+        if isinstance(coords, Coord):
+            coords = (coords,)
+        self._coords_info = CoordInfo(coords=coords)
+
+    def set_data_info(self) -> Data | tuple[Data, ...] | None:
+        """
+        Provide data_info metadata to the model.
+
+        Optional. Default implementation sets an empty DataInfo.
+        Child classes can override if the model requires exogenous data.
+        """
+        return
+
+    def _set_data_info(self) -> None:
+        data_info = self.set_data_info()
+        _validate_property(data_info, "data_info", Data)
+        if isinstance(data_info, Data):
+            data_info = (data_info,)
+        self._data_info = DataInfo(data=data_info)
+
     def _populate_prior_requirements(self) -> None:
         """
         Add requirements about priors needed for the model to a rich table, including their names,
         shapes, named dimensions, and any parameter constraints.
         """
-        # Check that the param_info class is implemented, and also that it's a dictionary. We can't proceed if either
-        # is not true.
-        try:
-            if not isinstance(self.param_info, dict):
-                return
-        except NotImplementedError:
+        if not self.param_info:
             return
 
         if self.requirement_table is None:
@@ -296,10 +420,7 @@ class PyMCStateSpace:
         """
         Add requirements about the data needed for the model, including their names, shapes, and named dimensions.
         """
-        try:
-            if not isinstance(self.data_info, dict):
-                return
-        except NotImplementedError:
+        if not self.data_info:
             return
 
         if self.requirement_table is None:
@@ -374,24 +495,24 @@ class PyMCStateSpace:
         return self.subbed_ssm
 
     @property
-    def param_names(self) -> list[str]:
+    def param_names(self) -> tuple[str, ...]:
         """
         Names of model parameters
 
         A list of all parameters expected by the model. Each parameter will be sought inside the active PyMC model
         context when ``build_statespace_graph`` is invoked.
         """
-        raise NotImplementedError("The param_names property has not been implemented!")
+        return self._param_info.names
 
     @property
-    def data_names(self) -> list[str]:
+    def data_names(self) -> tuple[str, ...]:
         """
         Names of data variables expected by the model.
 
         This does not include the observed data series, which is automatically handled by PyMC. This property only
         needs to be implemented for models that expect exogenous data.
         """
-        return []
+        return self._data_info.exogenous_names
 
     @property
     def param_info(self) -> dict[str, dict[str, Any]]:
@@ -406,7 +527,7 @@ class PyMCStateSpace:
               positive semi-definite, etc)
             * key: "dims", value: tuple of strings
         """
-        raise NotImplementedError("The params_info property has not been implemented!")
+        return self._param_info.to_dict()
 
     @property
     def data_info(self) -> dict[str, dict[str, Any]]:
@@ -419,31 +540,30 @@ class PyMCStateSpace:
             * key: "shape", value: a tuple of integers
             * key: "dims", value: tuple of strings
         """
-        raise NotImplementedError("The data_info property has not been implemented!")
+        return self._data_info.to_dict()
 
     @property
-    def state_names(self) -> list[str]:
+    def state_names(self) -> tuple[str, ...]:
         """
         A k_states length list of strings, associated with the model's hidden states
 
         """
-
-        raise NotImplementedError("The state_names property has not been implemented!")
+        return self._state_info.unobserved_state_names
 
     @property
-    def observed_states(self) -> list[str]:
+    def observed_states(self) -> tuple[str, ...]:
         """
         A k_endog length list of strings, associated with the model's observed states
         """
-        raise NotImplementedError("The observed_states property has not been implemented!")
+        return self._state_info.observed_state_names
 
     @property
-    def shock_names(self) -> list[str]:
+    def shock_names(self) -> tuple[str, ...]:
         """
         A k_posdef length list of strings, associated with the model's shock processes
 
         """
-        raise NotImplementedError("The shock_names property has not been implemented!")
+        return self._shock_info.names
 
     @property
     def default_priors(self) -> dict[str, Callable]:
@@ -464,10 +584,10 @@ class PyMCStateSpace:
         should come from the default names defined in ``statespace.utils.constants`` for them to be detected by
         sampling methods.
         """
-        raise NotImplementedError("The coords property has not been implemented!")
+        return self._coords_info.to_dict()
 
     @property
-    def param_dims(self) -> dict[str, Sequence[str]]:
+    def param_dims(self) -> dict[str, tuple[str, ...]]:
         """
         Dictionary of named dimensions for each model parameter
 
@@ -475,8 +595,17 @@ class PyMCStateSpace:
         PyMC random variable. Dimensions should come from the default names defined in ``statespace.utils.constants``
         for them to be detected by sampling methods.
 
+        Note: Scalar parameters (with dims=None) are not included in this dictionary.
         """
-        raise NotImplementedError("The param_dims property has not been implemented!")
+        return {param.name: param.dims for param in self._param_info if param.dims is not None}
+
+    @property
+    def _name_to_variable(self):
+        return self._tensor_variable_info.to_dict()
+
+    @property
+    def _name_to_data(self):
+        return self._tensor_data_info.to_dict()
 
     def add_default_priors(self) -> None:
         """
@@ -520,14 +649,15 @@ class PyMCStateSpace:
                 f"parameters."
             )
 
-        if name in self._name_to_variable.keys():
+        if name in self._tensor_variable_info:
             raise ValueError(
                 f"{name} is already a registered placeholder variable with shape "
-                f"{self._name_to_variable[name].type.shape}"
+                f"{self._tensor_variable_info[name].type.shape}"
             )
 
         placeholder = pt.tensor(name, shape=shape, dtype=dtype)
-        self._name_to_variable[name] = placeholder
+        tensor_var = SymbolicVariable(name=name, symbolic_variable=placeholder)
+        self._tensor_variable_info = self._tensor_variable_info.add(tensor_var)
         return placeholder
 
     def make_and_register_data(
@@ -559,14 +689,15 @@ class PyMCStateSpace:
                 f"parameters."
             )
 
-        if name in self._name_to_data.keys():
+        if name in self._tensor_data_info:
             raise ValueError(
                 f"{name} is already a registered placeholder variable with shape "
-                f"{self._name_to_data[name].type.shape}"
+                f"{self._tensor_data_info[name].type.shape}"
             )
 
         placeholder = pt.tensor(name, shape=shape, dtype=dtype)
-        self._name_to_data[name] = placeholder
+        tensor_data = SymbolicData(name=name, symbolic_data=placeholder)
+        self._tensor_data_info = self._tensor_data_info.add(tensor_data)
         return placeholder
 
     def make_symbolic_graph(self) -> None:
@@ -741,10 +872,8 @@ class PyMCStateSpace:
 
         Only used when models require exogenous data. The observed data is not added to the model using this method!
         """
-
-        try:
-            data_names = self.data_names
-        except NotImplementedError:
+        data_names = self.data_names
+        if not data_names:
             return
 
         pymc_model = modelcontext(None)
@@ -892,7 +1021,6 @@ class PyMCStateSpace:
             .. deprecated:: 0.2.5
                 The `mode` argument is deprecated and will be removed in a future version. Pass ``mode`` to the
                 model constructor, or manually specify ``compile_kwargs`` in sampling functions instead.
-
         """
         if mode is not None:
             warnings.warn(
@@ -1430,7 +1558,11 @@ class PyMCStateSpace:
         """
 
         return self._sample_conditional(
-            idata=idata, group="prior", random_seed=random_seed, mvn_method=mvn_method, **kwargs
+            idata=idata,
+            group="prior",
+            random_seed=random_seed,
+            mvn_method=mvn_method,
+            **kwargs,
         )
 
     def sample_conditional_posterior(
@@ -1473,7 +1605,11 @@ class PyMCStateSpace:
         """
 
         return self._sample_conditional(
-            idata=idata, group="posterior", random_seed=random_seed, mvn_method=mvn_method, **kwargs
+            idata=idata,
+            group="posterior",
+            random_seed=random_seed,
+            mvn_method=mvn_method,
+            **kwargs,
         )
 
     def sample_unconditional_prior(
@@ -2080,11 +2216,11 @@ class PyMCStateSpace:
         forecast_index: pd.RangeIndex | pd.DatetimeIndex,
         name=None,
     ):
-        try:
-            var_to_dims = {key: info["dims"][1:] for key, info in self.data_info.items()}
-        except NotImplementedError:
+        if not self.data_info:
             return scenario
 
+        var_to_dims = {key: info["dims"][1:] for key, info in self.data_info.items()}
+
         if any(len(dims) > 1 for dims in var_to_dims.values()):
             raise NotImplementedError(">2d exogenous data is not yet supported.")
         coords = {
@@ -2500,13 +2636,14 @@ class PyMCStateSpace:
                 next_x = c + T @ x + R @ shock
                 return next_x
 
-            irf, updates = pytensor.scan(
+            irf = pytensor.scan(
                 irf_step,
                 sequences=[shock_trajectory],
                 outputs_info=[x0],
                 non_sequences=[c, T, R],
                 n_steps=n_steps,
                 strict=True,
+                return_updates=False,
             )
 
             pm.Deterministic("irf", irf, dims=[TIME_DIM, ALL_STATE_DIM])

pymc_extras/statespace/filters/distributions.py

@@ -1,4 +1,3 @@
-import numpy as np
 import pymc as pm
 import pytensor
 import pytensor.tensor as pt
@@ -8,7 +7,9 @@ from pymc.distributions.dist_math import check_parameters
 from pymc.distributions.distribution import Continuous, SymbolicRandomVariable
 from pymc.distributions.shape_utils import get_support_shape_1d
 from pymc.logprob.abstract import _logprob
+from pymc.pytensorf import normalize_rng_param
 from pytensor.graph.basic import Node
+from pytensor.tensor.random import multivariate_normal
 
 floatX = pytensor.config.floatX
 COV_ZERO_TOL = 0
@@ -152,6 +153,7 @@ class _LinearGaussianStateSpace(Continuous):
         Q,
         steps,
         size=None,
+        rng=None,
         sequence_names=None,
         append_x0=True,
         method="svd",
@@ -178,7 +180,7 @@ class _LinearGaussianStateSpace(Continuous):
         ]
         non_sequences = [x for x in [c_, d_, T_, Z_, R_, H_, Q_] if x not in sequences]
 
-        rng = pytensor.shared(np.random.default_rng())
+        rng = normalize_rng_param(rng)
 
         def sort_args(args):
             sorted_args = []
@@ -197,10 +199,9 @@ class _LinearGaussianStateSpace(Continuous):
         n_seq = len(sequence_names)
 
         def step_fn(*args):
-            seqs, state, non_seqs = args[:n_seq], args[n_seq], args[n_seq + 1 :]
-            non_seqs, rng = non_seqs[:-1], non_seqs[-1]
+            seqs, (rng, state, *non_seqs) = args[:n_seq], args[n_seq:]
 
-            c, d, T, Z, R, H, Q = sort_args(seqs + non_seqs)
+            c, d, T, Z, R, H, Q = sort_args((*seqs, *non_seqs))
             k = T.shape[0]
             a = state[:k]
 
@@ -219,7 +220,7 @@ class _LinearGaussianStateSpace(Continuous):
 
             next_state = pt.concatenate([a_next, y_next], axis=0)
 
-            return next_state, {rng: next_rng}
+            return next_rng, next_state
 
         Z_init = Z_ if Z_ in non_sequences else Z_[0]
         H_init = H_ if H_ in non_sequences else H_[0]
@@ -229,13 +230,14 @@ class _LinearGaussianStateSpace(Continuous):
 
         init_dist_ = pt.concatenate([init_x_, init_y_], axis=0)
 
-        statespace, updates = pytensor.scan(
+        ss_rng, statespace = pytensor.scan(
             step_fn,
-            outputs_info=[init_dist_],
+            outputs_info=[rng, init_dist_],
             sequences=None if len(sequences) == 0 else sequences,
-            non_sequences=[*non_sequences, rng],
+            non_sequences=[*non_sequences],
             n_steps=steps,
             strict=True,
+            return_updates=False,
         )
 
         if append_x0:
@@ -245,7 +247,6 @@ class _LinearGaussianStateSpace(Continuous):
             statespace_ = statespace
             statespace_ = pt.specify_shape(statespace_, (steps, None))
 
-        (ss_rng,) = tuple(updates.values())
         linear_gaussian_ss_op = LinearGaussianStateSpaceRV(
             inputs=[a0_, P0_, c_, d_, T_, Z_, R_, H_, Q_, steps, rng],
             outputs=[ss_rng, statespace_],
@@ -368,41 +369,25 @@ class SequenceMvNormal(Continuous):
 
     @classmethod
     def dist(cls, mus, covs, logp, method="svd", **kwargs):
+        mus, covs, logp = map(pt.as_tensor_variable, (mus, covs, logp))
         return super().dist([mus, covs, logp], method=method, **kwargs)
 
     @classmethod
-    def rv_op(cls, mus, covs, logp, method="svd", size=None):
-        # Batch dimensions (if any) will be on the far left, but scan requires time to be there instead
-        if mus.ndim > 2:
-            mus = pt.moveaxis(mus, -2, 0)
-        if covs.ndim > 3:
-            covs = pt.moveaxis(covs, -3, 0)
-
-        mus_, covs_ = mus.type(), covs.type()
-
+    def rv_op(cls, mus, covs, logp, method="svd", size=None, rng=None):
+        rng = normalize_rng_param(rng)
         logp_ = logp.type()
-        rng = pytensor.shared(np.random.default_rng())
 
-        def step(mu, cov, rng):
-            new_rng, mvn = pm.MvNormal.dist(mu=mu, cov=cov, rng=rng, method=method).owner.outputs
-            return mvn, {rng: new_rng}
-
-        mvn_seq, updates = pytensor.scan(
-            step, sequences=[mus_, covs_], non_sequences=[rng], strict=True, n_steps=mus_.shape[0]
-        )
-        mvn_seq = pt.specify_shape(mvn_seq, mus.type.shape)
-
-        # Move time axis back to position -2 so batches are on the left
-        if mvn_seq.ndim > 2:
-            mvn_seq = pt.moveaxis(mvn_seq, 0, -2)
-
-        (seq_mvn_rng,) = tuple(updates.values())
+        mus_, covs_ = mus.type(), covs.type()
+        seq_mvn_rng, mvn_seq = multivariate_normal(
+            mean=mus_, cov=covs_, rng=rng, method=method
+        ).owner.outputs
 
         mvn_seq_op = KalmanFilterRV(
             inputs=[mus_, covs_, logp_, rng], outputs=[seq_mvn_rng, mvn_seq], ndim_supp=2
         )
 
         mvn_seq = mvn_seq_op(mus, covs, logp, rng)
+
         return mvn_seq
 
 

pymc_extras/statespace/filters/kalman_filter.py

@@ -8,7 +8,7 @@ from pymc.pytensorf import constant_fold
 from pytensor.graph.basic import Variable
 from pytensor.raise_op import Assert
 from pytensor.tensor import TensorVariable
-from pytensor.tensor.slinalg import solve_triangular
+from pytensor.tensor.linalg import solve_triangular
 
 from pymc_extras.statespace.filters.utilities import (
     quad_form_sym,
@@ -148,10 +148,9 @@ class BaseFilter(ABC):
         R,
         H,
         Q,
-        return_updates=False,
         missing_fill_value=None,
         cov_jitter=None,
-    ) -> list[TensorVariable] | tuple[list[TensorVariable], dict]:
+    ) -> list[TensorVariable]:
         """
         Construct the computation graph for the Kalman filter. See [1] for details.
 
@@ -211,20 +210,17 @@ class BaseFilter(ABC):
         if len(sequences) > 0:
             sequences = self.add_check_on_time_varying_shapes(data, sequences)
 
-        results, updates = pytensor.scan(
+        results = pytensor.scan(
             self.kalman_step,
             sequences=[data, *sequences],
             outputs_info=[None, a0, None, None, P0, None, None],
             non_sequences=non_sequences,
             name="forward_kalman_pass",
             strict=False,
+            return_updates=False,
         )
 
-        filter_results = self._postprocess_scan_results(results, a0, P0, n=data.type.shape[0])
-
-        if return_updates:
-            return filter_results, updates
-        return filter_results
+        return self._postprocess_scan_results(results, a0, P0, n=data.type.shape[0])
 
     def _postprocess_scan_results(self, results, a0, P0, n) -> list[TensorVariable]:
         """
@@ -652,7 +648,9 @@ class SquareRootFilter(BaseFilter):
         y_hat = Z.dot(a) + d
         v = y - y_hat
 
-        H_chol = pytensor.ifelse(pt.all(pt.eq(H, 0.0)), H, pt.linalg.cholesky(H, lower=True))
+        H_chol = pytensor.ifelse(
+            pt.all(pt.eq(H, 0.0)), H, pt.linalg.cholesky(H, lower=True, on_error="nan")
+        )
 
         # The following notation comes from https://ipnpr.jpl.nasa.gov/progress_report/42-233/42-233A.pdf
         # Construct upper-triangular block matrix A = [[chol(H), Z @ L_pred],
@@ -694,8 +692,10 @@ class SquareRootFilter(BaseFilter):
             """
             return [a, P_chol, pt.zeros(())]
 
+        degenerate = pt.eq(all_nan_flag, 1.0)
+        F_chol = pytensor.ifelse(degenerate, pt.eye(*F_chol.shape), F_chol)
         [a_filtered, P_chol_filtered, ll] = pytensor.ifelse(
-            pt.eq(all_nan_flag, 1.0),
+            degenerate,
             compute_degenerate(P_chol_filtered, F_chol, K_F_chol, v),
             compute_non_degenerate(P_chol_filtered, F_chol, K_F_chol, v),
         )
@@ -786,11 +786,12 @@ class UnivariateFilter(BaseFilter):
         H_masked = W.dot(H)
         y_masked = pt.set_subtensor(y[nan_mask], 0.0)
 
-        result, updates = pytensor.scan(
+        result = pytensor.scan(
             self._univariate_inner_filter_step,
             sequences=[y_masked, Z_masked, d, pt.diag(H_masked), nan_mask],
             outputs_info=[a, P, None, None, None],
             name="univariate_inner_scan",
+            return_updates=False,
         )
 
         a_filtered, P_filtered, obs_mu, obs_cov, ll_inner = result

pymc_extras/statespace/filters/kalman_smoother.py

@@ -76,16 +76,16 @@ class KalmanSmoother:
         self.seq_names = seq_names
         self.non_seq_names = non_seq_names
 
-        smoother_result, updates = pytensor.scan(
+        smoothed_states, smoothed_covariances = pytensor.scan(
             self.smoother_step,
             sequences=[filtered_states[:-1], filtered_covariances[:-1], *sequences],
             outputs_info=[a_last, P_last],
             non_sequences=non_sequences,
             go_backwards=True,
             name="kalman_smoother",
+            return_updates=False,
         )
 
-        smoothed_states, smoothed_covariances = smoother_result
         smoothed_states = pt.concatenate(
             [smoothed_states[::-1], pt.expand_dims(a_last, axis=(0,))], axis=0
         )