Upload bayesian_utils.py

bayesian_utils.py

@@ -53,22 +53,33 @@ def plot_trace(trace, var_names=['d', 'sigma']):
             n_warmup = warmup_data.shape[1]
             n_post = post_data.shape[1]
             
+            # 定義顏色,讓每條鏈用固定顏色
+            colors = plt.cm.tab10.colors  # 使用 matplotlib 的顏色循環
+            
             for chain_idx in range(warmup_data.shape[0]):
+                chain_color = colors[chain_idx % len(colors)]  # 每條鏈一個固定顏色
+                
                 # 繪 warmup 部分
                 x_warmup = np.arange(n_warmup)
                 axes[idx, 1].plot(x_warmup, warmup_data[chain_idx].flatten(), 
+                                color=chain_color,  # 👈 指定顏色
                                 alpha=0.7, linewidth=0.5,
                                 label=f'Chain {chain_idx+1}' if idx == 0 else '')
                 
-                # 繪 posterior 部分
+                # 繪 posterior 部分 (用同樣的顏色!)
                 x_post = np.arange(n_warmup, n_warmup + n_post)
                 axes[idx, 1].plot(x_post, post_data[chain_idx].flatten(), 
+                                color=chain_color,  # 👈 同一個顏色
                                 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:
             # 沒有 warmup: 只用 posterior
             post_data = trace.posterior[var_name].values