"""Interop with ``pymc_extras.statespace`` models.
Statespace structural models (level/trend, seasonality, cycles, AR, regression;
SARIMAX; VARMAX) cover the linear-Gaussian slice of what
:func:`~pymc_forecast.markov.markov_series` is used for — with
Kalman-filter marginalization instead of sampling per-step latents, which
usually gives better posteriors *and* faster sampling.
Their lifecycle differs from this package's model functions: components are
combined and ``.build()`` into a ``PyMCStateSpace``, parameter priors are
declared inside a ``pm.Model`` carrying the statespace coords, and
``build_statespace_graph`` inserts the Kalman-filter likelihood. The
:class:`StatespaceForecaster` adapter maps that lifecycle onto the package's
``fit + forecast(horizon)`` protocol, so a statespace model drops into
:func:`~pymc_forecast.evaluate.backtest` and the metrics layer wherever a
hand-written forecasting model is accepted.
Coords stay this package's responsibility: the statespace internals only ever
see positional (integer-indexed) data, and real time coordinates are stamped
onto the outputs — ``"time_future"`` on forecasts, ``"time"`` on in-sample
predictions — exactly like the core forecasters.
"""
import abc
import warnings
from collections.abc import Mapping, Sequence
import numpy as np
import pymc as pm
import xarray as xr
from pymc_forecast.data import (
FUTURE_DIM,
TIME_DIM,
_validate_covariate_structure,
as_dataarray,
concat_covariates,
concat_time_index,
extend_time_index,
null_covariates,
validate_alignment,
)
from pymc_forecast.exceptions import AlignmentError, HorizonError, OptionalDependencyError
from pymc_forecast.forecaster import HMCForecaster
from pymc_forecast.priors import PriorConfig
__all__ = ["StatespaceForecaster", "StatespaceModel"]
OBSERVED_STATE_DIM = "observed_state"
"""Dim name pymc-extras puts on the observed-series axis of its outputs."""
_NO_TIME_INDEX_MESSAGE = "No time index found on the supplied data"
"""pymc-extras warns when data has a plain range index — expected here, since
the adapter deliberately feeds positional data and stamps coords itself."""
[docs]
class StatespaceModel(PriorConfig, abc.ABC):
"""A pymc-extras statespace model definition for :class:`StatespaceForecaster`.
The statespace lifecycle is two-phase — the component graph must exist
before its coords can be registered on the ``pm.Model`` that holds the
parameter priors — so the definition is split accordingly. Example::
import pymc as pm
import pytensor.tensor as pt
from pymc_extras.statespace import structural as st
class LocalLinearTrend(StatespaceModel):
def statespace(self, data, covariates):
trend = st.LevelTrend(order=2, innovations_order=[0, 1])
return (trend + st.MeasurementError()).build(verbose=False)
def priors(self, ss_mod, data, covariates):
P0_diag = pm.Gamma("P0_diag", alpha=2, beta=5, dims="state")
pm.Deterministic("P0", pt.diag(P0_diag), dims=("state", "state_aux"))
pm.Normal("initial_level_trend", float(data[0]), 1, dims="state_level_trend")
pm.Gamma("sigma_level_trend", alpha=2, beta=10, dims="shock_level_trend")
pm.Gamma("sigma_MeasurementError", alpha=2, beta=10)
Both phases receive the normalized training ``data`` and ``covariates`` as
labeled DataArrays — ``statespace`` so the component graph can be sized
from them (series count, regression features), ``priors`` so priors can be
informed by them (as with ``initial_level_trend`` above). For a fit
without covariates, ``covariates`` is a zero-width DataArray with dims
``("time", "covariate")``, never ``None``. A model with a
``st.Regression`` component registers its feature matrix with
``pm.Data(name, ...)`` inside ``priors``, where ``name`` is the entry in
``ss_mod.data_names``; at forecast time the adapter feeds the future
covariate values through as the scenario.
Priors are user-injectable here too (see
:class:`~pymc_forecast.priors.PriorConfig`): declare named defaults in
``default_priors``, create them inside :meth:`priors` with
``self.create_prior(name)``, and let callers override any subset with
``priors=`` when constructing the model object.
Parameters
----------
priors : mapping, optional
Prior specs by name, merged over ``default_priors`` (see
:class:`~pymc_forecast.priors.PriorConfig`); ``None`` keeps the
defaults.
"""
[docs]
@abc.abstractmethod
def statespace(self, data: xr.DataArray, covariates: xr.DataArray):
"""Build and return the ``PyMCStateSpace`` (component sum + ``.build()``).
Parameters
----------
data : xarray.DataArray
Normalized training data, ``"time"`` first.
covariates : xarray.DataArray
Training-window covariates (zero-width for a covariate-free fit).
Returns
-------
object
A built ``pymc_extras`` ``PyMCStateSpace`` whose ``k_endog`` equals
the data width (1 for 1-D data, else the size of the non-time dim).
"""
[docs]
@abc.abstractmethod
def priors(self, ss_mod, data: xr.DataArray, covariates: xr.DataArray) -> None:
"""Declare the parameter priors listed by ``ss_mod.param_info``.
Called inside a ``pm.Model`` whose coords are ``ss_mod.coords``; the
adapter calls ``build_statespace_graph`` afterwards.
Parameters
----------
ss_mod : object
The ``PyMCStateSpace`` returned by :meth:`statespace`.
data : xarray.DataArray
Normalized training data, ``"time"`` first.
covariates : xarray.DataArray
Training-window covariates (zero-width for a covariate-free fit).
"""
def _observed_frame(data: xr.DataArray):
"""Convert normalized data to the positional DataFrame statespace expects.
The real time coords are dropped on purpose (statespace requires a
zero-based index and would regenerate one anyway); the adapter re-attaches
them to every output.
"""
import pandas as pd
values = data.transpose(TIME_DIM, ...).values
if values.ndim == 1:
values = values[:, None]
return pd.DataFrame(values)
_FIT_DATA_GROUPS = ("observed_data", "constant_data")
"""Fit-result groups pymc-extras reads back to rebuild the fit coords, observed
data, and exogenous inputs."""
def _posterior_like(fit_result, posterior: xr.Dataset):
"""Wrap a thinned posterior ``Dataset`` in the fit result's container type.
The container is a ``DataTree`` or ``arviz.InferenceData``, depending on the
pymc/arviz generation. By construction that is the idata flavor the installed
statespace methods accept. The fit's data groups travel along because recent
pymc-extras rebuilds the fit coords, observed data, and exogenous inputs
from them before forecasting.
"""
groups = {"posterior": posterior}
for name in _FIT_DATA_GROUPS:
if name in fit_result:
group = fit_result[name]
groups[name] = group.to_dataset() if isinstance(group, xr.DataTree) else group
if isinstance(fit_result, xr.DataTree):
return xr.DataTree.from_dict(groups)
return type(fit_result)(**groups)
def _predictive_dataset(result) -> xr.Dataset:
"""Extract the predictive ``Dataset`` from a statespace sampling result.
pymc-extras returns an ArviZ tree whose predictive variables live either at
the root or under a ``posterior_predictive`` child, depending on version.
"""
children = getattr(result, "children", None) or {}
if "posterior_predictive" in children:
result = result["posterior_predictive"]
elif hasattr(result, "posterior_predictive") and not hasattr(result, "data_vars"):
result = result.posterior_predictive
return result.to_dataset() if hasattr(result, "to_dataset") else result
[docs]
class StatespaceForecaster(HMCForecaster):
"""Fit and forecast a pymc-extras statespace model behind the forecaster protocol.
An :class:`~pymc_forecast.forecaster.HMCForecaster` whose training model
is built through the statespace lifecycle instead of
:func:`~pymc_forecast.model.build_model`: fit on construction with NUTS,
``draw_posterior``, :meth:`forecast` returning a labeled
``predictions`` group, :meth:`predict_in_sample` — so it drops into
:func:`~pymc_forecast.evaluate.backtest` via
``forecaster_cls=StatespaceForecaster``.
pymc-extras is imported on construction, so constructing this class is
the opt-in that requires it.
Parameters
----------
model_fn : StatespaceModel
The model definition, a :class:`StatespaceModel` (or any object with
its ``statespace`` / ``priors`` methods).
data : xarray.DataArray, pandas.Series, pandas.DataFrame or array_like, optional
Observed training series (univariate or 2-D with a named series dim),
normalized with :func:`~pymc_forecast.data.as_dataarray`, or ``None``
to construct unfitted and call ``fit`` later.
covariates : xarray.DataArray, pandas.Series, pandas.DataFrame or array_like, optional
Covariates on the same ``"time"`` coordinate as ``data``, covering (at
least) the training window, passed through to the model definition;
rows past the training window are dropped during fitting. ``None`` for
models without covariates.
draws : int, default 1000
Number of posterior draws per chain.
tune : int, default 1000
Number of tuning steps per chain.
chains : int, default 2
Number of chains.
nuts_sampler : {"pymc", "nutpie", "numpyro", "blackjax"}, default "pymc"
NUTS backend, forwarded to ``pm.sample``.
random_seed : int or numpy.random.Generator, optional
Seed for the fit, forwarded to ``pm.sample``; also the default seed
for later ``fit`` calls.
progressbar : bool, optional
Show the sampling progress bar; ``None`` means off. May instead be
given in ``sample_kwargs``, but not both.
sample_kwargs : mapping, optional
Extra keyword arguments for ``pm.sample``. ``progressbar`` is
accepted here for compatibility, but the direct argument is preferred
(passing both raises).
build_kwargs : mapping, optional
Extra keyword arguments forwarded to ``build_statespace_graph``; which
keywords exist depends on the installed pymc-extras release.
forecast_kwargs : mapping, optional
Extra keyword arguments forwarded to ``PyMCStateSpace.forecast``
(not applied to :meth:`predict_in_sample`). ``start``, ``periods``,
``end``, ``scenario``, ``random_seed``, ``verbose`` and
``progressbar`` are managed by the adapter and cannot be given here.
Attributes
----------
model_fn : StatespaceModel
The model definition passed to the constructor.
ss_mod : object
The built ``PyMCStateSpace`` (set by fitting).
model : pymc.Model or None
The ``pm.Model`` holding priors and the Kalman-filter likelihood;
``None`` before the first fit.
idata : xarray.DataTree or arviz.InferenceData
The full MCMC result, as returned by ``pm.sample`` (set by fitting).
approx : None
Always ``None`` (set by fitting).
losses : None
Always ``None`` (set by fitting).
Raises
------
pymc_forecast.exceptions.OptionalDependencyError
If ``pymc_extras.statespace`` cannot be imported.
ValueError
If ``forecast_kwargs`` contains an adapter-managed key, ``progressbar``
is given both directly and in ``sample_kwargs``, or ``covariates`` is
given without ``data``.
pymc_forecast.exceptions.AlignmentError
If (when fitting) ``data`` or ``covariates`` cannot be normalized, the
data has more than two dims, the covariates do not align with it, or
``ss_mod.k_endog`` does not match the data width.
"""
def __init__(
self,
model_fn,
data=None,
covariates=None,
*,
draws: int = 1000,
tune: int = 1000,
chains: int = 2,
nuts_sampler: str = "pymc",
random_seed=None,
progressbar: bool | None = None,
sample_kwargs: Mapping | None = None,
build_kwargs: Mapping | None = None,
forecast_kwargs: Mapping | None = None,
) -> None:
try:
import pymc_extras.statespace # noqa: F401
except ImportError as err:
raise OptionalDependencyError("pymc-extras", "extras", "StatespaceForecaster") from err
self._build_kwargs = dict(build_kwargs or {})
self._forecast_kwargs = dict(forecast_kwargs or {})
reserved = {"start", "periods", "end", "scenario", "random_seed", "verbose", "progressbar"}
overlap = reserved.intersection(self._forecast_kwargs)
if overlap:
msg = f"forecast_kwargs cannot override adapter-managed arguments: {sorted(overlap)}"
raise ValueError(msg)
super().__init__(
model_fn,
data,
covariates,
draws=draws,
tune=tune,
chains=chains,
nuts_sampler=nuts_sampler,
random_seed=random_seed,
progressbar=progressbar,
sample_kwargs=sample_kwargs,
)
def _build_model(self) -> pm.Model:
"""Drive the statespace lifecycle: components, priors, Kalman graph.
Validation happens here — after input normalization, before the
(expensive) fit.
"""
if self._data.ndim > 2:
msg = (
"pymc-extras statespace models take one- or two-dimensional "
f"data, got dims {self._data.dims}"
)
raise AlignmentError(msg)
validate_alignment(self._data, self._covariates)
self.ss_mod = self.model_fn.statespace(self._data, self._covariates)
data_width = 1 if self._data.ndim == 1 else self._data.sizes[self._series_dim]
if self.ss_mod.k_endog != data_width:
msg = (
"statespace observed width does not match the training data: "
f"k_endog={self.ss_mod.k_endog}, data width={data_width}"
)
raise AlignmentError(msg)
with pm.Model(coords=self.ss_mod.coords) as model:
self.model_fn.priors(self.ss_mod, self._data, self._covariates)
with warnings.catch_warnings():
warnings.filterwarnings("ignore", message=_NO_TIME_INDEX_MESSAGE)
self.ss_mod.build_statespace_graph(
_observed_frame(self._data), **self._build_kwargs
)
return model
@property
def _series_dim(self) -> str | None:
"""The data's non-time dim, or ``None`` for a univariate series."""
extra_dims = [d for d in self._data.dims if d != TIME_DIM]
return extra_dims[0] if extra_dims else None
def _relabel(self, da: xr.DataArray, time_dim: str, time_coords) -> xr.DataArray:
"""Stamp real coords on a statespace output: rename its positional
``"time"`` dim and map ``"observed_state"`` back to the data's series
dim (squeezed away for univariate data)."""
if "time" != time_dim:
da = da.rename({"time": time_dim})
da = da.assign_coords({time_dim: np.asarray(time_coords)})
series_dim = self._series_dim
if OBSERVED_STATE_DIM in da.dims:
if series_dim is None:
da = da.squeeze(OBSERVED_STATE_DIM, drop=True)
else:
da = da.rename({OBSERVED_STATE_DIM: series_dim})
da = da.assign_coords({series_dim: self._data[series_dim].values})
sample_dims = [d for d in ("chain", "draw") if d in da.dims]
return da.transpose(*sample_dims, time_dim, ...)
def _scenario(self, future_covariates: xr.DataArray):
"""Future exogenous values for ``PyMCStateSpace.forecast(scenario=...)``.
Returns ``None`` for models without exogenous data. A single
registered data variable (one ``st.Regression`` component) is fed the
future slice of the covariate matrix.
"""
data_names = tuple(self.ss_mod.data_names)
if not data_names:
return None
if len(data_names) > 1:
msg = (
"one covariate DataArray cannot populate a statespace model "
f"with multiple exogenous inputs {data_names}; combine them "
"into one Regression component"
)
raise NotImplementedError(msg)
feature_dims = [d for d in future_covariates.dims if d != TIME_DIM]
if len(feature_dims) != 1 or future_covariates.sizes[feature_dims[0]] == 0:
msg = (
f"the statespace model needs future values for {data_names[0]!r}: "
"pass full-horizon covariates instead of horizon=/future_index="
)
raise AlignmentError(msg)
return np.asarray(future_covariates.transpose(TIME_DIM, ...).values)
[docs]
def forecast(
self,
covariates=None,
num_samples: int | None = None,
*,
horizon: int | None = None,
future_index=None,
future_covariates=None,
posterior=None,
var_names: Sequence[str] | None = None,
random_seed=None,
progressbar: bool = False,
) -> xr.DataTree:
"""Sample forecasts beyond the training window.
The horizon is supplied at forecast time, in one of four mutually
exclusive ways: pass ``covariates`` spanning the training window plus
the forecast steps, ``future_covariates`` covering only the forecast
steps, or — for a model without exogenous inputs — pass ``horizon=N``
to forecast ``N`` steps past the training data (its time coord is
extended at the inferred spacing) or ``future_index=`` to forecast
over an arbitrary later time index (strictly increasing values lying
after the training window; the horizon length is derived from it).
Forecast steps are always iterated consecutively from the end of
training and labeled with the supplied coordinates.
The forecast draws the terminal state from its smoothed posterior and
iterates the statespace forward — the Kalman analogue of the core
mechanism of seeding ``*_future`` latents from the in-sample state.
For a model with a ``st.Regression`` component, the future slice of
the covariates is fed through as the scenario.
Parameters
----------
covariates : xarray.DataArray, pandas.Series, pandas.DataFrame or array_like, optional
Covariates spanning training window + forecast horizon (time
coords must extend the training data's), normalized with
:func:`~pymc_forecast.data.as_dataarray`; the steps past the
training window define the horizon. Non-time dimensions and
coordinate names/order must match the training covariates.
``None`` when the horizon comes from another argument.
num_samples : int, optional
Number of posterior draws (and forecast samples) to sample when
``posterior`` is not given; ``None`` means 100. Mutually exclusive
with ``posterior``.
horizon : int, optional
Number of steps to forecast past the training data (must be at
least 1). Rejected for models with exogenous inputs.
future_index : array_like, optional
Time coordinate values of the forecast horizon, supplied at
forecast time (rejected for models with exogenous inputs) — the
covariate-free half of the predict-time horizon capability;
``future_covariates`` is the with-exogenous half.
future_covariates : xarray.DataArray, pandas.DataFrame or array_like, optional
Covariates (any input :func:`~pymc_forecast.data.as_dataarray`
accepts, including a ``pandas.Series``) covering only the forecast
horizon, with a time index
lying after the training window; fed through as the forecast
scenario. Structure (dims, covariate names and order) must match
the training covariates.
posterior : xarray.Dataset, xarray.DataTree or arviz.InferenceData, optional
A fixed posterior to condition on: a posterior Dataset or any
object with a ``posterior`` group
(:func:`~pymc_forecast.prediction.posterior_dataset`), typically
from ``draw_posterior``; used as given, not thinned. Passing the
same posterior to ``predict_in_sample`` and ``forecast`` makes the
calls draw-coherent. Mutually exclusive with ``num_samples``.
var_names : sequence of str, optional
Subset of ``{"forecast", "forecast_latent"}`` to keep (a bare
string is not accepted as a single name); ``None`` keeps both.
random_seed : int or numpy.random.Generator, optional
Seeds both the posterior draws (when ``posterior`` is not given)
and ``PyMCStateSpace.forecast``.
progressbar : bool, default False
Show the sampling progress bar (forwarded to
``PyMCStateSpace.forecast``).
Returns
-------
xarray.DataTree
Always a DataTree, with a ``predictions`` group holding
``"forecast"`` (dims ``(chain, draw, time_future)``, plus the
series dim for 2-D data) and the latent state trajectories as
``"forecast_latent"`` (dims ``(chain, draw, time_future, state)``).
Raises
------
pymc_forecast.exceptions.NotFittedError
If the forecaster has not been fitted.
ValueError
If not exactly one of ``covariates``, ``horizon``, ``future_index``
and ``future_covariates`` is given, if both ``posterior`` and
``num_samples`` are given, or if a posterior draw count is not
positive.
pymc_forecast.exceptions.AlignmentError
If the horizon arguments or covariates are invalid or misaligned
(see ``extend_time_index``, ``concat_time_index``,
``concat_covariates``), or if ``horizon`` / ``future_index`` is
used with a model that has an exogenous input.
pymc_forecast.exceptions.HorizonError
If the resulting horizon is empty (e.g. ``horizon=0``).
NotImplementedError
If the model has more than one exogenous input.
KeyError
If ``var_names`` contains an unknown name.
TypeError
If ``posterior`` has no ``posterior`` group.
"""
self._require_fitted()
provided = sum(
arg is not None for arg in (covariates, horizon, future_index, future_covariates)
)
if provided != 1:
msg = "pass exactly one of covariates, horizon, future_index, or future_covariates"
raise ValueError(msg)
t_obs = self._data.sizes[TIME_DIM]
if horizon is not None:
full_index = extend_time_index(self._data.get_index(TIME_DIM), horizon)
cov = null_covariates(np.asarray(full_index))
elif future_index is not None:
full_index = concat_time_index(self._data[TIME_DIM].values, future_index)
cov = null_covariates(np.asarray(full_index))
elif future_covariates is not None:
cov = concat_covariates(self._covariates, future_covariates)
else:
cov = as_dataarray(covariates, role="covariates")
_validate_covariate_structure(self._covariates, cov)
validate_alignment(self._data, cov)
future_covariates = cov.isel({TIME_DIM: slice(t_obs, None)})
future_coords = future_covariates[TIME_DIM].values
if len(future_coords) == 0:
msg = (
"no forecast horizon: pass covariates longer than the data "
f"along '{TIME_DIM}', or horizon >= 1"
)
raise HorizonError(msg)
posterior = self._resolve_posterior(posterior, num_samples, random_seed)
with warnings.catch_warnings():
warnings.filterwarnings("ignore", message=_NO_TIME_INDEX_MESSAGE)
result = self.ss_mod.forecast(
_posterior_like(self.idata, posterior),
start=t_obs - 1,
periods=len(future_coords),
scenario=self._scenario(future_covariates),
random_seed=random_seed,
progressbar=progressbar,
verbose=False,
**self._forecast_kwargs,
)
ds = _predictive_dataset(result)
predictions = xr.Dataset(
{
"forecast": self._relabel(ds["forecast_observed"], FUTURE_DIM, future_coords),
"forecast_latent": self._relabel(ds["forecast_latent"], FUTURE_DIM, future_coords),
}
)
if var_names is not None:
names = list(var_names)
unknown = set(names).difference(predictions.data_vars)
if unknown:
msg = f"unknown statespace prediction variables: {sorted(unknown)}"
raise KeyError(msg)
predictions = predictions[names]
return xr.DataTree.from_dict({"predictions": predictions})
[docs]
def predict_in_sample(
self,
num_samples: int | None = None,
*,
posterior=None,
random_seed=None,
progressbar: bool = False,
) -> xr.DataTree:
"""Sample the in-sample predictive of the observed series.
Draws observations from the smoothed state posterior (conditioned on
the full training window) — the statespace analogue of replaying
in-sample latents and resampling the observation noise.
Parameters
----------
num_samples : int, optional
Number of posterior draws to sample when ``posterior`` is not
given; ``None`` means 100. Mutually exclusive with ``posterior``.
posterior : xarray.Dataset, xarray.DataTree or arviz.InferenceData, optional
A fixed posterior to condition on: a posterior Dataset or any
object with a ``posterior`` group
(:func:`~pymc_forecast.prediction.posterior_dataset`); used as
given, not thinned (see :meth:`forecast` for the draw-coherence
semantics). Mutually exclusive with ``num_samples``.
random_seed : int or numpy.random.Generator, optional
Seeds both the posterior draws (when ``posterior`` is not given)
and the conditional posterior sampling.
progressbar : bool, default False
Show the sampling progress bar.
Returns
-------
xarray.DataTree
Always a DataTree, with a ``posterior_predictive`` group holding
only ``"obs"`` (dims ``(chain, draw, time)``, plus the series dim
for 2-D data).
Raises
------
pymc_forecast.exceptions.NotFittedError
If the forecaster has not been fitted.
ValueError
If both ``posterior`` and ``num_samples`` are given, or a
posterior draw count is not positive.
TypeError
If ``posterior`` has no ``posterior`` group.
"""
self._require_fitted()
posterior = self._resolve_posterior(posterior, num_samples, random_seed)
with warnings.catch_warnings():
warnings.filterwarnings("ignore", message=_NO_TIME_INDEX_MESSAGE)
result = self.ss_mod.sample_conditional_posterior(
_posterior_like(self.idata, posterior),
random_seed=random_seed,
progressbar=progressbar,
)
ds = _predictive_dataset(result)
obs = self._relabel(
ds["smoothed_posterior_observed"], TIME_DIM, self._data[TIME_DIM].values
)
return xr.DataTree.from_dict({"posterior_predictive": xr.Dataset({"obs": obs})})