mdprep 0.2.0__py3-none-any.whl

mdprep/protonation/temporary_hydrogenation.py

@@ -0,0 +1,238 @@
+"""Temporary PDBFixer hydrogenation for xTB histidine-tautomer clusters.
+
+The generated hydrogens are environment atoms for HID/HIE ranking only.  They
+are never transferred to the protonation-stage output structure; the normal
+Amber/tleap hydrogenation remains authoritative for the prepared system.
+"""
+
+from __future__ import annotations
+
+from collections import OrderedDict
+from dataclasses import dataclass, replace
+from importlib.metadata import PackageNotFoundError, version
+from math import dist
+from pathlib import Path
+import random
+
+from mdprep.structure.models import AtomRecord, PdbStructure, ResidueId, ResidueRecord
+from mdprep.structure.pdb import read_pdb
+from mdprep.structure.writer import write_pdb
+
+
+class TemporaryHydrogenationError(ValueError):
+    """Raised when a safe temporary hydrogenated environment cannot be built."""
+
+
+@dataclass(frozen=True)
+class TemporaryHydrogenationResult:
+    backend: str
+    backend_version: str | None
+    platform: str
+    ph: float
+    random_seed: int
+    input_path: Path
+    raw_output_path: Path
+    restored_output_path: Path
+    structure: PdbStructure
+    added_hydrogen_count: int
+    maximum_original_atom_displacement_angstrom: float
+    final_prepared_pdb_modified: bool = False
+
+    def to_dict(self) -> dict[str, object]:
+        return {
+            "backend": self.backend,
+            "backend_version": self.backend_version,
+            "platform": self.platform,
+            "ph": self.ph,
+            "random_seed": self.random_seed,
+            "input_path": str(self.input_path),
+            "raw_output_path": str(self.raw_output_path),
+            "restored_output_path": str(self.restored_output_path),
+            "added_hydrogen_count": self.added_hydrogen_count,
+            "maximum_original_atom_displacement_angstrom": (
+                self.maximum_original_atom_displacement_angstrom
+            ),
+            "final_prepared_pdb_modified": self.final_prepared_pdb_modified,
+            "interpretation": (
+                "PDBFixer hydrogens complete the local xTB HID/HIE environment only. "
+                "They are not copied into the protonation-stage or final prepared PDB."
+            ),
+        }
+
+
+def add_temporary_protein_hydrogens(
+    structure: PdbStructure,
+    *,
+    ph: float,
+    work_dir: str | Path,
+    random_seed: int = 20260722,
+) -> TemporaryHydrogenationResult:
+    """Add a complete temporary H environment without changing input atoms."""
+
+    try:
+        from openmm import Platform
+        from openmm.app import PDBFile
+        from pdbfixer import PDBFixer
+    except Exception as exc:
+        raise TemporaryHydrogenationError(
+            "protonation.histidine.xtb.add_missing_protein_hydrogens was requested, "
+            "but PDBFixer and OpenMM are unavailable. Install PDBFixer or set the "
+            "option to false and provide a fully hydrogenated input model."
+        ) from exc
+
+    work = Path(work_dir)
+    work.mkdir(parents=True, exist_ok=True)
+    input_path = work / "histidine_xtb_environment_input.pdb"
+    raw_path = work / "histidine_xtb_environment.pdbfixer_raw.pdb"
+    restored_path = work / "histidine_xtb_environment.hydrogenated.pdb"
+    write_pdb(structure, input_path)
+
+    try:
+        platform_name = "Reference"
+        platform = Platform.getPlatformByName(platform_name)
+        fixer = PDBFixer(filename=str(input_path), platform=platform)
+        # OpenMM Modeller places hydrogens with Python's global random
+        # generator before minimizing them.  Restore the caller's RNG state
+        # after making this scientific input reproducible.
+        random_state = random.getstate()
+        random.seed(random_seed)
+        try:
+            fixer.addMissingHydrogens(float(ph))
+        finally:
+            random.setstate(random_state)
+        with raw_path.open("w", encoding="utf-8") as handle:
+            PDBFile.writeFile(fixer.topology, fixer.positions, handle, keepIds=True)
+    except Exception as exc:
+        raise TemporaryHydrogenationError(
+            f"PDBFixer could not build the temporary xTB hydrogen environment: {exc}"
+        ) from exc
+
+    raw = read_pdb(raw_path)
+    restored, maximum_displacement = _restore_and_validate(
+        structure,
+        raw,
+        output_path=restored_path,
+    )
+    write_pdb(restored, restored_path)
+    added = len(restored.atoms) - len(structure.atoms)
+    if added <= 0:
+        raise TemporaryHydrogenationError(
+            "PDBFixer did not add any temporary hydrogens to a dehydrogenated "
+            "histidine-ranking environment."
+        )
+    return TemporaryHydrogenationResult(
+        backend="pdbfixer",
+        backend_version=_package_version("pdbfixer"),
+        platform=platform_name,
+        ph=float(ph),
+        random_seed=random_seed,
+        input_path=input_path,
+        raw_output_path=raw_path,
+        restored_output_path=restored_path,
+        structure=restored,
+        added_hydrogen_count=added,
+        maximum_original_atom_displacement_angstrom=maximum_displacement,
+    )
+
+
+def _restore_and_validate(
+    reference: PdbStructure,
+    observed: PdbStructure,
+    *,
+    output_path: Path,
+) -> tuple[PdbStructure, float]:
+    if len(observed.residues) != len(reference.residues):
+        raise TemporaryHydrogenationError(
+            "Temporary PDBFixer hydrogenation changed the residue count: "
+            f"{len(reference.residues)} -> {len(observed.residues)}."
+        )
+    atoms: list[AtomRecord] = []
+    maximum_displacement = 0.0
+    serial = 0
+    for expected, actual in zip(reference.residues, observed.residues, strict=True):
+        expected_key = (expected.id.chain_id, expected.id.resid, expected.id.icode)
+        actual_key = (actual.id.chain_id, actual.id.resid, actual.id.icode)
+        if expected_key != actual_key or expected.id.resname != actual.id.resname:
+            raise TemporaryHydrogenationError(
+                "Temporary PDBFixer hydrogenation changed residue identity/order: "
+                f"expected {expected.id.display()}, found {actual.id.display()}."
+            )
+        expected_by_name = {atom.name: atom for atom in expected.atoms}
+        if len(expected_by_name) != len(expected.atoms):
+            raise TemporaryHydrogenationError(
+                f"Input residue {expected.id.display()} has duplicate atom names."
+            )
+        actual_names = {atom.name for atom in actual.atoms}
+        missing = sorted(set(expected_by_name) - actual_names)
+        if missing:
+            raise TemporaryHydrogenationError(
+                f"Temporary PDBFixer hydrogenation removed atoms from "
+                f"{expected.id.display()}: {missing}"
+            )
+        residue_record_name = "HETATM" if expected.record_names == {"HETATM"} else "ATOM"
+        for atom in actual.atoms:
+            serial += 1
+            original = expected_by_name.get(atom.name)
+            if original is not None:
+                displacement = dist(
+                    (original.x, original.y, original.z),
+                    (atom.x, atom.y, atom.z),
+                )
+                maximum_displacement = max(maximum_displacement, displacement)
+                if displacement > 0.002:
+                    raise TemporaryHydrogenationError(
+                        "Temporary PDBFixer hydrogenation moved an original atom by "
+                        f"{displacement:.6f} A: {original.atom_identity}."
+                    )
+            elif (atom.element or "").strip().upper() != "H":
+                raise TemporaryHydrogenationError(
+                    "Temporary PDBFixer hydrogenation added a non-hydrogen atom "
+                    f"to {expected.id.display()}: {atom.name}."
+                )
+            atoms.append(
+                replace(
+                    atom,
+                    serial=serial,
+                    resname=expected.id.resname,
+                    chain_id=expected.id.chain_id,
+                    resid=expected.id.resid,
+                    icode=expected.id.icode,
+                    record_name=(original.record_name if original else residue_record_name),
+                    occupancy=(original.occupancy if original else atom.occupancy),
+                    bfactor=(original.bfactor if original else atom.bfactor),
+                    element=(original.element if original else atom.element),
+                )
+            )
+    return (
+        PdbStructure(
+            path=output_path,
+            atoms=atoms,
+            residues=_build_residues(atoms),
+            model_count=1,
+            used_model=1,
+            warnings=list(reference.warnings),
+        ),
+        maximum_displacement,
+    )
+
+
+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=key[0], resname=key[1], resid=key[2], icode=key[3]),
+            atoms=residue_atoms,
+            record_names={atom.record_name for atom in residue_atoms},
+            original_index=index,
+        )
+        for index, (key, residue_atoms) in enumerate(grouped.items())
+    ]
+
+
+def _package_version(name: str) -> str | None:
+    try:
+        return version(name)
+    except PackageNotFoundError:
+        return None

