pymc-extras 0.2.7__py3-none-any.whl → 0.4.0__py3-none-any.whl

pymc_extras/inference/pathfinder/pathfinder.py

@@ -38,16 +38,15 @@ from pymc.blocking import DictToArrayBijection, RaveledVars
 from pymc.initial_point import make_initial_point_fn
 from pymc.model import modelcontext
 from pymc.model.core import Point
+from pymc.progress_bar import CustomProgress, default_progress_theme
 from pymc.pytensorf import (
     compile,
     find_rng_nodes,
     reseed_rngs,
 )
 from pymc.util import (
-    CustomProgress,
     RandomSeed,
     _get_seeds_per_chain,
-    default_progress_theme,
     get_default_varnames,
 )
 from pytensor.compile.function.types import Function
@@ -63,7 +62,7 @@ from rich.text import Text
 # TODO: change to typing.Self after Python versions greater than 3.10
 from typing_extensions import Self
 
-from pymc_extras.inference.laplace import add_data_to_inferencedata
+from pymc_extras.inference.laplace_approx.idata import add_data_to_inference_data
 from pymc_extras.inference.pathfinder.importance_sampling import (
     importance_sampling as _importance_sampling,
 )
@@ -1759,6 +1758,6 @@ def fit_pathfinder(
         importance_sampling=importance_sampling,
     )
 
-    idata = add_data_to_inferencedata(idata, progressbar, model, compile_kwargs)
+    idata = add_data_to_inference_data(idata, progressbar, model, compile_kwargs)
 
     return idata

pymc_extras/linearmodel.py

@@ -2,10 +2,12 @@ import numpy as np
 import pandas as pd
 import pymc as pm
 
+from sklearn.base import BaseEstimator
+
 from pymc_extras.model_builder import ModelBuilder
 
 
-class LinearModel(ModelBuilder):
+class LinearModel(ModelBuilder, BaseEstimator):
     def __init__(
         self, model_config: dict | None = None, sampler_config: dict | None = None, nsamples=100
     ):

pymc_extras/model/marginal/graph_analysis.py

@@ -5,6 +5,7 @@ from itertools import zip_longest
 
 from pymc import SymbolicRandomVariable
 from pymc.model.fgraph import ModelVar
+from pymc.variational.minibatch_rv import MinibatchRandomVariable
 from pytensor.graph import Variable, ancestors
 from pytensor.graph.basic import io_toposort
 from pytensor.tensor import TensorType, TensorVariable
@@ -313,6 +314,9 @@ def _subgraph_batch_dim_connection(var_dims: VAR_DIMS, input_vars, output_vars)
 
             var_dims[node.outputs[0]] = output_dims
 
+        elif isinstance(node.op, MinibatchRandomVariable):
+            var_dims[node.outputs[0]] = inputs_dims[0]
+
         else:
             raise NotImplementedError(f"Marginalization through operation {node} not supported.")
 

pymc_extras/prior.py

@@ -84,6 +84,7 @@ from __future__ import annotations
 import copy
 
 from collections.abc import Callable
+from functools import partial
 from inspect import signature
 from typing import Any, Protocol, runtime_checkable
 
@@ -278,7 +279,7 @@ class VariableFactory(Protocol):
 def sample_prior(
     factory: VariableFactory,
     coords=None,
-    name: str = "var",
+    name: str = "variable",
     wrap: bool = False,
     **sample_prior_predictive_kwargs,
 ) -> xr.Dataset:
