| """ |
| MolGate — Module 4: Pocket Definition (Engine C) |
| Stage 23: Pocket center calculation |
| |
| 輸入：stage_22_ligand.json + protein/cleaned.pdb + Master Index（可選 box_padding） |
| 輸出：stage_23_pocket_center.json、pocket_config.json（Module 5 / Vina 契約） |
| |
| 依賴：biopython、numpy；Index 載入與 Engine A/B 對齊。 |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import sys |
| from datetime import datetime |
| from pathlib import Path |
| from typing import Any |
|
|
| import numpy as np |
| from Bio.PDB import PDBParser |
|
|
| _PKG_ROOT = Path(__file__).resolve().parent |
| _UI_ROOT = _PKG_ROOT.parent |
| if str(_PKG_ROOT) not in sys.path: |
| sys.path.insert(0, str(_PKG_ROOT)) |
| if str(_UI_ROOT) not in sys.path: |
| sys.path.insert(0, str(_UI_ROOT)) |
|
|
| from molgate_module4_engineA import load_target_index_entry |
| from molgate_ui.console_utf8 import configure_console_utf8, should_prompt_interactive |
| from molgate_ui.protein_families import default_box_padding_for_family |
| from molgate_module4_engineB import load_prep_overlay |
|
|
| configure_console_utf8() |
|
|
| SCHEMA_VERSION = "0.2" |
| STAGE22_SCHEMA_VERSION = "0.2" |
| DEFAULT_PADDING = 10.0 |
| GPCR_DEFAULT_PADDING = 12.0 |
|
|
|
|
| def resolve_box_padding(entry: dict[str, Any]) -> tuple[float, str]: |
| """Index 條目或 override.box_padding；否則依 protein_families 登記之 family 預設。""" |
| ov = entry.get("override") if isinstance(entry.get("override"), dict) else {} |
| for label, src in (("Master Index", entry), ("override", ov)): |
| v = src.get("box_padding") |
| if v is not None: |
| return float(v), f"{label} box_padding" |
| return default_box_padding_for_family( |
| str(entry.get("family") or ""), |
| global_default=DEFAULT_PADDING, |
| gpcr_default=GPCR_DEFAULT_PADDING, |
| ) |
|
|
|
|
| |
| 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_stage22(session_dir: Path) -> dict[str, Any]: |
| path = session_dir / "stage_22_ligand.json" |
| if not path.exists(): |
| raise FileNotFoundError("找不到 stage_22_ligand.json，請先跑 Engine B") |
| data = json.loads(path.read_text(encoding="utf-8")) |
| if data.get("_schema_version") != STAGE22_SCHEMA_VERSION: |
| raise ValueError( |
| f"stage_22_ligand.json schema 須為 {STAGE22_SCHEMA_VERSION}，請重跑 Module 4 Engine B" |
| ) |
| if not data.get("ready_for_engine_c"): |
| raise ValueError("stage_22_ligand.json 未標記 ready_for_engine_c=True，請重跑 Engine B") |
| return data |
|
|
|
|
| def load_cleaned_pdb(session_dir: Path) -> Path: |
| p = session_dir / "protein" / "cleaned.pdb" |
| if not p.is_file(): |
| raise FileNotFoundError("找不到 protein/cleaned.pdb，請先跑 Module 3 Engine A") |
| return p |
|
|
|
|
| def get_ligand_coords( |
| structure, |
| het_id: str, |
| chain_id: str, |
| resseq: int, |
| ) -> np.ndarray: |
| """從 PDB 結構抓取指定配體所有重原子座標。""" |
| coords = [] |
| for model in structure: |
| for chain in model: |
| if chain.id != chain_id: |
| continue |
| for res in chain: |
| rn = res.get_resname().strip().upper() |
| rs = res.get_id()[1] |
| if rn == het_id.upper() and rs == resseq: |
| for atom in res.get_atoms(): |
| if atom.element and atom.element.upper() != "H": |
| coords.append(atom.get_vector().get_array()) |
| if not coords: |
| raise ValueError( |
| f"在 cleaned.pdb 找不到配體：{het_id} chain {chain_id} res {resseq}\n" |
| f" → 請確認 cleaned.pdb 是否保留了 HETATM，" |
| f"或 Stage 22 / Index het_id 與 resseq 是否正確" |
| ) |
| return np.array(coords) |
|
|
|
|
| def prompt_confirm(label: str, value: Any, *, non_interactive: bool = False) -> bool: |
| if not should_prompt_interactive(non_interactive=non_interactive): |
| print(f" [AUTO] 非互動：接受 {label}") |
| return True |
| print(f"\n ── CONFIRM：{label} ──") |
| if isinstance(value, dict): |
| for k, v in value.items(): |
| if not k.startswith("_"): |
| print(f" {k}: {v}") |
| else: |
| print(f" {value}") |
| while True: |
| try: |
| raw = input(" 接受此結果？[Y/n]：").strip().lower() |
| if raw in ("", "y", "yes"): |
| return True |
| if raw in ("n", "no"): |
| return False |
| except KeyboardInterrupt: |
| sys.exit(0) |
| print(" 請輸入 y 或 n") |
|
|
|
|
| def _update_manifest_module4_engine_c(session_dir: Path, pocket_config_path: Path) -> None: |
| try: |
| from molgate_manifest import read_manifest, rel_path, write_manifest |
| except ImportError: |
| return |
| mp = session_dir / "molgate_run_manifest.json" |
| if not mp.is_file(): |
| return |
| try: |
| m = read_manifest(session_dir) |
| paths = dict(m.get("paths") or {}) |
| rmap = dict(paths.get("relative") or {}) |
| rmap["pocket_config"] = rel_path(session_dir, pocket_config_path) |
| paths["relative"] = rmap |
| paths["pocket_config_absolute"] = str(pocket_config_path.resolve()) |
| ns = dict(paths.get("next_step") or {}) |
| ns["after_module4_engine_c"] = { |
| "read_pocket_config_from": rmap["pocket_config"], |
| "hint": "Module 5：AutoDock Vina 使用 center / box_size", |
| } |
| paths["next_step"] = ns |
| m["paths"] = paths |
| pl = dict(m.get("pipeline") or {}) |
| pl["module4_engine_c"] = { |
| "status": "complete", |
| "completed_at": datetime.now().isoformat(), |
| "stages": ["23"], |
| } |
| m["pipeline"] = pl |
| write_manifest(session_dir, m) |
| except Exception: |
| pass |
|
|
|
|
| |
| def stage23_pocket_center( |
| session_dir: Path, |
| entry: dict[str, Any], |
| *, |
| non_interactive: bool = False, |
| ) -> dict[str, Any]: |
| print("\n[Stage 23] 口袋中心計算...") |
|
|
| s22 = load_stage22(session_dir) |
| het_id = s22["selected_het_id"] |
| chain = s22["chain"] |
| resseq = s22["resseq"] |
| pocket_rule = (s22.get("pocket_rule") or "").strip() |
|
|
| padding, padding_src = resolve_box_padding(entry) |
|
|
| print(f" → 配體：{het_id} chain {chain} res {resseq}") |
| print(f" → Pocket rule：{pocket_rule}") |
| print(f" → Padding：{padding} Å（{padding_src}）") |
|
|
| cleaned_path = load_cleaned_pdb(session_dir) |
| parser = PDBParser(QUIET=True) |
| structure = parser.get_structure("protein", str(cleaned_path)) |
|
|
| try: |
| coords = get_ligand_coords(structure, het_id, chain, resseq) |
| except ValueError as e: |
| result = { |
| "status": "FAIL", |
| "reason": str(e), |
| } |
| save_stage(session_dir, 23, "pocket_center", result) |
| raise |
|
|
| atom_count = len(coords) |
| print(f" ✓ 配體重原子數：{atom_count}") |
|
|
| center = coords.mean(axis=0) |
| cx, cy, cz = float(center[0]), float(center[1]), float(center[2]) |
|
|
| min_coords = coords.min(axis=0) |
| max_coords = coords.max(axis=0) |
| span = max_coords - min_coords |
|
|
| box_x = round(float(span[0]) + padding * 2, 1) |
| box_y = round(float(span[1]) + padding * 2, 1) |
| box_z = round(float(span[2]) + padding * 2, 1) |
|
|
| box_x = max(box_x, 20.0) |
| box_y = max(box_y, 20.0) |
| box_z = max(box_z, 20.0) |
|
|
| print(f" ✓ 重心：({cx:.3f}, {cy:.3f}, {cz:.3f})") |
| print(f" ✓ Box size：{box_x} × {box_y} × {box_z} Å") |
|
|
| confirm_data = { |
| "center_x": round(cx, 3), |
| "center_y": round(cy, 3), |
| "center_z": round(cz, 3), |
| "box_size_x": box_x, |
| "box_size_y": box_y, |
| "box_size_z": box_z, |
| "padding": padding, |
| } |
|
|
| if non_interactive: |
| accepted = True |
| print("\n ✓ --non-interactive：自動接受計算之口袋座標與 Box Size") |
| else: |
| accepted = prompt_confirm( |
| "口袋座標與 Box Size", |
| confirm_data, |
| non_interactive=non_interactive, |
| ) |
|
|
| if not accepted: |
| print("\n 手動輸入座標（直接 Enter 保留計算值）：") |
|
|
| def ask_float(label: str, default: float) -> float: |
| try: |
| raw = input(f" {label} [{default}]：").strip() |
| return float(raw) if raw else default |
| except (ValueError, KeyboardInterrupt): |
| return default |
|
|
| cx = ask_float("center_x", round(cx, 3)) |
| cy = ask_float("center_y", round(cy, 3)) |
| cz = ask_float("center_z", round(cz, 3)) |
| box_x = ask_float("box_size_x", box_x) |
| box_y = ask_float("box_size_y", box_y) |
| box_z = ask_float("box_size_z", box_z) |
| padding = ask_float("padding", padding) |
| status = "MANUAL" |
| print(f" ✓ 手動座標接受") |
| else: |
| status = "CONFIRM" |
| print(f" ✓ 計算結果接受") |
|
|
| print(f" → Status: {status}") |
|
|
| pdb_id = (s22.get("pdb_id") or "").strip().upper() |
| result: dict[str, Any] = { |
| "status": status, |
| "het_id": het_id, |
| "chain": chain, |
| "resseq": resseq, |
| "atom_count": atom_count, |
| "center_x": round(cx, 3), |
| "center_y": round(cy, 3), |
| "center_z": round(cz, 3), |
| "box_size_x": box_x, |
| "box_size_y": box_y, |
| "box_size_z": box_z, |
| "padding": padding, |
| "padding_source": padding_src, |
| "pocket_rule": pocket_rule, |
| } |
| if pdb_id: |
| result["pdb_id"] = pdb_id |
| cn = (s22.get("common_name") or "").strip() |
| td = (s22.get("target_drug") or "").strip() |
| if cn: |
| result["common_name"] = cn |
| if td: |
| result["target_drug"] = td |
| return save_stage(session_dir, 23, "pocket_center", result) |
|
|
|
|
| def write_pocket_config( |
| session_dir: Path, |
| r23: dict[str, Any], |
| *, |
| target: str, |
| ) -> Path: |
| config: dict[str, Any] = { |
| "schema_version": SCHEMA_VERSION, |
| "target": target.strip().upper(), |
| "het_id": r23["het_id"], |
| "chain": r23["chain"], |
| "resseq": r23["resseq"], |
| "center_x": r23["center_x"], |
| "center_y": r23["center_y"], |
| "center_z": r23["center_z"], |
| "box_size_x": r23["box_size_x"], |
| "box_size_y": r23["box_size_y"], |
| "box_size_z": r23["box_size_z"], |
| "padding": r23["padding"], |
| "atom_count": r23["atom_count"], |
| "pocket_rule": r23.get("pocket_rule", ""), |
| "ready_for_docking": r23["status"] in ("CONFIRM", "MANUAL"), |
| "_timestamp": datetime.now().isoformat(), |
| } |
| if r23.get("pdb_id"): |
| config["pdb_id"] = r23["pdb_id"] |
| if r23.get("common_name"): |
| config["common_name"] = r23["common_name"] |
| if r23.get("target_drug"): |
| config["target_drug"] = r23["target_drug"] |
| path = session_dir / "pocket_config.json" |
| path.write_text(json.dumps(config, indent=2, ensure_ascii=False), encoding="utf-8") |
| return path |
|
|
|
|
| def run_engine_c( |
| session_dir: Path, |
| target: str, |
| *, |
| index_path: Path | None = None, |
| skip_manifest_target_check: bool = False, |
| skip_prep_target_check: bool = False, |
| non_interactive: bool = False, |
| ): |
| print("=" * 55) |
| print(" MolGate — Module 4 Engine C") |
| print(" Stage 23: 口袋中心計算 → pocket_config.json") |
| print("=" * 55) |
|
|
| if not skip_manifest_target_check: |
| try: |
| from molgate_manifest import validate_target_matches_manifest |
|
|
| validate_target_matches_manifest(session_dir, target) |
| except FileNotFoundError: |
| pass |
| except ValueError as e: |
| print(f"❌ {e}") |
| sys.exit(1) |
|
|
| prep = load_prep_overlay(session_dir, target) |
| if not skip_prep_target_check: |
| pt = (prep.get("target") or "").strip().upper() |
| if pt and pt != target.strip().upper(): |
| print( |
| f"❌ prep_decisions.target={prep.get('target')!r} 與 --target={target!r} 不一致" |
| ) |
| sys.exit(1) |
|
|
| entry = load_target_index_entry(session_dir, prep, index_path) |
| if not entry: |
| print(f"❌ 無法載入 Master Index 條目（--target 或 master_index.json）") |
| sys.exit(1) |
|
|
| print(f"\n Target：{target.upper()}") |
|
|
| try: |
| r23 = stage23_pocket_center(session_dir, entry, non_interactive=non_interactive) |
| 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) |
|
|
| if r23.get("status") == "FAIL": |
| print(f"\n❌ Stage 23 失敗：{r23.get('reason', '')}") |
| sys.exit(1) |
|
|
| config_path = write_pocket_config(session_dir, r23, target=target) |
| _update_manifest_module4_engine_c(session_dir, config_path) |
|
|
| print("\n" + "=" * 55) |
| print(" Engine C 完成 ✓") |
| print(f" 中心：({r23['center_x']}, {r23['center_y']}, {r23['center_z']})") |
| print(f" Box：{r23['box_size_x']} × {r23['box_size_y']} × {r23['box_size_z']} Å") |
| print(f" 契約：{config_path}") |
| print(f" ready_for_docking = {r23['status'] in ('CONFIRM', 'MANUAL')}") |
| print("=" * 55) |
|
|
| return r23 |
|
|
|
|
| if __name__ == "__main__": |
| ap = argparse.ArgumentParser(description="MolGate Module 4 Engine C") |
| ap.add_argument("--session-dir", required=True, help="Session 目錄路徑") |
| ap.add_argument("--target", required=True, help="Target 名稱（如 COX2）") |
| ap.add_argument( |
| "--index", |
| default=None, |
| help="master_index.json 路徑（預設：session_meta / manifest 或 MolGate 目錄）", |
| ) |
| ap.add_argument( |
| "--skip-manifest-target-check", |
| action="store_true", |
| help="不與 molgate_run_manifest.json 的 target 交叉驗證", |
| ) |
| ap.add_argument( |
| "--skip-prep-target-check", |
| action="store_true", |
| help="不檢查 prep_decisions.target 與 --target 一致", |
| ) |
| ap.add_argument( |
| "--non-interactive", |
| action="store_true", |
| help="Stage 23 自動接受計算之口袋與 box（不詢問 Y/n）", |
| ) |
| 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 |
| run_engine_c( |
| sd, |
| args.target, |
| index_path=idx, |
| skip_manifest_target_check=args.skip_manifest_target_check, |
| skip_prep_target_check=args.skip_prep_target_check, |
| non_interactive=args.non_interactive, |
| ) |
|
|