mdprep/protonation/xtb_parser.py

@@ -0,0 +1,79 @@
+"""xTB output energy parsing and tautomer selection helpers."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+import re
+
+
+HARTREE_TO_KCAL_MOL = 627.509474
+
+
+class XtbParseError(ValueError):
+    """Raised when xTB output does not contain a parseable energy."""
+
+
+ENERGY_PATTERNS = [
+    re.compile(r"TOTAL\s+ENERGY\s+(-?\d+(?:\.\d+)?)", re.IGNORECASE),
+    re.compile(r"total\s+energy\s+(-?\d+(?:\.\d+)?)", re.IGNORECASE),
+    re.compile(r"::\s*total\s+energy\s+(-?\d+(?:\.\d+)?)", re.IGNORECASE),
+    re.compile(r"total\s+E\s*=\s*(-?\d+(?:\.\d+)?)", re.IGNORECASE),
+]
+
+
+@dataclass(frozen=True)
+class XtbEnergyComparison:
+    hid_energy_hartree: float
+    hie_energy_hartree: float
+    delta_kcal_mol: float
+    selected_state: str
+    close_call: bool
+    warnings: list[str]
+
+
+def parse_xtb_energy_text(text: str, *, source: str = "xTB stdout") -> float:
+    energies: list[float] = []
+    for line in text.splitlines():
+        for pattern in ENERGY_PATTERNS:
+            match = pattern.search(line)
+            if match:
+                energies.append(float(match.group(1)))
+                break
+    if not energies:
+        raise XtbParseError(f"Could not parse xTB total energy from {source}")
+    return energies[-1]
+
+
+def parse_xtb_energy_file(path: str | Path) -> float:
+    path_obj = Path(path)
+    return parse_xtb_energy_text(path_obj.read_text(encoding="utf-8"), source=str(path_obj))
+
+
+def compare_hid_hie_energies(
+    *,
+    hid_energy_hartree: float,
+    hie_energy_hartree: float,
+    close_call_kcal_mol: float,
+) -> XtbEnergyComparison:
+    delta = (hid_energy_hartree - hie_energy_hartree) * HARTREE_TO_KCAL_MOL
+    warnings: list[str] = []
+    if hid_energy_hartree < hie_energy_hartree:
+        selected = "HID"
+    elif hie_energy_hartree < hid_energy_hartree:
+        selected = "HIE"
+    else:
+        selected = "HIE"
+        warnings.append("HID and HIE xTB energies are equal; selected HIE by deterministic tie-break.")
+    close_call = abs(delta) <= close_call_kcal_mol
+    if close_call:
+        warnings.append("HID/HIE xTB energy difference is within the configured close-call threshold.")
+    return XtbEnergyComparison(
+        hid_energy_hartree=hid_energy_hartree,
+        hie_energy_hartree=hie_energy_hartree,
+        delta_kcal_mol=delta,
+        selected_state=selected,
+        close_call=close_call,
+        warnings=warnings,
+    )
+

mdprep/protonation/xtb_runner.py

@@ -0,0 +1,108 @@
+"""xTB command construction and execution."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+
+from mdprep.config.models import HistidineXtbConfig
+from mdprep.external.discovery import which_executable
+from mdprep.external.runner import CommandResult, run_command
+
+
+class XtbExecutionError(RuntimeError):
+    """Raised when xTB is requested but unavailable or fails."""
+
+
+@dataclass(frozen=True)
+class XtbRunResult:
+    command_result: CommandResult
+    stdout_path: Path
+    stderr_path: Path
+
+
+def resolve_xtb_executable(config: HistidineXtbConfig) -> str:
+    configured = config.executable
+    if Path(configured).exists():
+        return configured
+    found = which_executable(configured)
+    if found:
+        return found
+    raise XtbExecutionError(f"xTB executable not found: {configured}")
+
+
+def build_xtb_command(
+    *,
+    config: HistidineXtbConfig,
+    xyz_path: str | Path,
+    cluster_charge: int,
+    executable: str | None = None,
+    input_path: str | Path | None = None,
+) -> list[str]:
+    exe = executable or config.executable
+    command = [exe, str(xyz_path)]
+    extra_args = list(config.extra_args)
+    if config.model == "gfn2":
+        command.extend(["--gfn", "2"])
+    elif "--gxtb" not in extra_args:
+        command.append("--gxtb")
+    if config.mode == "opt":
+        command.extend(["--opt", config.opt_level])
+    command.extend(["--chrg", str(cluster_charge), "--uhf", "0"])
+    command.extend(["--iterations", str(config.scf_iterations)])
+    if config.electronic_temperature_kelvin is not None:
+        command.extend(["--etemp", str(config.electronic_temperature_kelvin)])
+    if config.solvent is not None and not _extra_args_define_solvent(extra_args):
+        command.extend(["--alpb", config.solvent])
+    if input_path is not None:
+        command.extend(["--input", str(input_path)])
+    command.extend(extra_args)
+    return command
+
+
+def run_xtb(
+    *,
+    config: HistidineXtbConfig,
+    xyz_path: str | Path,
+    work_dir: str | Path,
+    cluster_charge: int,
+    stdout_path: str | Path,
+    stderr_path: str | Path,
+    input_path: str | Path | None = None,
+) -> XtbRunResult:
+    executable = resolve_xtb_executable(config)
+    command = build_xtb_command(
+        config=config,
+        xyz_path=Path(xyz_path).name,
+        cluster_charge=cluster_charge,
+        executable=executable,
+        input_path=Path(input_path).name if input_path is not None else None,
+    )
+    result = run_command(command, cwd=work_dir)
+    stdout = Path(stdout_path)
+    stderr = Path(stderr_path)
+    stdout.write_text(result.stdout, encoding="utf-8")
+    stderr.write_text(result.stderr, encoding="utf-8")
+    if result.returncode != 0:
+        raise XtbExecutionError(
+            "xTB failed with exit code "
+            f"{result.returncode}. Command: {' '.join(result.command)}. "
+            f"See {stdout} and {stderr}.\n"
+            f"stdout tail:\n{_tail(result.stdout)}\n"
+            f"stderr tail:\n{_tail(result.stderr)}"
+        )
+    return XtbRunResult(command_result=result, stdout_path=stdout, stderr_path=stderr)
+
+
+def _extra_args_define_solvent(extra_args: list[str]) -> bool:
+    return any(arg in {"--alpb", "--gbsa"} for arg in extra_args)
+
+
+def _tail(text: str, *, max_lines: int = 20, max_chars: int = 4000) -> str:
+    if not text:
+        return "<empty>"
+    lines = text.rstrip().splitlines()[-max_lines:]
+    tail = "\n".join(lines)
+    if len(tail) > max_chars:
+        return tail[-max_chars:]
+    return tail

