pymc-modeling

name

pymc-modeling

description

Bayesian statistical modeling with PyMC v5+. Use when building probabilistic models, specifying priors, running MCMC inference, diagnosing convergence, or comparing models. Covers PyMC, ArviZ, pymc-bart, pymc-extras, nutpie, and JAX/NumPyro backends. Triggers on tasks involving: Bayesian inference, posterior sampling, hierarchical/multilevel models, GLMs, time series, Gaussian processes, BART, mixture models, prior/posterior predictive checks, MCMC diagnostics, LOO-CV, WAIC, model comparison, or causal inference with do/observe.

PyMC Modeling

Bayesian modeling workflow for PyMC v5+ with modern API patterns.

Notebook preference: Use marimo for interactive modeling unless the project already uses Jupyter.

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()

Coords and Dims

Use coords/dims for interpretable InferenceData when 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

Default Sampling (nutpie)

Use nutpie as the default sampler—it's Rust-based and typically 2-5x faster:

with model:
    idata = pm.sample(
        draws=1000, tune=1000, chains=4,
        nuts_sampler="nutpie",
        random_seed=42,
    )

PyMC Native Sampling

Fall back to PyMC's NUTS when nutpie unavailable:

with model:
    idata = pm.sample(draws=1000, tune=1000, chains=4, random_seed=42)

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

Follow this systematic workflow after every sampling run:

Phase 1: Immediate Checks (Required)

# 1. Check for divergences (must be 0 or near 0)
n_div = idata.sample_stats["diverging"].sum().item()
print(f"Divergences: {n_div}")

# 2. Summary with convergence diagnostics
summary = az.summary(idata, var_names=["~offset"])  # exclude auxiliary
print(summary[["mean", "sd", "hdi_3%", "hdi_97%", "ess_bulk", "ess_tail", "r_hat"]])

# 3. Visual convergence check
az.plot_trace(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(idata, kind="evolution")

# 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:
    pm.sample_posterior_predictive(idata, extend_inferencedata=True)

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

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

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

Phase 4: Parameter Interpretation

# Posterior summaries
az.plot_posterior(idata, var_names=["beta"], ref_val=0)

# 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)

# Do prior predictions span reasonable outcome range?
az.plot_ppc(prior_pred, group="prior", kind="cumulative")

# Numerical sanity check
prior_y = prior_pred.prior_predictive["y"].values.flatten()
print(f"Prior predictive range: [{prior_y.min():.1f}, {prior_y.max():.1f}]")

Warning signs: Prior predictive covers implausible values (negative counts, probabilities > 1) or is extremely wide/narrow.

Posterior Predictive (After Fitting)

with model:
    pm.sample_posterior_predictive(idata, extend_inferencedata=True)

# Density comparison
az.plot_ppc(idata, kind="kde")

# Cumulative (better for systematic deviations)
az.plot_ppc(idata, kind="cumulative")

# Calibration diagnostic
az.plot_loo_pit(idata, y="y")

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

Model Debugging

Inspecting Model Structure

# Print model summary (variables, shapes, distributions)
print(model)

# Visualize model as directed graph
pm.model_to_graphviz(model)

Checking for Specification Errors

Before sampling, validate the model:

# Debug model: checks for common issues
model.debug()

# Check initial point log-probabilities
# Identifies which variables have invalid starting values
model.point_logps()

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
Poor out-of-sample prediction	Static data arrays	Use pm.Data containers with mutable=True

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

Debugging Divergences

# Identify where divergences occur in parameter space
az.plot_pair(idata, var_names=["alpha", "beta", "sigma"], divergences=True)

# Check if divergences cluster in specific regions
# Clustering suggests parameterization or prior issues

Profiling Slow Models

# Time individual operations in the log-probability computation
profile = model.profile(model.logp())
profile.summary()

# Identify bottlenecks in gradient computation
import pytensor
grad_profile = model.profile(pytensor.grad(model.logp(), model.continuous_value_vars))
grad_profile.summary()

See references/gotchas.md for additional troubleshooting.

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(idata)

Comparing Models

comparison = az.compare({
    "model_a": idata_a,
    "model_b": idata_b,
}, ic="loo")

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

Decision rule: If d_loo < 2*dse, models are effectively equivalent.

See references/arviz.md for detailed model comparison workflow.

Saving and Loading Results

InferenceData Persistence

Save sampling results for later analysis or sharing:

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

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

Compressed Storage

For large InferenceData objects (many draws, large posterior predictive):

# Compress with zlib (reduces file size 50-80%)
idata.to_netcdf(
    "results/model_v1.nc",
    engine="h5netcdf",
    encoding={var: {"zlib": True, "complevel": 4}
              for group in ["posterior", "posterior_predictive"]
              if hasattr(idata, group)
              for var in getattr(idata, group).data_vars}
)

What Gets Saved

InferenceData preserves the full Bayesian workflow:

• posterior: Parameter samples from MCMC
• prior, prior_predictive: Prior samples (if generated)
• posterior_predictive: Predictions (if generated)
• observed_data, constant_data: Data used in fitting
• sample_stats: Diagnostics (divergences, tree depth, energy)
• log_likelihood: Pointwise log-likelihood (for LOO-CV)
• All coordinates and dimensions

Workflow Pattern

# Save after each major step
with model:
    idata = pm.sample(nuts_sampler="nutpie")
idata.to_netcdf("results/step1_posterior.nc")

with model:
    pm.sample_posterior_predictive(idata, extend_inferencedata=True)
idata.to_netcdf("results/step2_with_ppc.nc")

# Resume later
idata = az.from_netcdf("results/step2_with_ppc.nc")
az.plot_ppc(idata)  # Continue analysis

Prior Selection

See references/priors.md for:

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

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")
    alpha = pm.Deterministic("alpha", mu_alpha + sigma_alpha * alpha_offset, dims="group")

    # Likelihood
    y = pm.Normal("y", alpha[group_idx], sigma, observed=y_obs, dims="obs")

GLMs

# Logistic regression
with pm.Model() as logistic:
    alpha = pm.Normal("alpha", 0, 2.5)  # intercept
    beta = pm.Normal("beta", 0, 2.5, dims="features")
    
    # Logit link
    logit_p = alpha + pm.math.dot(X, beta)
    p = pm.math.sigmoid(logit_p)
    
    y = pm.Bernoulli("y", p=p, observed=y_obs)

# Poisson regression
with pm.Model() as poisson:
    beta = pm.Normal("beta", 0, 1, dims="features")
    mu = pm.math.exp(pm.math.dot(X, beta))
    y = pm.Poisson("y", mu=mu, observed=y_obs)

Gaussian Processes