Baysian_Canonical_Identification/models_and_mcmc.py

import numpy as np
import jax
import jax.numpy as jnp
from jax.scipy.stats import multivariate_normal

import numpyro
import numpyro.distributions as dist
from numpyro.infer import MCMC, NUTS, Predictive
import numpyro.infer.initialization as init_strategy 

from functools import partial

from generateGroudTruth import generate_ground_truth_system
from generateSimData import simulate_lti_data

# =========================================================================
# JIT 编译卡尔曼滤波器似然函数 (JAX)
# (基于 Theorem B.1 )
# =========================================================================
@partial(jax.jit, static_argnums=(0, 1, 2, 3))
def kalman_likelihood(dx, du, dy, T, A, B, C, D, Q, R, u_data, y_data):
    """
    使用卡尔曼滤波器计算 LTI 系统的对数似然。
    
    参数:
        dx, du, dy, T: 系统的维度和时间
        A, B, C, D: LTI 系统矩阵 (JAX 数组)
        Q, R: 过程噪声和测量噪声的协方差 (JAX 数组)
        u_data, y_data: 输入和输出数据 (T x du) 和 (T x dy)
        
    返回:
        float: 总的对数似然
    """
    
    # 初始状态
    x_0_m1 = jnp.zeros(dx)      # x_{0|-1} (t=0 given t=-1)
    P_0_m1 = jnp.eye(dx) * 1.0  # P_{0|-1}
    I_dx = jnp.eye(dx)
    
    # 确保 u 和 y 是 (T, N, 1) 的形状以便于矩阵运算
    u_data = u_data.reshape((T, du, 1))
    y_data = y_data.reshape((T, dy, 1))

    def kalman_step(carry, t):
        # --- 0. 载入上一步的预测结果 ---
        # carry 是 (x_{t|t-1}, P_{t|t-1})
        x_t_tm1, P_t_tm1 = carry 
        
        # 获取当前数据
        u_t = u_data[t]
        y_t = y_data[t]
        
        # --- 1. 测量更新 (Measurement Update) ---
        # (使用 y_t 和 u_t 来从 x_{t|t-1} 得到 x_{t|t})
        
        # 确保 x_t_tm1 是列向量 (dx, 1)
        x_t_tm1_col = x_t_tm1.reshape((dx, 1))
        
        # 创新 nu_t = y_t - (C*x_t_tm1 + D*u_t)
        nu_t = y_t - (C @ x_t_tm1_col) - (D @ u_t)
        
        # 创新协方差 S_t = C*P_t_tm1*C^T + R
        S_t = C @ P_t_tm1 @ C.T + R
        
        # 计算对数似然 p(y_t | y_{t-1}, ...)
        sign, logdet_S_t = jnp.linalg.slogdet(S_t)
        S_inv_nu = jnp.linalg.solve(S_t, nu_t)
        quad_term = (nu_t.T @ S_inv_nu).squeeze()
        log_2pi = jnp.log(2.0 * jnp.pi)
        log_lik_t = -0.5 * dy * log_2pi - 0.5 * logdet_S_t - 0.5 * quad_term
        
        # 卡尔曼增益 K_t = P_t_tm1*C^T * S_t^{-1}
        K_t = jnp.linalg.solve(S_t, C @ P_t_tm1).T
        
        # 更新状态 x_{t|t} = x_t_tm1 + K_t * nu_t
        x_t_t = x_t_tm1_col + K_t @ nu_t
        
        # 更新协方差 P_{t|t} = (I - K_t*C)*P_t_tm1
        P_t_t = (I_dx - K_t @ C) @ P_t_tm1

        # --- 2. 时间预测 (Time Prediction) ---
        # (使用 u_t 来从 x_{t|t} 得到 x_{t+1|t})
        
        # 预测下一个状态 x_{t+1|t} = A*x_{t|t} + B*u_t
        x_tp1_t = (A @ x_t_t + B @ u_t).flatten()  # 确保输出是 (dx,) 形状
        
        # 预测下一个协方差 P_{t+1|t} = A*P_{t|t}*A^T + Q
        P_tp1_t = A @ P_t_t @ A.T + Q
        
        # 返回 (x_{t+1|t}, P_{t+1|t}) 作为下一次迭代的 carry
        return (x_tp1_t, P_tp1_t), log_lik_t

    # 运行 scan 循环
    # 初始 carry 是 (x_{0|-1}, P_{0|-1})
    initial_carry = (x_0_m1, P_0_m1)
    
    # jax.lax.scan 会在所有 t 上迭代 kalman_step
    (_, _), log_likelihoods = jax.lax.scan(kalman_step, initial_carry, jnp.arange(T))
    
    # 返回总的对数似然
    return jnp.sum(log_likelihoods)

