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()' 来获取样本。")