| """ |
| MolGate — Module 3: Protein Preparation (Engine A) |
| Stage 10: Family classification |
| Stage 11: State confirmation |
| Stage 12: Contaminant chain removal |
| Stage 13: Missing loop assessment |
| Stage 14: Water removal |
| |
| 全自動層，輸出 cleaned.pdb |
| Master Index 與 Module 1 共用 master_index.json（單一真相來源）。 |
| 可選合併 session 內 stage_03_index.json（與該次 run 對齊）。 |
| 依賴：biopython |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import sys |
| from datetime import datetime |
| from pathlib import Path |
| from typing import Any |
|
|
| from Bio.PDB import PDBIO, PDBParser, Select |
| from Bio.PDB.Polypeptide import is_aa |
|
|
| _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_ui.console_utf8 import configure_console_utf8 |
| from molgate_ui.protein_families import prep_profile_for_family |
|
|
| configure_console_utf8() |
|
|
| SCHEMA_VERSION = "0.2" |
| DEFAULT_INDEX_PATH = _PKG_ROOT / "master_index.json" |
| INDEX_SCHEMA_FALLBACK = "0.2" |
|
|
|
|
| def load_master_index_file(path: Path | None = None) -> dict[str, Any]: |
| """與 molgate_module1.load_master_index 對齊的純 JSON 載入（不依賴 RDKit）。""" |
| p = path or DEFAULT_INDEX_PATH |
| if not p.is_file(): |
| raise FileNotFoundError( |
| f"找不到 Master Index：{p}\n" |
| "請在 MolGate 目錄放置 master_index.json，或使用 --index 指定路徑。" |
| ) |
| raw = json.loads(p.read_text(encoding="utf-8")) |
| ver = raw.pop("_schema_version", None) |
| raw.pop("_comment", None) |
| out: dict[str, Any] = {} |
| for k, v in raw.items(): |
| if not isinstance(v, dict): |
| continue |
| entry = dict(v) |
| entry["_index_schema"] = ver or INDEX_SCHEMA_FALLBACK |
| out[k.upper()] = entry |
| return out |
|
|
|
|
| def load_stage_03(session_dir: Path) -> dict[str, Any] | None: |
| p = session_dir / "stage_03_index.json" |
| if not p.is_file(): |
| return None |
| return json.loads(p.read_text(encoding="utf-8")) |
|
|
|
|
| def load_session_index(session_dir: Path) -> dict[str, Any] | None: |
| """Module 1 UI 寫入的 session_index.json（與 master_index 同構）。""" |
| p = session_dir / "session_index.json" |
| if not p.is_file(): |
| return None |
| try: |
| from molgate_ui.index_normalize import normalize_index |
|
|
| raw = json.loads(p.read_text(encoding="utf-8")) |
| entry = normalize_index(raw) |
| for meta_key in ("_created", "_source"): |
| if meta_key in raw: |
| entry[meta_key] = raw[meta_key] |
| return entry |
| except (json.JSONDecodeError, OSError, TypeError): |
| return None |
|
|
|
|
| def build_engine_entry( |
| master: dict[str, Any], |
| target: str, |
| session_dir: Path, |
| ) -> dict[str, Any]: |
| """ |
| 優先 master_index[target]；若無則 fallback session_index.json（與 Module 1 UI 一致）。 |
| 再以 stage_03_index.json 覆寫本次 run 的 pdb_id / family 等。 |
| """ |
| key = target.strip().upper() |
| master_base = master.get(key) |
| if master_base is not None: |
| entry: dict[str, Any] = dict(master_base) |
| entry["_index_source"] = "master_index" |
| else: |
| session_entry = load_session_index(session_dir) |
| if session_entry is None: |
| raise KeyError(key) |
| entry = dict(session_entry) |
| entry["_index_source"] = "session_index" |
| print(f" → Master Index 無 {key!r}，使用 session_index.json") |
|
|
| entry.setdefault("allowed", True) |
|
|
| s3 = load_stage_03(session_dir) |
| if s3: |
| for k in ( |
| "pdb_id", |
| "state", |
| "family", |
| "notes", |
| "cocrystal_het", |
| "het_id", |
| "override", |
| "pocket_rule", |
| "target_drug", |
| "common_name", |
| "box_padding", |
| ): |
| if k in s3 and s3[k] is not None: |
| entry[k] = s3[k] |
|
|
| entry.setdefault("contaminant_chains", []) |
| entry.setdefault("keep_water", False) |
| base_for_defaults = master_base if master_base is not None else {} |
| if not entry["contaminant_chains"] and isinstance(entry.get("override"), dict): |
| oc = entry["override"].get("contaminant_chains") |
| if isinstance(oc, list): |
| entry["contaminant_chains"] = oc |
| if "keep_water" not in base_for_defaults and isinstance(entry.get("override"), dict): |
| kw = entry["override"].get("keep_water") |
| if isinstance(kw, bool): |
| entry["keep_water"] = kw |
|
|
| try: |
| from molgate_ui.anchor_routing import merge_session_route_overrides |
|
|
| entry = merge_session_route_overrides(entry, session_dir) |
| except ImportError: |
| pass |
| return entry |
|
|
| |
| def save_stage(session_dir: Path, stage: int, name: str, data: dict) -> dict: |
| 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 find_pdb_file(session_dir: Path, pdb_id: str) -> Path: |
| pdb_dir = session_dir / "pdb" |
| candidates = (list(pdb_dir.glob(f"*{pdb_id.lower()}*")) + |
| list(pdb_dir.glob(f"*{pdb_id.upper()}*"))) |
| if not candidates: |
| raise FileNotFoundError(f"找不到 PDB 檔案：{pdb_id}，請先跑 Module 1") |
| return candidates[0] |
|
|
| |
| def stage10_family(entry: dict, session_dir: Path) -> dict: |
| print("\n[Stage 10] 蛋白家族分類...") |
| family = entry.get("family", "other") or "other" |
| profile = prep_profile_for_family(str(family), lang="zh") |
| print(f" ✓ Family: {family}") |
| print(f" ✓ Prep profile: {profile}") |
| print(f" → Status: AUTO") |
| return save_stage(session_dir, 10, "family", { |
| "status": "AUTO", |
| "family": family, |
| "prep_profile": profile, |
| }) |
|
|
| |
| def stage11_state(entry: dict, session_dir: Path) -> dict: |
| print("\n[Stage 11] 結構態確認...") |
| state = entry.get("state", "unknown") |
| notes = entry.get("notes", "") |
| status = "WARN" if state == "unknown" else "AUTO" |
| if status == "WARN": |
| print(f" ⚠ WARN: state = unknown，結果可信度降低") |
| print(f" ✓ State: {state}") |
| if notes: |
| print(f" ✓ Notes: {notes}") |
| print(f" → Status: {status}") |
| return save_stage(session_dir, 11, "state", { |
| "status": status, |
| "state": state, |
| "notes": notes, |
| }) |
|
|
| |
| def stage12_chains(entry: dict, structure, session_dir: Path): |
| print("\n[Stage 12] 污染鏈移除...") |
| contaminant = entry.get("contaminant_chains", []) |
| all_chains = [c.id for m in structure for c in m] |
| kept = [c for c in all_chains if c not in contaminant] |
| if contaminant: |
| print(f" ⚠ 移除污染鏈：{contaminant}") |
| else: |
| print(f" ✓ 無污染鏈") |
| print(f" ✓ 保留 chains：{kept}") |
| print(f" → Status: AUTO") |
| r = save_stage(session_dir, 12, "chains", { |
| "status": "AUTO", |
| "all_chains": all_chains, |
| "contaminant_chains": contaminant, |
| "kept_chains": kept, |
| }) |
| return r, kept |
|
|
| |
| def stage13_loops(structure, kept_chains: list, session_dir: Path) -> dict: |
| print("\n[Stage 13] 缺失環評估...") |
| missing = [] |
| for model in structure: |
| for chain in model: |
| if chain.id not in kept_chains: |
| continue |
| residues = [r for r in chain if is_aa(r, standard=True)] |
| if not residues: |
| continue |
| seq_ids = sorted([r.get_id()[1] for r in residues]) |
| for i in range(len(seq_ids) - 1): |
| gap = seq_ids[i+1] - seq_ids[i] |
| if gap > 1: |
| missing.append({ |
| "chain": chain.id, |
| "after": seq_ids[i], |
| "before": seq_ids[i+1], |
| "gap_size": gap - 1, |
| }) |
| status = "WARN" if missing else "AUTO" |
| if missing: |
| print(f" ⚠ 偵測到 {len(missing)} 個缺失片段") |
| for seg in missing[:3]: |
| print(f" Chain {seg['chain']}: 殘基 {seg['after']}～{seg['before']} 缺失 {seg['gap_size']} 個") |
| if len(missing) > 3: |
| print(f" ...（共 {len(missing)} 個）") |
| else: |
| print(f" ✓ 無缺失環，結構完整") |
| print(f" → Status: {status}") |
| return save_stage(session_dir, 13, "loops", { |
| "status": status, |
| "missing_count": len(missing), |
| "missing_segments": missing[:20], |
| }) |
|
|
| |
| def stage14_water(entry: dict, structure, kept_chains: list, session_dir: Path) -> dict: |
| print("\n[Stage 14] 去結晶水...") |
| keep_water = entry.get("keep_water", False) |
| water_count = sum( |
| 1 for m in structure for c in m |
| if c.id in kept_chains |
| for r in c if r.get_resname().strip() in ["HOH", "WAT"] |
| ) |
| if keep_water: |
| status = "CONFIRM" |
| action = "keep_all" |
| print(f" ⚠ Master Index 標記保留水（共 {water_count} 個）") |
| else: |
| status = "AUTO" |
| action = "remove_all" |
| print(f" ✓ 移除全部結晶水（共 {water_count} 個）") |
| print(f" → Status: {status}") |
| return save_stage(session_dir, 14, "water", { |
| "status": status, |
| "action": action, |
| "water_count": water_count, |
| "keep_water": keep_water, |
| }) |
|
|
| |
| class CleanSelect(Select): |
| def __init__(self, kept_chains, remove_water): |
| self.kept_chains = kept_chains |
| self.remove_water = remove_water |
|
|
| def accept_chain(self, chain): |
| return chain.id in self.kept_chains |
|
|
| def accept_residue(self, residue): |
| if self.remove_water and residue.get_resname().strip() in ["HOH", "WAT"]: |
| return False |
| return True |
|
|
| def _update_manifest_engine_a( |
| session_dir: Path, |
| cleaned_path: Path, |
| flags: list[str], |
| ) -> None: |
| """若存在 molgate_run_manifest.json，寫入 cleaned.pdb 與 module3_engine_a 狀態。""" |
| 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["cleaned_pdb"] = rel_path(session_dir, cleaned_path) |
| paths["relative"] = rmap |
| paths["cleaned_pdb_absolute"] = str(cleaned_path.resolve()) |
| ns = dict(paths.get("next_step") or {}) |
| ns["after_module3_engine_a"] = { |
| "read_cleaned_pdb_from": rmap["cleaned_pdb"], |
| "hint": "下一步：Engine B（人機決策 15/17 + Flag 16/18 → prep_decisions）", |
| } |
| paths["next_step"] = ns |
| m["paths"] = paths |
| pl = dict(m.get("pipeline") or {}) |
| pl["module3_engine_a"] = { |
| "status": "complete", |
| "completed_at": datetime.now().isoformat(), |
| "stages": ["10", "11", "12", "13", "14"], |
| } |
| m["pipeline"] = pl |
| summ = dict(m.get("summary") or {}) |
| summ["module3_engine_a_flags"] = flags |
| summ["cleaned_pdb_relative"] = rmap["cleaned_pdb"] |
| m["summary"] = summ |
| write_manifest(session_dir, m) |
| except Exception: |
| pass |
|
|
|
|
| |
| def run_engine_a( |
| session_dir: Path, |
| target: str, |
| *, |
| index_path: Path | None = None, |
| skip_manifest_target_check: bool = False, |
| ): |
| print("=" * 55) |
| print(" MolGate — Module 3 Engine A（全自動層）") |
| print(" Stage 10/11/12/13/14 → cleaned.pdb") |
| print("=" * 55) |
|
|
| try: |
| master = load_master_index_file(index_path) |
| except FileNotFoundError as e: |
| print(f"❌ {e}") |
| sys.exit(1) |
|
|
| try: |
| entry = build_engine_entry(master, target, session_dir) |
| except KeyError: |
| print( |
| f"[ERROR] Target '{target}' 不在 Master Index，" |
| "且 session 內無 session_index.json" |
| ) |
| sys.exit(1) |
|
|
| if entry.get("allowed") is False: |
| print(f"❌ Target '{target}' 在 Master Index 標記為 allowed=false，已中止") |
| sys.exit(1) |
|
|
| 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) |
|
|
| pdb_id = (entry.get("pdb_id") or "").upper() |
| if not pdb_id: |
| print("❌ 條目缺少 pdb_id") |
| sys.exit(1) |
| print(f"\n Target：{target.upper()} | PDB：{pdb_id}") |
|
|
| try: |
| pdb_path = find_pdb_file(session_dir, pdb_id) |
| except FileNotFoundError as e: |
| print(f"❌ {e}") |
| sys.exit(1) |
|
|
| print(f" PDB：{pdb_path}") |
|
|
| parser = PDBParser(QUIET=True) |
| structure = parser.get_structure(pdb_id, str(pdb_path)) |
|
|
| flags = [] |
|
|
| try: |
| r10 = stage10_family(entry, session_dir) |
| r11 = stage11_state(entry, session_dir) |
| if r11["status"] == "WARN": |
| flags.append("WARN_state_unknown") |
|
|
| r12, kept_chains = stage12_chains(entry, structure, session_dir) |
| r13 = stage13_loops(structure, kept_chains, session_dir) |
| if r13["status"] == "WARN": |
| flags.append(f"WARN_missing_loops:{r13['missing_count']}") |
|
|
| r14 = stage14_water(entry, structure, kept_chains, session_dir) |
|
|
| except Exception as e: |
| print(f"\n❌ Pipeline stopped: {e}") |
| sys.exit(1) |
|
|
| |
| protein_dir = session_dir / "protein" |
| protein_dir.mkdir(exist_ok=True) |
| cleaned_path = protein_dir / "cleaned.pdb" |
|
|
| io = PDBIO() |
| io.set_structure(structure) |
| io.save(str(cleaned_path), CleanSelect(kept_chains, r14["action"] == "remove_all")) |
|
|
| with open(cleaned_path) as f: |
| lines = f.readlines() |
| atom_count = len([l for l in lines if l.startswith("ATOM")]) |
| hetatm_count = len([l for l in lines if l.startswith("HETATM")]) |
|
|
| print(f"\n ✓ cleaned.pdb 完成") |
| print(f" ✓ ATOM: {atom_count} | HETATM: {hetatm_count} | Chains: {kept_chains}") |
|
|
| _update_manifest_engine_a(session_dir, cleaned_path, flags) |
|
|
| print("\n" + "=" * 55) |
| print(" Engine A 完成 ✓") |
| if flags: |
| print(f" Flags：{' | '.join(flags)}") |
| print(f" 輸出：{cleaned_path}") |
| print(f" Session log：{session_dir}") |
| print("=" * 55) |
|
|
| return { |
| "stage10": r10, "stage11": r11, |
| "stage12": r12, "stage13": r13, "stage14": r14, |
| "cleaned_pdb": str(cleaned_path), |
| "kept_chains": kept_chains, |
| "flags": flags, |
| } |
|
|
| |
| if __name__ == "__main__": |
| parser = argparse.ArgumentParser(description="MolGate Module 3 Engine A") |
| parser.add_argument("--session-dir", required=True) |
| parser.add_argument("--target", required=True) |
| parser.add_argument( |
| "--index", |
| type=Path, |
| default=None, |
| help="Master Index JSON（預設：MolGate 目錄下 master_index.json）", |
| ) |
| parser.add_argument( |
| "--skip-manifest-target-check", |
| action="store_true", |
| help="不與 molgate_run_manifest.json 的 target 交叉驗證", |
| ) |
| args = parser.parse_args() |
|
|
| session_dir = Path(args.session_dir) |
| if not session_dir.exists(): |
| print(f"❌ Session 目錄不存在：{session_dir}") |
| sys.exit(1) |
|
|
| run_engine_a( |
| session_dir, |
| args.target, |
| index_path=args.index, |
| skip_manifest_target_check=args.skip_manifest_target_check, |
| ) |
|
|
