pymc-modeling

name

pymc-modeling

description

Load whenever the user is working on code that imports pymc, pytensor, or arviz, or asks about Bayesian modeling, MCMC, priors, posteriors, sampling, or model diagnostics. Covers PyMC 6+, PyTensor 3+, ArviZ 1.0+ (DataTree API), pymc-bart, pymc-extras, nutpie, and JAX/NumPyro backends. Use for building probabilistic models, specifying priors, running MCMC, diagnosing convergence, or comparing models. Triggers include: Bayesian inference, posterior sampling, hierarchical/multilevel models, GLMs, time series, Gaussian processes, HSGP, BART, mixture models, prior/posterior predictive checks, MCMC diagnostics, LOO-CV, model comparison, causal inference with do/observe, and any PyTensor Op or graph work.

PyMC Modeling

Modern Bayesian modeling with PyMC 6+ on the ArviZ 1.0 / PyTensor 3 stack. Key defaults: nutpie sampler (2-5x faster; PyMC 6 selects it automatically when installed — no nuts_sampler argument needed), non-centered parameterization for hierarchical models, HSGP over exact GPs, coords/dims for readable DataTree output, and save-early workflow to prevent data loss from late crashes.

pm.sample(...) returns an xarray.DataTree — the idata name is kept by convention, but it is a DataTree, not the old InferenceData. Access groups by bracket: idata["posterior"], idata["sample_stats"], etc.

Modeling strategy: Build models iteratively — start simple, check prior predictions, fit and diagnose, check posterior predictions, expand one piece at a time. See references/workflow.md for the full workflow.

Model Specification

Basic Structure

import pymc as pm
import arviz as az

with pm.Model(coords=coords) as model:
    # Data containers (for out-of-sample prediction)
    x = pm.Data("x", x_obs, dims="obs")

    # Priors
    beta = pm.Normal("beta", mu=0, sigma=1, dims="features")
    sigma = pm.HalfNormal("sigma", sigma=1)

    # Likelihood
    mu = pm.math.dot(x, beta)
    y = pm.Normal("y", mu=mu, sigma=sigma, observed=y_obs, dims="obs")

    # Inference
    idata = pm.sample(random_seed=42)  # PyMC 6 uses nutpie automatically when installed

Coords and Dims

Use coords/dims for an interpretable DataTree when the model has meaningful structure:

coords = {
    "obs": np.arange(n_obs),
    "features": ["intercept", "age", "income"],
    "group": group_labels,
}

Skip for simple models where overhead exceeds benefit.

Parameterization

Prefer non-centered parameterization for hierarchical models with weak data:

# Non-centered (better for divergences)
offset = pm.Normal("offset", 0, 1, dims="group")
alpha = mu_alpha + sigma_alpha * offset

# Centered (better with strong data)
alpha = pm.Normal("alpha", mu_alpha, sigma_alpha, dims="group")

Inference

Sampling budgets

Start with PyMC's default tuning budget: 400 iterations per chain. More iterations are not a substitute for a well-identified model with healthy geometry.

• Use 400 tuning iterations per chain by default. Use 1,000 only when model complexity, warmup diagnostics, or a justified adaptation need supports it. Treat 2,000 as an exceptional upper limit. Never exceed 2,000. If it is insufficient, repair parameterization, scaling, identifiability, priors, or the likelihood instead of increasing tuning.
• Choose a total posterior-draw budget from the required ESS and inferential precision, then distribute it across available compute. Prefer more independent chains and fewer draws per chain when sufficient cores are available; a large per-chain draw count is not a default.
• Before an expensive fit, state the available cores, chain count, tune per chain, draws per chain, total posterior draws, desired ESS/precision, and why that allocation is sufficient. These values depend on the model and machine, not a universal prescription.
• If diagnostics or precision are inadequate, diagnose geometry and parameterization before increasing the sampling budget.

For example, on a 32-core machine, 16 concurrent chains with roughly 200 draws each may be sensible for a 3,200-total-draw target—if that meets the required ESS and precision. This is an example, not a prescription.

Default Sampling (nutpie preferred)

In PyMC 6, pm.sample uses nutpie automatically whenever it is installed and the model can be compiled — do not pass nuts_sampler="nutpie" explicitly:

# Example only: 4 chains × 500 draws = 2,000 posterior draws total.
# Select chains and draws from available compute and required ESS/precision.
with model:
    idata = pm.sample(
        draws=500, tune=400, chains=4,
        random_seed=42,
    )
idata.to_netcdf("results.nc")  # Save immediately after sampling

