pymc-modeling

# 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](references/workflow.md) for the full workflow. ## Model Specification ### Basic Structure ```python 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: ```python 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: ```python # 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: ```python # 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: ```python 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"`: ```python 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](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) ```python # 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) ```python # 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) ```python # 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 ```python # 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](references/arviz.md) for comprehensive ArviZ usage. See [references/diagnostics.md](references/diagnostics.md) for troubleshooting. ## Prior and Posterior Predictive Checks ### Prior Predictive (Before Fitting) Always check prior implications before fitting: ```python 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) ```python 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](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](references/troubleshooting.md) § Cold Numba Kernel Cache | See [references/troubleshooting.md](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](references/diagnostics.md) § Divergence Troubleshooting. For profiling slow models, see [references/troubleshooting.md](references/troubleshooting.md) § Performance Issues. ## Model Comparison ### LOO-CV (Preferred) ```python # 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 ```python # 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](references/arviz.md) for detailed model comparison workflow. For detailed LOO-CV workflows, model stacking, and calibration diagnostics, see the [model-evaluation skill](../model-evaluation/SKILL.md). ### 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](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. ```python # 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](references/workflow.md). **Critical**: Save IMMEDIATELY after sampling — late crashes destroy valid results: ```python 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](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](../prior-elicitation/SKILL.md). ## Common Patterns ### Hierarchical/Multilevel ```python 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")