Pymc Bayesian Modeling

pip install pymc arviz numpy matplotlib
# Optional: JAX backend for GPU acceleration
pip install pymc[jax]

import pymc as pm
import arviz as az
import numpy as np

# Simulate data
np.random.seed(42)
X = np.random.randn(100)
y = 2.5 + 1.3 * X + np.random.randn(100) * 0.5

# Build and fit model
with pm.Model() as model:
    alpha = pm.Normal("alpha", mu=0, sigma=5)
    beta = pm.Normal("beta", mu=0, sigma=5)
    sigma = pm.HalfNormal("sigma", sigma=1)
    mu = alpha + beta * X
    y_obs = pm.Normal("y_obs", mu=mu, sigma=sigma, observed=y)
    idata = pm.sample(1000, tune=1000, chains=4, random_seed=42)

print(az.summary(idata, var_names=["alpha", "beta", "sigma"]))
# Expected: alpha ~ 2.5, beta ~ 1.3, sigma ~ 0.5

import pymc as pm
import arviz as az
import numpy as np

# Load data
X = np.random.randn(200, 3)  # 200 obs, 3 predictors
y = X @ np.array([1.0, -0.5, 0.3]) + np.random.randn(200) * 0.8

# Standardize predictors
X_mean, X_std = X.mean(axis=0), X.std(axis=0)
X_scaled = (X - X_mean) / X_std

# Define coordinates for named dimensions
coords = {
    "predictors": ["var1", "var2", "var3"],
    "obs_id": np.arange(len(y)),
}
print(f"Data shape: X={X_scaled.shape}, y={y.shape}")

with pm.Model(coords=coords) as model:
    # Priors — weakly informative, not flat
    alpha = pm.Normal("alpha", mu=0, sigma=1)
    beta = pm.Normal("beta", mu=0, sigma=1, dims="predictors")
    sigma = pm.HalfNormal("sigma", sigma=1)

    # Linear predictor
    mu = alpha + pm.math.dot(X_scaled, beta)

    # Likelihood
    y_obs = pm.Normal("y_obs", mu=mu, sigma=sigma, observed=y, dims="obs_id")

# Inspect model variables
print(model.basic_RVs)  # Lists: [alpha, beta, sigma, y_obs]

with model:
    prior_pred = pm.sample_prior_predictive(samples=1000, random_seed=42)

# Check prior-implied data range
prior_y = prior_pred.prior_predictive["y_obs"].values.flatten()
print(f"Prior predictive range: [{prior_y.min():.1f}, {prior_y.max():.1f}]")
print(f"Observed data range:    [{y.min():.1f}, {y.max():.1f}]")

az.plot_ppc(prior_pred, group="prior", num_pp_samples=100)

with model:
    idata = pm.sample(
        draws=2000,
        tune=1000,
        chains=4,
        target_accept=0.9,
        random_seed=42,
        idata_kwargs={"log_likelihood": True},
    )

print(f"Posterior shape: {idata.posterior['beta'].shape}")
# Expected: (4 chains, 2000 draws, 3 predictors)

# Summary with convergence diagnostics
summary = az.summary(idata, var_names=["alpha", "beta", "sigma"])
print(summary[["mean", "sd", "hdi_3%", "hdi_97%", "r_hat", "ess_bulk"]])

# R-hat convergence check
bad_rhat = summary[summary["r_hat"] > 1.01]
if len(bad_rhat) > 0:
    print(f"WARNING: {len(bad_rhat)} parameters with R-hat > 1.01")
    print(bad_rhat[["r_hat"]])

# Effective sample size check
low_ess = summary[summary["ess_bulk"] < 400]
if len(low_ess) > 0:
    print(f"WARNING: {len(low_ess)} parameters with ESS < 400")

# Divergence check
n_div = idata.sample_stats.diverging.sum().item()
total = len(idata.posterior.draw) * len(idata.posterior.chain)
print(f"Divergences: {n_div}/{total} ({n_div / total * 100:.2f}%)")

# Visual diagnostics — trace plots and rank plots
az.plot_trace(idata, var_names=["alpha", "beta", "sigma"])
az.plot_rank(idata, var_names=["alpha", "beta", "sigma"])

with model:
    pm.sample_posterior_predictive(idata, extend_inferencedata=True, random_seed=42)

az.plot_ppc(idata, num_pp_samples=100)
# Blue = observed data, grey = posterior simulations
# Systematic deviations indicate model misspecification

# Fit multiple models with log_likelihood=True, then compare
# Example: compare linear vs a second model
idatas = {"linear": idata}  # add more fitted models here

comparison = az.compare(idatas, ic="loo")
print(comparison[["rank", "elpd_loo", "p_loo", "d_loo", "weight"]])

