import numpy as np import scipy.linalg import matplotlib.pyplot as plt import control import jax import jax.numpy as jnp from jax.scipy.linalg import solve_triangular import numpyro import numpyro.distributions as dist from numpyro.infer import MCMC, NUTS import arviz as az plt.rcParams['font.sans-serif'] = ['SimHei'] # or 'Microsoft YaHei' etc. plt.rcParams['axes.unicode_minus'] = False DX, DU, DY = 2, 1, 1 T_sim = 400 SIGMA_W_STD = 0.3 SIGMA_Z_STD = 1e-12 # Nugget value instead of 0.0 # 定义 LTI 系统 a0_true,a1_true = 0.6, -1.5 b0_true,b1_true = 1.0, 0.5 d0_true = 0.0 Ac_true = jnp.array([[0.,1.], [-a0_true, -a1_true]]) Bc_true = jnp.array([[0.],[1.]]) Cc_true = jnp.array([[b0_true, b1_true]]) Dc_true = jnp.array([[d0_true]]) A_true, B_true, C_true, D_true = Ac_true, Bc_true, Cc_true, Dc_true Sigma_true = (SIGMA_W_STD**2)*jnp.eye(DX) Gamma_true = (SIGMA_Z_STD**2)*jnp.eye(DY) # 生成仿真数据 key = jax.random.PRNGKey(0) key_u, key_w, key_x0 = jax.random.split(key, 3) u_data = jax.random.normal(key_u, shape=(T_sim, DU)) x0 = jnp.zeros(DX) x_hist = np.zeros((T_sim + 1, DX)) y_hist = np.zeros((T_sim, DY)) x_hist[0] = x0 for t in range(T_sim): w_t = jax.random.multivariate_normal(key_w, jnp.zeros(DX), Sigma_true) key_w, _ = jax.random.split(key_w) # 更新随机数种子 x_next = A_true @ x_hist[t] + B_true @ u_data[t] + w_t y_t = C_true @ x_hist[t] + D_true @ u_data[t] # 理想观测，无观测噪声，不考虑直通项影响 y_t_observed = y_t + D_true @ u_data[t] # 理想观测，无观测噪声 x_hist[t+1] = x_next y_hist[t] = y_t_observed y_data = y_hist def kalman_log_likelihood(params_A, params_B, params_C, params_D, params_Sigma, params_Gamma, u, y): """ 使用 JAX 实现的卡尔曼滤波器计算状态空间模型的对数似然。 Args: params_A, params_B, params_C, params_D: 系统矩阵 (JAX arrays)。 params_Sigma: 过程噪声协方差矩阵 (JAX array)。 params_Gamma: 测量噪声协方差矩阵 (JAX array)。 u: 输入数据 (JAX array, shape (T_sim, DU))。 y: 输出数据 (JAX array, shape (T_sim, DY))。 Returns: 总对数似然 (JAX scalar)。 """ # 确保输入是 JAX 数组 (防御性编程) u = jnp.asarray(u) y = jnp.asarray(y) params_A = jnp.asarray(params_A) params_B = jnp.asarray(params_B) params_C = jnp.asarray(params_C) params_D = jnp.asarray(params_D) params_Sigma = jnp.asarray(params_Sigma) params_Gamma = jnp.asarray(params_Gamma) # 初始状态估计和协方差 # 假设初始状态均值为零，具有一定的先验不确定性 x_hat_0_0 = jnp.zeros(DX) P_0_0 = jnp.eye(DX) * 0.01 # 可以根据需要调整初始不确定性 def kf_step(carry, t): """ 卡尔曼滤波器的单步迭代函数，用于 jax.lax.scan。 Args: carry: 上一步的状态 (x_hat_{t-1|t-1}, P_{t-1|t-1})。 t: 当前时间步索引 (从 0 到 T_sim-1)。 Returns: 新的状态 (x_hat_{t|t}, P_{t|t}) 和当前步的对数似然 log_lik_t。 """ x_hat_t_t_prev, P_t_t_prev = carry # --- 获取当前和过去的输入 --- u_t_prev = jnp.where(t > 0, u[t - 1], jnp.zeros(DU)) u_t_obs = u[t] # --- 获取当前观测 --- y_t_obs = y[t] # --- 预测步骤 --- x_hat_t_tm1 = params_A @ x_hat_t_t_prev + params_B @ u_t_prev P_t_tm1 = params_A @ P_t_t_prev @ params_A.T + params_Sigma # --- 更新步骤 --- nu_t = y_t_obs - (params_C @ x_hat_t_tm1 + params_D @ u_t_obs) S_t = params_C @ P_t_tm1 @ params_C.T + params_Gamma try: # 尝试直接求逆 S_t_inv = jnp.linalg.inv(S_t) K_t = P_t_tm1 @ params_C.T @ S_t_inv except jnp.linalg.LinAlgError: S_t_inv = jnp.full_like(S_t, jnp.nan) K_t = jnp.full((DX, DY), jnp.nan) x_hat_t_t = x_hat_t_tm1 + K_t @ nu_t P_t_t = (jnp.eye(DX) - K_t @ params_C) @ P_t_tm1 sign, log_det_S = jnp.linalg.slogdet(S_t) # 确保 nu_t 是合适的形状进行二次型计算 nu_t_col = nu_t.reshape(-1, 1) if DY > 0 else nu_t quad_form = (nu_t_col.T @ S_t_inv @ nu_t_col).squeeze() log_lik_t = -0.5 * DY * jnp.log(2 * jnp.pi) - 0.5 * log_det_S - 0.5 * quad_form # 对于求逆失败的情况，返回一个非常差的似然 log_lik_t = jnp.where(jnp.isnan(K_t).any(), -jnp.inf, log_lik_t) # 返回新的状态和当前步的对数似然 (确保是标量) return (x_hat_t_t, P_t_t), jnp.asarray(log_lik_t).squeeze() initial_carry = (x_hat_0_0, P_0_0) time_indices = jnp.arange(T_sim) final_carry, log_lik_steps = jax.lax.scan(kf_step, initial_carry, time_indices) return jnp.sum(log_lik_steps) # 定义贝叶斯模型 def model_standard(u, y = None): # A 的先验 (DX*DX elements) A_s_flat = numpyro.sample("A_s_flat", dist.Normal(0., 1.).expand([DX*DX])) A_s = A_s_flat.reshape((DX, DX)) # B 的先验 (DX*DU elements) B_s_flag = numpyro.sample("B_s_flat", dist.Normal(0., 1.).expand([DX*DU])) B_s = B_s_flag.reshape((DX, DU)) # C 的先验 (DY*DX elements) C_s_flat = numpyro.sample("C_s_flat", dist.Normal(0., 1.).expand([DY*DX])) C_s = C_s_flat.reshape((DY, DX)) # D 的先验 (仍然采样一个标量) D_s_scalar = numpyro.sample("D_s", dist.Normal(0., 1.)) # 将标量 reshape 为 (DY, DU) = (1, 1) 矩阵 D_s = D_s_scalar.reshape((DY, DU)) # Sigma 的先验 (DX*DX elements) sigma_w = numpyro.sample("sigma_w", dist.HalfCauchy(1.)) Sigma_w = (sigma_w**2) * jnp.eye(DX) # Gamma 的先验 (DY*DY elements) Gamma_z = (SIGMA_Z_STD**2) * jnp.eye(DY) # 计算似然 log_lik = kalman_log_likelihood(A_s, B_s, C_s, D_s, Sigma_w, Gamma_z, u, y) # 观测数据的似然 numpyro.factor("obs", log_lik) # 运行 MCMC 进行推断 kernel_s = NUTS(model_standard) mcmc_s = MCMC(kernel_s, num_warmup=1000, num_samples=2000, num_chains=1) mcmc_s.run(jax.random.PRNGKey(2), u=u_data, y=y_data) samples_s = mcmc_s.get_samples() idata_s = az.from_numpyro(mcmc_s) def model_canonical(u, y=None): a0 = numpyro.sample("a0", dist.Normal(0., 1.)) a1 = numpyro.sample("a1", dist.Normal(0., 1.)) b0 = numpyro.sample("b0", dist.Normal(0., 1.)) b1 = numpyro.sample("b1", dist.Normal(0., 1.)) d0 = numpyro.sample("d0", dist.Normal(0., 1.)) A_c = jnp.array([[0., 1.], [-a0, -a1]]) B_c = jnp.array([[0.],[1.]]) C_c = jnp.array([[b0, b1]]) D_c = jnp.array([[d0]]) sigma_w = numpyro.sample("sigma_w", dist.HalfCauchy(1.)) Sigma_w = (sigma_w**2) * jnp.eye(DX) Gamma_z = (SIGMA_Z_STD**2) * jnp.eye(DY) log_lik = kalman_log_likelihood(A_c, B_c, C_c, D_c, Sigma_w, Gamma_z, u, y) numpyro.factor("obs", log_lik) # 运行 MCMC 进行推断 kernel_c = NUTS(model_canonical) mcmc_c = MCMC(kernel_c, num_warmup=1000, num_samples=2000, num_chains=1) mcmc_c.run(jax.random.PRNGKey(3), u=u_data, y=y_data) samples_c = mcmc_c.get_samples() idata_c = az.from_numpyro(mcmc_c) # # Plot pair plot for standard parameters (select a few, e.g., A_s elements) az.plot_pair(idata_s, var_names=['A_s_flat'], coords={'A_s_flat_dim_0': [0, 1, 2, 3]}) # Plot A11, A12, A21, A22 plt.suptitle("后验分布 (标准参数化 $\Theta_s$)") plt.tight_layout(rect=[0, 0.03, 1, 0.95]) # Adjust layout to prevent title overlap plt.show() # Plot pair plot for canonical parameters az.plot_pair(idata_c, var_names=['a0', 'a1', 'b0', 'b1']) plt.suptitle("后验分布 (规范参数化 $\Theta_c$)") plt.tight_layout(rect=[0, 0.03, 1, 0.95]) plt.show()