Important: In PyMC 6, pm.sample no longer computes the log-likelihood automatically — passing compute_log_likelihood=True emits a FutureWarning. Compute it explicitly after sampling whenever you plan to run LOO-CV, model comparison, or loo-pit checks:

pm.compute_log_likelihood(idata, model=model)

This applies to every sampler (nutpie, default NUTS, NumPyro) — not just nutpie.

When to Use PyMC's Default NUTS Instead

nutpie cannot handle discrete parameters or certain transforms (e.g., ordered transform with OrderedLogistic/OrderedProbit). PyMC 6 falls back automatically; to force the PyMC sampler explicitly, pass nuts_sampler="pymc":

idata = pm.sample(draws=500, tune=400, chains=4, nuts_sampler="pymc", random_seed=42)

Never change the model specification to work around sampler limitations.

If nutpie is not installed, install it (pip install nutpie) or fall back to nuts_sampler="numpyro".

Alternative MCMC Backends

See references/inference.md for:

• NumPyro/JAX: GPU acceleration, vectorized chains

Approximate Inference

For fast (but inexact) posterior approximations:

• ADVI/DADVI: Variational inference with Gaussian approximation
• Pathfinder: Quasi-Newton optimization for initialization or screening

Diagnostics and ArviZ Workflow

Minimum workflow checklist — every model script should include:

1. Prior predictive check (pm.sample_prior_predictive)
2. Save results immediately after sampling (idata.to_netcdf(...))
3. Divergence count + r_hat + ESS check
4. Posterior predictive check (pm.sample_posterior_predictive)

Follow this systematic workflow after every sampling run:

Phase 1: Immediate Checks (Required)

# 1. Check for divergences (must be 0 or near 0)
# idata is an xarray.DataTree; path-access gets a DataArray
n_div = idata["sample_stats"]["diverging"].sum().item()
print(f"Divergences: {n_div}")

# 2. Summary with convergence diagnostics
# Default CI is 0.89 ETI (equal-tailed) — bounds labelled eti_5.5% / eti_94.5%
summary = az.summary(idata, var_names=["~offset"])  # exclude auxiliary
print(summary[["mean", "sd", "eti_5.5%", "eti_94.5%", "ess_bulk", "ess_tail", "r_hat"]])

# 3. Visual convergence check
az.plot_trace_dist(idata, compact=True)
az.plot_rank(idata, var_names=["beta", "sigma"])

Pass criteria (all must pass before proceeding):

• Zero divergences (or < 0.1% and randomly scattered)
• r_hat < 1.01 for all parameters
• ess_bulk > 400 and ess_tail > 400
• Trace plots show good mixing (overlapping densities, fuzzy caterpillar)

Phase 2: Deep Convergence (If Phase 1 marginal)

# ESS evolution (should grow linearly)
az.plot_ess_evolution(idata)

# Energy diagnostic (HMC health)
az.plot_energy(idata)

# Autocorrelation (should decay rapidly)
az.plot_autocorr(idata, var_names=["beta"])

Phase 3: Model Criticism (Required)

# Generate posterior predictive
with model:
    idata.update(pm.sample_posterior_predictive(idata))

# Does the model capture the data?
az.plot_ppc_dist(idata, kind="ecdf")

# Calibration check
az.plot_loo_pit(idata, var_names=["y"])

Critical rule: Never interpret parameters until Phases 1-3 pass.

Phase 4: Parameter Interpretation

# Posterior summaries
az.plot_dist(idata, var_names=["beta"])

# Forest plots for hierarchical parameters
az.plot_forest(idata, var_names=["alpha"], combined=True)

# Parameter correlations (identify non-identifiability)
az.plot_pair(idata, var_names=["alpha", "beta", "sigma"])

See references/arviz.md for comprehensive ArviZ usage. See references/diagnostics.md for troubleshooting.

Prior and Posterior Predictive Checks

Prior Predictive (Before Fitting)

Always check prior implications before fitting:

with model:
    prior_pred = pm.sample_prior_predictive(draws=500)

az.plot_ppc_dist(prior_pred, group="prior_predictive", kind="ecdf")
prior_y = prior_pred["prior_predictive"]["y"].values.flatten()
print(f"Prior predictive range: [{prior_y.min():.1f}, {prior_y.max():.1f}]")

Rule: Run prior predictive checks before pm.sample() on any new model. If the range is implausible (negative counts, probabilities > 1), adjust priors before proceeding.

Posterior Predictive (After Fitting)

with model:
    idata.update(pm.sample_posterior_predictive(idata))

