diff --git a/__pycache__/generateGroudTruth.cpython-311.pyc b/__pycache__/generateGroudTruth.cpython-311.pyc new file mode 100644 index 0000000..15f3aae Binary files /dev/null and b/__pycache__/generateGroudTruth.cpython-311.pyc differ diff --git a/__pycache__/generateSimData.cpython-311.pyc b/__pycache__/generateSimData.cpython-311.pyc new file mode 100644 index 0000000..97b9437 Binary files /dev/null and b/__pycache__/generateSimData.cpython-311.pyc differ diff --git a/__pycache__/models_and_mcmc.cpython-311.pyc b/__pycache__/models_and_mcmc.cpython-311.pyc new file mode 100644 index 0000000..dc78ad5 Binary files /dev/null and b/__pycache__/models_and_mcmc.cpython-311.pyc differ diff --git a/demo.py b/demo.py index b74d4f1..96672a5 100644 --- a/demo.py +++ b/demo.py @@ -1,105 +1,99 @@ -## 实现一下重要抽样 import numpy as np import matplotlib.pyplot as plt -import scipy as stats -import seaborn as sns +from scipy.stats import beta as beta_dist # 导入beta分布用于绘图 -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 = β0 + β1*x + ε, ε ~ N(0, σ²) - = [1 x] * [β0 β1]' + ε + 计算逐步平均值 (累积平均值) + y_k_bar = (1/k) * sum(y_i for i=1 to k) + + 使用 np.cumsum() 可以高效实现 """ + 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. 定义关键数据 - x = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10]) - 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) +# 计算所有 Y 样本的逐步平均值 (包括预热期) +ergodic_mean_Y = calculate_ergodic_mean(samples_Y) - print(f"数据点数量：{n}") - print(f"x: {x}") - print(f"y: {y}") +# ------------------------------------------------------------------ +# 6. 绘制逐步平均值图 (实现您图片中的效果) +# ------------------------------------------------------------------ +print("正在绘制逐步平均值图...") - # 2. 定义先验分布 - # β0 ~ N(0, 10) - beta0_prior_mean = 0 - beta0_prior_var = 10 +plt.figure(figsize=(10, 6)) +plt.plot(np.arange(1, N_samples + 1), ergodic_mean_Y, label=r"逐步平均值 $\bar{y}_k$") +plt.axhline(theoretical_mean, color='red', linestyle='--', label=f"理论均值: {theoretical_mean:.4f}") - # β1 ~ N(1, 5) - beta1_prior_mean = 1 - beta1_prior_var = 5 +# 添加一个垂直线来标记我们估计的预热期 +plt.axvline(N_burn_in, color='gray', linestyle=':', label=f"估计的预热期 N1 = {N_burn_in}") - # σ² ~ Inverse-Gamma(2, 1) - sigma2_prior_a = 2 - sigma2_prior_b = 1 - - # 3. 定义重要抽样分布 - # 使用最小二乘进行估计 - X = np.column_stack([np.ones(n), x]) - 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() - - \ No newline at end of file +plt.title("使用逐步平均值图查看预热期") +plt.xlabel("迭代次数 (k)") +plt.ylabel(r"逐步平均值 $\bar{y}_k$") +plt.legend() +plt.grid(True) +plt.ylim(0, 1) # Y值在0到1之间 +plt.show() diff --git a/demo1.py b/demo1.py index d708046..a6d2f3d 100644 --- a/demo1.py +++ b/demo1.py @@ -1,3 +1,6 @@ +# JAX 入门示例代码 +# 演示 JAX 的基本用法，包括 jnp 数组、jit 加速 + import numpy as np import jax import jax.numpy as jnp diff --git a/generateGroudTruth.py b/generateGroudTruth.py new file mode 100644 index 0000000..9a1959d --- /dev/null +++ b/generateGroudTruth.py @@ -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) diff --git a/generateSimData.py b/generateSimData.py new file mode 100644 index 0000000..a1f2b89 --- /dev/null +++ b/generateSimData.py @@ -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() diff --git a/main.py b/main.py new file mode 100644 index 0000000..13dff69 --- /dev/null +++ b/main.py @@ -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程序执行完成!") + diff --git a/models_and_mcmc.py b/models_and_mcmc.py new file mode 100644 index 0000000..4e82f2f --- /dev/null +++ b/models_and_mcmc.py @@ -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()' 来获取样本。")