@@ -292,7 +293,7 @@ def sample_prior(
         The coordinates for the variable, by default None.
         Only required if the dims are specified.
     name : str, optional
-        The name of the variable, by default "var".
+        The name of the variable, by default "variable".
     wrap : bool, optional
         Whether to wrap the variable in a `pm.Deterministic` node, by default False.
     sample_prior_predictive_kwargs : dict
@@ -362,7 +363,7 @@ class Prior:
 
     - `preliz` attribute to get the equivalent distribution in `preliz`
     - `sample_prior` method to sample from the prior
-    - `graph` get a dummy model graph with the distribution
+    - `to_graph` get a dummy model graph with the distribution
     - `constrain` to shift the distribution to a different range
 
     Parameters
@@ -900,7 +901,7 @@ class Prior:
     def sample_prior(
         self,
         coords=None,
-        name: str = "var",
+        name: str = "variable",
         **sample_prior_predictive_kwargs,
     ) -> xr.Dataset:
         """Sample the prior distribution for the variable.
@@ -911,7 +912,7 @@ class Prior:
             The coordinates for the variable, by default None.
             Only required if the dims are specified.
         name : str, optional
-            The name of the variable, by default "var".
+            The name of the variable, by default "variable".
         sample_prior_predictive_kwargs : dict
             Additional arguments to pass to `pm.sample_prior_predictive`.
 
@@ -1175,7 +1176,7 @@ class Censored:
         """Create a censored distribution from a dictionary."""
         data = data["data"]
         return cls(  # type: ignore
-            distribution=Prior.from_dict(data["dist"]),
+            distribution=deserialize(data["dist"]),
             lower=data["lower"],
             upper=data["upper"],
         )
@@ -1354,3 +1355,34 @@ def _is_censored_type(data: dict) -> bool:
 
 register_deserialization(is_type=_is_prior_type, deserialize=Prior.from_dict)
 register_deserialization(is_type=_is_censored_type, deserialize=Censored.from_dict)
+
+
+def __getattr__(name: str):
+    """Get Prior class through the module.
+
+    Examples
+    --------
+    Create a normal distribution.
+
+    .. code-block:: python
+
+        from pymc_extras.prior import Normal
+
+        dist = Normal(mu=1, sigma=2)
+
+    Create a hierarchical normal distribution.
+
+    .. code-block:: python
+
+        import pymc_extras.prior as pr
+
+        dist = pr.Normal(mu=pr.Normal(), sigma=pr.HalfNormal(), dims="channel")
+        samples = dist.sample_prior(coords={"channel": ["C1", "C2", "C3"]})
+
+    """
+    # Protect against doctest
+    if name == "__wrapped__":
+        return
+
+    _get_pymc_distribution(name)
+    return partial(Prior, distribution=name)

pymc_extras/statespace/core/statespace.py

@@ -2047,6 +2047,69 @@ class PyMCStateSpace:
 
         return scenario
 
+    def _build_forecast_model(
+        self, time_index, t0, forecast_index, scenario, filter_output, mvn_method
+    ):
+        filter_time_dim = TIME_DIM
+        temp_coords = self._fit_coords.copy()
+
+        dims = None
+        if all([dim in temp_coords for dim in [filter_time_dim, ALL_STATE_DIM, OBS_STATE_DIM]]):
+            dims = [TIME_DIM, ALL_STATE_DIM, OBS_STATE_DIM]
+
+        t0_idx = np.flatnonzero(time_index == t0)[0]
+
+        temp_coords["data_time"] = time_index
+        temp_coords[TIME_DIM] = forecast_index
+
+        mu_dims, cov_dims = None, None
+        if all([dim in self._fit_coords for dim in [TIME_DIM, ALL_STATE_DIM, ALL_STATE_AUX_DIM]]):
+            mu_dims = ["data_time", ALL_STATE_DIM]
+            cov_dims = ["data_time", ALL_STATE_DIM, ALL_STATE_AUX_DIM]
+
+        with pm.Model(coords=temp_coords) as forecast_model:
+            (_, _, *matrices), grouped_outputs = self._kalman_filter_outputs_from_dummy_graph(
+                data_dims=["data_time", OBS_STATE_DIM],
+            )
+
+            group_idx = FILTER_OUTPUT_TYPES.index(filter_output)
+            mu, cov = grouped_outputs[group_idx]
+
+            sub_dict = {
+                data_var: pt.as_tensor_variable(data_var.get_value(), name="data")
+                for data_var in forecast_model.data_vars
+            }
+
+            missing_data_vars = np.setdiff1d(
+                ar1=[*self.data_names, "data"], ar2=[k.name for k, _ in sub_dict.items()]
+            )
+            if missing_data_vars.size > 0:
+                raise ValueError(f"{missing_data_vars} data used for fitting not found!")
+
+            mu_frozen, cov_frozen = graph_replace([mu, cov], replace=sub_dict, strict=True)
+
+            x0 = pm.Deterministic(
+                "x0_slice", mu_frozen[t0_idx], dims=mu_dims[1:] if mu_dims is not None else None
+            )
+            P0 = pm.Deterministic(
+                "P0_slice", cov_frozen[t0_idx], dims=cov_dims[1:] if cov_dims is not None else None
+            )
+
+            _ = LinearGaussianStateSpace(
+                "forecast",
+                x0,
+                P0,
+                *matrices,
+                steps=len(forecast_index),
+                dims=dims,
+                sequence_names=self.kalman_filter.seq_names,
+                k_endog=self.k_endog,
+                append_x0=False,
+                method=mvn_method,
+            )
+
+        return forecast_model
+
     def forecast(
         self,
         idata: InferenceData,
@@ -2139,8 +2202,6 @@ class PyMCStateSpace:
                   the latent state trajectories: `y[t] = Z @ x[t] + nu[t]`, where `nu ~ N(0, H)`.
 
         """
-        filter_time_dim = TIME_DIM
-
         _validate_filter_arg(filter_output)
 
         compile_kwargs = kwargs.pop("compile_kwargs", {})
@@ -2185,58 +2246,23 @@ class PyMCStateSpace:
             use_scenario_index=use_scenario_index,
         )
         scenario = self._finalize_scenario_initialization(scenario, forecast_index)
