Upload bayesian_core.py

bayesian_core.py

@@ -349,14 +349,129 @@ class BayesianHierarchicalAnalyzer:
             'heterogeneity': heterogeneity
         }
     
-    def get_model_graph(self):
-        """生成模型 DAG 圖（返回 graphviz 物件）"""
+    def get_model_graph(self, language='zh'):
+        """
+        生成模型 DAG 圖（返回 graphviz 物件）
+        
+        Args:
+            language: 'zh' 中文 | 'en' 英文
+            
+        Returns:
+            graphviz.Digraph 物件
+        """
         if self.model is None:
             raise ValueError("請先執行分析")
         
         try:
-            gv = pm.model_to_graphviz(self.model)
-            return gv
+            import graphviz
+            
+            # 獲取欄位名稱
+            control_prefix = getattr(self, 'col_control_win', 'control').replace('_win', '').replace('_battles', '').replace('_total', '')
+            treatment_prefix = getattr(self, 'col_treatment_win', 'treatment').replace('_win', '').replace('_battles', '').replace('_total', '')
+            control_win = getattr(self, 'col_control_win', 'control_win')
+            treatment_win = getattr(self, 'col_treatment_win', 'treatment_win')
+            
+            # 定義中英文對照
+            labels_zh = {
+                # 先驗參數
+                'd': f'整體效應 (d)\n{treatment_prefix} vs {control_prefix}',
+                'tau': '精度 (tau)',
+                'sigma': '配對間標準差 (σ)',
+                
+                # 階層參數
+                'mu': f'基線對數勝算 (μ[i])\n各道館基準',
+                'delta': f'配對特定效應 (δ[i])\n各道館{treatment_prefix}優勢',
+                'delta_new': f'新配對預測 (δ_new)',
+                
+                # 轉換參數
+                f'p_{control_prefix}': f'{control_prefix.capitalize()}勝率 (p_{control_prefix}[i])',
+                f'p_{treatment_prefix}': f'{treatment_prefix.capitalize()}勝率 (p_{treatment_prefix}[i])',
+                
+                # 觀測值
+                f'{control_win}_obs': f'{control_prefix.capitalize()}勝場觀測 ({control_win}_obs[i])',
+                f'{treatment_win}_obs': f'{treatment_prefix.capitalize()}勝場觀測 ({treatment_win}_obs[i])',
+                
+                # 其他
+                'or_speed': f'勝算比 (OR)\nexp(d)'
+            }
+            
+            labels_en = {
+                # Priors
+                'd': f'Overall Effect (d)\n{treatment_prefix} vs {control_prefix}',
+                'tau': 'Precision (tau)',
+                'sigma': 'Between-Pair SD (σ)',
+                
+                # Hierarchy
+                'mu': f'Baseline Log-Odds (μ[i])\nGym Baselines',
+                'delta': f'Pair-Specific Effect (δ[i])\n{treatment_prefix.capitalize()} Advantage per Gym',
+                'delta_new': f'New Pair Prediction (δ_new)',
+                
+                # Transformations
+                f'p_{control_prefix}': f'{control_prefix.capitalize()} Win Rate (p_{control_prefix}[i])',
+                f'p_{treatment_prefix}': f'{treatment_prefix.capitalize()} Win Rate (p_{treatment_prefix}[i])',
+                
+                # Observations
+                f'{control_win}_obs': f'{control_prefix.capitalize()} Wins Obs ({control_win}_obs[i])',
+                f'{treatment_win}_obs': f'{treatment_prefix.capitalize()} Wins Obs ({treatment_win}_obs[i])',
+                
+                # Others
+                'or_speed': f'Odds Ratio (OR)\nexp(d)'
+            }
+            
+            # 選擇語言
+            labels = labels_zh if language == 'zh' else labels_en
+            
+            # 使用 PyMC 生成基本圖
+            gv_original = pm.model_to_graphviz(self.model)
+            
+            # 創建新圖,使用自定義標籤
+            dot = graphviz.Digraph(comment='Bayesian Hierarchical Model')
+            dot.attr(rankdir='TB')  # 從上到下
+            dot.attr('node', shape='ellipse', style='filled', fontname='Arial')
+            
+            # 解析原始圖的結構
+            lines = gv_original.source.split('\n')
+            
+            # 重建圖,替換標籤
+            for line in lines:
+                line = line.strip()
+                
+                if not line or line.startswith('digraph') or line == '}':
+                    continue
+                
+                # 處理節點定義
+                if '[label=' in line:
+                    # 提取節點名稱
+                    node_name = line.split('[')[0].strip()
+                    
+                    # 使用自定義標籤
+                    custom_label = labels.get(node_name, node_name)
+                    
+                    # 設定節點樣式
+                    if 'obs' in node_name:
+                        # 觀測值 - 方框,淺紅色
+                        dot.node(node_name, custom_label, shape='box', fillcolor='lightcoral')
+                    elif node_name in ['sigma', f'p_{control_prefix}', f'p_{treatment_prefix}']:
+                        # 確定性節點 - 菱形,灰色
+                        dot.node(node_name, custom_label, shape='diamond', fillcolor='lightgray')
+                    elif node_name in ['mu', 'delta']:
+                        # 階層參數 - 橢圓,綠色
+                        dot.node(node_name, custom_label, fillcolor='lightgreen')
+                    else:
+                        # 先驗參數 - 橢圓,黃色
+                        dot.node(node_name, custom_label, fillcolor='lightyellow')
+                
+                # 處理邊
+                elif '->' in line:
+                    # 提取邊的定義
+                    edge_parts = line.replace(';', '').split('->')
+                    if len(edge_parts) == 2:
+                        from_node = edge_parts[0].strip()
+                        to_node = edge_parts[1].strip().split('[')[0].strip()
+                        dot.edge(from_node, to_node)
+            
+            return dot
+            
         except Exception as e:
             raise Exception(f"無法生成 DAG 圖: {str(e)}")
     
@@ -369,4 +484,4 @@ class BayesianHierarchicalAnalyzer:
     def clear_session_results(cls, session_id):
         """清除特定 session 的結果"""
         if session_id in cls._session_results:
-            del cls._session_results[session_id]
+            del cls._session_results[session_id]