| """ |
| MolGate — Module 4: Pocket Definition (Engine B) |
| Stage 22: Ligand identity confirmation (人機) |
| |
| 輸入：stage_21_hetatm.json + Master Index（het_id / cocrystal_het、pocket_rule、可選 common_name、target_drug） |
| 輸出：stage_22_ligand.json（選定配體） |
| 依賴：與 molgate_module4_engineA 共用 Index 解析（molgate_module3_engineA） |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import sys |
| from datetime import datetime |
| from pathlib import Path |
| from typing import Any |
|
|
| _PKG_ROOT = Path(__file__).resolve().parent |
| _UI_ROOT = _PKG_ROOT.parent |
| if str(_UI_ROOT) not in sys.path: |
| sys.path.insert(0, str(_UI_ROOT)) |
| if str(_PKG_ROOT) not in sys.path: |
| sys.path.insert(0, str(_PKG_ROOT)) |
|
|
| from molgate_ui.console_utf8 import configure_console_utf8, should_prompt_interactive |
|
|
| configure_console_utf8() |
|
|
| from molgate_module4_engineA import load_target_index_entry |
|
|
| SCHEMA_VERSION = "0.2" |
|
|
|
|
| def load_prep_overlay(session_dir: Path, target: str) -> dict[str, Any]: |
| """合併 prep_decisions.json 與 CLI target（CLI 優先）。""" |
| prep: dict[str, Any] = {"target": target.strip()} |
| p = session_dir / "prep_decisions.json" |
| if p.is_file(): |
| try: |
| data = json.loads(p.read_text(encoding="utf-8")) |
| if isinstance(data, dict): |
| for k, v in data.items(): |
| if k.startswith("_"): |
| continue |
| prep[k] = v |
| except (json.JSONDecodeError, OSError): |
| pass |
| prep["target"] = target.strip() |
| return prep |
|
|
|
|
| def expected_het_id(entry: dict[str, Any]) -> str: |
| """比對用單一代碼：優先 het_id，否則 cocrystal_het。""" |
| for key in ("het_id", "cocrystal_het"): |
| v = entry.get(key) |
| if isinstance(v, str) and v.strip(): |
| return v.strip().upper() |
| return "" |
|
|
|
|
| |
| def save_stage(session_dir: Path, stage: int, name: str, data: dict[str, Any]) -> dict[str, Any]: |
| data = dict(data) |
| data["_stage"] = stage |
| data["_name"] = name |
| data["_timestamp"] = datetime.now().isoformat() |
| data["_schema_version"] = SCHEMA_VERSION |
| path = session_dir / f"stage_{stage:02d}_{name}.json" |
| path.write_text(json.dumps(data, indent=2, ensure_ascii=False), encoding="utf-8") |
| return data |
|
|
|
|
| def load_stage21(session_dir: Path) -> dict[str, Any]: |
| path = session_dir / "stage_21_hetatm.json" |
| if not path.exists(): |
| raise FileNotFoundError("找不到 stage_21_hetatm.json，請先跑 Engine A") |
| data = json.loads(path.read_text(encoding="utf-8")) |
| if data.get("_schema_version") != SCHEMA_VERSION: |
| raise ValueError("Stage 21 schema version mismatch，請重跑 Engine A") |
| return data |
|
|
|
|
| def prompt_select( |
| candidates: list[dict], |
| label: str, |
| *, |
| non_interactive: bool = False, |
| ) -> dict: |
| if not candidates: |
| raise ValueError(f"無候選可選：{label}") |
| if not should_prompt_interactive(non_interactive=non_interactive): |
| pick = candidates[0] |
| print( |
| f" [AUTO] 非互動：{label} → [0] {pick.get('resname')} " |
| f"chain {pick.get('chain')} res {pick.get('resseq')}" |
| ) |
| return pick |
| print(f"\n ── {label} ──") |
| for i, c in enumerate(candidates): |
| tag = "【目標】 " if c.get("ligand_role") == "target" else "" |
| print( |
| f" [{i}] {tag}{c['resname']} chain {c['chain']} " |
| f"res {c['resseq']} (~{c['mw_estimate']:.0f} Da, " |
| f"{c['heavy_atom_count']} heavy atoms)" |
| ) |
| while True: |
| try: |
| raw = input(" 請輸入編號選擇：").strip() |
| idx = int(raw) |
| if 0 <= idx < len(candidates): |
| return candidates[idx] |
| except (ValueError, KeyboardInterrupt): |
| pass |
| print(" 請輸入有效編號") |
|
|
|
|
| |
| def stage22_ligand( |
| session_dir: Path, |
| target: str, |
| index_path: Path | None, |
| *, |
| non_interactive: bool = False, |
| selected: dict[str, Any] | None = None, |
| ) -> dict[str, Any]: |
| print("\n[Stage 22] 配體身份確認...") |
|
|
| prep = load_prep_overlay(session_dir, target) |
| entry = load_target_index_entry(session_dir, prep, index_path) |
| if not entry: |
| result = { |
| "status": "FAIL", |
| "reason": "無法載入 Master Index 條目（請確認 --target 與 master_index.json）", |
| } |
| save_stage(session_dir, 22, "ligand", result) |
| raise ValueError(f"Stage 22 FAIL: {result['reason']}") |
|
|
| s21 = load_stage21(session_dir) |
| candidates = s21.get("candidates", []) |
| route_fallback_status = "" |
|
|
| if not candidates: |
| try: |
| from molgate_ui.anchor_routing import m4_recover_empty_candidates |
|
|
| recovery = m4_recover_empty_candidates(session_dir, entry) |
| except ImportError: |
| recovery = None |
| if recovery: |
| if recovery.get("needs_rerun_from_m1"): |
| res_obj = recovery.get("resolution") |
| result = { |
| "status": "ROUTE_FALLBACK", |
| "reason": ( |
| "Direct 共晶 HET 不存在；已寫入 cross-PDB proxy 路由，" |
| "請以更新後的 pdb_id 自 Module 1 重跑" |
| ), |
| "route_resolution": res_obj.to_json() if hasattr(res_obj, "to_json") else res_obj, |
| } |
| save_stage(session_dir, 22, "ligand", result) |
| raise ValueError(f"Stage 22 ROUTE_FALLBACK: {result['reason']}") |
| entry = recovery.get("effective_entry") or entry |
| candidates = recovery.get("candidates") or [] |
| route_fallback_status = getattr(recovery.get("resolution"), "reason", "") or "route_fallback" |
| print(f" → Route fallback：{route_fallback_status}") |
|
|
| if not candidates: |
| result = { |
| "status": "FAIL", |
| "reason": "Stage 21 無候選配體，無法確認配體身份", |
| } |
| save_stage(session_dir, 22, "ligand", result) |
| raise ValueError(f"Stage 22 FAIL: {result['reason']}") |
|
|
| het_id = expected_het_id(entry) |
| pocket_rule = (entry.get("pocket_rule") or "").strip() |
| common_name = (entry.get("common_name") or "").strip() |
| target_drug = (entry.get("target_drug") or "").strip() |
| pdb_id = (entry.get("pdb_id") or prep.get("pdb_id") or "").strip().upper() |
|
|
| print(f" → Master Index expected het（het_id／cocrystal_het）：{het_id or '（未設定）'}") |
| if common_name: |
| print(f" → 共晶配體俗名（common_name）：{common_name}") |
| if target_drug: |
| print(f" → Docking 目標藥物（target_drug）：{target_drug}") |
| print(f" → Pocket rule: {pocket_rule}") |
| print(f" → 候選配體數：{len(candidates)}") |
|
|
| if selected is not None: |
| want = ( |
| str(selected.get("resname") or "").strip().upper(), |
| str(selected.get("chain") or "").strip(), |
| int(selected.get("resseq")), |
| ) |
| picked = None |
| for c in candidates: |
| if ( |
| str(c.get("resname") or "").strip().upper() == want[0] |
| and str(c.get("chain") or "").strip() == want[1] |
| and int(c.get("resseq")) == want[2] |
| ): |
| picked = c |
| break |
| if not picked: |
| raise ValueError( |
| f"UI 選定的配體 {want[0]} chain {want[1]} res {want[2]} 不在 Stage 21 候選列表中" |
| ) |
| selected = picked |
| status = "MANUAL" |
| print(f" ✓ UI 確認：{selected['resname']} chain {selected['chain']} res {selected['resseq']}") |
| elif not het_id: |
| print(f" ⚠ 無法自動比對（Index 缺 het_id / cocrystal_het），請手動選擇") |
| if non_interactive: |
| selected = candidates[0] |
| status = "AUTO_FORCED" |
| print(f" ✓ --non-interactive：採用候選 [0] {selected['resname']} chain {selected['chain']}") |
| else: |
| selected = prompt_select( |
| candidates, |
| "請選擇用於口袋定義的配體", |
| non_interactive=non_interactive, |
| ) |
| status = "MANUAL" |
| print(f" ✓ 選定：{selected['resname']} chain {selected['chain']}") |
| else: |
| matched = [c for c in candidates if c["resname"].upper() == het_id] |
|
|
| if len(matched) == 1: |
| selected = matched[0] |
| status = "AUTO" |
| print(f" ✓ 自動比對成功：{selected['resname']} chain {selected['chain']}") |
|
|
| elif len(matched) > 1: |
| print(f" ⚠ 多個 {het_id} 存在於不同 chain，需要確認") |
| if non_interactive: |
| selected = matched[0] |
| status = "AUTO_FORCED" |
| print( |
| f" ✓ --non-interactive：採用排序後第一個 {selected['resname']} " |
| f"chain {selected['chain']} res {selected['resseq']}" |
| ) |
| else: |
| selected = prompt_select( |
| matched, |
| f"選擇 {het_id} 所在 chain", |
| non_interactive=non_interactive, |
| ) |
| status = "MANUAL" |
| print(f" ✓ 選定：{selected['resname']} chain {selected['chain']}") |
|
|
| elif len(candidates) == 1: |
| print( |
| f" ⚠ WARN: 唯一候選 {candidates[0]['resname']} " |
| f"與 Index 期望 {het_id} 不符" |
| ) |
| print(f" → 可能是 PDB／CCD 代碼與 Index 不一致") |
| if non_interactive: |
| selected = candidates[0] |
| status = "AUTO_FORCED" |
| print(f" ✓ --non-interactive：採用唯一候選 {selected['resname']}") |
| else: |
| confirm = input( |
| f" 是否使用 {candidates[0]['resname']} 作為口袋定義配體？[y/N]：" |
| ).strip().lower() |
| if confirm == "y": |
| selected = candidates[0] |
| status = "MANUAL" |
| print(f" ✓ 手動確認：{selected['resname']}") |
| else: |
| result = { |
| "status": "FAIL", |
| "reason": f"用戶拒絕使用 {candidates[0]['resname']}，無法繼續", |
| } |
| save_stage(session_dir, 22, "ligand", result) |
| raise ValueError(f"Stage 22 FAIL: {result['reason']}") |
|
|
| else: |
| print(f" ⚠ 無法自動比對，候選列表如下：") |
| if non_interactive: |
| preferred = [c for c in candidates if c.get("ligand_role") == "target"] |
| selected = (preferred[0] if preferred else candidates[0]) |
| status = "AUTO_FORCED" |
| print( |
| f" ✓ --non-interactive：採用" |
| f"{'【目標】' if selected.get('ligand_role') == 'target' else ''}" |
| f"{selected['resname']} chain {selected['chain']}" |
| ) |
| else: |
| selected = prompt_select( |
| candidates, |
| "請選擇用於口袋定義的配體", |
| non_interactive=non_interactive, |
| ) |
| status = "MANUAL" |
| print(f" ✓ 選定：{selected['resname']} chain {selected['chain']}") |
|
|
| print(f" → Status: {status}") |
|
|
| result: dict[str, Any] = { |
| "status": status, |
| "selected_het_id": selected["resname"], |
| "chain": selected["chain"], |
| "resseq": selected["resseq"], |
| "mw_estimate": selected["mw_estimate"], |
| "heavy_atom_count": selected["heavy_atom_count"], |
| "master_index_expected_het": het_id, |
| "master_index_het_id": het_id, |
| "match": selected["resname"].upper() == het_id if het_id else None, |
| "pocket_rule": pocket_rule, |
| "ready_for_engine_c": True, |
| } |
| if route_fallback_status: |
| result["route_fallback"] = route_fallback_status |
| if pdb_id: |
| result["pdb_id"] = pdb_id |
| if common_name: |
| result["common_name"] = common_name |
| if target_drug: |
| result["target_drug"] = target_drug |
| return save_stage(session_dir, 22, "ligand", result) |
|
|
|
|
| |
| def run_engine_b( |
| session_dir: Path, |
| target: str, |
| index_path: Path | None = None, |
| *, |
| non_interactive: bool = False, |
| selected: dict[str, Any] | None = None, |
| ): |
| print("=" * 55) |
| print(" MolGate — Module 4 Engine B") |
| print(" Stage 22: 配體身份確認（人機）") |
| print("=" * 55) |
|
|
| prep = load_prep_overlay(session_dir, target) |
| entry = load_target_index_entry(session_dir, prep, index_path) |
| if not entry: |
| print(f"❌ Target '{target}' 無法從 Master Index 載入條目（檔案或鍵名錯誤）") |
| sys.exit(1) |
|
|
| pdb_id = (entry.get("pdb_id") or prep.get("pdb_id") or "?").strip() |
| het_hint = expected_het_id(entry) or "?" |
| print(f"\n Target：{target.upper()} | PDB：{pdb_id}") |
| print(f" Index 期望共晶代碼：{het_hint}") |
|
|
| try: |
| result = stage22_ligand( |
| session_dir, |
| target, |
| index_path, |
| non_interactive=non_interactive, |
| selected=selected, |
| ) |
| except FileNotFoundError as e: |
| print(f"❌ {e}") |
| sys.exit(1) |
| except ValueError as e: |
| print(f"\n❌ Pipeline stopped: {e}") |
| sys.exit(1) |
| except Exception as e: |
| print(f"\n❌ Pipeline stopped: {e}") |
| sys.exit(1) |
|
|
| print("\n" + "=" * 55) |
| print(" Engine B 完成 ✓") |
| print( |
| f" 選定配體：{result['selected_het_id']} " |
| f"chain {result['chain']} res {result['resseq']}" |
| ) |
| print(f" ready_for_engine_c = True") |
| print(f" 輸出：{session_dir / 'stage_22_ligand.json'}") |
| print("=" * 55) |
|
|
| return result |
|
|
|
|
| |
| if __name__ == "__main__": |
| ap = argparse.ArgumentParser(description="MolGate Module 4 Engine B") |
| ap.add_argument("--session-dir", required=True, help="Session 目錄路徑") |
| ap.add_argument("--target", required=True, help="Target 名稱（如 COX2，須對應 master_index.json）") |
| ap.add_argument( |
| "--index", |
| default=None, |
| help="master_index.json 路徑（預設：session_meta / manifest 或 MolGate 目錄下）", |
| ) |
| ap.add_argument( |
| "--non-interactive", |
| action="store_true", |
| help="Stage 22 無法唯一自動比對時，採用規則化預設（多命中取第一個／無命中取候選[0]或【目標】）", |
| ) |
| ap.add_argument( |
| "--selected-json", |
| default=None, |
| help="UI 人機選配體 JSON 檔（含 resname、chain、resseq）", |
| ) |
| args = ap.parse_args() |
|
|
| sd = Path(args.session_dir) |
| if not sd.is_dir(): |
| print(f"❌ Session 目錄不存在：{sd}") |
| sys.exit(1) |
|
|
| idx = Path(args.index) if args.index else None |
| ui_selected: dict[str, Any] | None = None |
| if args.selected_json: |
| pick_path = Path(args.selected_json) |
| if not pick_path.is_file(): |
| print(f"❌ 找不到 --selected-json：{pick_path}") |
| sys.exit(1) |
| try: |
| ui_selected = json.loads(pick_path.read_text(encoding="utf-8")) |
| except (json.JSONDecodeError, OSError) as e: |
| print(f"❌ 無法讀取 --selected-json：{e}") |
| sys.exit(1) |
| if not isinstance(ui_selected, dict): |
| print("❌ --selected-json 必須為 JSON 物件") |
| sys.exit(1) |
|
|
| run_engine_b( |
| sd, |
| args.target, |
| index_path=idx, |
| non_interactive=args.non_interactive, |
| selected=ui_selected, |
| ) |
|
|
