程序初版完成,目前收敛后结果不太对
This commit is contained in:
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -1,105 +1,99 @@
|
|||||||
## 实现一下重要抽样
|
|
||||||
import numpy as np
|
import numpy as np
|
||||||
import matplotlib.pyplot as plt
|
import matplotlib.pyplot as plt
|
||||||
import scipy as stats
|
from scipy.stats import beta as beta_dist # 导入beta分布用于绘图
|
||||||
import seaborn as sns
|
|
||||||
|
|
||||||
np.random.seed(42)
|
# ------------------------------------------------------------------
|
||||||
|
# 1. 设置 matplotlib 支持中文显示
|
||||||
|
# ------------------------------------------------------------------
|
||||||
|
try:
|
||||||
|
plt.rcParams['font.sans-serif'] = ['SimHei'] # Windows/Linux
|
||||||
|
plt.rcParams['axes.unicode_minus'] = False # 正常显示负号
|
||||||
|
except Exception:
|
||||||
|
try:
|
||||||
|
plt.rcParams['font.sans-serif'] = ['Arial Unicode MS'] # MacOS
|
||||||
|
plt.rcParams['axes.unicode_minus'] = False
|
||||||
|
except Exception:
|
||||||
|
print("未找到中文字体,绘图可能显示异常。请安装'SimHei'或'Arial Unicode MS'字体。")
|
||||||
|
|
||||||
def baysian_linear_regression_inportance_sampling():
|
# ------------------------------------------------------------------
|
||||||
|
# 2. 设定模型参数 (为了模拟您图中的效果)
|
||||||
|
# ------------------------------------------------------------------
|
||||||
|
n = 20 # X的试验总次数
|
||||||
|
# 更改 alpha 和 beta 以匹配图中 ~0.72 的均值
|
||||||
|
alpha = 8.0 # Beta分布的先验参数 alpha
|
||||||
|
beta = 3.0 # Beta分布的先验参数 beta
|
||||||
|
# 理论均值 E[Y] = alpha / (alpha + beta) = 8 / 11 ≈ 0.727
|
||||||
|
theoretical_mean = alpha / (alpha + beta)
|
||||||
|
|
||||||
|
# MCMC (Gibbs) 抽样参数
|
||||||
|
N_samples = 1000 # 总抽样量 N (同您图中的 N=1000)
|
||||||
|
N_burn_in = 200 # 预估的老化期(预热期) N1
|
||||||
|
|
||||||
|
print(f"模型参数: n={n}, alpha={alpha}, beta={beta}")
|
||||||
|
print(f"理论均值 E[Y]: {theoretical_mean:.4f}")
|
||||||
|
print(f"抽样设置: 总样本 N={N_samples}")
|
||||||
|
|
||||||
|
# ------------------------------------------------------------------
|
||||||
|
# 3. 初始化
|
||||||
|
# ------------------------------------------------------------------
|
||||||
|
# 创建数组来存储所有样本
|
||||||
|
samples_X = np.zeros(N_samples, dtype=int)
|
||||||
|
samples_Y = np.zeros(N_samples, dtype=float)
|
||||||
|
|
||||||
|
# 设定马尔可夫链的初始状态
|
||||||
|
# 故意设置一个远离均值(0.727)的初始值,以观察收敛
|
||||||
|
y_t = 0.1
|
||||||
|
|
||||||
|
# ------------------------------------------------------------------
|
||||||
|
# 4. 运行 Gibbs 抽样
|
||||||
|
# ------------------------------------------------------------------
|
||||||
|
print("开始Gibbs抽样...")
|
||||||
|
np.random.seed(101) # 使用和您图中一样的随机种子
|
||||||
|
|
||||||
|
for i in range(N_samples):
|
||||||
|
x_t = np.random.binomial(n, y_t)
|
||||||
|
y_t = np.random.beta(x_t + alpha, n - x_t + beta)
|
||||||
|
samples_Y[i] = y_t
|
||||||
|
samples_X[i] = x_t
|
||||||
|
print("抽样完成。")
|
||||||
|
|
||||||
|
# ------------------------------------------------------------------
|
||||||
|
# 5. 定义并计算逐步平均值 (Ergodic Mean)
|
||||||
|
# ------------------------------------------------------------------
|
||||||
|
def calculate_ergodic_mean(y_samples):
|
||||||
"""
|
"""
|
||||||
使用贝叶斯线性回归方法进行重要抽样
|
计算逐步平均值 (累积平均值)
|
||||||
|
y_k_bar = (1/k) * sum(y_i for i=1 to k)
|
||||||
本例需要估计一个简单的线性回归模型的后验分布
|
|
||||||
|
使用 np.cumsum() 可以高效实现
|
||||||
y = β0 + β1*x + ε, ε ~ N(0, σ²)
|
|
||||||
= [1 x] * [β0 β1]' + ε
|
|
||||||
"""
|
"""
|
||||||
|
n = len(y_samples)
|
||||||
|
# 1. 计算累积和 [y1, y1+y2, y1+y2+y3, ...]
|
||||||
|
s = np.cumsum(y_samples)
|
||||||
|
# 2. 创建 k 数组 [1, 2, 3, ...]
|
||||||
|
k_array = np.arange(1, n + 1)
|
||||||
|
# 3. 计算 avg[k] = s[k] / k
|
||||||
|
return s / k_array
|
||||||
|
|
||||||
# 1. 定义关键数据
|
# 计算所有 Y 样本的逐步平均值 (包括预热期)
|
||||||
x = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10])
|
ergodic_mean_Y = calculate_ergodic_mean(samples_Y)
|
||||||
y = np.array([2.1, 3.9, 6.2, 8.1, 9.8, 12.3, 13.9, 16.2, 17.8, 20.1])
|
|
||||||
n = len(x)
|
|
||||||
|
|
||||||
print(f"数据点数量:{n}")
|
# ------------------------------------------------------------------
|
||||||
print(f"x: {x}")
|
# 6. 绘制逐步平均值图 (实现您图片中的效果)
|
||||||
print(f"y: {y}")
|
# ------------------------------------------------------------------
|
||||||
|
print("正在绘制逐步平均值图...")
|
||||||
|
|
||||||
# 2. 定义先验分布
|
plt.figure(figsize=(10, 6))
|
||||||
# β0 ~ N(0, 10)
|
plt.plot(np.arange(1, N_samples + 1), ergodic_mean_Y, label=r"逐步平均值 $\bar{y}_k$")
|
||||||
beta0_prior_mean = 0
|
plt.axhline(theoretical_mean, color='red', linestyle='--', label=f"理论均值: {theoretical_mean:.4f}")
|
||||||
beta0_prior_var = 10
|
|
||||||
|
|
||||||
# β1 ~ N(1, 5)
|
# 添加一个垂直线来标记我们估计的预热期
|
||||||
beta1_prior_mean = 1
|
plt.axvline(N_burn_in, color='gray', linestyle=':', label=f"估计的预热期 N1 = {N_burn_in}")
|
||||||
beta1_prior_var = 5
|
|
||||||
|
|
||||||
# σ² ~ Inverse-Gamma(2, 1)
|
plt.title("使用逐步平均值图查看预热期")
|
||||||
sigma2_prior_a = 2
|
plt.xlabel("迭代次数 (k)")
|
||||||
sigma2_prior_b = 1
|
plt.ylabel(r"逐步平均值 $\bar{y}_k$")
|
||||||
|
plt.legend()
|
||||||
# 3. 定义重要抽样分布
|
plt.grid(True)
|
||||||
# 使用最小二乘进行估计
|
plt.ylim(0, 1) # Y值在0到1之间
|
||||||
X = np.column_stack([np.ones(n), x])
|
plt.show()
|
||||||
beta_estimate = np.linalg.inv(X.T @ X) @ X.T @ y
|
|
||||||
y_pred = X @ beta_estimate
|
|
||||||
residuals = y - y_pred
|
|
||||||
RSS = residuals.T @ residuals
|
|
||||||
sigma2_estimate = RSS / (n - 2)
|
|
||||||
|
|
||||||
print(f"最小二乘结果:β0 = {beta_estimate[0]}, β1 = {beta_estimate[1]}, σ² = {sigma2_estimate}")
|
|
||||||
print(f"残差平方和 RSS = {RSS}")
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
# 4. 定义联合后验分布(未归一化)
|
|
||||||
def unnormalized_log_posterior(param):
|
|
||||||
"""
|
|
||||||
计算后验概率密度的未归一化值
|
|
||||||
|
|
||||||
param: [β0, β1, log_sigma2)]
|
|
||||||
"""
|
|
||||||
|
|
||||||
beta0, beta1, log_sigma2 = param
|
|
||||||
|
|
||||||
sigma2 = np.exp(log_sigma2)
|
|
||||||
|
|
||||||
# 计算似然函数值
|
|
||||||
y_pred = beta0 + beta1 * x
|
|
||||||
residuals = y - y_pred
|
|
||||||
log_likelihood = -0.5 * n * np.log(2 * np.pi * sigma2) * -0.5 * np.sum(residuals**2) / sigma2
|
|
||||||
|
|
||||||
# 计算先验概率
|
|
||||||
log_prior_beta0 = stats.norm.logpdf(beta0, loc = beta0_prior_mean, scale = np.sqrt(beta0_prior_var))
|
|
||||||
log_prior_beta1 = stats.norm.logpdf(beta1, loc = beta1_prior_mean, scale = np.sqrt(beta1_prior_var))
|
|
||||||
log_prior_sigma2 = stats.invgamma.logpdf(sigma2, a = sigma2_prior_a, scale = sigma2_prior_b) + log_sigma2
|
|
||||||
|
|
||||||
return log_likelihood + log_prior_beta0 + log_prior_beta1 + log_prior_sigma2
|
|
||||||
|
|
||||||
# 5. 定义重要抽样分布
|
|
||||||
proposal_mean = np.array([beta_estimate[0], beta_estimate[1], np.log(sigma2_estimate)])
|
|
||||||
proposal_cov = np.diag([1.0, 1.0, 1.0])
|
|
||||||
|
|
||||||
def log_proposal_density(param):
|
|
||||||
return stats.multivariable_normal.logpdf(param, mean = proposal_mean, cov = proposal_cov)
|
|
||||||
|
|
||||||
# 6. 进行重要抽样
|
|
||||||
num_samples = 10000
|
|
||||||
samples = np.zeros((num_samples, 3))
|
|
||||||
weights = np.zeros(num_samples)
|
|
||||||
log_weights = np.zeros(num_samples)
|
|
||||||
|
|
||||||
for i in range(num_samples):
|
|
||||||
samples[i] = np.random.multivariate_normal(proposal_mean, proposal_cov)
|
|
||||||
unnormalized_log_posterior = unnormalized_log_posterior(samples[i])
|
|
||||||
log_proposal = log_proposal_density(samples[i])
|
|
||||||
log_weights[i] = unnormalized_log_posterior - log_proposal
|
|
||||||
|
|
||||||
weights = np.exp(log_weights)
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
|
||||||
baysian_linear_regression_inportance_sampling()
|
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -1,3 +1,6 @@
|
|||||||
|
# JAX 入门示例代码
|
||||||
|
# 演示 JAX 的基本用法,包括 jnp 数组、jit 加速
|
||||||
|
|
||||||
import numpy as np
|
import numpy as np
|
||||||
import jax
|
import jax
|
||||||
import jax.numpy as jnp
|
import jax.numpy as jnp
|
||||||
|
|||||||
@@ -0,0 +1,100 @@
|
|||||||
|
import numpy as np
|
||||||
|
from scipy.linalg import solve_discrete_lyapunov
|
||||||
|
import sys
|
||||||
|
|
||||||
|
def generate_ground_truth_system(dx=2, du=1, dy=1, rng_seed=None):
|
||||||
|
"""
|
||||||
|
生成一个 "真实" 的、稳定的、可控的、可观测的 LTI 系统。
|
||||||
|
|
||||||
|
该过程遵循论文 6.2 节中描述的方法。
|
||||||
|
|
||||||
|
参数:
|
||||||
|
dx (int): 状态维度 (state dimension)
|
||||||
|
du (int): 输入维度 (input dimension)
|
||||||
|
dy (int): 输出维度 (output dimension)
|
||||||
|
rng_seed (int, optional): 用于复现的随机种子
|
||||||
|
|
||||||
|
返回:
|
||||||
|
tuple: (A, B, C, D) 矩阵
|
||||||
|
"""
|
||||||
|
|
||||||
|
# 初始化随机数生成器
|
||||||
|
if rng_seed is None:
|
||||||
|
rng = np.random.default_rng()
|
||||||
|
else:
|
||||||
|
rng = np.random.default_rng(rng_seed)
|
||||||
|
|
||||||
|
# 尝试生成一个良态的系统,最多重试 100 次
|
||||||
|
max_retries = 100
|
||||||
|
for attempt in range(max_retries):
|
||||||
|
try:
|
||||||
|
# --- 1. 生成稳定的 A 矩阵 (dx x dx) ---
|
||||||
|
# 论文图 2(b) 显示了复特征值,遵循 6.2 节的极坐标法
|
||||||
|
|
||||||
|
# 在极坐标下采样一对共轭复特征值
|
||||||
|
r_squared = rng.uniform(0, 1.0) # 采样 r^2,确保在单位圆内
|
||||||
|
r = np.sqrt(r_squared)
|
||||||
|
theta = rng.uniform(0, np.pi) # 仅在上半平面采样角度
|
||||||
|
|
||||||
|
lambda1 = r * (np.cos(theta) + 1j * np.sin(theta))
|
||||||
|
lambda2 = np.conjugate(lambda1)
|
||||||
|
|
||||||
|
# 创建对应的实数块对角矩阵
|
||||||
|
alpha, beta = lambda1.real, lambda1.imag
|
||||||
|
lambda_block = np.array([[alpha, beta],
|
||||||
|
[-beta, alpha]])
|
||||||
|
|
||||||
|
# 生成一个随机正交矩阵 V (dx x dx) [cite: 511, 584]
|
||||||
|
Z = rng.standard_normal(size=(dx, dx))
|
||||||
|
V, _ = np.linalg.qr(Z)
|
||||||
|
|
||||||
|
# 组装 A = V * Lambda_block * V^T
|
||||||
|
A = V @ lambda_block @ V.T
|
||||||
|
|
||||||
|
# --- 2. 生成 B (dx x du) 和 C (dy x dx) 矩阵 ---
|
||||||
|
# 元素从 N(0, 1) 独立采样
|
||||||
|
B = rng.standard_normal(size=(dx, du))
|
||||||
|
C = rng.standard_normal(size=(dy, dx))
|
||||||
|
|
||||||
|
# --- 3. 检查可控性和可观测性 ---
|
||||||
|
# 求解离散时间李雅普诺夫方程 (Lyapunov equation)
|
||||||
|
Wc = solve_discrete_lyapunov(A, B @ B.T) # Controllability Gramian
|
||||||
|
Wo = solve_discrete_lyapunov(A.T, C.T @ C) # Observability Gramian
|
||||||
|
|
||||||
|
# 检查 Gramian 矩阵的条件
|
||||||
|
# "拒绝主要特征值占总能量 99% 以上的系统"
|
||||||
|
eig_Wc = np.linalg.eigvalsh(Wc)
|
||||||
|
eig_Wo = np.linalg.eigvalsh(Wo)
|
||||||
|
|
||||||
|
cond_c = np.max(eig_Wc) / np.sum(eig_Wc)
|
||||||
|
cond_o = np.max(eig_Wo) / np.sum(eig_Wo)
|
||||||
|
|
||||||
|
# 如果系统是良态的,则跳出循环
|
||||||
|
if cond_c < 0.99 and cond_o < 0.99:
|
||||||
|
|
||||||
|
# --- 4. 定义 D 矩阵 (dy x du) ---
|
||||||
|
# 算例 6.3 中 D=0
|
||||||
|
D = np.zeros((dy, du))
|
||||||
|
|
||||||
|
print(f"--- 成功生成 Ground Truth 系统 (尝试次数: {attempt + 1}) ---")
|
||||||
|
print(f"特征值: {lambda1:.4f}, {lambda2:.4f}")
|
||||||
|
print(f"Gramian 能量占比: Wc={cond_c:.4f}, Wo={cond_o:.4f}")
|
||||||
|
print("A = \n", A)
|
||||||
|
print("B = \n", B)
|
||||||
|
print("C = \n", C)
|
||||||
|
print("D = \n", D)
|
||||||
|
|
||||||
|
return A, B, C, D
|
||||||
|
|
||||||
|
except np.linalg.LinAlgError:
|
||||||
|
# 李雅普诺夫方程求解器可能失败(例如,A 矩阵数值上不稳定)
|
||||||
|
print(f"尝试 {attempt + 1} 失败 (LinAlgError)。正在重试...", file=sys.stderr)
|
||||||
|
continue
|
||||||
|
|
||||||
|
# 如果循环结束仍未成功
|
||||||
|
raise RuntimeError(f"在 {max_retries} 次尝试后未能生成一个良态的系统。")
|
||||||
|
|
||||||
|
if __name__ == '__main__':
|
||||||
|
# --- 运行示例 ---
|
||||||
|
# 设置一个随机种子,以便每次运行时都能得到相同的结果
|
||||||
|
A_true, B_true, C_true, D_true = generate_ground_truth_system(rng_seed=42)
|
||||||
@@ -0,0 +1,126 @@
|
|||||||
|
import numpy as np
|
||||||
|
from scipy.linalg import solve_discrete_lyapunov
|
||||||
|
import sys
|
||||||
|
import matplotlib.pyplot as plt
|
||||||
|
from generateGroudTruth import generate_ground_truth_system
|
||||||
|
|
||||||
|
|
||||||
|
def simulate_lti_data(A, B, C, D, T, sigma_process, sigma_measurement, rng_seed=None):
|
||||||
|
"""
|
||||||
|
使用 LTI 系统 (A, B, C, D) 仿真生成数据。
|
||||||
|
|
||||||
|
参数:
|
||||||
|
A (np.ndarray): 状态矩阵 (dx x dx)
|
||||||
|
B (np.ndarray): 输入矩阵 (dx x du)
|
||||||
|
C (np.ndarray): 观测矩阵 (dy x dx)
|
||||||
|
D (np.ndarray): 前馈矩阵 (dy x du)
|
||||||
|
T (int): 轨迹长度 (时间步数)
|
||||||
|
sigma_process (float): 过程噪声的标准差 (sigma_Sigma)
|
||||||
|
sigma_measurement (float): 测量噪声的标准差 (sigma_Gamma)
|
||||||
|
rng_seed (int, optional): 用于复现的随机种子
|
||||||
|
|
||||||
|
返回:
|
||||||
|
tuple: (u_data, y_data)
|
||||||
|
u_data (np.ndarray): 输入轨迹 (T x du)
|
||||||
|
y_data (np.ndarray): 输出轨迹 (T x dy)
|
||||||
|
"""
|
||||||
|
|
||||||
|
# 初始化随机数生成器
|
||||||
|
if rng_seed is None:
|
||||||
|
rng = np.random.default_rng()
|
||||||
|
else:
|
||||||
|
rng = np.random.default_rng(rng_seed)
|
||||||
|
|
||||||
|
# 从矩阵形状获取维度
|
||||||
|
dx = A.shape[0]
|
||||||
|
du = B.shape[1]
|
||||||
|
dy = C.shape[0]
|
||||||
|
|
||||||
|
# 初始化状态向量 x_0 = 0
|
||||||
|
x = np.zeros((dx, 1))
|
||||||
|
|
||||||
|
# 初始化用于存储历史数据的列表
|
||||||
|
x_history = []
|
||||||
|
y_history = []
|
||||||
|
u_history = []
|
||||||
|
|
||||||
|
# 生成噪声序列
|
||||||
|
# 输入 u_t ~ N(0, I)
|
||||||
|
u_data_gen = rng.standard_normal(size=(T, du, 1))
|
||||||
|
|
||||||
|
# 过程噪声 w_t ~ N(0, sigma_process^2 * I) [cite: 523]
|
||||||
|
w_data_gen = rng.normal(scale=sigma_process, size=(T, dx, 1))
|
||||||
|
|
||||||
|
# 测量噪声 z_t ~ N(0, sigma_measurement^2 * I) [cite: 523]
|
||||||
|
z_data_gen = rng.normal(scale=sigma_measurement, size=(T, dy, 1))
|
||||||
|
|
||||||
|
print(f"\n--- 开始仿真数据 (T={T}) ---")
|
||||||
|
print(f"过程噪声 (sigma_Sigma): {sigma_process}")
|
||||||
|
print(f"测量噪声 (sigma_Gamma): {sigma_measurement}")
|
||||||
|
|
||||||
|
# 循环 T 个时间步
|
||||||
|
for t in range(T):
|
||||||
|
u = u_data_gen[t]
|
||||||
|
w = w_data_gen[t]
|
||||||
|
z = z_data_gen[t]
|
||||||
|
|
||||||
|
# 1. 计算当前输出 y_t = C*x_t + D*u_t + z_t
|
||||||
|
y = C @ x + D @ u + z
|
||||||
|
|
||||||
|
# 2. 计算下一个状态 x_{t+1} = A*x_t + B*u_t + w_t
|
||||||
|
x_next = A @ x + B @ u + w
|
||||||
|
|
||||||
|
# 存储数据
|
||||||
|
u_history.append(u.squeeze())
|
||||||
|
y_history.append(y.squeeze())
|
||||||
|
x_history.append(x.squeeze())
|
||||||
|
|
||||||
|
# 更新状态
|
||||||
|
x = x_next
|
||||||
|
|
||||||
|
print("仿真完成。")
|
||||||
|
|
||||||
|
# 将列表转换为 numpy 数组
|
||||||
|
# 我们需要 (T, du) 和 (T, dy) 的形状
|
||||||
|
# 使用 .reshape(-1, du) 和 .reshape(-1, dy) 来处理 du/dy=1 的情况
|
||||||
|
u_data = np.array(u_history).reshape(T, du)
|
||||||
|
y_data = np.array(y_history).reshape(T, dy)
|
||||||
|
|
||||||
|
return u_data, y_data
|
||||||
|
|
||||||
|
if __name__ == '__main__':
|
||||||
|
# --- 运行示例 ---
|
||||||
|
|
||||||
|
# 1. 生成 Ground Truth 系统
|
||||||
|
# 使用与第一步相同的种子,确保系统一致
|
||||||
|
A_true, B_true, C_true, D_true = generate_ground_truth_system(rng_seed=42)
|
||||||
|
|
||||||
|
# 2. 仿真数据
|
||||||
|
# 根据 6.3 节的设置
|
||||||
|
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=123
|
||||||
|
)
|
||||||
|
|
||||||
|
# 打印结果形状
|
||||||
|
print(f"\n--- 仿真结果 ---")
|
||||||
|
print(f"输入数据 u_data 形状: {u_data.shape}")
|
||||||
|
print(f"输出数据 y_data 形状: {y_data.shape}")
|
||||||
|
|
||||||
|
# 可视化检查
|
||||||
|
plt.figure(figsize=(12, 4))
|
||||||
|
plt.plot(y_data, label=f'Simulated Output $y_t$ (noise $\sigma_w={sigma_proc}$)')
|
||||||
|
plt.title('Simulated Data Trajectory (First Output Channel)')
|
||||||
|
plt.xlabel('Time Step $t$')
|
||||||
|
plt.ylabel('Output $y_t$')
|
||||||
|
plt.legend()
|
||||||
|
plt.grid(True)
|
||||||
|
plt.show()
|
||||||
@@ -0,0 +1,313 @@
|
|||||||
|
import numpy as np
|
||||||
|
import jax
|
||||||
|
import jax.numpy as jnp
|
||||||
|
import matplotlib.pyplot as plt
|
||||||
|
from matplotlib.patches import Circle
|
||||||
|
import seaborn as sns
|
||||||
|
import arviz as az
|
||||||
|
import numpyro
|
||||||
|
|
||||||
|
from generateGroudTruth import generate_ground_truth_system
|
||||||
|
from generateSimData import simulate_lti_data
|
||||||
|
# 确保 models_and_mcmc.py 中的 run_mcmc 接受 init_params 参数
|
||||||
|
from models_and_mcmc import run_mcmc, model_canonical, model_standard
|
||||||
|
|
||||||
|
# =========================================================================
|
||||||
|
# 配置中文字体
|
||||||
|
# =========================================================================
|
||||||
|
try:
|
||||||
|
# Windows 系统优先尝试这些字体
|
||||||
|
plt.rcParams['font.sans-serif'] = ['Microsoft YaHei', 'SimHei', 'SimSun', 'KaiTi', 'FangSong', 'Arial Unicode MS']
|
||||||
|
plt.rcParams['axes.unicode_minus'] = False # 正常显示负号
|
||||||
|
print("✓ 中文字体配置成功")
|
||||||
|
except Exception as e:
|
||||||
|
print(f"⚠ 字体配置警告: {e}")
|
||||||
|
print(" 如果图表中文显示异常,请运行 check_fonts.py 查看可用字体")
|
||||||
|
|
||||||
|
# =========================================================================
|
||||||
|
# 实现标准型到规范型的转换
|
||||||
|
# =========================================================================
|
||||||
|
def standard_to_canonical(A_s, B_s, C_s):
|
||||||
|
"""
|
||||||
|
将一个 2x2 的标准状态空间系统 (A_s, B_s, C_s) 转换为控制器规范型参数。
|
||||||
|
|
||||||
|
返回:
|
||||||
|
dict: 包含 'a0', 'a1', 'b0', 'b1' 的字典
|
||||||
|
"""
|
||||||
|
dx = A_s.shape[0]
|
||||||
|
if dx != 2:
|
||||||
|
raise ValueError("此转换函数仅为 dx=2 的情况实现。")
|
||||||
|
|
||||||
|
# 1. 计算特征多项式系数: p(λ) = λ^2 + a1*λ + a0
|
||||||
|
a1_true = -np.trace(A_s)
|
||||||
|
a0_true = np.linalg.det(A_s)
|
||||||
|
|
||||||
|
# 2. 构造转换矩阵 T_c
|
||||||
|
I = np.eye(dx)
|
||||||
|
f1 = (A_s + a1_true * I) @ B_s
|
||||||
|
f2 = B_s
|
||||||
|
|
||||||
|
Tc_inv = np.hstack([f1, f2])
|
||||||
|
|
||||||
|
if np.linalg.matrix_rank(Tc_inv) < dx:
|
||||||
|
raise np.linalg.LinAlgError("系统 (A_s, B_s) 不是可控的,无法转换为控制器规范型。")
|
||||||
|
|
||||||
|
Tc = np.linalg.inv(Tc_inv)
|
||||||
|
|
||||||
|
# 3. 转换 C 矩阵: C_c = C_s * T_c
|
||||||
|
C_c = C_s @ Tc
|
||||||
|
b0_true = C_c[0, 0]
|
||||||
|
b1_true = C_c[0, 1]
|
||||||
|
|
||||||
|
true_params = {'a0': a0_true, 'a1': a1_true, 'b0': b0_true, 'b1': b1_true}
|
||||||
|
|
||||||
|
print("\n--- Ground Truth 转换结果 ---")
|
||||||
|
print(f"真实规范型参数: {true_params}")
|
||||||
|
|
||||||
|
return true_params
|
||||||
|
|
||||||
|
# =========================================================================
|
||||||
|
# 可视化函数
|
||||||
|
# =========================================================================
|
||||||
|
|
||||||
|
def plot_canonical_results(mcmc_canonical, true_params_c):
|
||||||
|
"""
|
||||||
|
可视化规范型模型的 MCMC 结果,重现图 2。
|
||||||
|
"""
|
||||||
|
print("\n--- 正在生成规范型模型的后验分布图 (图 2)... ---")
|
||||||
|
|
||||||
|
idata_c = az.from_numpyro(mcmc_canonical)
|
||||||
|
samples_c = mcmc_canonical.get_samples()
|
||||||
|
|
||||||
|
# 图 2(a): 参数的配对图
|
||||||
|
az.plot_pair(
|
||||||
|
idata_c,
|
||||||
|
var_names=['a0', 'a1', 'b0', 'b1'],
|
||||||
|
kind='kde',
|
||||||
|
marginals=True,
|
||||||
|
point_estimate='mean',
|
||||||
|
reference_values=true_params_c,
|
||||||
|
figsize=(10, 10)
|
||||||
|
)
|
||||||
|
plt.suptitle("图 2(a) 复现: 规范型参数的后验分布", y=1.02, fontsize=16)
|
||||||
|
|
||||||
|
# 图 2(b): 特征值在复平面上的分布
|
||||||
|
true_eigenvalues = np.roots([1, true_params_c['a1'], true_params_c['a0']])
|
||||||
|
|
||||||
|
posterior_eigenvalues = []
|
||||||
|
# 避免使用过多样本导致计算缓慢,可以对样本进行降采样
|
||||||
|
num_plot_samples = min(5000, len(samples_c['a0']))
|
||||||
|
plot_indices = np.random.choice(len(samples_c['a0']), num_plot_samples, replace=False)
|
||||||
|
|
||||||
|
for i in plot_indices:
|
||||||
|
a0 = samples_c['a0'][i]
|
||||||
|
a1 = samples_c['a1'][i]
|
||||||
|
posterior_eigenvalues.extend(np.roots([1, a1, a0]))
|
||||||
|
posterior_eigenvalues = np.array(posterior_eigenvalues, dtype=np.complex128)
|
||||||
|
|
||||||
|
mean_params = {k: np.mean(v) for k, v in samples_c.items()}
|
||||||
|
map_eigenvalues = np.roots([1, mean_params['a1'], mean_params['a0']])
|
||||||
|
|
||||||
|
plt.figure(figsize=(8, 8))
|
||||||
|
sns.kdeplot(x=posterior_eigenvalues.real, y=posterior_eigenvalues.imag,
|
||||||
|
fill=True, cmap="Blues", levels=10, alpha=0.7) # 添加透明度
|
||||||
|
plt.plot(true_eigenvalues.real, true_eigenvalues.imag, 'ro', markersize=10,
|
||||||
|
label=f'真实特征值: {true_eigenvalues[0]:.3f}')
|
||||||
|
plt.plot(map_eigenvalues.real, map_eigenvalues.imag, 'gs', markersize=10,
|
||||||
|
label=f'后验均值估计: {map_eigenvalues[0]:.3f}')
|
||||||
|
circle = Circle((0, 0), 1, color='gray', fill=False, linestyle='--')
|
||||||
|
plt.gca().add_artist(circle)
|
||||||
|
|
||||||
|
plt.title('图 2(b) 复现: 主特征值的后验分布', fontsize=16)
|
||||||
|
plt.xlabel('Re(λ)')
|
||||||
|
plt.ylabel('Im(λ)')
|
||||||
|
plt.legend()
|
||||||
|
plt.axis('equal')
|
||||||
|
plt.grid(True)
|
||||||
|
|
||||||
|
|
||||||
|
def plot_standard_results(mcmc_standard, true_params_s):
|
||||||
|
"""
|
||||||
|
可视化标准型模型的 MCMC 结果,重现图 3。
|
||||||
|
"""
|
||||||
|
print("\n--- 正在生成标准型模型的后验分布图 (图 3)... ---")
|
||||||
|
|
||||||
|
idata_s = az.from_numpyro(mcmc_standard)
|
||||||
|
|
||||||
|
# 为了清晰起见,只选择论文图 3a 中显示的几个参数子集
|
||||||
|
plot_vars = ['A11', 'A12', 'A21', 'A22', 'B1', 'B2', 'C1', 'C2']
|
||||||
|
# 限制绘制的参数数量,避免图像过于拥挤
|
||||||
|
az.plot_pair(
|
||||||
|
idata_s,
|
||||||
|
var_names=['A11', 'A12', 'A21','A22', 'B1', 'B2', 'C1','C2'], # 选择部分参数展示
|
||||||
|
kind='kde',
|
||||||
|
marginals=True,
|
||||||
|
# point_estimate='mean', # 对于多峰分布,均值可能误导,不显示
|
||||||
|
reference_values={k: v for k, v in true_params_s.items() if k in plot_vars},
|
||||||
|
figsize=(12, 12) # 稍微增大图像尺寸
|
||||||
|
)
|
||||||
|
plt.suptitle("图 3(a) 复现: 标准型参数的后验分布 (部分)", y=1.02, fontsize=16)
|
||||||
|
|
||||||
|
def get_user_choice():
|
||||||
|
"""
|
||||||
|
获取用户选择运行哪个模型
|
||||||
|
"""
|
||||||
|
print("\n" + "="*60)
|
||||||
|
print("请选择要运行的模型类型:")
|
||||||
|
print("1. 只运行规范型模型 (Canonical)")
|
||||||
|
print("2. 只运行标准型模型 (Standard ABCD)")
|
||||||
|
print("3. 两个模型都运行")
|
||||||
|
print("="*60)
|
||||||
|
|
||||||
|
while True:
|
||||||
|
try:
|
||||||
|
choice = input("请输入选择 (1/2/3): ").strip()
|
||||||
|
if choice in ['1', '2', '3']:
|
||||||
|
return int(choice)
|
||||||
|
else:
|
||||||
|
print("请输入有效的选择 (1, 2 或 3)")
|
||||||
|
except KeyboardInterrupt:
|
||||||
|
print("\n程序被用户中断")
|
||||||
|
exit()
|
||||||
|
except:
|
||||||
|
print("请输入有效的选择 (1, 2 或 3)")
|
||||||
|
|
||||||
|
# =========================================================================
|
||||||
|
# 主程序
|
||||||
|
# =========================================================================
|
||||||
|
if __name__ == '__main__':
|
||||||
|
# 获取用户选择
|
||||||
|
choice = get_user_choice()
|
||||||
|
|
||||||
|
# --- 0. 设置 ---
|
||||||
|
numpyro.set_host_device_count(4)
|
||||||
|
main_rng_key = jax.random.PRNGKey(42)
|
||||||
|
|
||||||
|
# --- 1. 生成 Ground Truth 系统 ---
|
||||||
|
print("\n--- (步骤 1) 生成真实系统 ---")
|
||||||
|
gt_key, sim_key, mcmc_key, init_noise_key = jax.random.split(main_rng_key, 4) # 多分配一个 key
|
||||||
|
A_true, B_true, C_true, D_true = generate_ground_truth_system(
|
||||||
|
rng_seed=int(gt_key[0])
|
||||||
|
)
|
||||||
|
|
||||||
|
# --- 2. 仿真数据 ---
|
||||||
|
print("\n--- (步骤 2) 生成仿真数据 ---")
|
||||||
|
T_steps, sigma_proc, sigma_meas = 400, 0.3, 0.5
|
||||||
|
u_data, y_data = simulate_lti_data(
|
||||||
|
A_true, B_true, C_true, D_true, T_steps, sigma_proc, sigma_meas,
|
||||||
|
rng_seed=int(sim_key[0])
|
||||||
|
)
|
||||||
|
u_data_jax, y_data_jax = jnp.array(u_data), jnp.array(y_data)
|
||||||
|
|
||||||
|
# --- 步骤 2.5: 计算 MCMC 初始值 ---
|
||||||
|
true_params_c = standard_to_canonical(A_true, B_true, C_true)
|
||||||
|
true_params_s = {
|
||||||
|
'A11': A_true[0, 0], 'A12': A_true[0, 1],
|
||||||
|
'A21': A_true[1, 0], 'A22': A_true[1, 1],
|
||||||
|
'B1': B_true[0, 0], 'B2': B_true[1, 0],
|
||||||
|
'C1': C_true[0, 0], 'C2': C_true[0, 1]
|
||||||
|
}
|
||||||
|
|
||||||
|
# --- 3 & 4. 运行 MCMC ---
|
||||||
|
mcmc_canonical = None
|
||||||
|
mcmc_standard = None
|
||||||
|
|
||||||
|
# 分配用于初始值噪声的 key
|
||||||
|
init_noise_key_c, init_noise_key_s = jax.random.split(init_noise_key)
|
||||||
|
|
||||||
|
if choice in [1, 3]: # 运行规范型
|
||||||
|
print("\n--- 运行规范型模型 MCMC (在真实值附近初始化) ---")
|
||||||
|
mcmc_key_c, mcmc_key = jax.random.split(mcmc_key)
|
||||||
|
|
||||||
|
# --- (修改) 在真实值附近添加小的随机扰动 (±10%) ---
|
||||||
|
init_key = init_noise_key_c # 使用独立的 key
|
||||||
|
noise_scale = 0.1 # 10% 的扰动
|
||||||
|
init_params_c_noisy = {}
|
||||||
|
for param_name, true_value in true_params_c.items():
|
||||||
|
# 确保即使 true_value 为 0 也有扰动,添加一个小的基准值
|
||||||
|
base_value = abs(true_value) if abs(true_value) > 1e-6 else 1.0
|
||||||
|
noise = jax.random.normal(init_key, shape=()) * base_value * noise_scale
|
||||||
|
# 确保 a0, a1 扰动后仍在有效范围内
|
||||||
|
if param_name == 'a0':
|
||||||
|
noisy_val = np.clip(true_value + float(noise), -0.99, 0.99) # 限制在 (-1, 1) 内
|
||||||
|
elif param_name == 'a1':
|
||||||
|
# 需要知道 a0 的扰动值来确定 a1 的范围
|
||||||
|
a0_noisy = init_params_c_noisy.get('a0', true_params_c['a0']) # 获取已扰动的 a0
|
||||||
|
low_bound = -1 - a0_noisy + 1e-6 # 加一点边界防止卡住
|
||||||
|
high_bound = 1 + a0_noisy - 1e-6
|
||||||
|
noisy_val = np.clip(true_value + float(noise), low_bound, high_bound)
|
||||||
|
else:
|
||||||
|
noisy_val = true_value + float(noise)
|
||||||
|
|
||||||
|
init_params_c_noisy[param_name] = noisy_val
|
||||||
|
init_key, _ = jax.random.split(init_key) # 更新 key
|
||||||
|
|
||||||
|
print(f"规范型初始参数 (带扰动): {init_params_c_noisy}")
|
||||||
|
|
||||||
|
mcmc_canonical = run_mcmc(
|
||||||
|
model_canonical,
|
||||||
|
mcmc_key_c,
|
||||||
|
u_data_jax,
|
||||||
|
y_data_jax,
|
||||||
|
sigma_proc,
|
||||||
|
init_params=init_params_c_noisy # <-- 使用带扰动的初始值
|
||||||
|
)
|
||||||
|
|
||||||
|
if choice in [2, 3]: # 运行标准型
|
||||||
|
print("\n--- 运行标准型模型 MCMC (在真实值附近初始化) ---")
|
||||||
|
mcmc_key_s, mcmc_key = jax.random.split(mcmc_key)
|
||||||
|
|
||||||
|
# --- (修改) 在真实值附近添加小的随机扰动 (±10%) ---
|
||||||
|
init_key = init_noise_key_s # 使用独立的 key
|
||||||
|
noise_scale = 0.1 # 10% 的扰动
|
||||||
|
init_params_s_noisy = {}
|
||||||
|
for param_name, true_value in true_params_s.items():
|
||||||
|
# 确保即使 true_value 为 0 也有扰动
|
||||||
|
base_value = abs(true_value) if abs(true_value) > 1e-6 else 1.0
|
||||||
|
noise = jax.random.normal(init_key, shape=()) * base_value * noise_scale
|
||||||
|
init_params_s_noisy[param_name] = true_value + float(noise)
|
||||||
|
init_key, _ = jax.random.split(init_key) # 更新 key
|
||||||
|
|
||||||
|
print(f"标准型初始参数 (带扰动): {init_params_s_noisy}")
|
||||||
|
|
||||||
|
mcmc_standard = run_mcmc(
|
||||||
|
model_standard,
|
||||||
|
mcmc_key_s,
|
||||||
|
u_data_jax,
|
||||||
|
y_data_jax,
|
||||||
|
sigma_proc,
|
||||||
|
init_params=init_params_s_noisy # <-- 使用带扰动的初始值
|
||||||
|
)
|
||||||
|
|
||||||
|
# --- 5. 分析与可视化 ---
|
||||||
|
|
||||||
|
# (计算真实参数的代码已移至 步骤 2.5)
|
||||||
|
|
||||||
|
# 根据选择显示结果
|
||||||
|
if choice in [1, 3] and mcmc_canonical is not None:
|
||||||
|
plot_canonical_results(mcmc_canonical, true_params_c)
|
||||||
|
|
||||||
|
print("\n" + "="*50)
|
||||||
|
print("规范型模型分析:")
|
||||||
|
print("1. 后验分布应为单峰 (unimodal) 且近似高斯。")
|
||||||
|
print("2. 参数解释简单,后验均值是好的点估计。")
|
||||||
|
print("3. 特征值分布应集中围绕真实值。")
|
||||||
|
print("="*50)
|
||||||
|
|
||||||
|
if choice in [2, 3] and mcmc_standard is not None:
|
||||||
|
plot_standard_results(mcmc_standard, true_params_s)
|
||||||
|
|
||||||
|
print("\n" + "="*50)
|
||||||
|
print("标准型模型分析:")
|
||||||
|
print("1. 后验分布可能呈现复杂的多峰 (multi-modal) 和强相关性。")
|
||||||
|
print("2. 这是由参数的非唯一性 (non-identifiability) 造成的。")
|
||||||
|
print("3. MCMC 采样效率可能较低,点估计(如均值)可能无意义。")
|
||||||
|
print("="*50)
|
||||||
|
|
||||||
|
# 显示所有图像
|
||||||
|
if choice != 0:
|
||||||
|
print("\n显示所有图像...")
|
||||||
|
plt.show()
|
||||||
|
|
||||||
|
print("\n程序执行完成!")
|
||||||
|
|
||||||
@@ -0,0 +1,283 @@
|
|||||||
|
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}, ...)
|
||||||
|
log_lik_t = multivariate_normal.logpdf(nu_t.squeeze(), mean=jnp.zeros(dy), cov=S_t)
|
||||||
|
|
||||||
|
# 卡尔曼增益 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 [cite: 21, 1042]
|
||||||
|
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, nugget=1e-12):
|
||||||
|
"""
|
||||||
|
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, 1))
|
||||||
|
b1 = numpyro.sample("b1", dist.Normal(0, 1))
|
||||||
|
|
||||||
|
# --- 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)
|
||||||
|
# 为数值稳定性添加 "nugget" [cite: 533]
|
||||||
|
R = jnp.eye(dy) * nugget
|
||||||
|
|
||||||
|
# --- 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, nugget=1e-12):
|
||||||
|
"""
|
||||||
|
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, 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))
|
||||||
|
|
||||||
|
# 输入矩阵 B (dx*du = 2 个参数)
|
||||||
|
B1 = numpyro.sample("B1", dist.Normal(0, 1))
|
||||||
|
B2 = numpyro.sample("B2", dist.Normal(0, 1))
|
||||||
|
|
||||||
|
# 观测矩阵 C (dy*dx = 2 个参数)
|
||||||
|
C1 = numpyro.sample("C1", dist.Normal(0, 1))
|
||||||
|
C2 = numpyro.sample("C2", dist.Normal(0, 1))
|
||||||
|
|
||||||
|
# --- 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) * nugget
|
||||||
|
|
||||||
|
# --- 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, init_params=None):
|
||||||
|
"""辅助函数,用于运行 NUTS 采样器"""
|
||||||
|
|
||||||
|
print(f"\n--- 开始为模型 {model.__name__} 运行 MCMC ---")
|
||||||
|
|
||||||
|
# 论文中的 MCMC 设置
|
||||||
|
num_warmup = 5000
|
||||||
|
num_samples = 20000
|
||||||
|
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)
|
||||||
|
|
||||||
|
# 打印总结
|
||||||
|
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
|
||||||
|
)
|
||||||
|
|
||||||
|
# (模型 2: 标准型)
|
||||||
|
mcmc_standard = run_mcmc(
|
||||||
|
model_standard,
|
||||||
|
mcmc_key_s,
|
||||||
|
u_data_jax,
|
||||||
|
y_data_jax,
|
||||||
|
sigma_proc
|
||||||
|
)
|
||||||
|
|
||||||
|
print("\n--- MCMC 运行完成 ---")
|
||||||
|
print("下一步是分析和可视化后验分布 (步骤 5)。")
|
||||||
|
print("例如,使用 'mcmc_canonical.get_samples()' 来获取样本。")
|
||||||
Reference in New Issue
Block a user