Use posterior draws#

Extract draws, compute predictions, and save a finished run.

from pathlib import Path
from tempfile import TemporaryDirectory

import jax
import jax.numpy as jnp
import numpy as np
import pandas as pd
import tensorflow_probability.substrates.jax.bijectors as tfb
import tensorflow_probability.substrates.jax.distributions as tfd

import liesel.goose as gs
import liesel.model as lsl

Prepare the example#

These examples require the regression model and sampling results from Sample your first posterior, with sigma_sq included in the stored draws.

Extract the posterior#

samples = results.get_posterior_samples()
sample_shapes = pd.DataFrame(
    {
        "Variable": ["beta", "sigma_sq"],
        "Shape": [samples["beta"].shape, samples["sigma_sq"].shape],
    },
)

sample_shapes

	Variable	Shape
0	beta	(4, 1000, 2)
1	sigma_sq	(4, 1000)

The shapes are (4, 1000, 2) and (4, 1000) in the tutorial: chains, draws, then any parameter dimensions. Preserve the first two axes when computing diagnostics. get_samples() also includes initial values and warmup; use get_posterior_samples() for posterior inference.

Record derived values#

Goose records variables updated by kernels. The tutorial also sets positions_included=["sigma_sq"] to record the derived variance.

Here NUTS samples log_sigma_sq; sigma_sq is its positive transformation. positions_excluded overrides included positions and affects storage only. It does not remove a kernel. Excluding sampled parameters can prevent later predictions or diagnostics that need those draws.

For large intermediate arrays, retain the underlying parameter draws and recompute quantities as needed instead:

variance = model.predict(samples, predict=["sigma_sq"])

gs.SamplesSummary(variance).to_dataframe()[["mean", "sd"]].round(3)

	mean	sd
variable
sigma_sq	1.039	0.065

Predict new responses#

For the tutorial’s design matrix variable X, construct a grid with the same two columns: intercept and covariate.

x_grid = jnp.linspace(0.0, 1.0, 50)
X_grid = jnp.column_stack([jnp.ones_like(x_grid), x_grid])

mean = model.predict(
    samples,
    predict=["mu"],
    newdata={"X": X_grid},
)
mean_summary = gs.SamplesSummary(mean, which=["quantiles"]).to_dataframe()
mean_summary["x"] = x_grid

mean_summary.iloc[[0, -1]][["x", "q_0.05", "q_0.5", "q_0.95"]].round(3)

	x	q_0.05	q_0.5	q_0.95
variable
mu	0.0	0.839	0.992	1.140
mu	1.0	2.741	2.888	3.041

mean["mu"] contains the regression mean at each grid point for every draw. The displayed endpoints show how the fitted mean rises across the grid. The 5th and 95th percentiles form a pointwise 90% credible interval for the mean, as plotted in the first tutorial.

To include observation noise, simulate responses conditional on those draws:

parameter_draws = {name: samples[name] for name in model.parameters}
predictive = model.sample(
    sample_shape=(),
    seed=jax.random.key(2027),
    posterior_samples=parameter_draws,
    newdata={"X": X_grid, "y": jnp.zeros_like(x_grid)},
)
y_predictive = predictive["y"]

y_predictive.shape

(4, 1000, 50)

Pass draws for the writable parameters to model.sample. The recorded sigma_sq is a derived value in this model; it is recomputed from log_sigma_sq rather than set directly.

The placeholder y gives the response the new grid length while the model updates its state. These values are replaced by simulated responses; they are not additional observations used for fitting.

With shape=(), there is one simulated response vector per posterior draw. Predictive intervals also include observation variation and are generally wider. See predict() and sample() for new-data and shape rules. Additional sample dimensions, if requested, precede the posterior chain/draw dimensions.

Save a finished run#

with TemporaryDirectory() as directory:
    path = Path(directory) / "regression-samples.pkl"
    results.pkl_save(path)
    restored = gs.SamplingResults.pkl_load(path)

This example uses a temporary directory; choose a persistent path to keep your run. Load only trusted pickle files. Keep the model definition, data, and environment information with the result.

Alternatively, pass save_path="regression-samples.pkl" to run_for_epochs. If that path already exists, Goose loads it and skips sampling, even when you changed the model, seed, or sampling settings. Use a fresh path for a new fit. A saved SamplingResults object contains results; it is not an engine checkpoint for resuming a chain.

Use posterior draws

Contents