mdprep 0.2.0__py3-none-any.whl

mdprep/refinement/coordinates.py

@@ -0,0 +1,210 @@
+"""Transfer optimized Amber coordinates back to stable input residue identities."""
+
+from __future__ import annotations
+
+from collections import OrderedDict
+from dataclasses import dataclass, replace
+from pathlib import Path
+
+from mdprep.structure.classify import is_standard_protein_residue
+from mdprep.refinement.mapping import (
+    RefinementMappingError,
+    compatible_resnames,
+    is_virtual_site,
+    map_residues_by_coordinates,
+)
+from mdprep.structure.models import AtomRecord, PdbStructure, ResidueId, ResidueRecord
+
+
+class CoordinateTransferError(ValueError):
+    """Raised when optimized coordinates cannot be mapped without guessing."""
+
+
+@dataclass(frozen=True)
+class CoordinateTransferResult:
+    structure: PdbStructure
+    restored_residue_names: tuple[dict[str, object], ...]
+    generated_hydrogen_count: int
+    excluded_virtual_site_count: int
+
+    def to_dict(self) -> dict[str, object]:
+        return {
+            "restored_residue_names": list(self.restored_residue_names),
+            "generated_hydrogen_count": self.generated_hydrogen_count,
+            "excluded_virtual_site_count": self.excluded_virtual_site_count,
+            "output_atom_count": len(self.structure.atoms),
+        }
+
+
+def transfer_optimized_coordinates(
+    *,
+    reference_structure: PdbStructure,
+    leap_input_structure: PdbStructure,
+    topology_structure: PdbStructure,
+    coordinates_angstrom: tuple[tuple[float, float, float], ...],
+    output_path: str | Path,
+) -> CoordinateTransferResult:
+    """Return a hydrogenated optimized structure with pre-protonation names.
+
+    tleap-generated hydrogens are retained here. With
+    ``post_refinement_protonation: reuse_initial`` they remain the authoritative
+    optimized coordinates for subsequent MCPB/charge/final-tleap stages. With
+    the legacy ``rerun`` policy, the normal protonation stage may remove them
+    before final parameterization.
+    """
+
+    if len(coordinates_angstrom) != len(topology_structure.atoms):
+        raise CoordinateTransferError(
+            "Optimized coordinate count does not match provisional tleap PDB atom count: "
+            f"{len(coordinates_angstrom)} != {len(topology_structure.atoms)}"
+        )
+    try:
+        reference_to_leap = map_residues_by_coordinates(
+            reference_structure,
+            leap_input_structure,
+        )
+        leap_to_topology = map_residues_by_coordinates(
+            leap_input_structure,
+            topology_structure,
+        )
+    except RefinementMappingError as exc:
+        raise CoordinateTransferError(str(exc)) from exc
+    leap_pair_by_input_index = {
+        pair.input_index: pair for pair in leap_to_topology
+    }
+
+    topology_atom_index = {
+        id(atom): index for index, atom in enumerate(topology_structure.atoms)
+    }
+    new_atoms: list[AtomRecord] = []
+    old_to_new_serial: dict[int, int] = {}
+    residue_last_serial: dict[int, int] = {}
+    restored_names: list[dict[str, object]] = []
+    generated_hydrogens = 0
+    excluded_virtual_sites = 0
+
+    for residue_index, reference_pair in enumerate(reference_to_leap):
+        reference = reference_pair.input_residue
+        leap_input = reference_pair.topology_residue
+        topology_pair = leap_pair_by_input_index[reference_pair.topology_index]
+        topology = topology_pair.topology_residue
+        if not compatible_resnames(reference.id.resname, leap_input.id.resname):
+            raise CoordinateTransferError(
+                f"Residue {residue_index + 1} changed identity before tleap: "
+                f"{reference.id.display()} -> {leap_input.id.display()}"
+            )
+        if not compatible_resnames(leap_input.id.resname, topology.id.resname):
+            raise CoordinateTransferError(
+                f"Residue {residue_index + 1} changed identity in tleap: "
+                f"{leap_input.id.display()} -> {topology.id.display()}"
+            )
+        if reference.id.resname != leap_input.id.resname:
+            restored_names.append(
+                {
+                    "chain_id": reference.id.chain_id,
+                    "resid": reference.id.resid,
+                    "icode": reference.id.icode,
+                    "optimized_provisional_resname": leap_input.id.resname,
+                    "restored_pre_protonation_resname": reference.id.resname,
+                    "reason": "rerun requested protonation protocol after QM/MM refinement",
+                }
+            )
+        reference_by_name = {atom.name: atom for atom in reference.atoms}
+        if len(reference_by_name) != len(reference.atoms):
+            raise CoordinateTransferError(
+                f"Reference residue {reference.id.display()} has duplicate atom names"
+            )
+        record_name = "ATOM" if is_standard_protein_residue(reference) else next(
+            iter(reference.record_names), "HETATM"
+        )
+        for topology_atom in topology.atoms:
+            atom_index = topology_atom_index[id(topology_atom)]
+            if is_virtual_site(topology_atom):
+                excluded_virtual_sites += 1
+                continue
+            x, y, z = coordinates_angstrom[atom_index]
+            reference_atom = reference_by_name.get(topology_atom.name)
+            serial = len(new_atoms) + 1
+            if reference_atom is not None:
+                atom_record_name = reference_atom.record_name
+                occupancy = reference_atom.occupancy
+                bfactor = reference_atom.bfactor
+                original_line = reference_atom.original_line
+                if reference_atom.serial is not None:
+                    old_to_new_serial[reference_atom.serial] = serial
+            else:
+                atom_record_name = record_name
+                occupancy = 1.0
+                bfactor = 0.0
+                original_line = ""
+                if (topology_atom.element or "").strip().upper() == "H":
+                    generated_hydrogens += 1
+            new_atoms.append(
+                replace(
+                    topology_atom,
+                    serial=serial,
+                    resname=reference.id.resname,
+                    chain_id=reference.id.chain_id,
+                    resid=reference.id.resid,
+                    icode=reference.id.icode,
+                    x=x,
+                    y=y,
+                    z=z,
+                    occupancy=occupancy,
+                    bfactor=bfactor,
+                    record_name=atom_record_name,  # type: ignore[arg-type]
+                    original_line=original_line,
+                )
+            )
+            residue_last_serial[residue_index] = serial
+
+    conect_bonds = {
+        (min(old_to_new_serial[left], old_to_new_serial[right]),
+         max(old_to_new_serial[left], old_to_new_serial[right]))
+        for left, right in reference_structure.conect_bonds
+        if left in old_to_new_serial and right in old_to_new_serial
+    }
+    reference_residue_index_by_serial = {
+        atom.serial: index
+        for index, residue in enumerate(reference_structure.residues)
+        for atom in residue.atoms
+        if atom.serial is not None
+    }
+    ter_after_serials = {
+        residue_last_serial[reference_residue_index_by_serial[serial]]
+        for serial in reference_structure.ter_after_serials
+        if serial in reference_residue_index_by_serial
+        and reference_residue_index_by_serial[serial] in residue_last_serial
+    }
+    output = Path(output_path)
+    structure = PdbStructure(
+        path=output,
+        atoms=new_atoms,
+        residues=_build_residues(new_atoms),
+        model_count=1,
+        used_model=1,
+        warnings=list(reference_structure.warnings),
+        conect_bonds=conect_bonds,
+        ter_after_serials=ter_after_serials,
+    )
+    return CoordinateTransferResult(
+        structure=structure,
+        restored_residue_names=tuple(restored_names),
+        generated_hydrogen_count=generated_hydrogens,
+        excluded_virtual_site_count=excluded_virtual_sites,
+    )
+
+
+def _build_residues(atoms: list[AtomRecord]) -> list[ResidueRecord]:
+    grouped: "OrderedDict[tuple[str, str, int, str | None], list[AtomRecord]]" = OrderedDict()
+    for atom in atoms:
+        grouped.setdefault(atom.residue_key, []).append(atom)
+    return [
+        ResidueRecord(
+            id=ResidueId(chain_id=chain, resname=resname, resid=resid, icode=icode),
+            atoms=residue_atoms,
+            record_names={atom.record_name for atom in residue_atoms},
+            original_index=index,
+        )
+        for index, ((chain, resname, resid, icode), residue_atoms) in enumerate(grouped.items())
+    ]

