pymc-modeling

PyMC 建模

使用 PyMC v5+ 和现代 API 模式的贝叶斯建模工作流。

笔记本偏好：除非项目已使用 Jupyter，否则请使用 marimo 进行交互式建模。

模型规范

基本结构

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/dims 以便更好地解读 InferenceData：

coords = {
    "obs": np.arange(n_obs),
    "features": ["intercept", "age", "income"],
    "group": group_labels,
}

对于简单模型，如果开销大于收益，则跳过。

参数化

对于数据较少的层级模型，优先使用非中心化参数化：

# 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")

推断

默认采样（nutpie）

默认使用 nutpie 作为采样器——它基于 Rust，通常快 2-5 倍：

with model:
    idata = pm.sample(
        draws=1000, tune=1000, chains=4,
        nuts_sampler="nutpie",
        random_seed=42,
    )

PyMC 原生采样

nutpie 不可用时，回退到 PyMC 的 NUTS：

with model:
    idata = pm.sample(draws=1000, tune=1000, chains=4, random_seed=42)

其他 MCMC 后端

参见 references/inference.md：

• NumPyro/JAX：GPU 加速、向量化链

近似推断

用于快速（但不精确）的后验近似：

• ADVI/DADVI：使用高斯近似的变分推断
• Pathfinder：用于初始化或筛选的拟牛顿优化

诊断和 ArviZ 工作流

每次采样后都应遵循以下系统化工作流：

阶段 1：即时检查（必需）

# 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"])

通过标准（继续之前必须全部满足）：

• 无发散（或发散率 < 0.1% 且随机分布）
• 所有参数的 r_hat < 1.01
• ess_bulk > 400 且 ess_tail > 400
• 轨迹图显示混合良好（密度重叠，呈模糊的毛毛虫状）

阶段 2：深入检查收敛性（阶段 1 结果不明确时）

# 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"])

阶段 3：模型评估（必需）

# 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")

重要规则：阶段 1-3 全部通过前，切勿解读参数。

阶段 4：参数解读

# 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"])

参见 references/arviz.md 了解 ArviZ 的完整用法。 参见 references/diagnostics.md 了解故障排查方法。

先验和后验预测检查

先验预测（拟合前）

拟合前务必检查先验的影响：

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}]")

警示信号：先验预测覆盖了不合理的值（负计数、概率 > 1），或范围极宽/极窄。

后验预测（拟合后）

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")

解读：观测数据（深色线）应落在后验预测分布（浅色线）范围内。详细解读请参见 references/arviz.md。

模型调试

检查模型结构

# Print model summary (variables, shapes, distributions)
print(model)

# Visualize model as directed graph
pm.model_to_graphviz(model)

检查规范错误

采样前验证模型：

# Debug model: checks for common issues
model.debug()

# Check initial point log-probabilities
# Identifies which variables have invalid starting values
model.point_logps()

常见问题

症状	可能原因	修复方法
ValueError: Shape mismatch	参数与观测值的维度不匹配	使用索引向量：alpha[group_idx]
Initial evaluation failed	数据超出分布支持范围	检查边界；使用 init="adapt_diag"
Mass matrix contains zeros	预测变量未缩放或先验平坦	标准化特征；使用弱信息先验
发散次数较多	漏斗几何结构	使用非中心化参数化
对数概率中出现 NaN	参数组合无效	检查参数约束，添加边界
对数概率为 -inf	观测值超出似然支持范围	确认数据符合分布定义域
离散采样缓慢	NUTS 与离散变量不兼容	边缘化离散变量
样本外预测效果不佳	数据数组静态	使用 mutable=True 的 pm.Data 容器

参见 references/troubleshooting.md 获取完整的问题与解决方案指南。

调试发散

# 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

分析缓慢模型的性能

# 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()

其他故障排查方法请参见 references/gotchas.md。

模型比较

LOO-CV（优先）

# 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)

比较模型

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)

判定规则：如果 d_loo < 2*dse，则模型实际上等效。

有关详细的模型比较工作流，请参阅 references/arviz.md。

保存和加载结果

InferenceData 持久化

保存采样结果，以供日后分析或分享：

