pymc-modeling

# PyMC Modeling Modern Bayesian modeling with PyMC 6+ on the ArviZ 1.1 / 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 ### 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 with model: idata = pm.sample( draws=1000, tune=1000, chains=4, random_seed=42, ) idata.to_netcdf("results.nc") # Save immediately after sampling ``` **Important**: For LOO-CV, model comparison, or LOO-PIT checks, ensure the `log_likelihood` group exists. In PyMC 6, do not pass a top-level `compute_log_likelihood=` argument to `pm.sample`. Either request it during conversion with `idata_kwargs={"log_likelihood": True}` or compute it explicitly after sampling: ```python idata = pm.sample(idata_kwargs={"log_likelihood": True}, random_seed=42) # or, after an existing sample: pm.compute_log_likelihood(idata, model=model) ``` This applies to every sampler (nutpie, PyMC 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=1000, tune=1000, 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 eti89_lb / eti89_ub summary = az.summary(idata, var_names=["~offset"]) # exclude auxiliary print(summary[["mean", "sd", "eti89_lb", "eti89_ub", "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 | 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:.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.1 — only loo is supported (waic was removed) comparison = az.compare({ "model_a": idata_a, "model_b": idata_b, }) print(comparison[["rank", "elpd", "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") 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 ```python # Logistic regression with pm.Model() as logistic: alpha = pm.Normal("alpha", 0, 2.5) beta = pm.Normal("beta", 0, 2.5, dims="features")