Upload bayesian_utils.py

bayesian_utils.py

@@ -13,6 +13,7 @@ from PIL import Image
 def plot_trace(trace, var_names=['d', 'sigma']):
     """
     繪製 Trace Plot（MCMC 收斂診斷）
+    包含完整的 warmup + posterior
     
     Args:
         trace: ArviZ InferenceData 物件
@@ -24,21 +25,65 @@ def plot_trace(trace, var_names=['d', 'sigma']):
     fig, axes = plt.subplots(len(var_names), 2, figsize=(14, 4 * len(var_names)))
     if len(var_names) == 1:
         axes = axes.reshape(1, -1)
-       
-    az.plot_trace(trace, var_names=var_names, axes=axes, 
-                  combined=False,
-                  compact=False)
-
-    # 加上 Tune 結束的垂直線
-    for i in range(len(var_names)):
-        if i == 0:
-            axes[i, 1].axvline(x=1000, color='red', linestyle='--', 
-                               linewidth=2, alpha=0.7, label='Tune結束')
-            axes[i, 1].legend(loc='upper right', fontsize=9)
+    
+    # 檢查是否有 warmup_posterior
+    has_warmup = hasattr(trace, 'warmup_posterior') and trace.warmup_posterior is not None
+    
+    for idx, var_name in enumerate(var_names):
+        # 左圖: KDE 密度圖（只用 posterior, 不用 warmup）
+        post_data = trace.posterior[var_name].values
+        for chain_idx in range(post_data.shape[0]):
+            from scipy import stats
+            data = post_data[chain_idx].flatten()
+            density = stats.gaussian_kde(data)
+            xs = np.linspace(data.min(), data.max(), 200)
+            axes[idx, 0].plot(xs, density(xs), alpha=0.8, label=f'Chain {chain_idx+1}')
+        axes[idx, 0].set_xlabel(var_name, fontsize=12)
+        axes[idx, 0].set_ylabel('Density', fontsize=12)
+        axes[idx, 0].set_title(f'{var_name}', fontsize=13, fontweight='bold')
+        if idx == 0:
+            axes[idx, 0].legend()
+        
+        # 右圖: Trace 圖（完整 warmup + posterior）
+        if has_warmup:
+            # 有 warmup: 合併繪製
+            warmup_data = trace.warmup_posterior[var_name].values
+            post_data = trace.posterior[var_name].values
+            
+            n_warmup = warmup_data.shape[1]
+            n_post = post_data.shape[1]
+            
+            for chain_idx in range(warmup_data.shape[0]):
+                # 繪 warmup 部分
+                x_warmup = np.arange(n_warmup)
+                axes[idx, 1].plot(x_warmup, warmup_data[chain_idx].flatten(), 
+                                alpha=0.7, linewidth=0.5,
+                                label=f'Chain {chain_idx+1}' if idx == 0 else '')
+                
+                # 繪 posterior 部分
+                x_post = np.arange(n_warmup, n_warmup + n_post)
+                axes[idx, 1].plot(x_post, post_data[chain_idx].flatten(), 
+                                alpha=0.7, linewidth=0.5)
+            
+            # 加 Tune 結束的紅線
+            axes[idx, 1].axvline(x=n_warmup, color='red', linestyle='--', 
+                               linewidth=2, alpha=0.7, 
+                               label='Tune結束' if idx == 0 else '')
         else:
-            axes[i, 1].axvline(x=1000, color='red', linestyle='--', 
-                               linewidth=2, alpha=0.7)
-
+            # 沒有 warmup: 只用 posterior
+            post_data = trace.posterior[var_name].values
+            for chain_idx in range(post_data.shape[0]):
+                axes[idx, 1].plot(post_data[chain_idx].flatten(), 
+                                alpha=0.7, linewidth=0.5,
+                                label=f'Chain {chain_idx+1}' if idx == 0 else '')
+        
+        axes[idx, 1].set_xlabel('Iteration', fontsize=12)
+        axes[idx, 1].set_ylabel(var_name, fontsize=12)
+        axes[idx, 1].set_title(f'{var_name} trace', fontsize=13, fontweight='bold')
+        if idx == 0:
+            axes[idx, 1].legend(loc='upper right', fontsize=9)
+        axes[idx, 1].grid(alpha=0.3)
+    
     plt.tight_layout()
     
     # 轉換為圖片
@@ -50,6 +95,13 @@ def plot_trace(trace, var_names=['d', 'sigma']):
     
     return img
 
+
+# ============================================
+# 替換說明:
+# 在 bayesian_utils.py 中,把第 13-51 行的整個 plot_trace 函數
+# 替換成上面這個版本
+# ============================================
+
 def plot_posterior(trace, var_names=['d', 'sigma', 'or_speed'], hdi_prob=0.95):
     """
     繪製後驗分佈圖