mdprep 0.2.0__py3-none-any.whl

mdprep/external/__init__.py

@@ -0,0 +1,14 @@
+"""External command helpers."""
+
+from __future__ import annotations
+
+from mdprep.external.discovery import optional_executable_report, which_executable
+from mdprep.external.runner import CommandResult, run_command
+
+__all__ = [
+    "CommandResult",
+    "optional_executable_report",
+    "run_command",
+    "which_executable",
+]
+

mdprep/external/discovery.py

@@ -0,0 +1,14 @@
+"""Optional external executable discovery."""
+
+from __future__ import annotations
+
+from shutil import which
+
+
+def which_executable(name: str) -> str | None:
+    return which(name)
+
+
+def optional_executable_report(names: list[str]) -> dict[str, str | None]:
+    return {name: which_executable(name) for name in names}
+

mdprep/external/runner.py

@@ -0,0 +1,55 @@
+"""Small subprocess wrapper with reproducibility metadata."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+from time import perf_counter
+from typing import Mapping, Sequence
+import subprocess
+
+
+@dataclass(frozen=True)
+class CommandResult:
+    command: tuple[str, ...]
+    cwd: str
+    returncode: int
+    stdout: str
+    stderr: str
+    runtime_seconds: float
+
+
+def run_command(
+    command: Sequence[str],
+    *,
+    cwd: str | Path | None = None,
+    env: Mapping[str, str] | None = None,
+    timeout: float | None = None,
+    check: bool = False,
+) -> CommandResult:
+    workdir = Path(cwd) if cwd is not None else Path.cwd()
+    started = perf_counter()
+    completed = subprocess.run(
+        list(command),
+        cwd=workdir,
+        env=dict(env) if env is not None else None,
+        timeout=timeout,
+        text=True,
+        capture_output=True,
+        check=False,
+    )
+    runtime = perf_counter() - started
+    result = CommandResult(
+        command=tuple(str(part) for part in command),
+        cwd=str(workdir),
+        returncode=completed.returncode,
+        stdout=completed.stdout,
+        stderr=completed.stderr,
+        runtime_seconds=runtime,
+    )
+    if check and result.returncode != 0:
+        raise RuntimeError(
+            f"Command failed with exit code {result.returncode}: {' '.join(result.command)}"
+        )
+    return result
+

mdprep/leap/__init__.py

@@ -0,0 +1,2 @@
+"""tleap system-build helpers."""
+

mdprep/leap/builder.py

@@ -0,0 +1,476 @@
+"""Build and run tleap scripts for final Amber outputs."""
+
+from __future__ import annotations
+
+import shutil
+from dataclasses import dataclass
+import os
+from pathlib import Path
+
+from mdprep.config.models import ManifestConfig
+from mdprep.leap.forcefields import ForceFieldError, ForceFieldSources, forcefield_sources
+from mdprep.leap.ions import IonPlan, IonPlanError, amber_inpcrd_box_volume, build_ion_plan
+from mdprep.leap.log_parser import LeapLogError, assert_tleap_success
+from mdprep.leap.residues import (
+    DisulfideBondCommand,
+    LeapInputResult,
+    LeapResidueError,
+    LigandParameterFiles,
+    append_disulfide_conect_records,
+    disulfide_bond_commands,
+    prepare_leap_input_pdb,
+    validate_ligand_parameter_files,
+    validate_tleap_ligand_coordinates,
+)
+from mdprep.leap.runner import TLeapRun, TLeapRunError, run_tleap
+from mdprep.ligands.workflow import LigandStageResult
+from mdprep.metals.workflow import MetalStageResult
+from mdprep.metals.mcpb import McpbRefittedLigand
+from mdprep.metals.c4 import C4PostprocessError, C4PostprocessRun, apply_12_6_4_c4
+from mdprep.protonation.apply import ProtonationResult
+from mdprep.structure.models import PdbStructure
+
+
+class TLeapBuildError(ValueError):
+    """Raised when final tleap system generation cannot complete safely."""
+
+
+@dataclass(frozen=True)
+class TLeapOutputs:
+    prmtop: Path
+    inpcrd: Path
+    pdb: Path
+
+    def to_dict(self) -> dict[str, object]:
+        return {
+            "prmtop": str(self.prmtop),
+            "inpcrd": str(self.inpcrd),
+            "pdb": str(self.pdb),
+        }
+
+
+@dataclass(frozen=True)
+class TLeapStageResult:
+    forcefields: ForceFieldSources
+    leap_input: LeapInputResult
+    ligands: list[LigandParameterFiles]
+    mcpb_resp_ligands: list[McpbRefittedLigand]
+    disulfide_bonds: list[DisulfideBondCommand]
+    dry_run: TLeapRun
+    dry_outputs: TLeapOutputs
+    final_run: TLeapRun | None
+    final_outputs: TLeapOutputs
+    solvation_enabled: bool
+    solvation_box_type: str
+    buffer_angstrom: float
+    neutralization_requested: bool
+    salt_concentration_molar: float
+    ion_plan: IonPlan | None
+    salt_volume_a3: float | None
+    c4_postprocessing: list[C4PostprocessRun]
+    warnings: list[str]
+
+    def to_report_dict(self) -> dict[str, object]:
+        return {
+            "force_fields_sourced": self.forcefields.all_sources,
+            "force_field_warnings": self.forcefields.warnings,
+            "water_box": self.forcefields.water_box,
+            "ligands": [ligand.to_dict() for ligand in self.ligands],
+            "mcpb_resp_ligands": [
+                ligand.to_dict() for ligand in self.mcpb_resp_ligands
+            ],
+            "disulfide_bond_commands": [bond.to_dict() for bond in self.disulfide_bonds],
+            "leap_input": self.leap_input.to_dict(),
+            "dry_tleap": self.dry_run.to_dict(),
+            "dry_outputs": self.dry_outputs.to_dict(),
+            "dry_system_charge": self.dry_run.summary.total_charge,
+            "solvation_enabled": self.solvation_enabled,
+            "solvation_box_type": self.solvation_box_type,
+            "buffer_angstrom": self.buffer_angstrom,
+            "neutralization_requested": self.neutralization_requested,
+            "neutralizing_ions_added": (
+                self.ion_plan.neutralizing_count if self.ion_plan is not None else 0
+            ),
+            "salt_concentration_molar": self.salt_concentration_molar,
+            "salt_ion_pairs_requested": self.ion_plan.salt_pairs if self.ion_plan is not None else 0,
+            "salt_volume_a3": self.salt_volume_a3,
+            "c4_postprocessing": [run.to_dict() for run in self.c4_postprocessing],
+            "final_tleap": self.final_run.to_dict() if self.final_run is not None else None,
+            "final_outputs": self.final_outputs.to_dict(),
+            "warnings": self.warnings,
+        }
+
+
+def run_tleap_stage(
+    *,
+    structure: PdbStructure,
+    manifest: ManifestConfig,
+    output_dir: str | Path,
+    protonation_result: ProtonationResult,
+    ligand_result: LigandStageResult,
+    metal_result: MetalStageResult | None = None,
+) -> TLeapStageResult:
+    output = Path(output_dir)
+    input_dir = output / "leap" / "input"
+    dry_dir = output / "leap" / "dry"
+    final_dir = output / "final"
+    input_dir.mkdir(parents=True, exist_ok=True)
+    dry_dir.mkdir(parents=True, exist_ok=True)
+    final_dir.mkdir(parents=True, exist_ok=True)
+
+    try:
+        sources = forcefield_sources(
+            protein_forcefield=manifest.protein.forcefield,
+            water_model=manifest.protein.water_model,
+            ligands=manifest.ligands,
+        )
+        metal_structure = (
+            metal_result.structure
+            if metal_result is not None and metal_result.structure is not None
+            else structure
+        )
+        bonded_mcpb_complete = bool(
+            metal_result is not None
+            and metal_result.mcpb_site is not None
+            and metal_result.mcpb_site.complete
+        )
+        mcpb_resp_ligand_ids = {
+            ligand.id
+            for ligand in manifest.ligands
+            if ligand.charge_method == "mcpb_resp_pyscf"
+        }
+        refitted_ligands = (
+            list(metal_result.mcpb_site.refitted_ligands)
+            if bonded_mcpb_complete
+            and metal_result is not None
+            and metal_result.mcpb_site is not None
+            else []
+        )
+        if {item.ligand_id for item in refitted_ligands} != mcpb_resp_ligand_ids:
+            raise LeapResidueError(
+                "Final tLEAP build requires one MCPB RESP-fitted mol2 for every "
+                "mcpb_resp_pyscf ligand; expected "
+                f"{sorted(mcpb_resp_ligand_ids)}, obtained "
+                f"{sorted(item.ligand_id for item in refitted_ligands)}."
+            )
+        leap_input = prepare_leap_input_pdb(
+            metal_structure,
+            input_dir / "system.leap_input.pdb",
+            manifest=None if bonded_mcpb_complete else manifest,
+            ligand_result=None if bonded_mcpb_complete else ligand_result,
+            # MCPB.py writes bond commands against its own deterministic,
+            # one-based residue order. Original PDB residue numbers can be
+            # negative, zero, non-contiguous, or insertion-coded, and LEaP
+            # otherwise remaps those numbers independently. Keep atom order
+            # and coordinates unchanged while making the input numbering
+            # match MCPB's command indices exactly.
+            renumber_residues_sequentially=bonded_mcpb_complete,
+        )
+        ligand_files = validate_ligand_parameter_files(
+            manifest=manifest,
+            structure=structure,
+            ligand_result=ligand_result,
+            excluded_ligand_ids=mcpb_resp_ligand_ids,
+        )
+        disulfide_bonds = (
+            []
+            if bonded_mcpb_complete
+            else disulfide_bond_commands(
+                structure=leap_input.structure,
+                protonation_result=protonation_result,
+            )
+        )
+        conect_records = append_disulfide_conect_records(leap_input.path, disulfide_bonds)
+        dry_outputs = TLeapOutputs(
+            prmtop=dry_dir / "system.dry.prmtop",
+            inpcrd=dry_dir / "system.dry.inpcrd",
+            pdb=dry_dir / "system.dry.pdb",
+        )
+        dry_script = build_tleap_script(
+            sources=sources,
+            ligands=ligand_files,
+            input_pdb=leap_input.path,
+            disulfide_bonds=disulfide_bonds,
+            outputs=dry_outputs,
+            work_dir=dry_dir,
+            setup_commands=metal_result.leap_setup_commands if metal_result else None,
+            extra_bond_commands=metal_result.leap_bond_commands if metal_result else None,
+        )
+        dry_input = dry_dir / "tleap.in"
+        dry_input.write_text(dry_script, encoding="utf-8")
+        dry_run = run_tleap(dry_input, work_dir=dry_dir)
+        assert_tleap_success(
+            dry_run.summary,
+            fail_on_warnings=manifest.validation.fail_on_warnings,
+            context="dry",
+        )
+        _ensure_outputs(dry_outputs)
+        c4_runs: list[C4PostprocessRun] = []
+        if metal_result is not None and metal_result.c4_atom_types:
+            c4_runs.append(
+                apply_12_6_4_c4(
+                    dry_outputs.prmtop,
+                    atom_types=metal_result.c4_atom_types,
+                    water_model=manifest.protein.water_model,
+                    work_dir=dry_dir / "c4",
+                )
+            )
+        dry_coordinate_checks = validate_tleap_ligand_coordinates(
+            manifest=manifest,
+            reference_structure=structure,
+            output_pdb=dry_outputs.pdb,
+            stage="dry",
+            output_resnames_by_ligand_id={
+                item.ligand_id: item.final_resname
+                for item in refitted_ligands
+            },
+        )
+
+        warnings = list(sources.warnings)
+        if conect_records:
+            warnings.append(
+                "Disulfide SG-SG bonds were encoded as PDB CONECT records in the tleap input PDB."
+            )
+        warnings.extend(leap_input.structure.warnings)
+        warnings.extend(warning for warning in dry_run.summary.warnings)
+        warnings.extend(
+            f"Ligand {check['ligand_id']} dry tleap coordinate max deviation: "
+            f"{check['max_coordinate_deviation_angstrom']:.3f} A"
+            for check in dry_coordinate_checks
+        )
+        ion_plan: IonPlan | None = None
+        salt_volume_a3: float | None = None
+        final_run: TLeapRun | None = None
+        if not manifest.solvation.enabled:
+            final_outputs = _copy_outputs(dry_outputs, final_dir)
+        else:
+            final_outputs, final_run, ion_plan, salt_volume_a3, final_c4 = _run_solvated_build(
+                manifest=manifest,
+                sources=sources,
+                ligands=ligand_files,
+                input_pdb=leap_input.path,
+                disulfide_bonds=disulfide_bonds,
+                dry_charge=dry_run.summary.total_charge,
+                output_dir=output,
+                metal_result=metal_result,
+            )
+            if final_c4 is not None:
+                c4_runs.append(final_c4)
+            warnings.extend(final_run.summary.warnings)
+
+    except (
+        ForceFieldError,
+        IonPlanError,
+        LeapLogError,
+        LeapResidueError,
+        TLeapRunError,
+        C4PostprocessError,
+        FileNotFoundError,
+    ) as exc:
+        raise TLeapBuildError(str(exc)) from exc
+
+    return TLeapStageResult(
+        forcefields=sources,
+        leap_input=leap_input,
+        ligands=ligand_files,
+        mcpb_resp_ligands=refitted_ligands,
+        disulfide_bonds=disulfide_bonds,
+        dry_run=dry_run,
+        dry_outputs=dry_outputs,
+        final_run=final_run,
+        final_outputs=final_outputs,
+        solvation_enabled=manifest.solvation.enabled,
+        solvation_box_type=manifest.solvation.box,
+        buffer_angstrom=manifest.solvation.buffer_angstrom,
+        neutralization_requested=manifest.solvation.neutralize,
+        salt_concentration_molar=manifest.solvation.salt_concentration_molar,
+        ion_plan=ion_plan,
+        salt_volume_a3=salt_volume_a3,
+        c4_postprocessing=c4_runs,
+        warnings=warnings,
+    )
+
+
+def build_tleap_script(
+    *,
+    sources: ForceFieldSources,
+    ligands: list[LigandParameterFiles],
+    input_pdb: str | Path,
+    disulfide_bonds: list[DisulfideBondCommand],
+    outputs: TLeapOutputs,
+    work_dir: str | Path | None = None,
+    solvation_command: str | None = None,
+    ion_commands: list[str] | None = None,
+    setup_commands: list[str] | None = None,
+    extra_bond_commands: list[str] | None = None,
+) -> str:
+    lines: list[str] = []
+    for source in sources.all_sources:
+        lines.append(f"source {source}")
+    for ligand in ligands:
+        mol2_path = _tleap_path(ligand.final_mol2_path, work_dir)
+        frcmod_path = _tleap_path(ligand.final_frcmod_path, work_dir)
+        lines.append(f"{ligand.variable_name} = loadmol2 {mol2_path}")
+        if ligand.variable_name != ligand.residue_name:
+            lines.append(f"{ligand.residue_name} = {ligand.variable_name}")
+        lines.append(f"loadamberparams {frcmod_path}")
+    lines.extend(setup_commands or [])
+    lines.append(f"system = loadpdb {_tleap_path(input_pdb, work_dir)}")
+    lines.extend(bond.command for bond in disulfide_bonds if bond.command.startswith("bond "))
+    lines.extend(command for command in (extra_bond_commands or []) if command.startswith("bond "))
+    if solvation_command is not None:
+        lines.append(solvation_command)
+    lines.extend(ion_commands or [])
+    lines.extend(
+        [
+            "check system",
+            "charge system",
+            f"savepdb system {outputs.pdb.name}",
+            f"saveamberparm system {outputs.prmtop.name} {outputs.inpcrd.name}",
+            "quit",
+            "",
+        ]
+    )
+    return "\n".join(lines)
+
+
+def _tleap_path(path: str | Path, work_dir: str | Path | None) -> str:
+    """Return a path that tleap can resolve from its execution directory."""
+    path_obj = Path(path)
+    if work_dir is None:
+        return str(path_obj)
+    base = Path(work_dir)
+    base_abs = base if base.is_absolute() else Path.cwd() / base
+    target_abs = path_obj if path_obj.is_absolute() else Path.cwd() / path_obj
+    return os.path.relpath(target_abs.resolve(), base_abs.resolve())
+
+
+def solvation_command(*, box: str, water_box: str, buffer_angstrom: float) -> str:
+    if box == "truncated_octahedron":
+        return f"solvateOct system {water_box} {buffer_angstrom:.3f}"
+    if box == "rectangular":
+        return f"solvateBox system {water_box} {buffer_angstrom:.3f}"
+    raise TLeapBuildError(f"Unsupported solvation box: {box}")
+
+
+def _run_solvated_build(
+    *,
+    manifest: ManifestConfig,
+    sources: ForceFieldSources,
+    ligands: list[LigandParameterFiles],
+    input_pdb: Path,
+    disulfide_bonds: list[DisulfideBondCommand],
+    dry_charge: float | None,
+    output_dir: Path,
+    metal_result: MetalStageResult | None,
+) -> tuple[TLeapOutputs, TLeapRun, IonPlan | None, float | None, C4PostprocessRun | None]:
+    if dry_charge is None and manifest.solvation.neutralize:
+        raise TLeapBuildError("tleap dry-system charge was not parsed; cannot neutralize safely.")
+    solv_dir = output_dir / "leap" / "solvated"
+    solv_dir.mkdir(parents=True, exist_ok=True)
+    final_outputs = TLeapOutputs(
+        prmtop=solv_dir / "system.solvated.prmtop",
+        inpcrd=solv_dir / "system.solvated.inpcrd",
+        pdb=solv_dir / "system.solvated.pdb",
+    )
+    solvate = solvation_command(
+        box=manifest.solvation.box,
+        water_box=sources.water_box,
+        buffer_angstrom=manifest.solvation.buffer_angstrom,
+    )
+
+    volume_a3: float | None = None
+    if manifest.solvation.salt_concentration_molar > 0:
+        pre_outputs = TLeapOutputs(
+            prmtop=solv_dir / "system.presalt.prmtop",
+            inpcrd=solv_dir / "system.presalt.inpcrd",
+            pdb=solv_dir / "system.presalt.pdb",
+        )
+        pre_input = solv_dir / "tleap.presalt.in"
+        pre_input.write_text(
+            build_tleap_script(
+                sources=sources,
+                ligands=ligands,
+                input_pdb=input_pdb,
+                disulfide_bonds=disulfide_bonds,
+                outputs=pre_outputs,
+                work_dir=solv_dir,
+                solvation_command=solvate,
+                setup_commands=metal_result.leap_setup_commands if metal_result else None,
+                extra_bond_commands=metal_result.leap_bond_commands if metal_result else None,
+            ),
+            encoding="utf-8",
+        )
+        pre_run = run_tleap(pre_input, work_dir=solv_dir)
+        assert_tleap_success(
+            pre_run.summary,
+            fail_on_warnings=manifest.validation.fail_on_warnings,
+            context="pre-salt solvated",
+        )
+        _ensure_outputs(pre_outputs)
+        volume_a3 = amber_inpcrd_box_volume(str(pre_outputs.inpcrd))
+        if volume_a3 is None:
+            raise TLeapBuildError(
+                "Could not determine solvated box volume; cannot add configured salt concentration."
+            )
+
+    ion_plan = build_ion_plan(
+        total_charge=dry_charge or 0.0,
+        neutralize=manifest.solvation.neutralize,
+        positive_ion=manifest.solvation.positive_ion,
+        negative_ion=manifest.solvation.negative_ion,
+        salt_concentration_molar=manifest.solvation.salt_concentration_molar,
+        volume_a3=volume_a3,
+    )
+    script = build_tleap_script(
+        sources=sources,
+        ligands=ligands,
+        input_pdb=input_pdb,
+        disulfide_bonds=disulfide_bonds,
+        outputs=final_outputs,
+        work_dir=solv_dir,
+        solvation_command=solvate,
+        ion_commands=ion_plan.commands,
+        setup_commands=metal_result.leap_setup_commands if metal_result else None,
+        extra_bond_commands=metal_result.leap_bond_commands if metal_result else None,
+    )
+    script_path = solv_dir / "tleap.in"
+    script_path.write_text(script, encoding="utf-8")
+    final_run = run_tleap(script_path, work_dir=solv_dir)
+    assert_tleap_success(
+        final_run.summary,
+        fail_on_warnings=manifest.validation.fail_on_warnings,
+        context="solvated",
+    )
+    _ensure_outputs(final_outputs)
+    c4_run: C4PostprocessRun | None = None
+    if metal_result is not None and metal_result.c4_atom_types:
+        c4_run = apply_12_6_4_c4(
+            final_outputs.prmtop,
+            atom_types=metal_result.c4_atom_types,
+            water_model=manifest.protein.water_model,
+            work_dir=solv_dir / "c4",
+        )
+    copied = _copy_outputs(final_outputs, output_dir / "final")
+    return copied, final_run, ion_plan, volume_a3, c4_run
+
+
+def _copy_outputs(outputs: TLeapOutputs, final_dir: Path) -> TLeapOutputs:
+    final_dir.mkdir(parents=True, exist_ok=True)
+    copied = TLeapOutputs(
+        prmtop=final_dir / "system.prmtop",
+        inpcrd=final_dir / "system.inpcrd",
+        pdb=final_dir / "system.pdb",
+    )
+    shutil.copyfile(outputs.prmtop, copied.prmtop)
+    shutil.copyfile(outputs.inpcrd, copied.inpcrd)
+    shutil.copyfile(outputs.pdb, copied.pdb)
+    return copied
+
+
+def _ensure_outputs(outputs: TLeapOutputs) -> None:
+    for path in [outputs.prmtop, outputs.inpcrd, outputs.pdb]:
+        if not path.exists():
+            raise FileNotFoundError(f"tleap did not produce expected output: {path}")
+        if path.stat().st_size == 0:
+            raise TLeapBuildError(f"tleap produced an empty output file: {path}")

mdprep/leap/forcefields.py

@@ -0,0 +1,93 @@
+"""Amber force-field mapping helpers."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from mdprep.config.models import LigandConfig
+
+
+class ForceFieldError(ValueError):
+    """Raised when a manifest requests an unsupported force-field setting."""
+
+
+PROTEIN_LEAPRC = {
+    "ff14SB": "leaprc.protein.ff14SB",
+    "ff19SB": "leaprc.protein.ff19SB",
+}
+
+WATER_LEAPRC = {
+    "TIP3P": "leaprc.water.tip3p",
+    "OPC": "leaprc.water.opc",
+}
+
+WATER_BOX = {
+    "TIP3P": "TIP3PBOX",
+    "OPC": "OPCBOX",
+}
+
+LIGAND_LEAPRC = {
+    "gaff": "leaprc.gaff",
+    "gaff2": "leaprc.gaff2",
+}
+
+
+@dataclass(frozen=True)
+class ForceFieldSources:
+    protein: str
+    water: str
+    ligand: list[str]
+    water_box: str
+    warnings: list[str]
+
+    @property
+    def all_sources(self) -> list[str]:
+        return [self.protein, self.water, *self.ligand]
+
+
+def protein_leaprc(name: str) -> str:
+    try:
+        return PROTEIN_LEAPRC[name]
+    except KeyError as exc:
+        raise ForceFieldError(f"Unsupported protein force field: {name}") from exc
+
+
+def water_leaprc(name: str) -> str:
+    try:
+        return WATER_LEAPRC[name]
+    except KeyError as exc:
+        raise ForceFieldError(f"Unsupported water model: {name}") from exc
+
+
+def water_box(name: str) -> str:
+    try:
+        return WATER_BOX[name]
+    except KeyError as exc:
+        raise ForceFieldError(f"Unsupported water model: {name}") from exc
+
+
+def ligand_leaprc(name: str) -> str:
+    try:
+        return LIGAND_LEAPRC[name]
+    except KeyError as exc:
+        raise ForceFieldError(f"Unsupported ligand atom type family: {name}") from exc
+
+
+def forcefield_sources(
+    *,
+    protein_forcefield: str,
+    water_model: str,
+    ligands: list[LigandConfig],
+) -> ForceFieldSources:
+    ligand_families = sorted({ligand.atom_types for ligand in ligands})
+    ligand_sources = [ligand_leaprc(family) for family in ligand_families]
+    warnings: list[str] = []
+    if len(ligand_families) > 1:
+        warnings.append("Both GAFF and GAFF2 ligand force fields are sourced.")
+    return ForceFieldSources(
+        protein=protein_leaprc(protein_forcefield),
+        water=water_leaprc(water_model),
+        ligand=ligand_sources,
+        water_box=water_box(water_model),
+        warnings=warnings,
+    )