pymc-modeling

# 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 ```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() ``` ### Coords and Dims Use coords/dims for interpretable InferenceData when 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) Use nutpie as the default sampler—it's Rust-based and typically 2-5x faster: ```python 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: ```python with model: idata = pm.sample(draws=1000, tune=1000, chains=4, random_seed=42) ``` ### 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 Follow this systematic workflow after every sampling run: ### Phase 1: Immediate Checks (Required) ```python # 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) ```python # 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) ```python # 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 ```python # 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](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) # 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) ```python 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](references/arviz.md) for detailed interpretation. ## Model Debugging ### Inspecting Model Structure ```python # 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: ```python # 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](references/troubleshooting.md) for comprehensive problem-solution guide. ### Debugging Divergences ```python # 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 ```python # 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](references/gotchas.md) for additional troubleshooting. ## 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(idata) ``` ### Comparing Models ```python 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](references/arviz.md) for detailed model comparison workflow. ## Saving and Loading Results ### InferenceData Persistence Save sampling results for later analysis or sharing: ```python # 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): ```python # 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 ```python # 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](references/priors.md) for: - Weakly informative defaults by distribution type - Prior predictive checking workflow - Domain-specific recommendations ## 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) # 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