Configure sampling#

An iteration applies every configured kernel once. A chain repeats those iterations; an epoch is a run of iterations with the same phase and storage settings.

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, including its transformation and inference specifications, from Sample your first posterior.

Set the phase lengths#

For an existing model with complete inference specifications:

results = gs.LieselMCMC(model).run_for_epochs(
    seed=2026,
    num_chains=4,
    adaptation=1000,
    burnin=200,
    posterior=2000,
    show_progress=False,
)

liesel.goose.engine - INFO - Initializing kernels...

liesel.goose.engine - INFO - Done

liesel.goose.engine - INFO - Finished warmup

These counts apply to each chain. Adaptation tunes kernel settings. Optional burnin lets the chains continue with tuning finished. Both belong to warmup. get_posterior_samples() returns only the posterior phase: here, 2,000 draws from each of four chains. The additional recorded initial state is not a posterior draw.

Initial
state

Adaptation
Fast → expanding slow epochs → fast

Optional
burnin

Posterior sampling
Draws used for inference

Warmup includes adaptation and burnin. Box widths do not represent recommended durations.

During adaptation, kernels can tune within iterations and between epochs. NUTS and HMC tune step size and update their inverse mass matrix after slow epochs. RW and IWLS tune step size; Gibbs has nothing to tune. See Use the engine to customize the schedule.

The example lengths are a starting point. Use diagnostics to assess the result. More posterior iterations do not repair invalid starting values or a problematic model.

Thin stored draws#

Keep posterior_thinning=1 unless storing the draws is too expensive. For example, requesting 2,000 posterior iterations with posterior_thinning=2 retains 1,000 draws per chain; all 2,000 transitions still run. Thinning discards information and does not fix poor mixing.

Avoid adaptation thinning for kernels that rely on sample history for tuning. See add_adaptation() for constraints.

Inspect tuning#

Kernel states must be recorded during sampling. For the single joint NUTS kernel in Sample your first posterior, run:

recorded = gs.LieselMCMC(model).run_for_epochs(
    seed=2026,
    num_chains=4,
    adaptation=1000,
    posterior=1000,
    show_progress=False,
    store_kernel_states=True,
)

states = recorded.get_warmup_kernel_states()
kernel_id = recorded.get_kernels_by_pos_key()["beta"]
step_size = states[kernel_id]["step_size"]

liesel.goose.engine - INFO - Initializing kernels...

liesel.goose.engine - INFO - Done

liesel.goose.engine - INFO - Finished warmup

step_size.shape

(4, 1000)

pd.DataFrame(
    {
        "Chain": range(step_size.shape[0]),
        "Last recorded step size": np.asarray(step_size[:, -1]),
    },
).round(3)

	Chain	Last recorded step size
0	0	0.568
1	1	0.317
2	2	0.577
3	3	0.334

step_size has shape (chains, warmup iterations). It contains the adapting step sizes. The table shows the last recorded step size for each chain, not convergence diagnostics. To investigate changes over warmup, compare earlier entries along the second axis. Values may change at epoch boundaries; the final dual-averaging update can also change the step size used for posterior sampling. Recording kernel states uses additional memory, especially for mass matrices. It cannot be enabled retroactively for an existing result.

Allow compilation time#

Goose uses JAX to compile the sampling calculations. Compilation can make the first part of a run slower than later iterations. Multiple chains have separate states and random keys; they do not imply one physical CPU core per chain. Record the backend and numerical environment for reproducibility.

Configure sampling

Contents