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_time_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)[source]#

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 lazily, so constructing this class is the opt-in that requires it.

Parameters:

  • model_fn – The model definition, a StatespaceModel (or any object with its statespace / priors methods).

  • data – Observed training data (univariate series or 2-d with a named series dim), or None to construct unfitted and call fit() later.

  • covariates – Covariates covering (at least) the training window, passed through to the model definition; surplus future steps are ignored during fitting. None for models without covariates.

  • draws – MCMC schedule (defaults 1000 / 1000 / 2).

  • tune – MCMC schedule (defaults 1000 / 1000 / 2).

  • chains – MCMC schedule (defaults 1000 / 1000 / 2).

  • nuts_sampler – NUTS backend: "pymc" (default), "nutpie", "numpyro", or "blackjax".

  • random_seed – Seed for the fit.

  • progressbar – Show the sampling progress bar.

  • sample_kwargs – Extra keyword arguments for pm.sample. progressbar is accepted here for compatibility, but the direct argument is preferred (passing both raises).

  • build_kwargs – Extra keyword arguments for build_statespace_graph, such as mvn_method for the likelihood decomposition.

  • forecast_kwargs – Extra keyword arguments for PyMCStateSpace.forecast, such as filter_output or mvn_method. Horizon, scenario, seed, verbosity, and progress are managed by the adapter and cannot be overridden here.

ss_mod#

The built PyMCStateSpace.

model#

The pm.Model holding priors and the Kalman-filter likelihood.

idata#

The full MCMC result.

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

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 – Covariates spanning training window + forecast horizon (time coords must extend the training data’s).

  • num_samples – Number of posterior draws (and forecast samples); default 100. Mutually exclusive with posterior.

  • horizon – Number of steps to forecast past the training data.

  • future_index – Time coordinate values of the forecast horizon, supplied at forecast time (models without exogenous inputs only) — the covariate-free half of the predict-time horizon capability; future_covariates is the with-exogenous half.

  • future_covariates – Covariates 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 – A fixed posterior to condition on (any shape posterior_dataset() accepts, typically from draw_posterior); passing the same posterior to predict_in_sample and forecast makes the calls draw-coherent.

  • var_names – Subset of prediction variables to keep ("forecast", "forecast_latent"). Default: both.

  • random_seed – Passed through to PyMCStateSpace.forecast.

  • progressbar – Passed through to PyMCStateSpace.forecast.

Returns:

DataTree – With a predictions group holding "forecast" (dims (chain, draw, time_future, ...)) and the latent state trajectories as "forecast_latent".

predict_in_sample(num_samples=None, *, posterior=None, random_seed=None, progressbar=False)[source]#

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.

num_samples defaults to 100 and is mutually exclusive with posterior, a fixed posterior to condition on (see forecast() for the draw-coherence semantics).

Returns:

DataTree – With a posterior_predictive group holding "obs" (dims (chain, draw, time, ...)).

class pymc_forecast.statespace.StatespaceModel(priors=None)[source]#

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

abstractmethod statespace(data, covariates)[source]#

Build and return the PyMCStateSpace (component sum + .build()).

abstractmethod priors(ss_mod, data, covariates)[source]#

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.