mdprep/refinement/mapping.py

@@ -0,0 +1,163 @@
+"""Coordinate-anchored residue mapping across tleap reordering."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from math import sqrt
+
+from mdprep.structure.models import AtomRecord, PdbStructure, ResidueRecord
+
+
+class RefinementMappingError(ValueError):
+    """Raised when two structure representations cannot be mapped uniquely."""
+
+
+@dataclass(frozen=True)
+class MappedResidue:
+    input_index: int
+    topology_index: int
+    input_residue: ResidueRecord
+    topology_residue: ResidueRecord
+    anchor_rmsd_angstrom: float
+
+
+_PROTONATION_FAMILIES = (
+    frozenset({"HIS", "HID", "HIE", "HIP"}),
+    frozenset({"ASP", "ASH"}),
+    frozenset({"GLU", "GLH"}),
+    frozenset({"CYS", "CYM", "CYX"}),
+    frozenset({"LYS", "LYN"}),
+    frozenset({"HOH", "WAT", "H2O", "TIP3", "OPC"}),
+)
+
+
+def map_residues_by_coordinates(
+    input_structure: PdbStructure,
+    topology_structure: PdbStructure,
+    *,
+    max_anchor_rmsd_angstrom: float = 0.25,
+    ambiguity_tolerance_angstrom: float = 1.0e-5,
+) -> list[MappedResidue]:
+    """Map residues after tleap without relying on output residue order.
+
+    tleap commonly moves loaded ligand templates ahead of waters. Heavy-atom
+    coordinates are preserved in the provisional build, so compatible residue
+    names plus same-named atom anchors provide a deterministic mapping. This
+    also distinguishes hundreds of otherwise identical water residues.
+    """
+
+    if len(input_structure.residues) != len(topology_structure.residues):
+        raise RefinementMappingError(
+            "Provisional tleap changed the residue count, so its atom ordering cannot be "
+            f"mapped safely ({len(input_structure.residues)} input residues, "
+            f"{len(topology_structure.residues)} output residues)"
+        )
+    unused = set(range(len(topology_structure.residues)))
+    result: list[MappedResidue] = []
+    for input_index, input_residue in enumerate(input_structure.residues):
+        candidates: list[tuple[float, int, ResidueRecord]] = []
+        for topology_index in unused:
+            topology_residue = topology_structure.residues[topology_index]
+            if not compatible_resnames(
+                input_residue.id.resname,
+                topology_residue.id.resname,
+            ):
+                continue
+            score = _anchor_rmsd(input_residue, topology_residue)
+            if score is not None and score <= max_anchor_rmsd_angstrom:
+                candidates.append((score, topology_index, topology_residue))
+        if not candidates:
+            raise RefinementMappingError(
+                "No coordinate-preserving provisional residue matched "
+                f"{input_residue.id.display()} within {max_anchor_rmsd_angstrom:.3f} A"
+            )
+        candidates.sort(key=lambda item: (item[0], item[1]))
+        best_score, topology_index, topology_residue = candidates[0]
+        if (
+            len(candidates) > 1
+            and abs(candidates[1][0] - best_score) <= ambiguity_tolerance_angstrom
+        ):
+            raise RefinementMappingError(
+                "Coordinate mapping is ambiguous for "
+                f"{input_residue.id.display()}: provisional residues "
+                f"{topology_index + 1} and {candidates[1][1] + 1} have indistinguishable "
+                "anchor coordinates"
+            )
+        unused.remove(topology_index)
+        result.append(
+            MappedResidue(
+                input_index=input_index,
+                topology_index=topology_index,
+                input_residue=input_residue,
+                topology_residue=topology_residue,
+                anchor_rmsd_angstrom=best_score,
+            )
+        )
+    if unused:
+        raise RefinementMappingError(
+            f"{len(unused)} provisional residues remained unmapped after coordinate matching"
+        )
+    return result
+
+
+def compatible_resnames(left: str, right: str) -> bool:
+    if left == right:
+        return True
+    return any(left in family and right in family for family in _PROTONATION_FAMILIES)
+
+
+def is_virtual_site(atom: AtomRecord) -> bool:
+    element = (atom.element or "").strip().upper()
+    name = atom.name.strip().upper()
+    return element in {"EP", "LP", "X"} or name.startswith("EP") or name.startswith("LP")
+
+
+def _anchor_rmsd(
+    input_residue: ResidueRecord,
+    topology_residue: ResidueRecord,
+) -> float | None:
+    input_by_name = _unique_physical_atoms(input_residue)
+    topology_by_name = _unique_physical_atoms(topology_residue)
+    common_heavy = sorted(
+        name
+        for name in input_by_name.keys() & topology_by_name.keys()
+        if not _is_hydrogen(input_by_name[name])
+        and not _is_hydrogen(topology_by_name[name])
+    )
+    names = common_heavy or sorted(input_by_name.keys() & topology_by_name.keys())
+    if not names:
+        return None
+    squared = 0.0
+    for name in names:
+        left = input_by_name[name]
+        right = topology_by_name[name]
+        if _element(left) != _element(right):
+            return None
+        squared += (
+            (left.x - right.x) ** 2
+            + (left.y - right.y) ** 2
+            + (left.z - right.z) ** 2
+        )
+    return sqrt(squared / len(names))
+
+
+def _unique_physical_atoms(residue: ResidueRecord) -> dict[str, AtomRecord]:
+    result: dict[str, AtomRecord] = {}
+    duplicates: set[str] = set()
+    for atom in residue.atoms:
+        if is_virtual_site(atom):
+            continue
+        if atom.name in result:
+            duplicates.add(atom.name)
+        result[atom.name] = atom
+    for name in duplicates:
+        result.pop(name, None)
+    return result
+
+
+def _element(atom: AtomRecord) -> str:
+    return (atom.element or "").strip().upper()
+
+
+def _is_hydrogen(atom: AtomRecord) -> bool:
+    return _element(atom) == "H"

mdprep/refinement/report.py

@@ -0,0 +1,66 @@
+"""JSON and Markdown reporting for QM/MM refinement."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+from mdprep.refinement.workflow import RefinementResult
+
+
+def write_refinement_reports(
+    result: RefinementResult,
+    *,
+    json_path: str | Path,
+    markdown_path: str | Path,
+) -> dict[str, object]:
+    report = result.to_report_dict()
+    json_output = Path(json_path)
+    markdown_output = Path(markdown_path)
+    json_output.parent.mkdir(parents=True, exist_ok=True)
+    markdown_output.parent.mkdir(parents=True, exist_ok=True)
+    json_output.write_text(
+        json.dumps(report, indent=2, sort_keys=True) + "\n",
+        encoding="utf-8",
+    )
+    selection = report["selection"]
+    assert isinstance(selection, dict)
+    ash = report["ash"]
+    assert isinstance(ash, dict)
+    lines = [
+        "# QM/MM refinement report",
+        "",
+        "- Status: complete",
+        "- Backend: ASH",
+        f"- QM method: {report['qm_method']}",
+        "- MM theory: provisional Amber/OpenMM",
+        f"- Embedding: {report['embedding']}",
+        "- Provisional ligand/cofactor charges: AM1-BCC",
+        f"- QM atoms: {len(selection['qm_atom_indices_zero_based'])}",
+        f"- Active atoms: {len(selection['active_atom_indices_zero_based'])}",
+        f"- Movable-atom policy: {selection['movable_atom_policy']}",
+        f"- Active waters: {report['active_water_residue_count']}",
+        f"- Non-water active-region cutoff (Å): {selection['active_region_cutoff_angstrom']}",
+        f"- Active-water cutoff (Å): {selection['active_water_cutoff_angstrom']}",
+        f"- Total QM charge: {selection['total_qm_charge']}",
+        f"- Total QM multiplicity: {selection['total_qm_multiplicity']}",
+        f"- Final energy (Eh): {ash['final_energy_hartree']}",
+        f"- Refined hydrogenated PDB: `{report['refined_hydrogenated_pdb_path']}`",
+        "",
+        "## QM charge derivation",
+        "",
+    ]
+    for component in selection["charge_components"]:
+        lines.append(
+            f"- {component['label']}: {component['charge']:+d} ({component['kind']})"
+        )
+    lines.extend(["", "## External command", ""])
+    lines.append(f"- Command: `{' '.join(ash['command'])}`")
+    lines.append(f"- Working directory: `{ash['cwd']}`")
+    lines.append(f"- Return code: {ash['returncode']}")
+    lines.append(f"- Runtime (s): {ash['runtime_seconds']}")
+    lines.extend(["", "## Notes", ""])
+    for warning in report["warnings"]:
+        lines.append(f"- {warning}")
+    markdown_output.write_text("\n".join(lines) + "\n", encoding="utf-8")
+    return report