-        temp_coords = self._fit_coords.copy()
-
-        dims = None
-        if all([dim in temp_coords for dim in [filter_time_dim, ALL_STATE_DIM, OBS_STATE_DIM]]):
-            dims = [TIME_DIM, ALL_STATE_DIM, OBS_STATE_DIM]
-
-        t0_idx = np.flatnonzero(time_index == t0)[0]
-
-        temp_coords["data_time"] = time_index
-        temp_coords[TIME_DIM] = forecast_index
-
-        mu_dims, cov_dims = None, None
-        if all([dim in self._fit_coords for dim in [TIME_DIM, ALL_STATE_DIM, ALL_STATE_AUX_DIM]]):
-            mu_dims = ["data_time", ALL_STATE_DIM]
-            cov_dims = ["data_time", ALL_STATE_DIM, ALL_STATE_AUX_DIM]
-
-        with pm.Model(coords=temp_coords) as forecast_model:
-            (_, _, *matrices), grouped_outputs = self._kalman_filter_outputs_from_dummy_graph(
-                scenario=scenario,
-                data_dims=["data_time", OBS_STATE_DIM],
-            )
-
-            for name in self.data_names:
-                if name in scenario.keys():
-                    pm.set_data(
-                        {"data": np.zeros((len(forecast_index), self.k_endog))},
-                        coords={"data_time": np.arange(len(forecast_index))},
-                    )
-                    break
 
-            group_idx = FILTER_OUTPUT_TYPES.index(filter_output)
-            mu, cov = grouped_outputs[group_idx]
-
-            x0 = pm.Deterministic(
-                "x0_slice", mu[t0_idx], dims=mu_dims[1:] if mu_dims is not None else None
-            )
-            P0 = pm.Deterministic(
-                "P0_slice", cov[t0_idx], dims=cov_dims[1:] if cov_dims is not None else None
-            )
+        forecast_model = self._build_forecast_model(
+            time_index=time_index,
+            t0=t0,
+            forecast_index=forecast_index,
+            scenario=scenario,
+            filter_output=filter_output,
+            mvn_method=mvn_method,
+        )
 
