Prediction output schema#

The dimension, coordinate, group, and variable names on prediction outputs are public API. Downstream code (plotting, metrics, adapters such as CausalPy) can rely on the names below; any change to them is a breaking change, made only in a minor release and called out in the changelog.

Groups#

Prediction results are ArviZ trees (DataTree / InferenceData) with exactly one predictive group:

Producer	Group	Contents
`forecast(...)`	`predictions`	out-of-sample forecasts
`predict_in_sample(...)`	`posterior_predictive`	in-sample predictive of the observed variable, plus its registered noise-free predictors

The group names follow the ArviZ convention (out-of-sample predictions live in predictions). pymc_forecast.prediction_samples() extracts the samples Dataset from either group, so adapters don’t need to branch on the name.

Variables#

Name	Constant	Where	Meaning
`obs`	`pymc_forecast.OBS_VAR`	`posterior_predictive`	the observed (in-sample) variable
`mu`	`pymc_forecast.MU_VAR`	`posterior_predictive`	noise-free latent predictor over the observed window
`expected_observation`	`pymc_forecast.EXPECTED_OBSERVATION_VAR`	`posterior_predictive`	optional conditional expected observation over the observed window
`forecast`	`pymc_forecast.FORECAST_VAR`	`predictions`	the forecast-horizon variable
`mu_future`	`pymc_forecast.MU_FORECAST_VAR`	`predictions`	noise-free latent predictor over the forecast horizon
`expected_observation_future`	`pymc_forecast.EXPECTED_OBSERVATION_FORECAST_VAR`	`predictions`	optional conditional expected observation over the forecast horizon
`{name}_future`	—	`predictions`	forecast-horizon slice of each per-step latent registered by `innovations` (and by `markov_series` / `ssoe` for their own future names)

mu / mu_future carry the full draw-level samples of the latent passed to pymc_forecast.predict() — before observation noise, so they separate parameter and latent uncertainty from observation-level noise (e.g. for plotting the expected trajectory against the predictive draws). For GLM-style models this is the linear predictor, not the distribution mean — under a non-identity link mu is on the link scale, so consumers that need outcome units (causal-impact expectations, for instance) should read expected_observation / expected_observation_future instead, and use mu only for identity-link models. The names are reserved: a model body must not define its own mu or mu_future variable when calling predict. Models that register observations directly can explicitly register these Deterministics, as in the pymc_forecast.ssoe() examples. The predict_mvn helper and statespace adapter do not currently expose them.

expected_observation / expected_observation_future are opt-in outputs: model authors supply expected_observation= to pymc_forecast.predict() explicitly. They represent E[Y | parameters, latent state, covariates] in observed outcome units and exclude observation-level sampling noise. For example, a Poisson log-link model passes latent=eta and expected_observation=exp(eta), retaining eta as mu / mu_future while exposing expected counts separately. The library does not infer inverse links or distribution means from the observation factory. These names are reserved just like mu / mu_future — a model body must not define its own expected_observation or expected_observation_future variable, whether or not it passes expected_observation=. Models that omit the argument retain the previous output schema.

The statespace adapter additionally exposes the latent state trajectories as forecast_latent in its predictions group.

Dimensions and coordinates#

Every predictive variable carries, in order:

Sample dims ("chain", "draw") — always both, always leading, holding the full draw-level posterior-predictive samples (no reduction to means or quantiles happens on the default path). Exported as pymc_forecast.CHAIN_DIM, pymc_forecast.DRAW_DIM, and the ordered pair pymc_forecast.SAMPLE_DIMS; the metrics layer recognizes sample dims through the same constants.
A time dim — "time" (pymc_forecast.TIME_DIM) on in-sample variables, "time_future" (pymc_forecast.FUTURE_DIM) on forecast variables. Its coordinate values are real: whatever time index the inputs carried (a DatetimeIndex, periods, or the integer-range fallback). time_future carries the horizon index supplied at forecast time (the surplus covariate steps, or the index extended by horizon=).
Batch dims carried over from the data — e.g. "series" for 2-d data normalized by pymc_forecast.as_dataarray(), or any named dims of an xarray input.

So a univariate forecast is (chain, draw, time_future) and a hierarchical one (chain, draw, time_future, series), with real coordinate values on the time and series axes.

Draw coherence across predictive calls#

Each predictive call draws its own posterior subsample by default, so draw i of a predict_in_sample(...) result and draw i of a forecast(...) result do not come from the same parameter draw. To make several calls draw-coherent — required for any draw-level quantity that spans groups, such as joint pre/post summaries — draw one posterior and pass it to every call:

posterior = fc.draw_posterior(500, random_seed=42)
pre = fc.predict_in_sample(posterior=posterior)
post = fc.forecast(future_covariates=future_cov, posterior=posterior)
# draw i in `pre` and `post` now share parameters; chain/draw sizes match

With posterior= given, the chain/draw sizes and coordinates of every result equal the posterior’s, and num_samples must not be passed. This includes the optional expected-observation outputs, so each expected observation and latent predictor at draw i use the same parameter draw. The same coordinates are retained when batch_size= splits predictive sampling into blocks.

Mapping onto downstream coordinates#

Because the names are fixed and the coordinates are real, remapping is a one-liner, e.g.:

samples = pymc_forecast.prediction_samples(result)["forecast"]
samples = samples.rename({"time_future": "obs_ind", "series": "treated_units"})

Observation-driven recursion outputs#

The ssoe helper registers only its future error variable. Example models register obs, mu, mu_future and forecast explicitly using the same schema. forecast already includes observation noise; adding a further likelihood draw would count it twice. mu_future excludes the current observation error, but previous simulated future errors can still influence its state. It represents a conditional one-step mean, not an unconditional multi-step expectation. In-sample mu and posterior predictive observations condition on the actual training history; they are not simulations of a new training trajectory.