pymc_forecast.statespace

Interop with pymc_extras.statespace models.

Statespace structural models (level/trend, seasonality, cycles, AR, regression; SARIMAX; VARMAX) cover the linear-Gaussian slice of what 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 StatespaceForecaster adapter maps that lifecycle onto the package’s fit + forecast(horizon) protocol, so a statespace model drops into 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.

class pymc_forecast.statespace.StatespaceForecaster(model_fn, data=None, covariates=None, *, draws=1000, tune=1000, chains=2, nuts_sampler='pymc', random_seed=None, progressbar=None, sample_kwargs=None, build_kwargs=None, forecast_kwargs=None)

Bases: HMCForecaster

Fit and forecast a pymc-extras statespace model behind the forecaster protocol.

An HMCForecaster whose training model is built through the statespace lifecycle instead of build_model(): fit on construction with NUTS, draw_posterior, forecast() returning a labeled predictions group, predict_in_sample() — so it drops into 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 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 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 predict_in_sample()). start, periods, end, scenario, random_seed, verbose and progressbar are managed by the adapter and cannot be given here.

model_fn

The model definition passed to the constructor.

Type:

StatespaceModel

ss_mod

The built PyMCStateSpace (set by fitting).

Type:

object

model

The pm.Model holding priors and the Kalman-filter likelihood; None before the first fit.

Type:

pymc.Model or None

idata

The full MCMC result, as returned by pm.sample (set by fitting).

Type:

xarray.DataTree or arviz.InferenceData

approx

Always None (set by fitting).

Type:

None

losses

Always None (set by fitting).

Type:

None

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.

forecast(covariates=None, num_samples=None, *, horizon=None, future_index=None, future_covariates=None, posterior=None, var_names=None, random_seed=None, progressbar=False)

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 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 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 (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.

predict_in_sample(num_samples=None, *, posterior=None, random_seed=None, progressbar=False)

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 (posterior_dataset()); used as given, not thinned (see 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.

class pymc_forecast.statespace.StatespaceModel(priors=None)

Bases: PriorConfig, ABC

A pymc-extras statespace model definition for 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 PriorConfig): declare named defaults in default_priors, create them inside 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 PriorConfig); None keeps the defaults.

abstractmethod statespace(data, covariates)

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).

abstractmethod priors(ss_mod, data, covariates)

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 statespace().

• data (xarray.DataArray) – Normalized training data, "time" first.

• covariates (xarray.DataArray) – Training-window covariates (zero-width for a covariate-free fit).