# Save to NetCDF (recommended format)
idata.to_netcdf("results/model_v1.nc")

# Load
idata = az.from_netcdf("results/model_v1.nc")

压缩存储

对于大型 InferenceData 对象（抽样次数多、后验预测数据量大）：

# 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}
)

保存的内容

InferenceData 会保留完整的贝叶斯工作流：

• posterior：MCMC 的参数样本
• prior、prior_predictive：先验样本（如果已生成）
• posterior_predictive：预测结果（如果已生成）
• observed_data、constant_data：拟合时使用的数据
• sample_stats：诊断信息（散度、树深度、能量）
• log_likelihood：逐点对数似然（用于 LOO-CV）
• 所有坐标和维度

工作流模式

# 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

先验选择

请参阅 references/priors.md，了解：

• 按分布类型提供的弱信息默认值
• 先验预测检验工作流
• 特定领域的建议

常见模式

分层／多层模型

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")

GLM

# 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)

高斯过程

对于大多数 GP 问题（n > 500，输入为 1–3 维），默认使用 HSGP。它的复杂度为 O(nm)，而非 O(n³)：

with pm.Model() as gp_model:
    # Hyperparameters
    ell = pm.InverseGamma("ell", alpha=5, beta=5)
    eta = pm.HalfNormal("eta", sigma=2)
    sigma = pm.HalfNormal("sigma", sigma=0.5)

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

    # HSGP approximation
    gp = pm.gp.HSGP(m=[20], c=1.5, cov_func=cov)
    f = gp.prior("f", X=X[:, None])  # X must be 2D

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

对于周期性模式，请使用 pm.gp.HSGPPeriodic。对于小型数据集（n < 500），请使用 pm.gp.Marginal 或 pm.gp.Latent。

请参阅 references/gp.md，了解：

• HSGP 参数选择（选择 m 和 c、自动启发式方法）
• 用于季节性／周期性模式的 HSGPPeriodic
• 近似质量诊断
• 协方差函数和先验
• 常见模式（趋势 + 季节性、分类、异方差）

时间序列

with pm.Model(coords={"time": range(T)}) as ar_model:
    rho = pm.Uniform("rho", -1, 1)
    sigma = pm.HalfNormal("sigma", sigma=1)

    y = pm.AR("y", rho=[rho], sigma=sigma, constant=True,
              observed=y_obs, dims="time")

请参阅 references/timeseries.md，了解：

• 自回归模型（AR、ARMA）
• 随机游走和局部水平模型
• 结构时间序列（趋势 + 季节性）
• 状态空间模型
• 用于时间序列的 GP
• 处理多个季节性
• 预测模式

BART（贝叶斯加性回归树）

import pymc_bart as pmb

with pm.Model() as bart_model:
    mu = pmb.BART("mu", X=X, Y=y, m=50)
    sigma = pm.HalfNormal("sigma", 1)
    y_obs = pm.Normal("y_obs", mu=mu, sigma=sigma, observed=y)

请参阅 references/bart.md，了解：

• 回归和分类
• 变量重要性和偏依赖
• 将 BART 与参数化组件结合
• 配置（树的数量、深度先验）

混合模型

import numpy as np

coords = {"component": range(K)}

with pm.Model(coords=coords) as gmm:
    # Mixture weights
    w = pm.Dirichlet("w", a=np.ones(K), dims="component")

    # Component parameters (with ordering to avoid label switching)
    mu = pm.Normal("mu", mu=0, sigma=10, dims="component",
                   transform=pm.distributions.transforms.ordered)
    sigma = pm.HalfNormal("sigma", sigma=2, dims="component")

    # Mixture likelihood
    y = pm.NormalMixture("y", w=w, mu=mu, sigma=sigma, observed=y_obs)

请参阅 references/mixtures.md，了解：

• 有限混合模型和混合回归
• 标签切换问题及解决方案（顺序约束、重新标记）
• 边缘化混合模型（pymc-extras）
• 混合模型诊断

专用似然

# Zero-Inflated Poisson (excess zeros)
with pm.Model() as zip_model:
    psi = pm.Beta("psi", alpha=2, beta=2)  # P(structural zero)
    mu = pm.Exponential("mu", lam=1)
    y = pm.ZeroInflatedPoisson("y", psi=psi, mu=mu, observed=y_obs)