# Check LOO reliability via Pareto-k diagnostics
loo_result = az.loo(idata, pointwise=True)
high_k = (loo_result.pareto_k > 0.7).sum().item()
print(f"Observations with Pareto-k > 0.7: {high_k}")
# Interpretation: Dloo < 2 = similar models; Dloo > 10 = strong evidence

az.plot_compare(comparison)

X_new = np.array([[0.5, -1.0, 0.2]])
X_new_scaled = (X_new - X_mean) / X_std

with model:
    pm.set_data({"X_scaled": X_new_scaled})
    post_pred = pm.sample_posterior_predictive(
        idata.posterior, var_names=["y_obs"], random_seed=42
    )

y_pred = post_pred.posterior_predictive["y_obs"]
print(f"Predicted mean: {y_pred.mean().item():.3f}")
print(f"94% HDI: {az.hdi(y_pred, hdi_prob=0.94).values}")

import pymc as pm
import arviz as az
import numpy as np

n_groups, n_per_group = 5, 30
group_idx = np.repeat(np.arange(n_groups), n_per_group)
group_names = [f"group_{i}" for i in range(n_groups)]

# Simulated grouped data
true_alphas = np.random.normal(3.0, 1.5, n_groups)
y_obs = np.random.normal(true_alphas[group_idx], 0.5)

with pm.Model(coords={"groups": group_names}) as hierarchical_model:
    # Hyperpriors (population level)
    mu_alpha = pm.Normal("mu_alpha", mu=0, sigma=5)
    sigma_alpha = pm.HalfNormal("sigma_alpha", sigma=2)

    # Non-centered parameterization
    alpha_offset = pm.Normal("alpha_offset", mu=0, sigma=1, dims="groups")
    alpha = pm.Deterministic(
        "alpha", mu_alpha + sigma_alpha * alpha_offset, dims="groups"
    )

    # Observation level
    sigma = pm.HalfNormal("sigma", sigma=1)
    y = pm.Normal("y", mu=alpha[group_idx], sigma=sigma, observed=y_obs)

    idata_hier = pm.sample(2000, tune=1000, target_accept=0.95, random_seed=42)

print(az.summary(idata_hier, var_names=["mu_alpha", "sigma_alpha", "alpha", "sigma"]))

import pymc as pm
import arviz as az
import numpy as np

# Simulated binary outcome
np.random.seed(42)
n, p = 200, 3
X_lr = np.random.randn(n, p)
true_beta = np.array([1.0, -0.5, 0.3])
prob = 1 / (1 + np.exp(-(0.2 + X_lr @ true_beta)))
y_binary = np.random.binomial(1, prob)

with pm.Model() as logistic_model:
    alpha = pm.Normal("alpha", mu=0, sigma=2)
    beta = pm.Normal("beta", mu=0, sigma=2, shape=p)
    logit_p = alpha + pm.math.dot(X_lr, beta)
    y = pm.Bernoulli("y", logit_p=logit_p, observed=y_binary)

    idata_logit = pm.sample(2000, tune=1000, target_accept=0.9, random_seed=42)

summary = az.summary(idata_logit, var_names=["alpha", "beta"])
print(summary[["mean", "sd", "hdi_3%", "hdi_97%"]])
# Odds ratios
beta_samples = idata_logit.posterior["beta"].values.reshape(-1, p)
print(f"Odds ratios (median): {np.median(np.exp(beta_samples), axis=0)}")

import pymc as pm
import arviz as az
import numpy as np

# Simulated nonlinear data
np.random.seed(42)
X_gp = np.sort(np.random.uniform(0, 10, 60))[:, None]
y_gp = np.sin(X_gp[:, 0]) + np.random.randn(60) * 0.3

with pm.Model() as gp_model:
    # GP hyperparameters
    ls = pm.Gamma("lengthscale", alpha=2, beta=1)
    eta = pm.HalfNormal("amplitude", sigma=2)
    sigma_noise = pm.HalfNormal("noise", sigma=0.5)

    # Covariance function
    cov = eta**2 * pm.gp.cov.Matern52(1, ls=ls)
    gp = pm.gp.Marginal(cov_func=cov)

    # Marginal likelihood
    y_ = gp.marginal_likelihood("y", X=X_gp, y=y_gp, sigma=sigma_noise)
    idata_gp = pm.sample(1000, tune=1000, target_accept=0.9, random_seed=42)

# Predict on new points
X_new_gp = np.linspace(0, 10, 100)[:, None]
with gp_model:
    f_pred = gp.conditional("f_pred", X_new_gp)
    pred_samples = pm.sample_posterior_predictive(idata_gp, var_names=["f_pred"])

print(f"GP predictions shape: {pred_samples.posterior_predictive['f_pred'].shape}")