pymc

# PyMC & NumPyro — Probabilistic Programming PyMC is the leading Python library for Bayesian statistical modeling. You declare priors and a likelihood — the sampler computes the posterior automatically. NumPyro provides the identical paradigm on JAX for GPU-accelerated inference. Both integrate with ArviZ for diagnostics. **Core value:** Instead of "is this significant?" (p-value), Bayesian methods answer "what is the full probability distribution over possible values?" — a fundamentally richer answer for every scientific and business question. ## When to Use - Quantifying uncertainty: credible intervals, prediction intervals, error propagation through calculations. - Hierarchical models: data naturally nested (students in schools, patients in hospitals, stores in regions). - A/B testing where "P(B > A) = 0.95" is more actionable than "p = 0.03". - Scientific measurement where uncertainty must propagate through a pipeline. - Model comparison: which explanation fits better? (WAIC, LOO-CV) - Small samples: Bayesian priors regularize and prevent overfitting where MLE fails. - Any problem where you can articulate the generative story (how data was produced). **When NOT to use:** Simple hypothesis tests on huge datasets where a p-value suffices. Pure prediction without uncertainty (use sklearn). Real-time inference (MCMC is too slow). ## Reference Documentation **PyMC docs**: https://pymc.io/en/stable/ **NumPyro docs**: https://num.pyro.ai/en/stable/ **ArviZ docs**: https://arviz.org/en/stable/ **GitHub**: https://github.com/pymc-devs/pymc **Search patterns**: `pm.Model`, `pm.sample`, `pm.Normal`, `az.summary`, `az.plot_trace` ## Core Principles ### Bayes' Theorem **posterior ∝ prior × likelihood** Prior = belief before data. Likelihood = how probable is this data given the parameter? Posterior = updated belief after seeing data. PyMC specifies prior + likelihood; MCMC computes the posterior by sampling. ### MCMC — Markov Chain Monte Carlo The posterior is rarely analytically tractable. MCMC draws samples that, in the limit, are distributed exactly as the posterior. PyMC uses NUTS (No-U-Turn Sampler) — state of the art. Output: thousands of parameter values drawn from the posterior. ### Posterior Predictive Check (PPC) Generate synthetic data from the fitted model. Compare to real data. If they match visually, the model is reasonable. This catches systematic failures that summary statistics miss — the single most important validation step. ### Credible Intervals A 95% credible interval [a, b] means: "Given the data, there is a 95% probability the true parameter lies in [a, b]." This is the intuitive interpretation people actually want. (Frequentist confidence intervals do NOT mean this.) ## Quick Reference ### Installation ```bash pip install pymc arviz # PyMC + diagnostics pip install numpyro jax jaxlib # NumPyro (JAX backend, GPU-accelerated) ``` ### Standard Imports ```python import pymc as pm import arviz as az import numpy as np import pandas as pd ``` ### Basic Pattern — Coin Flip (Simplest Bayesian Model) ```python import pymc as pm import arviz as az # Observed: 15 heads out of 50 flips n_flips, n_heads = 50, 15 with pm.Model() as model: # Prior: belief about P(heads) before seeing data p = pm.Beta('p', alpha=2, beta=2) # Weakly informative, centered at 0.5 # Likelihood: how many heads given p heads = pm.Binomial('heads', n=n_flips, p=p, observed=n_heads) # Sample the posterior trace = pm.sample(4000, cores=2, return_inferencedata=True, random_seed=42) # Interpret posterior az.summary(trace, var_names=['p']) # → mean ≈ 0.31, HDI[94%] = [0.19, 0.43] # Translation: "P(heads) is 31%; 94% credible that it's between 19% and 43%" az.plot_posterior(trace, var_names=['p']) ``` ### Basic Pattern — NumPyro Equivalent ```python import numpyro import numpyro.distributions as dist from numpyro.infer import MCMC, NUTS from jax import random import jax.numpy as jnp n_flips, n_heads = 50, 15 def model(): p = numpyro.sample('p', dist.Beta(2, 2)) numpyro.sample('heads', dist.Binomial(n_flips, p), obs=n_heads) kernel = NUTS(model) mcmc = MCMC(kernel, num_warmup=1000, num_samples=4000, num_chains=2) mcmc.run(random.PRNGKey(0)) samples = mcmc.get_samples() print(f"P(heads): {samples['p'].mean():.3f} " f"HDI=[{jnp.percentile(samples['p'], 3):.3f}, {jnp.percentile(samples['p'], 97):.3f}]") ``` ## Critical Rules ### ✅ DO - **Check convergence ALWAYS** — `az.summary()` shows Rhat. Rhat > 1.01 means not converged. Never trust results without this. - **Use ≥ 2 chains, prefer 4** — Cannot assess convergence with 1 chain. - **Do posterior predictive checks** — `pm.sample_posterior_predictive()`. Non-negotiable validation. - **Use informative priors where possible** — Even weakly informative (`HalfNormal(1)` for σ) beats flat priors. Flat priors on unbounded scales cause convergence failures. - **Use `HalfNormal` or `Gamma` for scale parameters** — σ must be positive. `pm.HalfNormal('sigma', sigma=1)` is the default choice. - **Wrap derived quantities in `pm.Deterministic`** — Only way to track computed posteriors (odds ratios, differences, lifts). - **Set `random_seed`** — Reproducibility. - **Use `return_inferencedata=True`** — Returns ArviZ InferenceData. Legacy dict format is deprecated. ### ❌ DON'T - **Don't ignore divergences** — Divergence = sampler lost. Requires reparameterization or prior adjustment. - **Don't use `Uniform` on scale parameters** — `pm.Uniform('sigma', 0, 100)` causes sampler to struggle near 0. Use `HalfNormal`. - **Don't confuse `observed` and unobserved** — `observed=data` binds to data. Without it, the variable is sampled from the prior. - **Don't run only 1 chain** — Cannot diagnose convergence. - **Don't use PyMC for >100k rows without considering NumPyro** — CPU-bound PyMC becomes very slow. NumPyro + JAX + GPU is 10–100x faster. - **Don't mix up loss functions / label formats** — In Bayesian context: if labels are integers use `Categorical`; if binary use `Bernoulli`; if proportions use `Beta`. ## Anti-Patterns (NEVER) ```python import pymc as pm import numpy as np # ❌ BAD: Flat prior on scale → sampler struggles near zero with pm.Model(): sigma = pm.Uniform('sigma', lower=0, upper=100) # Divergences guaranteed # ✅ GOOD: HalfNormal is the standard for positive scales with pm.Model(): sigma = pm.HalfNormal('sigma', sigma=1.0) # ❌ BAD: Ignoring convergence — trusting undiagnosed results with pm.Model(): mu = pm.Normal('mu', mu=0, sigma=10) obs = pm.Normal('obs', mu=mu, sigma=1, observed=data) trace = pm.sample(100, cores=1) # Too few draws, 1 chain print(trace.posterior['mu'].mean().values) # Meaningless if not converged # ✅ GOOD: Proper setup + convergence gate with pm.Model(): mu = pm.Normal('mu', mu=0, sigma=10) obs = pm.Normal('obs', mu=mu, sigma=1, observed=data) trace = pm.sample(4000, cores=4, return_inferencedata=True, random_seed=42) import arviz as az summary = az.summary(trace) assert summary['r_hat'].max() < 1.01, f"Not converged: Rhat={summary['r_hat'].max():.3f}" # ❌ BAD: Derived quantity not tracked — lost after sampling with pm.Model(): p_a = pm.Beta('p_a', 1, 1) p_b = pm.Beta('p_b', 1, 1) # Difference p_b - p_a is never recorded in the trace # ✅ GOOD: Deterministic captures derived posteriors with pm.Model(): p_a = pm.Beta('p_a', 1, 1) p_b = pm.Beta('p_b', 1, 1) diff = pm.Deterministic('diff', p_b - p_a) # Now in trace.posterior['diff'] # ❌ BAD: No posterior predictive check — model could be completely wrong with pm.Model() as model: mu = pm.Normal('mu', mu=0, sigma=10) sigma = pm.HalfNormal('sigma', sigma=1) obs = pm.Normal('obs', mu=mu, sigma=sigma, observed=data) trace = pm.sample(4000, return_inferencedata=True) # "Converged" ≠ "correct". Convergence just means the sampler found A distribution. # ✅ GOOD: PPC validates whether the model generates plausible data with pm.Model() as model: mu = pm.Normal('mu', mu=0, sigma=10) sigma = pm.HalfNormal('sigma', sigma=1) obs = pm.Normal('obs', mu=mu, sigma=sigma, observed=data) trace = pm.sample(4000, return_inferencedata=True) ppc = pm.sample_posterior_predictive(trace, model=model) az.plot_ppc(ppc, var_names=['obs']) # Simulated vs real — must look similar ``` ## Prior Selection Guide ``` PARAMETER | RECOMMENDED PRIOR | NOTES ---------------------------|--------------------------------|-------------------------------------- Mean / intercept | Normal(domain_mean, wide_sd) | Weakly informative; scale to data Positive scale (σ, τ) | HalfNormal(1) | Default; adjust sigma to data scale Rate / proportion (0–1) | Beta(1, 1) | Uniform on [0,1]; Beta(2,5) if informed Count rate (λ) | Gamma(2, 1) or Exponential(1) | Positive, right-skewed Regression slope | Normal(0, 2–5) | Weakly regularizing toward zero Odds ratio | Lognormal(0, 1) | Positive, multiplicative scale Correlation | LKJCorrelation(1) | PyMC-specific; valid [-1, 1] RULES OF THUMB: → Know the rough scale? Use Normal(known_mean, 2–10 × known_std). → Know nothing? Wide but bounded priors, NOT Uniform(-∞, ∞). → Scale parameters (σ, τ) ALWAYS positive → HalfNormal, Gamma, or Exponential. → More data = prior matters less. With n > 100, even bad priors get washed out. ``` ## Model Specification — PyMC Patterns ### Basic Structure ```python import pymc as pm with pm.Model() as model: # 1. PRIORS — parameters to learn mu = pm.Normal('mu', mu=0, sigma=10) sigma = pm.HalfNormal('sigma', sigma=1) # 2. LIKELIHOOD — data generation mechanism obs = pm.Normal('obs', mu=mu, sigma=sigma, observed=data) # 3. DERIVED QUANTITIES (optional) # pm.Deterministic('name', expression) # 4. SAMPLE trace = pm.sample(4000, cores=4, return_inferencedata=True, random_seed=42) ``` ### Linear Regression ```python import pymc as pm import numpy as np # X shape: (n,) single predictor or (n, p) multiple # y shape: (n,) with pm.Model(): intercept = pm.Normal('intercept', mu=0, sigma=5) slope = pm.Normal('slope', mu=0, sigma=5) # shape=(p,) for multivariate sigma = pm.HalfNormal('sigma', sigma=2) mu = pm.Deterministic('mu', intercept + slope * X) # Track mean in trace y_obs = pm.Normal('y_obs', mu=mu, sigma=sigma, observed=y) trace = pm.sample(4000, cores=4, return_inferencedata=True, random_seed=42) ``` ### Multivariate Regression ```python import pymc as pm import numpy as np X = np.array([...]) # Shape: (n, p) y = np.array([...]) n, p = X.shape with pm.Model(): intercept = pm.Normal('intercept', mu=0, sigma=5) slopes = pm.Normal('slopes', mu=0, sigma=5, shape=p) # One per feature sigma = pm.HalfNormal('sigma', sigma=2) mu = pm.Deterministic('mu', intercept + X @ slopes) # Matrix multiply y_obs = pm.Normal('y_obs', mu=mu, sigma=sigma, observed=y) trace = pm.sample(4000, cores=4, return_inferencedata=True) ``` ## Sampling and Diagnostics ### Sampling Parameters ```python import pymc as pm with pm.Model() as model: # ... model ... trace = pm.sample( draws=4000, # Posterior samples per chain (after warmup) tune=2000, # Warmup draws — discarded; sampler adapts here cores=4, # Parallel chains = CPU cores chains=4, # Independent chains for convergence check return_inferencedata=True, # ArviZ InferenceData format random_seed=42, # Reproducibility ) ``` ### Diagnostics Checklist ```python import pymc as pm import arviz as az # 1. SUMMARY — Rhat + ESS at a glance # Rhat ≈ 1.00 → converged | ESS > 100 → enough effective samples summary = az.summary(trace) print(summary[['mean', 'sd', 'hdi_3%', 'hdi_97%', 'r_hat', 'ess_bulk']]) # 2. TRACE PLOT — visual convergence # Chains should overlap, look like random noise. No trends or drift. az.plot_trace(trace, var_names=['mu', 'sigma']) # 3. POSTERIOR az.plot_posterior(trace, var_names=['mu', 'sigma']) # 4. ENERGY — divergence check az.plot_energy(trace) # Two distributions should overlap heavily # 5. PPC — model validation ppc = pm.sample_posterior_predictive(trace, model=model) az.plot_ppc(ppc, var_names=['y_obs']) # ─── Convergence decision tree ─── # Rhat > 1.05 → STOP. Increase tune, add chains, or reparameterize. # Rhat 1.01–1.05 → WARNING. Increase draws and tune. # Rhat < 1.01 → OK. Verify ESS > 100. # Many divergences → Reparameterize or tighten priors.