Baysian_Canonical_Identification/demo.py

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