# =========================================================================
# 步骤 3.1: 定义模型一 (规范型, Canonical)
# =========================================================================
def model_canonical(u_data, y_data, sigma_process, sigma_measure):
    """
    NumPyro 模型 - 规范型 (Canonical Form)
    """
    dx, du = 2, 1
    dy = 1
    T = u_data.shape[0]

    # --- 1. 采样先验 (Priors) ---
    
    # 状态矩阵 A 的先验
    # 使用 Lemma 4.3  的稳定先验：|a0| < 1 和 |a1| < 1 + a0
    a0 = numpyro.sample("a0", dist.Uniform(-1, 1)) # type: ignore
    # a1 的范围取决于 a0
    a1 = numpyro.sample("a1", dist.Uniform(-1 - a0, 1 + a0)) # type: ignore
    
    # 观测矩阵 C 的先验
    b0 = numpyro.sample("b0", dist.Normal(0, 2))
    b1 = numpyro.sample("b1", dist.Normal(0, 2))

    # --- 2. 构造系统矩阵 ---
    A = jnp.array([[0.0, 1.0], [-a0, -a1]]) # type: ignore
    B = jnp.array([[0.0], [1.0]])
    C = jnp.array([[b0, b1]])
    D = jnp.zeros((dy, du))
    
    # --- 3. 构造噪声协方差 ---
    # 噪声是固定的 (已知的)，如 6.3 节算例所述
    Q = jnp.eye(dx) * (sigma_process ** 2)
    R = jnp.eye(dy) * (sigma_measure ** 2) 
    
    # --- 4. 计算总似然 ---
    log_lik_total = kalman_likelihood(dx, du, dy, T, A, B, C, D, Q, R, u_data, y_data)
    
    # 将似然注册到 NumPyro
    numpyro.factor("log_likelihood", log_lik_total)

# =========================================================================
# 步骤 3.2: 定义模型二 (标准型, Standard)
# =========================================================================
def model_standard(u_data, y_data, sigma_process, sigma_measure):
    """
    NumPyro 模型 - 标准型 (Standard Form)
    """
    dx, du = 2, 1
    dy = 1
    T = u_data.shape[0]

    # --- 1. 采样先验 (Priors) ---
    # 所有系数都是 N(0, 1) 
    
    # 状态矩阵 A (dx*dx = 4 个参数)
    A11 = numpyro.sample("A11", dist.Normal(0, 2))
    A12 = numpyro.sample("A12", dist.Normal(0, 2))
    A21 = numpyro.sample("A21", dist.Normal(0, 2))
    A22 = numpyro.sample("A22", dist.Normal(0, 2))
    
    # 输入矩阵 B (dx*du = 2 个参数)
    B1 = numpyro.sample("B1", dist.Normal(0, 2))
    B2 = numpyro.sample("B2", dist.Normal(0, 2))

    # 观测矩阵 C (dy*dx = 2 个参数)
    C1 = numpyro.sample("C1", dist.Normal(0, 2))
    C2 = numpyro.sample("C2", dist.Normal(0, 2))

    # --- 2. 构造系统矩阵 ---
    A = jnp.array([[A11, A12], [A21, A22]])
    B = jnp.array([[B1], [B2]])
    C = jnp.array([[C1, C2]])
    D = jnp.zeros((dy, du))

    # --- 3. 构造噪声协方差 ---
    Q = jnp.eye(dx) * (sigma_process ** 2)
    R = jnp.eye(dy) * (sigma_measure ** 2)
    
    # --- 4. 计算总似然 ---
    log_lik_total = kalman_likelihood(dx, du, dy, T, A, B, C, D, Q, R, u_data, y_data)
    
    # 将似然注册到 NumPyro
    numpyro.factor("log_likelihood", log_lik_total)

