Choose kernels#

A kernel updates one block of parameters conditional on the current values of the other blocks. Goose applies the kernels in sequence within each iteration. You can use different methods for different blocks without changing the model’s priors or likelihood.

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.

Choose an update#

Kernel	Use it for	What you supply or check
`NUTSKernel`	Continuous parameters with gradients.	Parameters on an unconstrained scale; trajectory length is selected automatically.
`HMCKernel`	Hamiltonian updates with a chosen trajectory length.	Gradients and the number of integration steps.
`IWLSKernel`	Blocks where local curvature gives a useful Gaussian proposal.	A differentiable density with suitable curvature; larger blocks make matrix operations more expensive.
`RWKernel`	A simple Gaussian random-walk update.	A useful parameter scale; mixing can be slow for strongly dependent parameters.
`GibbsKernel`	A block whose full conditional you can sample directly.	A function drawing from that full conditional.
`MHKernel`	A custom Metropolis-Hastings proposal.	A proposal function and the backward-minus-forward log proposal correction.

See Sample your first posterior for a positive parameter sampled on a log scale, and Combine NUTS and Gibbs for an exact Gibbs update. For custom proposals, start with Define a custom kernel.

Check kernel assignments#

For the regression model from Sample your first posterior:

mcmc = gs.LieselMCMC(model)
kernel_table = pd.DataFrame(
    [
        {"Kernel": type(kernel).__name__, "Parameters": kernel.position_keys}
        for kernel in mcmc.get_kernel_list()
    ]
)

kernel_table

	Kernel	Parameters
0	NUTSKernel	(beta, log_sigma_sq)

Check that every parameter you intend to estimate appears in a kernel. A parameter without an inference specification stays at its current value unless you add a kernel manually through an engine builder. Declaring a prior and setting the parameter flag does not assign a kernel. If you sample regression coefficients but leave the variance without a kernel, you obtain a posterior conditional on that fixed variance, not uncertainty about both. See Change a model for deliberate model changes.

Group parameters#

For the model from Sample your first posterior, give the coefficients and log variance separate NUTS kernels:

beta = model.vars["beta"]
log_variance = model.vars["log_sigma_sq"]

beta.inference = gs.MCMCSpec(gs.NUTSKernel, order=1)
log_variance.inference = gs.MCMCSpec(gs.NUTSKernel, order=2)

One iteration first updates beta given the old log variance, then updates the log variance given the new beta. The vector beta is already one block: both coefficients are proposed together.

To update both variables jointly, give them an explicit shared group:

joint = gs.MCMCSpec(
    gs.NUTSKernel,
    kernel_group="regression",
    kernel_kwargs={"da_target_accept": 0.9},
    jitter_dist=tfd.Normal(0.0, 0.2),
)
beta.inference = joint
log_variance.inference = joint

The shared specification keeps the group’s kernel, arguments, and order consistent. Reusing a specification without kernel_group does not combine variables. Assigning a new specification replaces the previous one, including its jitter settings.

Joint updates can help with posterior dependence but cost more per update. Inspect diagnostics after changing the grouping. A higher da_target_accept generally adapts toward smaller steps; it is not a substitute for a well-scaled model.

Order the updates#

Smaller order values run first. With equal values, the high-level helper uses reverse graph order, starting nearer the responses. Inspect get_kernel_list() if the order matters, and use the same order for every member of a joint group. See MCMCSpec and get_kernel_groups() for the exact grouping rules.

Current β and σ² Update β given σ² Update σ² given new β Store the iteration

Each block sees updates already made in this iteration. Separate chains run their own sequences.

Choose kernels

Contents