313 lines
12 KiB
Python
313 lines
12 KiB
Python
import numpy as np
|
||
import matplotlib.pyplot as plt
|
||
import scipy.stats as stats
|
||
from scipy.special import gammaln # 用于计算 log(Γ(x))
|
||
|
||
# --- 1. 定义先验和似然函数 ---
|
||
|
||
# 定义先验参数
|
||
# π(λ) ~ Gamma(α, β) (注意：scipy.stats.gamma用 a=shape, scale=1/rate)
|
||
# 我们使用 α=2, β=1 (rate=1) 作为先验
|
||
ALPHA_LAM = 2
|
||
BETA_LAM = 1 # 这是 rate (或 1/scale)
|
||
|
||
# π(r) ~ InverseGamma(α, β) (注意：scipy.stats.invgamma用 a=shape, scale=scale)
|
||
# 我们使用 α=2, β=1 (scale=1) 作为先验
|
||
ALPHA_R = 2
|
||
BETA_R = 1 # 这是 scale
|
||
|
||
# 模型先验
|
||
LOG_PRIOR_K1 = np.log(0.5)
|
||
LOG_PRIOR_K2 = np.log(0.5)
|
||
|
||
# 模型跳跃提议概率 q(k'|k)
|
||
# q(1|1)=0.5, q(2|1)=0.5, q(1|2)=0.5, q(2|2)=0.5
|
||
LOG_Q_1_GIVEN_1 = np.log(0.5)
|
||
LOG_Q_2_GIVEN_1 = np.log(0.5)
|
||
LOG_Q_1_GIVEN_2 = np.log(0.5)
|
||
LOG_Q_2_GIVEN_2 = np.log(0.5)
|
||
|
||
|
||
def get_log_prior(k, params):
|
||
"""计算参数的对数先验概率"""
|
||
if k == 1:
|
||
lam = params[0]
|
||
if lam <= 0:
|
||
return -np.inf
|
||
# π(λ)
|
||
return stats.gamma.logpdf(lam, a=ALPHA_LAM, scale=1.0/BETA_LAM)
|
||
|
||
elif k == 2:
|
||
lam, r = params
|
||
if lam <= 0 or r <= 0:
|
||
return -np.inf
|
||
# π(λ, r) = π(λ) * π(r) (假设先验独立)
|
||
log_p_lam = stats.gamma.logpdf(lam, a=ALPHA_LAM, scale=1.0/BETA_LAM)
|
||
log_p_r = stats.invgamma.logpdf(r, a=ALPHA_R, scale=BETA_R)
|
||
return log_p_lam + log_p_r
|
||
|
||
def get_log_likelihood(k, params, data):
|
||
"""计算数据的对数似然"""
|
||
if k == 1:
|
||
lam = params[0]
|
||
if lam <= 0:
|
||
return -np.inf
|
||
# Model 1: Poisson(λ)
|
||
return stats.poisson.logpmf(data, lam).sum()
|
||
|
||
elif k == 2:
|
||
lam, r = params
|
||
if lam <= 0 or r <= 0:
|
||
return -np.inf
|
||
# Model 2: Negative Binomial
|
||
# 使用 p = r / (λ + r) 的参数化
|
||
p = r / (lam + r)
|
||
# 必须检查 p 是否在 (0, 1] 范围内
|
||
if p <= 0 or p > 1:
|
||
return -np.inf
|
||
return stats.nbinom.logpmf(data, n=r, p=p).sum()
|
||
|
||
def get_log_posterior(k, params, data):
|
||
"""计算完整的对数后验（正比于）"""
|
||
log_prior = get_log_prior(k, params)
|
||
if log_prior == -np.inf:
|
||
return -np.inf
|
||
|
||
log_lik = get_log_likelihood(k, params, data)
|
||
if log_lik == -np.inf:
|
||
return -np.inf
|
||
|
||
log_model_prior = LOG_PRIOR_K1 if k == 1 else LOG_PRIOR_K2
|
||
|
||
return log_lik + log_prior + log_model_prior
|
||
|
||
# --- 2. 生成模拟数据 ---
|
||
|
||
# 我们故意从一个过度离散的负二项分布生成数据
|
||
# 泊松分布：均值=方差。 负二项：方差 > 均值。
|
||
TRUE_LAMBDA = 5.0
|
||
TRUE_R = 10.0 # R 值变大，方差接近均值 (方差 = 5 + 25/10 = 7.5)
|
||
TRUE_P = TRUE_R / (TRUE_LAMBDA + TRUE_R)
|
||
np.random.seed(42)
|
||
N_data = 5000 # <--- 之前缺失的行
|
||
data = stats.nbinom.rvs(n=TRUE_R, p=TRUE_P, size=N_data)
|
||
|
||
print(f"模拟数据均值: {data.mean():.2f} (真实均值 = {TRUE_LAMBDA})")
|
||
print(f"模拟数据方差: {data.var():.2f} (泊松模型的方差应为 {data.mean():.2f})")
|
||
|
||
# --- 3. RJMCMC 主函数 ---
|
||
|
||
def run_rjmcmc(data, n_iter=50000, burn_in=10000):
|
||
|
||
# 初始化
|
||
# 从模型1开始，λ 使用数据的均值
|
||
current_k = 1
|
||
current_lambda = data.mean()
|
||
current_params = [current_lambda]
|
||
|
||
# 存储轨迹
|
||
trace_k = np.zeros(n_iter, dtype=int)
|
||
trace_lambda = np.zeros(n_iter)
|
||
trace_r = np.full(n_iter, np.nan) # 仅当 k=2 时有值
|
||
|
||
# 接受计数器
|
||
acceptance = {
|
||
"1_to_1": 0, "2_to_2": 0, "1_to_2": 0, "2_to_1": 0
|
||
}
|
||
attempts = {
|
||
"1_to_1": 0, "2_to_2": 0, "1_to_2": 0, "2_to_1": 0
|
||
}
|
||
|
||
for i in range(n_iter):
|
||
# 1. 提议一个目标模型 k_prop
|
||
# 无论当前 k 是多少，都以 50/50 的概率提议 k=1 或 k=2
|
||
k_prop = np.random.choice([1, 2])
|
||
|
||
# 获取当前的对数后验
|
||
current_log_post = get_log_posterior(current_k, current_params, data)
|
||
|
||
# ----------------------------------
|
||
# 情况 A: 模型内移动 (k_prop == current_k)
|
||
# ----------------------------------
|
||
if k_prop == current_k:
|
||
if current_k == 1:
|
||
# --- Model 1 -> Model 1 ---
|
||
attempts["1_to_1"] += 1
|
||
|
||
# 提议一个新的 λ (使用正态分布随机游走)
|
||
lambda_prop = current_params[0] + np.random.normal(0, 0.5)
|
||
prop_params = [lambda_prop]
|
||
|
||
# 计算接受率
|
||
prop_log_post = get_log_posterior(1, prop_params, data)
|
||
log_alpha = prop_log_post - current_log_post
|
||
# (提议分布是对称的, q(λ'|λ) = q(λ|λ'))
|
||
|
||
if np.log(np.random.rand()) < log_alpha:
|
||
current_params = prop_params
|
||
acceptance["1_to_1"] += 1
|
||
|
||
elif current_k == 2:
|
||
# --- Model 2 -> Model 2 ---
|
||
attempts["2_to_2"] += 1
|
||
|
||
# 提议新的 (λ, r)
|
||
lambda_prop = current_params[0] + np.random.normal(0, 0.5)
|
||
r_prop = current_params[1] + np.random.normal(0, 0.5)
|
||
prop_params = [lambda_prop, r_prop]
|
||
|
||
# 计算接受率
|
||
prop_log_post = get_log_posterior(2, prop_params, data)
|
||
log_alpha = prop_log_post - current_log_post
|
||
|
||
if np.log(np.random.rand()) < log_alpha:
|
||
current_params = prop_params
|
||
acceptance["2_to_2"] += 1
|
||
|
||
# ----------------------------------
|
||
# 情况 B: 跨模型移动 (k_prop != current_k)
|
||
# ----------------------------------
|
||
else:
|
||
if current_k == 1 and k_prop == 2:
|
||
# --- Model 1 -> Model 2 (诞生) ---
|
||
attempts["1_to_2"] += 1
|
||
|
||
# 1. 抽取辅助变量 w
|
||
w = np.random.uniform(0, 1)
|
||
log_g_w = stats.uniform.logpdf(w, 0, 1) # 这是 log(1) = 0
|
||
|
||
# 2. 应用映射
|
||
lambda_prop = current_params[0]
|
||
r_prop = -np.log(w)
|
||
prop_params = [lambda_prop, r_prop]
|
||
|
||
# 3. 计算雅可比项 |J| = 1/w
|
||
log_jacobian = np.log(1.0 / w)
|
||
|
||
# 4. 计算接受率
|
||
prop_log_post = get_log_posterior(2, prop_params, data)
|
||
|
||
# log_alpha = (log_post_prop + log_q_backward) - (log_post_curr + log_q_forward + log_g_w) + log_jacobian
|
||
log_alpha = (prop_log_post + LOG_Q_1_GIVEN_2) - \
|
||
(current_log_post + LOG_Q_2_GIVEN_1 + log_g_w) + \
|
||
log_jacobian
|
||
|
||
if np.log(np.random.rand()) < log_alpha:
|
||
current_k = 2
|
||
current_params = prop_params
|
||
acceptance["1_to_2"] += 1
|
||
|
||
elif current_k == 2 and k_prop == 1:
|
||
# --- Model 2 -> Model 1 (死亡) ---
|
||
attempts["2_to_1"] += 1
|
||
|
||
# 1. 这是一个确定性映射（h' 的逆）
|
||
current_lambda, current_r = current_params
|
||
|
||
# 2. 应用逆映射
|
||
lambda_prop = current_lambda
|
||
w_prime = np.exp(-current_r) # 这就是辅助变量 w'
|
||
prop_params = [lambda_prop]
|
||
|
||
# 3. 计算雅可比项 |J'| = e^(-r)
|
||
log_jacobian_prime = np.log(np.exp(-current_r)) # 即 -current_r
|
||
|
||
# 4. 计算 g(w')，这是 *正向* 移动 (1->2) 中 w 的密度
|
||
# 正向移动是 w ~ U(0, 1)，所以 g(w') = 1 (只要 0 < w' < 1)
|
||
# 因为 r > 0, 所以 w' = e^(-r) 总是在 (0, 1) 区间内
|
||
log_g_w_prime = stats.uniform.logpdf(w_prime, 0, 1) # log(1) = 0
|
||
|
||
# 5. 计算接受率
|
||
prop_log_post = get_log_posterior(1, prop_params, data)
|
||
|
||
# log_alpha = (log_post_prop + log_q_backward + log_g_w_prime) - (log_post_curr + log_q_forward) + log_jacobian
|
||
log_alpha = (prop_log_post + LOG_Q_2_GIVEN_1 + log_g_w_prime) - \
|
||
(current_log_post + LOG_Q_1_GIVEN_2) + \
|
||
log_jacobian_prime
|
||
|
||
if np.log(np.random.rand()) < log_alpha:
|
||
current_k = 1
|
||
current_params = prop_params
|
||
acceptance["2_to_1"] += 1
|
||
|
||
# 存储当前状态
|
||
trace_k[i] = current_k
|
||
trace_lambda[i] = current_params[0]
|
||
if current_k == 2:
|
||
trace_r[i] = current_params[1]
|
||
else:
|
||
trace_r[i] = np.nan
|
||
|
||
# 打印接受率
|
||
print("\n--- 接受率 ---")
|
||
for move, count in attempts.items():
|
||
if count > 0:
|
||
rate = acceptance[move] / count
|
||
print(f"{move}: {acceptance[move]}/{count} ({rate:.2%})")
|
||
|
||
# 丢弃 Burn-in
|
||
trace_k_burned = trace_k[burn_in:]
|
||
trace_lambda_burned = trace_lambda[burn_in:]
|
||
trace_r_burned = trace_r[burn_in:]
|
||
|
||
return trace_k_burned, trace_lambda_burned, trace_r_burned
|
||
|
||
# --- 4. 运行和绘图 ---
|
||
|
||
N_ITER = 20000
|
||
BURN_IN = 5000
|
||
trace_k, trace_lambda, trace_r = run_rjmcmc(data, n_iter=N_ITER, burn_in=BURN_IN)
|
||
|
||
# --- 绘图 ---
|
||
plt.rcParams['font.sans-serif'] = ['SimHei'] # 用来正常显示中文标签
|
||
plt.rcParams['axes.unicode_minus'] = False # 用来正常显示负号
|
||
|
||
# 图 1: 模型后验概率
|
||
plt.figure(figsize=(12, 10))
|
||
|
||
prob_k1 = np.mean(trace_k == 1)
|
||
prob_k2 = np.mean(trace_k == 2)
|
||
|
||
ax1 = plt.subplot(3, 1, 1)
|
||
bars = plt.bar([1, 2], [prob_k1, prob_k2], color=["#227dbe", "#f97b0df9"], tick_label=['模型 1 (泊松)', '模型 2 (负二项)'])
|
||
plt.title('模型后验概率', fontsize=16)
|
||
plt.ylabel('P(k | data)', fontsize=12)
|
||
ax1.bar_label(bars, fmt='{:.2%}', fontsize=12)
|
||
ax1.set_ylim(0, 1)
|
||
|
||
# 图 2: 模型跳跃轨迹 (仅显示前 2000 步，看得更清楚)
|
||
ax2 = plt.subplot(3, 1, 2)
|
||
ax2.plot(trace_k[:2000], 'k.', markersize=2, alpha=0.5)
|
||
ax2.set_yticks([1, 2])
|
||
ax2.set_yticklabels(['模型 1 (泊松)', '模型 2 (负二项)'])
|
||
ax2.set_title('模型空间轨迹 (前2000次迭代)', fontsize=16)
|
||
ax2.set_xlabel('迭代次数', fontsize=12)
|
||
|
||
# 图 3: 参数后验分布
|
||
# λ (lambda) 的后验
|
||
ax3 = plt.subplot(3, 2, 5)
|
||
ax3.hist(trace_lambda, bins=50, density=True, color='#1f77b4', alpha=0.7, label='$\lambda$ 的后验分布')
|
||
ax3.axvline(data.mean(), color='red', linestyle='--', label=f'数据均值 ({data.mean():.2f})')
|
||
ax3.axvline(TRUE_LAMBDA, color='black', linestyle=':', label=f'真实 $\lambda$ ({TRUE_LAMBDA})')
|
||
ax3.set_title('参数 $\lambda$ 的后验分布', fontsize=14)
|
||
ax3.set_xlabel('$\lambda$ 值', fontsize=12)
|
||
ax3.set_ylabel('密度', fontsize=12)
|
||
ax3.legend()
|
||
|
||
# r 的后验
|
||
ax4 = plt.subplot(3, 2, 6)
|
||
# 仅使用 k=2 时的 r 值
|
||
trace_r_k2 = trace_r[~np.isnan(trace_r)]
|
||
if len(trace_r_k2) > 0:
|
||
ax4.hist(trace_r_k2, bins=50, density=True, color='#ff7f0e', alpha=0.7, label='$r$ 的后验分布 (当 k=2)')
|
||
ax4.axvline(TRUE_R, color='black', linestyle=':', label=f'真实 $r$ ({TRUE_R})')
|
||
ax4.set_title('参数 $r$ 的后验分布 (仅当k=2)', fontsize=14)
|
||
ax4.set_xlabel('$r$ 值', fontsize=12)
|
||
ax4.legend()
|
||
else:
|
||
ax4.set_title('参数 $r$ 的后验分布 (未采样到)', fontsize=14)
|
||
ax4.text(0.5, 0.5, '从未接受过模型 2', horizontalalignment='center', verticalalignment='center', transform=ax4.transAxes)
|
||
|
||
plt.tight_layout()
|
||
plt.show()
|