-            _ = LinearGaussianStateSpace(
-                "forecast",
-                x0,
-                P0,
-                *matrices,
-                steps=len(forecast_index),
-                dims=dims,
-                sequence_names=self.kalman_filter.seq_names,
-                k_endog=self.k_endog,
-                append_x0=False,
-                method=mvn_method,
-            )
+        with forecast_model:
+            if scenario is not None:
+                dummy_obs_data = np.zeros((len(forecast_index), self.k_endog))
+                pm.set_data(
+                    scenario | {"data": dummy_obs_data},
+                    coords={"data_time": np.arange(len(forecast_index))},
+                )
 
         forecast_model.rvs_to_initial_values = {
             k: None for k in forecast_model.rvs_to_initial_values.keys()

pymc_extras/statespace/filters/kalman_smoother.py

@@ -105,7 +105,7 @@ class KalmanSmoother:
         a_hat, P_hat = self.predict(a, P, T, R, Q)
 
         # Use pinv, otherwise P_hat is singular when there is missing data
-        smoother_gain = matrix_dot(pt.linalg.pinv(P_hat), T, P).T
+        smoother_gain = matrix_dot(pt.linalg.pinv(P_hat, hermitian=True), T, P).T
         a_smooth_next = a + smoother_gain @ (a_smooth - a_hat)
 
         P_smooth_next = P + quad_form_sym(smoother_gain, P_smooth - P_hat)

pymc_extras/statespace/models/structural/__init__.py

@@ -0,0 +1,21 @@
+from pymc_extras.statespace.models.structural.components.autoregressive import (
+    AutoregressiveComponent,
+)
+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.measurement_error import MeasurementError
+from pymc_extras.statespace.models.structural.components.regression import RegressionComponent
+from pymc_extras.statespace.models.structural.components.seasonality import (
+    FrequencySeasonality,
+    TimeSeasonality,
+)
+
+__all__ = [
+    "LevelTrendComponent",
+    "MeasurementError",
+    "AutoregressiveComponent",
+    "TimeSeasonality",
+    "FrequencySeasonality",
+    "RegressionComponent",
+    "CycleComponent",
+]

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

@@ -0,0 +1,188 @@
+import numpy as np
+import pytensor.tensor as pt
+
+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):
+    r"""
+    Autoregressive timeseries component
+
+    Parameters
+    ----------
+    order: int or sequence of int
+
+        If int, the number of lags to include in the model.
+        If a sequence, an array-like of zeros and ones indicating which lags to include in the model.
+
+    name: str, default "auto_regressive"
+        A name for this autoregressive component. Used to label dimensions and coordinates.
+
+    observed_state_names: list[str] | None, default None
+        List of strings for observed state labels. If None, defaults to ["data"].
+
+    Notes
+    -----
+    An autoregressive component can be thought of as a way o introducing serially correlated errors into the model.
+    The process is modeled:
+
+    .. math::
+        x_t = \sum_{i=1}^p \rho_i x_{t-i}
+
+    Where ``p``, the number of autoregressive terms to model, is the order of the process. By default, all lags up to
+    ``p`` are included in the model. To disable lags, pass a list of zeros and ones to the ``order`` argumnet. For
+    example, ``order=[1, 1, 0, 1]`` would become:
+
+    .. math::
+        x_t = \rho_1 x_{t-1} + \rho_2 x_{t-1} + \rho_4 x_{t-1}
+
+    The coefficient :math:`\rho_3` has been constrained to zero.
+
+    .. warning:: This class is meant to be used as a component in a structural time series model. For modeling of
+              stationary processes with ARIMA, use ``statespace.BayesianSARIMA``.
+
+    Examples
+    --------
+    Model a timeseries as an AR(2) process with non-zero mean:
+
+    .. code:: python
+
+        from pymc_extras.statespace import structural as st
+        import pymc as pm
+        import pytensor.tensor as pt
+
+        trend = st.LevelTrendComponent(order=1, innovations_order=0)
+        ar = st.AutoregressiveComponent(2)
+        ss_mod = (trend + ar).build()
+
+        with pm.Model(coords=ss_mod.coords) as model:
+            P0 = pm.Deterministic('P0', pt.eye(ss_mod.k_states) * 10, dims=ss_mod.param_dims['P0'])
+            intitial_trend = pm.Normal('initial_trend', sigma=10, dims=ss_mod.param_dims['initial_trend'])
+            ar_params = pm.Normal('ar_params', dims=ss_mod.param_dims['ar_params'])
+            sigma_ar = pm.Exponential('sigma_ar', 1, dims=ss_mod.param_dims['sigma_ar'])
+
+            ss_mod.build_statespace_graph(data)
+            idata = pm.sample(nuts_sampler='numpyro')
+
+    """
+
+    def __init__(
+        self,
+        order: int = 1,
+        name: str = "auto_regressive",
+        observed_state_names: list[str] | None = None,
+    ):
+        if observed_state_names is None:
+            observed_state_names = ["data"]
+
+        k_posdef = k_endog = len(observed_state_names)
+
+        order = order_to_mask(order)
+        ar_lags = np.flatnonzero(order).ravel().astype(int) + 1
+        k_states = len(order)
+
+        self.order = order
+        self.ar_lags = ar_lags
+
+        super().__init__(
+            name=name,
+            k_endog=k_endog,
+            k_states=k_states * k_endog,
+            k_posdef=k_posdef,
+            measurement_error=True,
+            combine_hidden_states=True,
+            observed_state_names=observed_state_names,
+            obs_state_idxs=np.tile(np.r_[[1.0], np.zeros(k_states - 1)], k_endog),
+        )
+
+    def populate_component_properties(self):
+        k_states = self.k_states // self.k_endog  # this is also the number of AR lags
+
+        self.state_names = [
+            f"L{i + 1}[{state_name}]"
+            for state_name in self.observed_state_names
+            for i in range(k_states)
+        ]
+
+        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()}
+
+        if self.k_endog > 1:
+            self.param_dims[f"params_{self.name}"] = (
+                f"endog_{self.name}",
+                AR_PARAM_DIM,
+            )
+            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_states,) if self.k_endog == 1 else (self.k_endog, k_states),
+                "constraints": None,
+                "dims": (AR_PARAM_DIM,)
+                if self.k_endog == 1
+                else (
+                    f"endog_{self.name}",
+                    f"lag_{self.name}",
+                ),
+            },
+            f"sigma_{self.name}": {
+                "shape": () if self.k_endog == 1 else (self.k_endog,),
+                "constraints": "Positive",
+                "dims": None if self.k_endog == 1 else (f"endog_{self.name}",),
+            },
+        }
+
+    def make_symbolic_graph(self) -> None:
+        k_endog = self.k_endog
+        k_states = self.k_states // k_endog
+        k_posdef = self.k_posdef
+
+        k_nonzero = int(sum(self.order))
+        ar_params = self.make_and_register_variable(
+            f"params_{self.name}", shape=(k_nonzero,) if k_endog == 1 else (k_endog, k_nonzero)
+        )
+        sigma_ar = self.make_and_register_variable(
+            f"sigma_{self.name}", shape=() if k_endog == 1 else (k_endog,)
+        )
+
+        if k_endog == 1:
+            T = pt.eye(k_states, k=-1)
+            ar_idx = (np.zeros(k_nonzero, dtype="int"), np.nonzero(self.order)[0])
+            T = T[ar_idx].set(ar_params)
+
+        else:
+            transition_matrices = []
+
+            for i in range(k_endog):
+                T = pt.eye(k_states, k=-1)
+                ar_idx = (np.zeros(k_nonzero, dtype="int"), np.nonzero(self.order)[0])
+                T = T[ar_idx].set(ar_params[i])
+                transition_matrices.append(T)
+            T = pt.specify_shape(
+                pt.linalg.block_diag(*transition_matrices), (self.k_states, self.k_states)
+            )
+
+        self.ssm["transition", :, :] = T
+
+        R = np.eye(k_states)
+        R_mask = np.full((k_states), False)
+        R_mask[0] = True
+        R = R[:, R_mask]
+
+        self.ssm["selection", :, :] = pt.specify_shape(
+            pt.linalg.block_diag(*[R for _ in range(k_endog)]), (self.k_states, self.k_posdef)
+        )
+
+        Z = pt.zeros((1, k_states))[0, 0].set(1.0)
+        self.ssm["design", :, :] = pt.specify_shape(
+            pt.linalg.block_diag(*[Z for _ in range(k_endog)]), (self.k_endog, self.k_states)
+        )
+
+        cov_idx = ("state_cov", *np.diag_indices(k_posdef))
+        self.ssm[cov_idx] = sigma_ar**2