pymc-fundamentals

# PyMC 5 Fundamentals ## When to Use This Skill - Writing new PyMC models in Python - Understanding PyMC syntax and API - Converting models from Stan/JAGS to PyMC - Diagnosing sampling issues with ArviZ ## Model Structure ```python import pymc as pm import numpy as np import arviz as az with pm.Model() as model: # 1. Priors mu = pm.Normal("mu", mu=0, sigma=10) sigma = pm.HalfNormal("sigma", sigma=1) # 2. Likelihood y_obs = pm.Normal("y_obs", mu=mu, sigma=sigma, observed=y_data) # 3. Sample trace = pm.sample(1000, tune=1000, return_inferencedata=True) # 4. Diagnostics az.summary(trace) ``` ## CRITICAL: SD Parameterization **PyMC uses SD (like Stan), NOT precision (like BUGS):** ```python # PyMC (SD) pm.Normal("x", mu=0, sigma=1) # sigma is SD # BUGS equivalent would be tau = 1/sigma² = 1 ``` ## Distribution Quick Reference ### Continuous ```python pm.Normal("x", mu=0, sigma=1) # Normal pm.HalfNormal("x", sigma=1) # Half-normal (>0) pm.HalfCauchy("x", beta=2.5) # Half-Cauchy (>0) pm.Exponential("x", lam=1) # Exponential pm.Uniform("x", lower=0, upper=1) # Uniform pm.Beta("x", alpha=1, beta=1) # Beta pm.Gamma("x", alpha=2, beta=1) # Gamma pm.StudentT("x", nu=3, mu=0, sigma=1) # Student-t pm.LogNormal("x", mu=0, sigma=1) # Log-normal pm.TruncatedNormal("x", mu=0, sigma=1, lower=0) # Truncated ``` ### Discrete ```python pm.Bernoulli("x", p=0.5) # Bernoulli pm.Binomial("x", n=10, p=0.5) # Binomial pm.Poisson("x", mu=5) # Poisson pm.NegativeBinomial("x", mu=5, alpha=1) # Negative binomial pm.Categorical("x", p=[0.3, 0.5, 0.2]) # Categorical ``` ### Multivariate ```python pm.MvNormal("x", mu=np.zeros(K), cov=np.eye(K)) pm.Dirichlet("x", a=np.ones(K)) pm.LKJCholeskyCov("chol", n=K, eta=2, sd_dist=pm.Exponential.dist(1)) ``` ## Sampling ```python # Standard NUTS trace = pm.sample( draws=1000, # Samples per chain tune=1000, # Warmup chains=4, cores=4, target_accept=0.8, # Increase for divergences random_seed=42, return_inferencedata=True ) # Variational inference (fast) approx = pm.fit(n=30000, method="advi") trace = approx.sample(1000) # Predictive sampling prior_pred = pm.sample_prior_predictive(500) post_pred = pm.sample_posterior_predictive(trace) ``` ## Bayesian Workflow (Statistical Rethinking) ### 1. Prior Predictive Check ```python with model: prior_pred = pm.sample_prior_predictive(500, random_seed=42) az.plot_ppc(prior_pred, group="prior") ``` ### 2. Fit Model ```python with model: trace = pm.sample(1000, tune=1000, target_accept=0.9, return_inferencedata=True) ``` ### 3. Diagnostics ```python az.summary(trace, hdi_prob=0.89) az.plot_trace(trace) az.plot_rank_hist(trace) # Ranked histograms (preferred) ``` ### 4. Posterior Predictive Check ```python with model: post_pred = pm.sample_posterior_predictive(trace) az.plot_ppc(post_pred, num_pp_samples=100) ``` ### 5. Model Comparison ```python loo1 = az.loo(trace1) loo2 = az.loo(trace2) az.compare({"m1": trace1, "m2": trace2}) az.plot_khat(loo1) # k > 0.7 is problematic ``` ## pm.Deterministic for Tracking **Always track mu for plotting:** ```python # Inside model mu = pm.Deterministic("mu", alpha + pm.math.dot(X, beta)) # Access later trace.posterior["mu"] # All samples of mu ``` ## Data Extraction Patterns ```python # Extract to DataFrame trace_df = az.extract_dataset(trace).to_dataframe() # Access specific parameters post = az.extract_dataset(trace["posterior"]) mu_samples = post["mu"].values # Get numpy arrays alpha_values = trace.posterior["alpha"].values # (chains, draws) ``` ## HDI Visualization ```python # Compute mu at new x values x_seq = np.linspace(x.min(), x.max(), 100) mu_pred = post["alpha"] + post["beta"] * x_seq[:, None] # Plot HDI bands az.plot_hdi(x_seq, mu_pred.T, hdi_prob=0.89) plt.scatter(x, y) ``` ## ArviZ Diagnostics ```python import arviz as az # Configure defaults az.rcParams["stats.hdi_prob"] = 0.89 # Summary table summary = az.summary(trace, hdi_prob=0.89) # Key metrics max_rhat = summary["r_hat"].max() # Should be < 1.01 min_ess = summary["ess_bulk"].min() # Should be > 400 # Plots az.plot_trace(trace) # Trace plots az.plot_rank_hist(trace) # Ranked histograms (preferred!) az.plot_posterior(trace) # Posteriors az.plot_forest(trace) # Forest plot az.plot_pair(trace) # Pairs plot # Model comparison az.loo(trace) # LOO-CV az.waic(trace) # WAIC az.compare({"m1": trace1, "m2": trace2}) ``` ## Diagnostic Checklist - [ ] Rhat < 1.01 for all parameters - [ ] ESS_bulk > 400 - [ ] ESS_tail > 400 - [ ] Prior predictive produces sensible values - [ ] Posterior predictive matches data pattern - [ ] Pareto k < 0.7 for LOO ## Non-Centered Parameterization For hierarchical models: ```python # Centered (may have divergences) theta = pm.Normal("theta", mu=mu, sigma=tau, shape=J) # Non-centered (recommended) theta_raw = pm.Normal("theta_raw", mu=0, sigma=1, shape=J) theta = pm.Deterministic("theta", mu + tau * theta_raw) ``` ## PyTensor Math Operations Inside `with pm.Model()`, use `pm.math` not `np`: ```python # Correct mu = pm.math.dot(X, beta) p = pm.math.sigmoid(eta) log_x = pm.math.log(x) # Wrong (will fail) mu = np.dot(X, beta) # Don't use numpy inside model ``` ## Common Priors ```python # Intercept alpha = pm.Normal("alpha", mu=0, sigma=10) # Coefficients beta = pm.Normal("beta", mu=0, sigma=2.5, shape=K) # Scale (SD) sigma = pm.HalfNormal("sigma", sigma=1) sigma = pm.HalfCauchy("sigma", beta=2.5) sigma = pm.Exponential("sigma", lam=1) # Hierarchical SD tau = pm.HalfCauchy("tau", beta=2.5) # Correlation matrix chol, corr, stds = pm.LKJCholeskyCov("chol", n=K, eta=2, sd_dist=pm.Exponential.dist(1)) ``` ## Key Differences from Stan | Feature | PyMC | Stan | |---------|------|------| | Syntax | Python | DSL | | Arrays | `shape=K` | `array[K]` | | Math | `pm.math.dot()` | `*` operator | | Blocks | Single context | 7 blocks | | Output | InferenceData | CmdStanMCMC | ## Troubleshooting | Issue | Solution | |-------|----------| | Divergences | Increase `target_accept` to 0.9-0.99 | | Low ESS | Run longer chains, reparameterize | | Shape errors | Check `shape=` parameter | | Slow | Use ADVI for quick approximation | | Memory | Reduce chains or use mini-batch |