# =========================================================================
# (步骤 4: 运行 MCMC - 作为本脚本的 main)
# =========================================================================
def run_mcmc(model, rng_key, u_data, y_data, sigma_process, sigma_meas, init_params=None):
    """辅助函数，用于运行 NUTS 采样器"""
    
    print(f"\n--- 开始为模型 {model.__name__} 运行 MCMC ---")
    
    # 论文中的 MCMC 设置
    num_warmup = 20000
    num_samples = 40000
    num_chains = 4
    
    # 使用 NUTS 内核
    if init_params:
        print("--- MCMC 正在使用 'init_to_value' 策略 (从真实值开始) ---")
        # 确保所有值都是 JAX 数组
        init_values = {k: jnp.array(v) for k, v in init_params.items()}
        # 使用 init_to_value 策略
        strategy = init_strategy.init_to_value(values=init_values) 
        kernel = NUTS(model, init_strategy=strategy) 
    else:
        print("--- MCMC 正在使用默认初始化策略 ---")
        kernel = NUTS(model) # <-- 无需策略
    
    # 配置 MCMC
    mcmc = MCMC(
        kernel,
        num_warmup=num_warmup,
        num_samples=num_samples,
        num_chains=num_chains,
        progress_bar=True
    )
    
    # 运行
    mcmc.run(rng_key, u_data, y_data, sigma_process=sigma_process, sigma_measure=sigma_meas)
    
    # 打印总结
    print(f"\n--- MCMC 总结: {model.__name__} ---")
    mcmc.print_summary()
    
    return mcmc

if __name__ == '__main__':
    # --- 0. 设置 JAX 和 NumPyro ---
    # numpyro.set_platform("cpu") # 或者 "gpu"
    numpyro.set_host_device_count(4) # 使用 4 个 CPU 核心 (对应 4 条链)
    
    # JAX 随机种子
    main_rng_key = jax.random.PRNGKey(42)
    
    # --- 1. 生成 Ground Truth 系统 ---
    print("--- (步骤 1) ---")
    gt_key, sim_key, mcmc_key = jax.random.split(main_rng_key, 3)
    A_true, B_true, C_true, D_true = generate_ground_truth_system(
        rng_seed=int(gt_key[0])
    )
    
    # --- 2. 仿真数据 ---
    print("\n--- (步骤 2) ---")
    T_steps = 400
    sigma_proc = 0.3
    sigma_meas = 0.0
    
    u_data, y_data = simulate_lti_data(
        A_true, B_true, C_true, D_true,
        T=T_steps,
        sigma_process=sigma_proc,
        sigma_measurement=sigma_meas,
        rng_seed=int(sim_key[0])
    )
    
    # 将数据转换为 JAX 数组
    u_data_jax = jnp.array(u_data)
    y_data_jax = jnp.array(y_data)

    # --- 3 & 4. 运行两个模型的 MCMC ---
    
    # (模型 1: 规范型)
    mcmc_key_c, mcmc_key_s = jax.random.split(mcmc_key)
    mcmc_canonical = run_mcmc(
        model_canonical, 
        mcmc_key_c, 
        u_data_jax, 
        y_data_jax,
        sigma_proc,
        sigma_meas
    )
    
    # (模型 2: 标准型)
    mcmc_standard = run_mcmc(
        model_standard, 
        mcmc_key_s, 
        u_data_jax, 
        y_data_jax,
        sigma_proc,
        sigma_meas
    )
    
    print("\n--- MCMC 运行完成 ---")
    print("下一步是分析和可视化后验分布 (步骤 5)。")
    print("例如，使用 'mcmc_canonical.get_samples()' 来获取样本。")