# Censored data (e.g., right-censored survival)
with pm.Model() as censored_model:
    mu = pm.Normal("mu", mu=0, sigma=10)
    sigma = pm.HalfNormal("sigma", sigma=5)
    y = pm.Censored("y", dist=pm.Normal.dist(mu=mu, sigma=sigma),
                    lower=None, upper=censoring_time, observed=y_obs)

# Ordinal regression
with pm.Model() as ordinal:
    beta = pm.Normal("beta", mu=0, sigma=2, dims="features")
    cutpoints = pm.Normal("cutpoints", mu=0, sigma=2,
                          transform=pm.distributions.transforms.ordered,
                          shape=n_categories - 1)
    y = pm.OrderedLogistic("y", eta=pm.math.dot(X, beta),
                           cutpoints=cutpoints, observed=y_obs)

请参阅 references/specialized_likelihoods.md，了解：

• 零膨胀模型（Poisson、Negative Binomial、Binomial）
• 计数数据的门槛模型
• 删失和截断数据
• 有序回归
• 使用 Student-t 似然的稳健回归

常见陷阱

请参阅 references/gotchas.md，了解：

• 居中与非居中参数化
• 尺度参数的先验
• 混合模型中的标签切换
• 性能问题（GP、大型 Deterministics）
• Python 条件语句和硬裁剪
• 分层模型中的冗余截距

请参阅 references/troubleshooting.md，获取全面的问题解决指南，涵盖：

• 形状和维度错误
• 初始化失败
• 质量矩阵和数值问题
• 离散变量难题
• 数据容器和预测问题

因果推断操作

pm.do（干预）

应用 do-calculus 干预，将变量设为固定值：

with pm.Model() as causal_model:
    x = pm.Normal("x", 0, 1)
    y = pm.Normal("y", x, 1)
    z = pm.Normal("z", y, 1)

# Intervene: set x = 2 (breaks incoming edges to x)
with pm.do(causal_model, {"x": 2}) as intervention_model:
    idata = pm.sample_prior_predictive()
    # Samples from P(y, z | do(x=2))

pm.observe（条件化）

在不进行干预的情况下，根据观测值进行条件化：

# Condition: observe y = 1 (doesn't break causal structure)
with pm.observe(causal_model, {"y": 1}) as conditioned_model:
    idata = pm.sample()
    # Samples from P(x, z | y=1)

结合 do 和 observe

# Intervention + observation for causal queries
with pm.do(causal_model, {"x": 2}) as m1:
    with pm.observe(m1, {"z": 0}) as m2:
        idata = pm.sample()
        # P(y | do(x=2), z=0)

pymc-extras

用于专用模型和推断：

import pymc_extras as pmx

# Marginalize discrete parameters from a model
model = pmx.marginalize(model, ["discrete_var"])

# R2D2 prior for regression (requires output_sigma and input_sigma)
residual_sigma, beta = pmx.R2D2M2CP(
    "r2d2",
    output_sigma=y.std(),
    input_sigma=X.std(axis=0),
    dims="features",
    r2=0.5,
)

# Laplace approximation for fast inference
idata = pmx.fit_laplace(model)

自定义分布和模型组件

要扩展 PyMC，超越内置分布：

import pymc as pm
import pytensor.tensor as pt

# Custom likelihood via DensityDist
def custom_logp(value, mu, sigma):
    return pm.logp(pm.Normal.dist(mu=mu, sigma=sigma), value)

with pm.Model() as model:
    mu = pm.Normal("mu", 0, 1)
    y = pm.DensityDist("y", mu, 1.0, logp=custom_logp, observed=y_obs)

# Soft constraints via Potential
with pm.Model() as model:
    alpha = pm.Normal("alpha", 0, 1, dims="group")
    pm.Potential("sum_to_zero", -100 * pt.sqr(alpha.sum()))

请参阅 references/custom_models.md，了解：

• pm.DensityDist，用于自定义似然
• pm.Potential，用于软约束和雅可比调整
• pm.Simulator，用于基于模拟的推断（ABC）
• pm.CustomDist，用于自定义先验分布