mdprep/qm/__init__.py

@@ -0,0 +1,2 @@
+"""QM backend helpers."""
+

mdprep/qm/point_charges.py

@@ -0,0 +1,233 @@
+"""MM point-charge extraction for QMMESP electrostatic embedding."""
+
+from __future__ import annotations
+
+import csv
+import json
+from dataclasses import dataclass
+from pathlib import Path
+
+import numpy as np
+
+from mdprep.config.models import LigandConfig, ManifestConfig
+from mdprep.structure.classify import STANDARD_PROTEIN_RESIDUES, WATER_RESIDUES
+
+
+class PointChargeError(ValueError):
+    """Raised when QMMESP point charges cannot be extracted safely."""
+
+
+@dataclass(frozen=True)
+class PointCharge:
+    x: float
+    y: float
+    z: float
+    charge: float
+    residue_name: str
+    residue_number: int | None
+    atom_name: str
+    category: str
+
+    def to_dict(self) -> dict[str, object]:
+        return {
+            "x": self.x,
+            "y": self.y,
+            "z": self.z,
+            "charge": self.charge,
+            "residue_name": self.residue_name,
+            "residue_number": self.residue_number,
+            "atom_name": self.atom_name,
+            "category": self.category,
+        }
+
+
+@dataclass(frozen=True)
+class PointChargeSelection:
+    target_atom_indices: list[int]
+    point_charges: list[PointCharge]
+    total_before_cutoff: int
+    total_after_cutoff: int
+    net_embedding_charge: float
+    min_distance: float | None
+    max_distance: float | None
+    categories: dict[str, int]
+
+    def to_dict(self) -> dict[str, object]:
+        return {
+            "target_atom_indices": self.target_atom_indices,
+            "target_atom_count": len(self.target_atom_indices),
+            "target_ligand_excluded_from_embedding": True,
+            "point_charge_count_before_cutoff": self.total_before_cutoff,
+            "point_charge_count_after_cutoff": self.total_after_cutoff,
+            "net_embedding_charge": self.net_embedding_charge,
+            "min_distance": self.min_distance,
+            "max_distance": self.max_distance,
+            "categories": self.categories,
+            "qmmesp_interpretation": (
+                "The target-ligand SCF uses electrostatic QM/MM embedding through "
+                "pyscf.qmmm.mm_charge. MM point charges polarize the target-ligand QM density; "
+                "they are not fitted and are not written to the ligand mol2."
+            ),
+        }
+
+    @property
+    def charge_array(self) -> np.ndarray:
+        return np.asarray([charge.charge for charge in self.point_charges], dtype=float)
+
+    @property
+    def coordinate_array(self) -> np.ndarray:
+        return np.asarray([[charge.x, charge.y, charge.z] for charge in self.point_charges], dtype=float)
+
+
+def parmed_available() -> bool:
+    try:
+        import parmed  # noqa: F401
+    except Exception:
+        return False
+    return True
+
+
+def extract_point_charges_from_prmtop(
+    *,
+    prmtop: str | Path,
+    inpcrd: str | Path,
+    ligand: LigandConfig,
+    manifest: ManifestConfig,
+    target_coordinates: np.ndarray,
+    target_atom_names: list[str] | None = None,
+) -> PointChargeSelection:
+    try:
+        import parmed
+    except Exception as exc:
+        raise PointChargeError("ParmEd is required for qmmesp_pyscf point-charge extraction.") from exc
+    try:
+        structure = parmed.load_file(str(prmtop), str(inpcrd))
+    except Exception as exc:
+        raise PointChargeError(f"Could not load provisional Amber topology with ParmEd: {exc}") from exc
+    coordinates = np.asarray(structure.coordinates, dtype=float)
+    target_indices = _find_target_ligand_atoms(
+        structure,
+        ligand,
+        topology_coordinates=coordinates,
+        target_coordinates=target_coordinates,
+        target_atom_names=target_atom_names,
+    )
+    target_set = set(target_indices)
+    target_coords = np.asarray(target_coordinates, dtype=float)
+    charges: list[PointCharge] = []
+    before_cutoff = 0
+    distances_kept: list[float] = []
+    categories: dict[str, int] = {}
+    for index, atom in enumerate(structure.atoms):
+        if index in target_set:
+            continue
+        category = _category(atom)
+        env = ligand.qmmesp.environment if ligand.qmmesp is not None else None
+        if env is not None:
+            if category == "protein" and not env.include_protein:
+                continue
+            if category == "water" and not env.include_waters:
+                continue
+            if category == "ligand" and not env.include_other_ligands:
+                continue
+        before_cutoff += 1
+        distance = float(np.min(np.linalg.norm(target_coords - coordinates[index], axis=1)))
+        cutoff = ligand.qmmesp.embedding_cutoff_angstrom if ligand.qmmesp is not None else None
+        if cutoff is not None and distance > cutoff:
+            continue
+        residue = atom.residue
+        charge = PointCharge(
+            x=float(coordinates[index][0]),
+            y=float(coordinates[index][1]),
+            z=float(coordinates[index][2]),
+            charge=float(atom.charge),
+            residue_name=str(residue.name),
+            residue_number=getattr(residue, "number", None),
+            atom_name=str(atom.name),
+            category=category,
+        )
+        charges.append(charge)
+        distances_kept.append(distance)
+        categories[category] = categories.get(category, 0) + 1
+    return PointChargeSelection(
+        target_atom_indices=target_indices,
+        point_charges=charges,
+        total_before_cutoff=before_cutoff,
+        total_after_cutoff=len(charges),
+        net_embedding_charge=float(sum(charge.charge for charge in charges)),
+        min_distance=min(distances_kept) if distances_kept else None,
+        max_distance=max(distances_kept) if distances_kept else None,
+        categories=categories,
+    )
+
+
+def write_point_charge_files(selection: PointChargeSelection, *, csv_path: str | Path, xyz_path: str | Path, summary_path: str | Path) -> None:
+    csv_output = Path(csv_path)
+    csv_output.parent.mkdir(parents=True, exist_ok=True)
+    with csv_output.open("w", encoding="utf-8", newline="") as handle:
+        writer = csv.DictWriter(handle, fieldnames=["x", "y", "z", "charge", "residue_name", "residue_number", "atom_name", "category"])
+        writer.writeheader()
+        for charge in selection.point_charges:
+            writer.writerow(charge.to_dict())
+    xyz_output = Path(xyz_path)
+    lines = [str(len(selection.point_charges)), "mdprep MM point charges"]
+    for charge in selection.point_charges:
+        lines.append(f"X {charge.x:.8f} {charge.y:.8f} {charge.z:.8f} {charge.charge:.8f}")
+    xyz_output.write_text("\n".join(lines) + "\n", encoding="utf-8")
+    Path(summary_path).write_text(json.dumps(selection.to_dict(), indent=2, sort_keys=True) + "\n", encoding="utf-8")
+
+
+def _find_target_ligand_atoms(
+    structure: object,
+    ligand: LigandConfig,
+    *,
+    topology_coordinates: np.ndarray,
+    target_coordinates: np.ndarray,
+    target_atom_names: list[str] | None,
+    coordinate_tolerance_angstrom: float = 0.1,
+) -> list[int]:
+    target = np.asarray(target_coordinates, dtype=float)
+    if target.ndim != 2 or target.shape[1] != 3 or len(target) == 0:
+        raise PointChargeError("Target-ligand coordinates must have shape (n_atoms, 3) and be non-empty.")
+    if target_atom_names is not None and len(target_atom_names) != len(target):
+        raise PointChargeError("Target-ligand atom-name and coordinate counts do not match.")
+
+    matches: list[tuple[list[int], float]] = []
+    rejected: list[str] = []
+    for residue in structure.residues:
+        if residue.name != ligand.selector.resname:
+            continue
+        indices = [atom.idx for atom in residue.atoms]
+        if len(indices) != len(target):
+            rejected.append(f"residue {getattr(residue, 'number', '?')}: {len(indices)} atoms")
+            continue
+        names = [str(atom.name) for atom in residue.atoms]
+        if target_atom_names is not None and names != target_atom_names:
+            rejected.append(f"residue {getattr(residue, 'number', '?')}: atom names {names}")
+            continue
+        candidate = topology_coordinates[indices]
+        max_deviation = float(np.max(np.linalg.norm(candidate - target, axis=1)))
+        if max_deviation <= coordinate_tolerance_angstrom:
+            matches.append((indices, max_deviation))
+        else:
+            rejected.append(
+                f"residue {getattr(residue, 'number', '?')}: coordinate max deviation "
+                f"{max_deviation:.3f} A"
+            )
+    if len(matches) != 1:
+        details = "; ".join(rejected) if rejected else "no residue-name candidates"
+        raise PointChargeError(
+            f"Could not map ligand {ligand.id} uniquely in provisional topology; "
+            "mapping requires residue name, atom count/order/names, and coordinates to agree. "
+            f"Candidates: {details}."
+        )
+    return matches[0][0]
+
+
+def _category(atom: object) -> str:
+    name = str(atom.residue.name).upper()
+    if name in WATER_RESIDUES:
+        return "water"
+    if name in STANDARD_PROTEIN_RESIDUES:
+        return "protein"
+    return "ligand"