API Reference

Model building blocks

Plain model functions that register the train/forecast sites for you.

models.Horizon: The train/forecast split for a single model call.
models.Transition: (carry, x_t) -> (dist_t, carry_fn) where carry_fn(z_t) builds the next
models.innovations(): Sample conditionally iid per-step innovations over the full horizon.
models.markov_series(): Sample a Markov (state-space) latent over the full horizon.
models.ssoe(): Run a single-source-of-error recursion over the full horizon.
models.SSOEStep: (carry, x_t) -> (mu_t, carry_fn) where mu_t is the one-step-ahead mean
models.SSOEResult: The means and sampled future values produced by ssoe().
models.predict(): Register the observation and forecast sites for the model.

Time-axis operations on observation distributions, extensible via singledispatch.

surgery.shift_loc(): Re-center a zero-centered noise distribution at loc.
surgery.slice_time(): Slice an elementwise distribution along the time axis -2.
surgery.prefix_condition(): Condition a (t+f)-length distribution on a t-length data prefix.
surgery.register_elementwise(): Declare a distribution family elementwise (usable as a decorator).

Drawing posterior samples and generating forecasts and in-sample predictions.

predictive.draw_posterior(): Draw num_samples posterior samples of the latent sites from a fitted guide.
predictive.forecast(): Sample forecasts for the steps in [t, duration) from a posterior.
predictive.predict_in_sample(): Sample the in-sample posterior predictive of the obs site.

Rolling-window backtesting and forecast metrics.

evaluate.backtest(): Backtest a forecasting model on a moving window of (train, test) data.
evaluate.backtest_vectorized(): Rolling-window backtest with all windows fitted in one vmapped SVI run.
evaluate.BacktestResult: Per-window result of a backtest() run.
evaluate.VectorizedBacktestResult: Result of a backtest_vectorized() run (all windows at once).
evaluate.evaluate_forecast(): Evaluate forecast samples against ground truth for several metrics at once.
evaluate.results_to_dataframe(): Flatten backtest results into a tidy one-row-per-window DataFrame.
evaluate.eval_crps(): Empirical CRPS averaged over all data elements.
evaluate.eval_mae(): Mean absolute error using the forecast sample median as point estimate.
evaluate.eval_rmse(): Root mean squared error using the forecast sample mean as point estimate.
evaluate.eval_coverage(): Empirical coverage of the central alpha prediction interval.
metrics.crps_empirical(): Compute the empirical Continuous Ranked Probability Score (CRPS).
metrics.eval_pinball(): Mean pinball (quantile) loss of the forecast quantile.
metrics.eval_interval_score(): Mean Winkler interval score for the central alpha prediction interval.
metrics.make_mase(): Build a Mean Absolute Scaled Error metric scaled by train_data.

Convert posteriors into ArviZ-schema xarray DataTrees for diagnostics and plotting.

convert.to_datatree(): Convert an already-drawn posterior into an ArviZ-schema xarray.DataTree.
convert.add_forecast_groups(): Attach out-of-sample forecast groups to a copy of tree.
convert.predictions_to_datatree(): Pack prediction draws into a DataTree laid out for per-series plot_lm faceting.

Optional backends behind pyproject extras (never imported by default).

contrib.blackjax.BlackjaxNUTSKernel: BlackJAX NUTS with Stan-style window adaptation.
contrib.blackjax.BlackjaxMCLMCKernel: BlackJAX Microcanonical Langevin Monte Carlo (MCLMC).
contrib.blackjax.BlackjaxCustomKernel: Adapt an arbitrary BlackJAX sampler via a user-supplied build_fn.
contrib.blackjax.PathfinderFit: The result of fitting a forecasting model with BlackJAX Pathfinder.
contrib.blackjax.fit_pathfinder(): Fit a forecasting model with BlackJAX Pathfinder variational inference.
contrib.blackjax.pathfinder_samples(): Draw num_samples posterior samples from a fitted Pathfinder approximation.
contrib.blackjax.MultiPathfinderFit: The result of fitting a forecasting model with multi-path BlackJAX Pathfinder.
contrib.blackjax.fit_multipathfinder(): Fit a forecasting model with multi-path BlackJAX Pathfinder and PSIS resampling.
contrib.blackjax.multipathfinder_samples(): Draw num_samples posterior samples from a fitted multipath Pathfinder fit.

Public type contracts.

typing.ForecastModel: A NumPyro forecasting model: a callable (covariates, data=None) -> None.
typing.ForecastFn: A closure that fits a model on a training window and forecasts its test horizon.
typing.Guide: A NumPyro guide for a ForecastModel: a callable with the model’s signature.
typing.InSampleFn: A closure that fits a model on a training window and scores its in-sample fit.
typing.Metric: A metric maps (pred, truth) forecast samples and ground truth to a scalar array.
typing.ModelFactory: A zero-argument callable returning a fresh ForecastModel instance.

Batched autocorrelation and partial autocorrelation diagnostics.

acf.acf(): Compute the empirical autocorrelation function up to max_lag.
acf.pacf(): Compute the empirical partial autocorrelation function up to max_lag.

Fourier design matrices and seasonal tiling.

features.fourier_features(): Build a Fourier seasonality design matrix.
features.periodic_repeat(): Tile a seasonal pattern to cover duration time steps.

Time-axis array shaping for the train/forecast split.

arrays.zero_data_like(): Return zeros shaped like data but extended to the covariate duration.
arrays.concat_future(): Concatenate in-sample and forecast-horizon arrays along the time axis.
arrays.pad_future(): Append future rows filled with value along the time axis.

Example datasets used in the tutorials.

datasets.load_bart_weekly(): Load total weekly BART ridership (log scale) for the univariate example.
datasets.load_bart_hierarchical(): Load the windowed hierarchical BART panel for the hierarchical example.
datasets.load_victoria_electricity(): Load hourly Victoria (Australia) electricity demand and temperature.
datasets.bart_available(): Return whether the BART dataset can be loaded (download succeeds).

Lazy imports behind pyproject extras.

optional.require(): Import an optional dependency, or raise a targeted ImportError.

Package exception hierarchy raised at validation boundaries.

exceptions.NumpyroForecastError: Base class for all deliberate numpyro_forecast errors.
exceptions.BacktestWindowError: A backtest window configuration is invalid.
exceptions.VectorizedMetricError: A metric is not vmappable in the vectorized backtest.
exceptions.KernelConfigError: A contrib.blackjax kernel is run unbound or misconfigured.
exceptions.CovariateDimsError: Covariate dimension names are inconsistent or malformed.
exceptions.MVNLayoutError: A MultivariateNormal layout is unsupported for time-axis surgery.
exceptions.DeviceMemoryError: A memory pool ran out during posterior or predictive sampling.
exceptions.HostMemoryKindError: A device exposes no host memory kind for device="pinned_host".
exceptions.DevicePlatformError: A device platform name has no initialized JAX backend.