修改代码及README

This commit is contained in:
2025-11-02 20:47:28 +08:00
parent 64ca5e257a
commit b6a9feddc7
14 changed files with 1785 additions and 236 deletions
+28 -23
View File
@@ -63,7 +63,11 @@ def kalman_likelihood(dx, du, dy, T, A, B, C, D, Q, R, u_data, y_data):
S_t = C @ P_t_tm1 @ C.T + R
# 计算对数似然 p(y_t | y_{t-1}, ...)
log_lik_t = multivariate_normal.logpdf(nu_t.squeeze(), mean=jnp.zeros(dy), cov=S_t)
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
@@ -77,7 +81,7 @@ def kalman_likelihood(dx, du, dy, T, A, B, C, D, Q, R, u_data, y_data):
# --- 2. 时间预测 (Time Prediction) ---
# (使用 u_t 来从 x_{t|t} 得到 x_{t+1|t})
# 预测下一个状态 x_{t+1|t} = A*x_{t|t} + B*u_t [cite: 21, 1042]
# 预测下一个状态 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
@@ -99,7 +103,7 @@ def kalman_likelihood(dx, du, dy, T, A, B, C, D, Q, R, u_data, y_data):
# =========================================================================
# 步骤 3.1: 定义模型一 (规范型, Canonical)
# =========================================================================
def model_canonical(u_data, y_data, sigma_process, nugget=1e-12):
def model_canonical(u_data, y_data, sigma_process, sigma_measure):
"""
NumPyro 模型 - 规范型 (Canonical Form)
"""
@@ -116,8 +120,8 @@ def model_canonical(u_data, y_data, sigma_process, nugget=1e-12):
a1 = numpyro.sample("a1", dist.Uniform(-1 - a0, 1 + a0)) # type: ignore
# 观测矩阵 C 的先验
b0 = numpyro.sample("b0", dist.Normal(0, 1))
b1 = numpyro.sample("b1", dist.Normal(0, 1))
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
@@ -128,8 +132,7 @@ def model_canonical(u_data, y_data, sigma_process, nugget=1e-12):
# --- 3. 构造噪声协方差 ---
# 噪声是固定的 (已知的),如 6.3 节算例所述
Q = jnp.eye(dx) * (sigma_process ** 2)
# 为数值稳定性添加 "nugget" [cite: 533]
R = jnp.eye(dy) * nugget
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)
@@ -140,7 +143,7 @@ def model_canonical(u_data, y_data, sigma_process, nugget=1e-12):
# =========================================================================
# 步骤 3.2: 定义模型二 (标准型, Standard)
# =========================================================================
def model_standard(u_data, y_data, sigma_process, nugget=1e-12):
def model_standard(u_data, y_data, sigma_process, sigma_measure):
"""
NumPyro 模型 - 标准型 (Standard Form)
"""
@@ -152,18 +155,18 @@ def model_standard(u_data, y_data, sigma_process, nugget=1e-12):
# 所有系数都是 N(0, 1)
# 状态矩阵 A (dx*dx = 4 个参数)
A11 = numpyro.sample("A11", dist.Normal(0, 1))
A12 = numpyro.sample("A12", dist.Normal(0, 1))
A21 = numpyro.sample("A21", dist.Normal(0, 1))
A22 = numpyro.sample("A22", dist.Normal(0, 1))
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, 1))
B2 = numpyro.sample("B2", dist.Normal(0, 1))
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, 1))
C2 = numpyro.sample("C2", dist.Normal(0, 1))
C1 = numpyro.sample("C1", dist.Normal(0, 2))
C2 = numpyro.sample("C2", dist.Normal(0, 2))
# --- 2. 构造系统矩阵 ---
A = jnp.array([[A11, A12], [A21, A22]])
@@ -173,7 +176,7 @@ def model_standard(u_data, y_data, sigma_process, nugget=1e-12):
# --- 3. 构造噪声协方差 ---
Q = jnp.eye(dx) * (sigma_process ** 2)
R = jnp.eye(dy) * nugget
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)
@@ -184,14 +187,14 @@ def model_standard(u_data, y_data, sigma_process, nugget=1e-12):
# =========================================================================
# (步骤 4: 运行 MCMC - 作为本脚本的 main)
# =========================================================================
def run_mcmc(model, rng_key, u_data, y_data, sigma_process, init_params=None):
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 = 5000
num_samples = 20000
num_warmup = 20000
num_samples = 40000
num_chains = 4
# 使用 NUTS 内核
@@ -216,7 +219,7 @@ def run_mcmc(model, rng_key, u_data, y_data, sigma_process, init_params=None):
)
# 运行
mcmc.run(rng_key, u_data, y_data, sigma_process=sigma_process)
mcmc.run(rng_key, u_data, y_data, sigma_process=sigma_process, sigma_measure=sigma_meas)
# 打印总结
print(f"\n--- MCMC 总结: {model.__name__} ---")
@@ -266,7 +269,8 @@ if __name__ == '__main__':
mcmc_key_c,
u_data_jax,
y_data_jax,
sigma_proc
sigma_proc,
sigma_meas
)
# (模型 2: 标准型)
@@ -275,7 +279,8 @@ if __name__ == '__main__':
mcmc_key_s,
u_data_jax,
y_data_jax,
sigma_proc
sigma_proc,
sigma_meas
)
print("\n--- MCMC 运行完成 ---")