az.plot_ppc_dist(idata, kind="ecdf")
az.plot_loo_pit(idata, var_names=["y"])

Observed data (dark line) should fall within posterior predictive distribution. See references/arviz.md for detailed interpretation.

Model Debugging

Before sampling, validate the model with model.debug() and model.point_logps(). Use print(model) for structure and pm.model_to_graphviz(model) for a DAG visualization.

Common Issues

Symptom	Likely Cause	Fix
ValueError: Shape mismatch	Parameter vs observation dimensions	Use index vectors: alpha[group_idx]
Initial evaluation failed	Data outside distribution support	Check bounds; use init="adapt_diag"
Mass matrix contains zeros	Unscaled predictors or flat priors	Standardize features; use weakly informative priors
High divergence count	Funnel geometry	Non-centered parameterization
NaN in log-probability	Invalid parameter combinations	Check parameter constraints, add bounds
-inf log-probability	Observations outside likelihood support	Verify data matches distribution domain
Slow discrete sampling	NUTS incompatible with discrete	Marginalize discrete variables
Fresh env: sampling never starts, one core pinned in numba typing for hours	Cold ~/.pytensor/numba kernel cache (invalidated by pytensor version change)	Warm cache with a small-data fit first; never delete the cache — see references/troubleshooting.md § Cold Numba Kernel Cache

See references/troubleshooting.md for comprehensive problem-solution guide.

For debugging divergences, use az.plot_pair(idata, divergences=True) to locate clusters. See references/diagnostics.md § Divergence Troubleshooting.

For profiling slow models, see references/troubleshooting.md § Performance Issues.

Model Comparison

LOO-CV (Preferred)

# Compute LOO with pointwise diagnostics
loo = az.loo(idata, pointwise=True)
print(f"ELPD: {loo.elpd_loo:.1f} ± {loo.se:.1f}")

# Check Pareto k values (must be < 0.7 for reliable LOO)
print(f"Bad k (>0.7): {(loo.pareto_k > 0.7).sum().item()}")
az.plot_khat(loo)

Comparing Models

# PyMC 6 requires an explicit log-likelihood compute before LOO
pm.compute_log_likelihood(idata_a, model=model_a)
pm.compute_log_likelihood(idata_b, model=model_b)

# ArviZ 1.0 — only loo is supported (waic was removed)
comparison = az.compare({
    "model_a": idata_a,
    "model_b": idata_b,
})

print(comparison[["rank", "elpd_loo", "elpd_diff", "weight"]])
az.plot_compare(comparison)

Decision rule: If two models have similar stacking weights, they are effectively equivalent.

See references/arviz.md for detailed model comparison workflow. For detailed LOO-CV workflows, model stacking, and calibration diagnostics, see the model-evaluation skill.

Iterative Model Building

Build complexity incrementally: fit the simplest plausible model first, diagnose it, check posterior predictions, then add ONE piece of complexity at a time. Compare each expansion via LOO. If stacking weights are similar, the models are effectively equivalent. See references/workflow.md for the full iterative workflow.

Saving and Loading Results

DataTree Persistence

pm.sample() returns an xarray.DataTree. Persist with NetCDF; the idata name is convention.

# Save to NetCDF (recommended format)
idata.to_netcdf("results/model_v1.nc")

# Load
idata = az.from_netcdf("results/model_v1.nc")

For compressed storage of large DataTree objects, see references/workflow.md.

Critical: Save IMMEDIATELY after sampling — late crashes destroy valid results:

with model:
    idata = pm.sample()  # nutpie by default in PyMC 6; returns a DataTree
idata.to_netcdf("results.nc")  # Save before any post-processing!

with model:
    idata.update(pm.sample_posterior_predictive(idata))  # .update() merges the new group in place
idata.to_netcdf("results.nc")  # Update with posterior predictive

Note: Use .update({...}) or direct assignment (idata["posterior_predictive"] = ppd_ds) to add groups.

Prior Selection

See references/priors.md for:

• Weakly informative defaults by distribution type
• Prior predictive checking workflow
• Domain-specific recommendations

For constrained priors, expert elicitation workflows, and PreliZ integration, see the prior-elicitation skill.

Common Patterns

Hierarchical/Multilevel

with pm.Model(coords={"group": groups, "obs": obs_idx}) as hierarchical:
    # Hyperpriors
    mu_alpha = pm.Normal("mu_alpha", 0, 1)
    sigma_alpha = pm.HalfNormal("sigma_alpha", 1)

    # Group-level (non-centered)
    alpha_offset = pm.Normal("alpha_offset", 0, 1, dims="group")