From da1f0f52a1c5e7f5cb1f8a4fe3bb14f4a8ebc4a5 Mon Sep 17 00:00:00 2001 From: bielj Date: Tue, 7 Jul 2026 14:47:12 -0400 Subject: [PATCH] Refactor bash MD scripts into typed Python core + Typer CLI Restructure the repo so the existing bash MD-prep workflows become resilient, dynamic Python workflows organized in layers. At the base, reusable GROMACS primitives (grompp+mdrun, solvate, energy, WAT-topology management) wrap single operations with structured errors and reproducibility metadata; on top, composable typed step functions and hardened pipelines build on those primitives, including the adaptive pressure-driven resolvation loop. This layering decouples what a step does from how it is invoked, so the same core supports multiple front ends: a Typer CLI (subcommand per step + run-pipeline), an SDK public API (run_standard_md), and eventual Prefect orchestration without rewrites. Per-system behavior is driven by a Pydantic config (YAML/TOML + flag overrides) with per-invocation GROMACS run profiles, replacing hard-coded, typo-prone shell scripts. Adopt taylor as the canonical variant and add gmx-free contract/CLI tests wired into CI. --- .github/workflows/ci.yml | 8 +- Dockerfile.actl | 2 +- README.md | 63 ++- md_workflows/__init__.py | 14 +- md_workflows/cli.py | 368 ++++++++++++++---- md_workflows/core/__init__.py | 6 + md_workflows/core/config.py | 232 +++++++++++ md_workflows/core/exceptions.py | 72 ++++ md_workflows/core/gmx.py | 347 +++++++++++++++++ md_workflows/core/paths.py | 16 + md_workflows/core/pipelines/__init__.py | 5 + md_workflows/core/pipelines/standard_md.py | 90 +++++ md_workflows/core/results.py | 71 ++++ md_workflows/core/steps/__init__.py | 11 + md_workflows/core/steps/equilibrate.py | 282 ++++++++++++++ md_workflows/core/steps/make_crystal.py | 178 +++++++++ md_workflows/core/steps/make_waterbox.py | 208 ++++++++++ md_workflows/core/steps/minimize.py | 96 +++++ md_workflows/core/steps/param_prot.py | 196 ++++++++++ .../core/steps/pressure_interpolate.py | 154 ++++++++ md_workflows/core/steps/resolvate.py | 299 ++++++++++++++ .../core/steps/run_params_gaussian.py | 317 +++++++++++++++ md_workflows/core/steps/solvate.py | 291 ++++++++++++++ md_workflows/equilibrate.py | 166 -------- md_workflows/make_crystal.py | 124 ------ md_workflows/make_waterbox.py | 202 ---------- md_workflows/minimize.py | 48 --- md_workflows/param_prot.py | 144 ------- md_workflows/resolvate.py | 151 ------- md_workflows/run_params_gaussian.py | 316 --------------- md_workflows/sdk/__init__.py | 74 ++++ md_workflows/sdk/convenience.py | 28 ++ md_workflows/solvate.py | 174 --------- md_workflows/workflows/__init__.py | 1 - md_workflows/workflows/mdmx.py | 100 ----- pyproject.toml | 32 +- tests/test_cli.py | 50 +++ tests/test_config.py | 84 ++++ tests/test_gmx.py | 72 ++++ tests/test_pressure_interpolate.py | 47 +++ tests/test_steps_contract.py | 83 ++++ 41 files changed, 3692 insertions(+), 1530 deletions(-) create mode 100644 md_workflows/core/__init__.py create mode 100644 md_workflows/core/config.py create mode 100644 md_workflows/core/exceptions.py create mode 100644 md_workflows/core/gmx.py create mode 100644 md_workflows/core/paths.py create mode 100644 md_workflows/core/pipelines/__init__.py create mode 100644 md_workflows/core/pipelines/standard_md.py create mode 100644 md_workflows/core/results.py create mode 100644 md_workflows/core/steps/__init__.py create mode 100644 md_workflows/core/steps/equilibrate.py create mode 100644 md_workflows/core/steps/make_crystal.py create mode 100644 md_workflows/core/steps/make_waterbox.py create mode 100644 md_workflows/core/steps/minimize.py create mode 100644 md_workflows/core/steps/param_prot.py create mode 100644 md_workflows/core/steps/pressure_interpolate.py create mode 100644 md_workflows/core/steps/resolvate.py create mode 100644 md_workflows/core/steps/run_params_gaussian.py create mode 100644 md_workflows/core/steps/solvate.py delete mode 100644 md_workflows/equilibrate.py delete mode 100644 md_workflows/make_crystal.py delete mode 100644 md_workflows/make_waterbox.py delete mode 100644 md_workflows/minimize.py delete mode 100644 md_workflows/param_prot.py delete mode 100644 md_workflows/resolvate.py delete mode 100644 md_workflows/run_params_gaussian.py create mode 100644 md_workflows/sdk/__init__.py create mode 100644 md_workflows/sdk/convenience.py delete mode 100644 md_workflows/solvate.py delete mode 100644 md_workflows/workflows/__init__.py delete mode 100644 md_workflows/workflows/mdmx.py create mode 100644 tests/test_cli.py create mode 100644 tests/test_config.py create mode 100644 tests/test_gmx.py create mode 100644 tests/test_pressure_interpolate.py create mode 100644 tests/test_steps_contract.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 52127c1..42030c6 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -40,10 +40,14 @@ jobs: pip install '.[dev]' - name: Run ruff lint check - run: ruff check md_workflows + run: ruff check md_workflows tests - name: Run ruff format check - run: ruff format --check md_workflows + run: ruff format --check md_workflows tests + + - name: Run tests + # Contract + CLI tests are gmx-free, so they run on the hosted runner. + run: pytest -q build-check: name: Dockerfile build checks (no push) diff --git a/Dockerfile.actl b/Dockerfile.actl index 38dc8c9..97d8511 100644 --- a/Dockerfile.actl +++ b/Dockerfile.actl @@ -83,7 +83,7 @@ RUN chmod 0644 /usr/local/share/actl-md-workflows-shell-init.sh \ && printf '\n# Astera md-workflows environment\n[ -r /usr/local/share/actl-md-workflows-shell-init.sh ] && . /usr/local/share/actl-md-workflows-shell-init.sh\n' >> /etc/bash.bashrc \ && printf '\n# Astera md-workflows environment\n[ -r /usr/local/share/actl-md-workflows-shell-init.sh ] && . /usr/local/share/actl-md-workflows-shell-init.sh\n' >> /etc/zsh/zshrc \ && for cmd in \ - md_workflows.mdmx gmx micromamba curl wget rsync tini vim nano emacs git zsh htop tmux ncdu ping dig; do \ + md-workflows gmx micromamba curl wget rsync tini vim nano emacs git zsh htop tmux ncdu ping dig; do \ command -v "${cmd}" >/dev/null; \ done diff --git a/README.md b/README.md index d392d4a..5215dc0 100644 --- a/README.md +++ b/README.md @@ -51,25 +51,66 @@ docker run --rm -it \ bash ``` -This registers the CLI entry points from `pyproject.toml`, including `md_workflows.mdmx`. +This registers the single `md-workflows` CLI entry point from `pyproject.toml`. -## 3) Run the full workflow command +## 3) Run the workflow -Inside the container shell: +The CLI is one Typer app with a subcommand per step plus `run-pipeline`. Global options +`--workdir/-w` (run directory), `--config/-c` (YAML/TOML), and `--resume/--force` come +before the subcommand. + +Run the full prep pipeline (`param_prot → make_crystal → make_waterbox → solvate → +minimize → equilibrate → resolvate`, ending with the adaptive pressure loop): + +```bash +md-workflows --workdir . run-pipeline --pdb-id 4LZT --ix 5 +``` + +Run a single step (each resolves its inputs from `--workdir` and fails loudly if any are +missing): ```bash -md_workflows.mdmx \ - --param-pdb-id 6B8X \ - --ix 1 \ - --ntomp 26 \ - --resolv-ntmpi 8 \ - --resolv-ntomp 1 +md-workflows -w . minimize +md-workflows -w . resolvate --target-bar 1.0 +``` + +Per-system settings live in a config file; flags override individual values. The +dominant cross-machine knob is the per-invocation GROMACS `mdrun` profile +(`run_profiles`, e.g. GPU offload + thread counts): + +```yaml +# config.yaml +pdb_id: 4LZT +crystal: { ix: 5 } +waterbox: { nc_scale: 5, conc: 60.0 } +solvate: { ionic_strength: 0.1 } +run_profiles: + equil: { ntomp: 16, nb: gpu, pme: gpu, bonded: gpu, tunepme: false } + min: { ntomp: 16, nb: gpu, pme: cpu, bonded: cpu, tunepme: false } +``` + +```bash +md-workflows -w run_dir -c config.yaml run-pipeline +md-workflows -w run_dir -c config.yaml --resume run-pipeline # skip completed steps +``` + +`--resume` skips a step whose durable outputs already exist. Exit codes: `0` ok, +`1` domain error, `2` missing input, `3` external tool failed. + +### Python / SDK + +The same logic is importable: + +```python +from md_workflows import run_standard_md +result = run_standard_md("4LZT", "run_dir", crystal={"ix": 5}, resume=True) ``` -To see all available flags: +To see all commands and flags: ```bash -md_workflows.mdmx --help +md-workflows --help +md-workflows run-pipeline --help ``` ## Development diff --git a/md_workflows/__init__.py b/md_workflows/__init__.py index 87cd617..63d6e9c 100644 --- a/md_workflows/__init__.py +++ b/md_workflows/__init__.py @@ -1 +1,13 @@ -"""md_workflows package.""" +"""md_workflows — crystalline molecular-dynamics prep workflows. + +The public API is re-exported from :mod:`md_workflows.sdk`; ``import md_workflows`` gives +access to the blessed steps, the standard pipeline, config/result types, and the +``run_standard_md`` convenience wrapper. +""" + +from __future__ import annotations + +from .sdk import * # noqa: F401,F403 +from .sdk import __all__ as _sdk_all + +__all__ = list(_sdk_all) diff --git a/md_workflows/cli.py b/md_workflows/cli.py index 6250c9b..ab9d6e9 100644 --- a/md_workflows/cli.py +++ b/md_workflows/cli.py @@ -1,99 +1,313 @@ -"""CLI entrypoints for workflow commands.""" +"""Typer CLI: a subcommand per step plus ``run-pipeline``. + +The CLI is a thin caller of ``core`` (and the SDK): it resolves ``--workdir`` + +``--config`` (+ flag overrides) into typed inputs, runs the step, and maps structured +exceptions to process exit codes. It never imports orchestration code. + +Exit codes: 0 ok · 1 domain error · 2 missing input · 3 external tool failed. +""" from __future__ import annotations -import argparse -import sys - -from . import ( - equilibrate, - make_crystal, - make_waterbox, - minimize, - param_prot, - resolvate, - run_params_gaussian, - solvate, +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import typer + +from .core.config import RUN_PROFILE_KEYS, SystemConfig +from .core.exceptions import MDWorkflowError, MissingInputError, StepToolError +from .core.pipelines.standard_md import standard_md_pipeline +from .core.results import StepResult, StepStatus +from .core.steps import ( + equilibrate as equilibrate_step, +) +from .core.steps import ( + make_crystal as make_crystal_step, +) +from .core.steps import ( + make_waterbox as make_waterbox_step, +) +from .core.steps import ( + minimize as minimize_step, +) +from .core.steps import ( + param_prot as param_prot_step, +) +from .core.steps import ( + resolvate as resolvate_step, +) +from .core.steps import ( + run_params_gaussian as gaussian_step, +) +from .core.steps import ( + solvate as solvate_step, +) + +app = typer.Typer( + no_args_is_help=True, + add_completion=False, + help="Crystalline molecular-dynamics prep workflows.", ) -def _single_command_cli(command: str) -> None: - parser = argparse.ArgumentParser(description=f"Run {command} workflow") - - if command == "param_prot": - parser.add_argument("--pdb-id", default="6B8X") - elif command == "make_crystal": - parser.add_argument("--ix", type=int, default=1) - parser.add_argument("--iy", type=int, default=None) - parser.add_argument("--iz", type=int, default=None) - parser.add_argument("--chimerax-exec", default="chimerax") - elif command == "make_waterbox": - parser.add_argument("--ntomp", type=int, default=26) - elif command == "solvate": - pass - elif command == "minimize": - parser.add_argument("--ntomp", type=int, default=26) - elif command == "equilibrate": - parser.add_argument("--ntomp", type=int, default=26) - elif command == "resolvate": - parser.add_argument("--ntmpi", type=int, default=8) - parser.add_argument("--ntomp", type=int, default=1) - elif command == "run_params_gaussian": - parser.add_argument("--g16root", default="/Users/mewall/packages") - parser.add_argument("--nproc", type=int, default=8) - else: - raise ValueError(f"Unknown command: {command}") - - args = parser.parse_args(sys.argv[1:]) - - if command == "param_prot": - param_prot.run(pdb_id=args.pdb_id) - elif command == "make_crystal": - make_crystal.run(ix=args.ix, iy=args.iy, iz=args.iz, chimerax_exec=args.chimerax_exec) - elif command == "make_waterbox": - make_waterbox.run(ntomp=args.ntomp) - elif command == "solvate": - solvate.run() - elif command == "minimize": - minimize.run(ntomp=args.ntomp) - elif command == "equilibrate": - equilibrate.run(ntomp=args.ntomp) - elif command == "resolvate": - resolvate.run(ntmpi=args.ntmpi, ntomp=args.ntomp) - elif command == "run_params_gaussian": - run_params_gaussian.run( - g16root=args.g16root, - nproc=args.nproc, - ) +@dataclass +class Common: + workdir: Path + config: Path | None + resume: bool + + +@app.callback() +def _main( + ctx: typer.Context, + workdir: Path = typer.Option(Path("."), "--workdir", "-w", help="Run directory."), + config: Path | None = typer.Option(None, "--config", "-c", help="YAML/TOML config file."), + resume: bool = typer.Option( + False, "--resume/--force", help="Skip steps whose outputs already exist." + ), +) -> None: + ctx.obj = Common(workdir=workdir, config=config, resume=resume) -def param_prot_cli() -> None: - _single_command_cli("param_prot") +# --------------------------------------------------------------------------- # +# helpers +# --------------------------------------------------------------------------- # +def _load_cfg(common: Common, overrides: dict[str, Any] | None = None) -> SystemConfig: + try: + return SystemConfig.load(common.config, overrides=overrides or {}) + except (FileNotFoundError, ValueError) as exc: + typer.secho(f"config error: {exc}", fg=typer.colors.RED, err=True) + raise typer.Exit(1) from exc -def make_crystal_cli() -> None: - _single_command_cli("make_crystal") +def _ntomp_override(ntomp: int | None) -> dict[str, Any]: + """Apply a quick ntomp tweak across every mdrun profile.""" + if ntomp is None: + return {} + return {"run_profiles": {k: {"ntomp": ntomp} for k in RUN_PROFILE_KEYS}} -def make_waterbox_cli() -> None: - _single_command_cli("make_waterbox") +def _emit(result: StepResult) -> None: + color = { + StepStatus.COMPLETED: typer.colors.GREEN, + StepStatus.SKIPPED: typer.colors.YELLOW, + StepStatus.FAILED: typer.colors.RED, + }[result.status] + typer.secho(f"{result.step}: {result.status.value}", fg=color) + for name, path in result.outputs.items(): + typer.echo(f" {name}: {path}") -def solvate_cli() -> None: - _single_command_cli("solvate") +def _run(fn) -> None: + """Execute a step callable, mapping structured exceptions to exit codes.""" + try: + result = fn() + except MissingInputError as exc: + typer.secho(str(exc), fg=typer.colors.RED, err=True) + raise typer.Exit(2) from exc + except StepToolError as exc: + typer.secho(str(exc), fg=typer.colors.RED, err=True) + raise typer.Exit(3) from exc + except MDWorkflowError as exc: + typer.secho(str(exc), fg=typer.colors.RED, err=True) + raise typer.Exit(1) from exc + if isinstance(result, StepResult): + _emit(result) + else: # pipeline result + for step in result.steps: + _emit(step) -def minimize_cli() -> None: - _single_command_cli("minimize") + +# --------------------------------------------------------------------------- # +# per-step commands +# --------------------------------------------------------------------------- # +@app.command("param-prot") +def param_prot_cmd( + ctx: typer.Context, + pdb_id: str | None = typer.Option(None, "--pdb-id", help="PDB id to parameterize."), +) -> None: + """Parameterize the protein (clean, tleap, Amber->GROMACS).""" + common: Common = ctx.obj + cfg = _load_cfg(common, {"pdb_id": pdb_id} if pdb_id else None) + inputs = param_prot_step.resolve_inputs(common.workdir, cfg) + _run(lambda: param_prot_step.param_prot(inputs, resume=common.resume)) + + +@app.command("make-crystal") +def make_crystal_cmd( + ctx: typer.Context, + ix: int | None = typer.Option(None, "--ix", help="Supercell replication in x (and y/z)."), + iy: int | None = typer.Option(None, "--iy"), + iz: int | None = typer.Option(None, "--iz"), + chimerax_exec: str | None = typer.Option(None, "--chimerax-exec"), +) -> None: + """Build the crystal supercell.""" + common: Common = ctx.obj + crystal = { + k: v + for k, v in {"ix": ix, "iy": iy, "iz": iz, "chimerax_exec": chimerax_exec}.items() + if v is not None + } + cfg = _load_cfg(common, {"crystal": crystal} if crystal else None) + inputs = make_crystal_step.resolve_inputs(common.workdir, cfg) + _run(lambda: make_crystal_step.make_crystal(inputs, cfg.crystal, resume=common.resume)) + + +@app.command("make-waterbox") +def make_waterbox_cmd( + ctx: typer.Context, + nc_scale: int | None = typer.Option(None, "--nc-scale"), + conc: float | None = typer.Option(None, "--conc"), + ntomp: int | None = typer.Option(None, "--ntomp", help="Override mdrun OpenMP threads."), +) -> None: + """Build the equilibrated bulk-water reservoir.""" + common: Common = ctx.obj + wb = {k: v for k, v in {"nc_scale": nc_scale, "conc": conc}.items() if v is not None} + overrides = _ntomp_override(ntomp) + if wb: + overrides["waterbox"] = wb + cfg = _load_cfg(common, overrides) + inputs = make_waterbox_step.resolve_inputs(common.workdir, cfg) + _run( + lambda: make_waterbox_step.make_waterbox( + inputs, + cfg.waterbox, + cfg.profile("waterbox_min"), + cfg.profile("waterbox_equil"), + resume=common.resume, + ) + ) + + +@app.command("solvate") +def solvate_cmd( + ctx: typer.Context, + ionic_strength: float | None = typer.Option(None, "--ionic-strength", help="Target salt M."), +) -> None: + """Solvate the crystal and add ions.""" + common: Common = ctx.obj + ov = {"solvate": {"ionic_strength": ionic_strength}} if ionic_strength is not None else None + cfg = _load_cfg(common, ov) + inputs = solvate_step.resolve_inputs(common.workdir, cfg) + _run(lambda: solvate_step.solvate(inputs, cfg.solvate, resume=common.resume)) + + +@app.command("minimize") +def minimize_cmd( + ctx: typer.Context, + ntomp: int | None = typer.Option(None, "--ntomp", help="Override mdrun OpenMP threads."), +) -> None: + """Energy-minimize the solvated model.""" + common: Common = ctx.obj + cfg = _load_cfg(common, _ntomp_override(ntomp)) + inputs = minimize_step.resolve_inputs(common.workdir, cfg) + _run(lambda: minimize_step.minimize(inputs, cfg.profile("min"), resume=common.resume)) + + +@app.command("equilibrate") +def equilibrate_cmd( + ctx: typer.Context, + ntomp: int | None = typer.Option(None, "--ntomp"), + ligand_resname: str | None = typer.Option(None, "--ligand-resname"), + handle_ligand_restraint: bool | None = typer.Option( + None, "--handle-ligand-restraint/--no-ligand-restraint" + ), + chain_split_mode: str | None = typer.Option( + None, "--chain-split-mode", help="ignore_ligand | split_intelligently" + ), +) -> None: + """Set up restraints and run the NPT equilibration.""" + common: Common = ctx.obj + eq = { + k: v + for k, v in { + "ligand_resname": ligand_resname, + "handle_ligand_restraint": handle_ligand_restraint, + "chain_split_mode": chain_split_mode, + }.items() + if v is not None + } + overrides = _ntomp_override(ntomp) + if eq: + overrides["equilibrate"] = eq + cfg = _load_cfg(common, overrides) + inputs = equilibrate_step.resolve_inputs(common.workdir, cfg) + _run( + lambda: equilibrate_step.equilibrate( + inputs, cfg.equilibrate, cfg.profile("equil"), resume=common.resume + ) + ) + + +@app.command("resolvate") +def resolvate_cmd( + ctx: typer.Context, + target_bar: float | None = typer.Option(None, "--target-bar", help="Target mean pressure."), + trial_fraction: float | None = typer.Option(None, "--trial-fraction"), + ntomp: int | None = typer.Option(None, "--ntomp"), +) -> None: + """Adaptive resolvation to the target pressure (run 1 -> branch -> final).""" + common: Common = ctx.obj + rv = { + k: v + for k, v in {"target_bar": target_bar, "trial_fraction": trial_fraction}.items() + if v is not None + } + overrides = _ntomp_override(ntomp) + if rv: + overrides["resolvate"] = rv + cfg = _load_cfg(common, overrides) + inputs = resolvate_step.resolve_inputs(common.workdir, cfg) + _run( + lambda: resolvate_step.resolvate( + inputs, + cfg.resolvate, + cfg.profile("resolv_min"), + cfg.profile("resolv_equil"), + resume=common.resume, + ) + ) -def equilibrate_cli() -> None: - _single_command_cli("equilibrate") +@app.command("run-params-gaussian") +def run_params_gaussian_cmd( + ctx: typer.Context, + pdb_id: str | None = typer.Option(None, "--pdb-id"), + g16root: str | None = typer.Option(None, "--g16root"), + nproc: int | None = typer.Option(None, "--nproc"), +) -> None: + """Ligand parameterization with Gaussian + AmberTools (standalone).""" + common: Common = ctx.obj + gaussian = {k: v for k, v in {"g16root": g16root, "nproc": nproc}.items() if v is not None} + overrides: dict[str, Any] = {} + if pdb_id: + overrides["pdb_id"] = pdb_id + if gaussian: + overrides["gaussian"] = gaussian + cfg = _load_cfg(common, overrides) + inputs = gaussian_step.resolve_inputs(common.workdir, cfg) + _run(lambda: gaussian_step.run_params_gaussian(inputs, cfg.gaussian, resume=common.resume)) -def resolvate_cli() -> None: - _single_command_cli("resolvate") +@app.command("run-pipeline") +def run_pipeline_cmd( + ctx: typer.Context, + pdb_id: str | None = typer.Option(None, "--pdb-id"), + ix: int | None = typer.Option(None, "--ix", help="Crystal supercell replication."), + ntomp: int | None = typer.Option(None, "--ntomp", help="Override mdrun OpenMP threads."), +) -> None: + """Run the full prep pipeline (param_prot -> ... -> resolvate).""" + common: Common = ctx.obj + overrides = _ntomp_override(ntomp) + if pdb_id: + overrides["pdb_id"] = pdb_id + if ix is not None: + overrides["crystal"] = {"ix": ix} + cfg = _load_cfg(common, overrides) + _run(lambda: standard_md_pipeline(common.workdir, cfg, resume=common.resume)) -def run_params_gaussian_cli() -> None: - _single_command_cli("run_params_gaussian") +if __name__ == "__main__": + app() diff --git a/md_workflows/core/__init__.py b/md_workflows/core/__init__.py new file mode 100644 index 0000000..9aa6681 --- /dev/null +++ b/md_workflows/core/__init__.py @@ -0,0 +1,6 @@ +"""Pure, orchestrator-agnostic MD workflow logic. + +Everything under ``core`` is a plain, typed Python function with no CLI parsing and +no orchestration-framework dependency. The CLI, the SDK, and (later) a Prefect layer +are all just different callers of these functions. +""" diff --git a/md_workflows/core/config.py b/md_workflows/core/config.py new file mode 100644 index 0000000..acb9b96 --- /dev/null +++ b/md_workflows/core/config.py @@ -0,0 +1,232 @@ +"""Typed configuration for the MD prep pipeline. + +A :class:`SystemConfig` captures everything that varies between systems and machines: +the PDB id, crystal/water-box/solvation parameters, ligand-handling flags, and — the +dominant cross-machine knob — the per-invocation GROMACS ``mdrun`` execution profiles +(GPU offload + thread counts). It loads from YAML/TOML and accepts CLI-flag overrides. + +Defaults follow the canonical taylor scripts +(``taylor_scripts_4lzt_5x5x5``): GPU ``mdrun`` with ``-pme cpu`` on minimizations and +``-pme gpu`` on equilibrations. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any, Literal + +import tomllib +import yaml +from pydantic import BaseModel, ConfigDict, Field + +Offload = Literal["auto", "cpu", "gpu"] + +# Keys for the per-invocation mdrun profiles. Each GROMACS mdrun call in the pipeline +# looks up its profile by one of these names. +RUN_PROFILE_KEYS = ( + "waterbox_min", + "waterbox_equil", + "min", + "equil", + "resolv_min", + "resolv_equil", +) + + +class GromacsRunProfile(BaseModel): + """Execution flags for a single ``gmx mdrun`` invocation. + + These are hardware-dependent and differ per call (e.g. ``-pme cpu`` for + minimization vs ``-pme gpu`` for equilibration), so they are configured rather than + hard-coded. ``to_mdrun_flags`` renders only the set (non-``None``) options. + """ + + model_config = ConfigDict(extra="forbid") + + gmx_bin: str = "gmx" + ntmpi: int | None = 1 + ntomp: int | None = 16 + nb: Offload | None = "gpu" + pme: Offload | None = "gpu" + bonded: Offload | None = "gpu" + update: Offload | None = None + pin: Literal["auto", "on", "off"] | None = None + dlb: Literal["auto", "yes", "no"] | None = None + tunepme: bool = True # emit -notunepme when False + extra_mdrun_args: list[str] = Field(default_factory=list) + + def to_mdrun_flags(self) -> list[str]: + """Render the mdrun option list (excludes ``-s``/``-deffnm``/``-v``).""" + flags: list[str] = [] + if self.ntmpi is not None: + flags += ["-ntmpi", str(self.ntmpi)] + if self.ntomp is not None: + flags += ["-ntomp", str(self.ntomp)] + if self.nb is not None: + flags += ["-nb", self.nb] + if self.pme is not None: + flags += ["-pme", self.pme] + if self.bonded is not None: + flags += ["-bonded", self.bonded] + if self.update is not None: + flags += ["-update", self.update] + if self.pin is not None: + flags += ["-pin", self.pin] + if self.dlb is not None: + flags += ["-dlb", self.dlb] + if not self.tunepme: + flags += ["-notunepme"] + flags += list(self.extra_mdrun_args) + return flags + + +def _min_profile() -> GromacsRunProfile: + # Minimizations: GPU nonbonded, CPU PME/bonded, no PME tuning (taylor canonical). + return GromacsRunProfile(nb="gpu", pme="cpu", bonded="cpu", tunepme=False) + + +def _equil_profile() -> GromacsRunProfile: + # Equilibrations: full GPU offload, no PME tuning (taylor canonical). + return GromacsRunProfile(nb="gpu", pme="gpu", bonded="gpu", tunepme=False) + + +def _default_run_profiles() -> dict[str, GromacsRunProfile]: + return { + "waterbox_min": _min_profile(), + "waterbox_equil": _min_profile(), # taylor equilibrates the water box on CPU PME + "min": _min_profile(), + "equil": _equil_profile(), + "resolv_min": _min_profile(), + "resolv_equil": _equil_profile(), + } + + +class CrystalParams(BaseModel): + model_config = ConfigDict(extra="forbid") + ix: int = 1 + iy: int | None = None # falls back to ix + iz: int | None = None # falls back to ix + chimerax_exec: str = "chimerax" + + +class WaterboxParams(BaseModel): + model_config = ConfigDict(extra="forbid") + nc_scale: int = 5 # cell subdivision / reservoir tiling factor (taylor) + conc: float = 60.0 # gmx insert-molecules water concentration (mol/L) + min_mdp: str = "min_water.mdp" + equil_mdp: str = "equil_water.mdp" + + +class SolvateParams(BaseModel): + model_config = ConfigDict(extra="forbid") + ionic_strength: float = 0.1 # target molarity for added ions + water_molarity: float = 55.0 # water molarity used in the ion-count formula + + +class MinimizeParams(BaseModel): + model_config = ConfigDict(extra="forbid") + min_mdp: str = "min.mdp" + + +class EquilibrateParams(BaseModel): + model_config = ConfigDict(extra="forbid") + equil_mdp: str = "equil.mdp" + restraint_fc: float = 209.2 # kJ/mol/nm^2 position-restraint force constant + restraint_group: str = "Protein-H" # gmx genrestr selection group + ligand_resname: str | None = None # None => auto-detect via pdb_file_processing + handle_ligand_restraint: bool = False # OFF until the full ligand feature lands + chain_split_mode: Literal["ignore_ligand", "split_intelligently"] = "ignore_ligand" + + +class ResolvateParams(BaseModel): + model_config = ConfigDict(extra="forbid") + trial_fraction: float = 0.25 # fraction of trial-solvation waters used in run 1 + scale_add: float = 1.5 # run-2 probe multiplier on run-1 water count + target_bar: float = 1.0 # target mean pressure + pressure_tol: float = 100.0 # bar; within this of target counts as "converged" + max_dNw_factor: float = 1.0 # sanity bound on |dNw| as a multiple of current waters + + +class GaussianParams(BaseModel): + model_config = ConfigDict(extra="forbid") + g16root: str = "" # required only when running the standalone gaussian utility + nproc: int = 8 + net_charge: int = -2 + method: str = "B3LYP/6-31+G(d,p)" + + +class SystemConfig(BaseModel): + """Top-level configuration for one MD prep run.""" + + model_config = ConfigDict(extra="forbid") + + pdb_id: str = "6B8X" + variant: Literal["original", "taylor"] = "taylor" + mdp_dir: Path = Path("artifacts") + + crystal: CrystalParams = Field(default_factory=CrystalParams) + waterbox: WaterboxParams = Field(default_factory=WaterboxParams) + solvate: SolvateParams = Field(default_factory=SolvateParams) + minimize: MinimizeParams = Field(default_factory=MinimizeParams) + equilibrate: EquilibrateParams = Field(default_factory=EquilibrateParams) + resolvate: ResolvateParams = Field(default_factory=ResolvateParams) + gaussian: GaussianParams = Field(default_factory=GaussianParams) + + run_profiles: dict[str, GromacsRunProfile] = Field(default_factory=_default_run_profiles) + + def profile(self, key: str) -> GromacsRunProfile: + """Return the mdrun profile for an invocation key, defaulting if unset.""" + return self.run_profiles.get(key, GromacsRunProfile()) + + @classmethod + def load( + cls, + path: str | Path | None = None, + *, + overrides: dict[str, Any] | None = None, + ) -> SystemConfig: + """Build a config from an optional YAML/TOML file plus optional overrides. + + ``overrides`` is a (possibly nested) dict deep-merged over the file/defaults — + this is how CLI flags win over the config file. + + The merge is seeded from the full default config so a partial override (e.g. + one field of one run profile) merges into the defaults instead of replacing the + whole ``run_profiles`` mapping. + """ + data: dict[str, Any] = cls().model_dump(mode="python") + if path is not None: + data = _deep_merge(data, _read_config_file(Path(path))) + if overrides: + data = _deep_merge(data, overrides) + return cls.model_validate(data) + + +def _read_config_file(path: Path) -> dict[str, Any]: + if not path.exists(): + raise FileNotFoundError(f"config file not found: {path}") + suffix = path.suffix.lower() + try: + if suffix in {".yaml", ".yml"}: + loaded = yaml.safe_load(path.read_text()) or {} + elif suffix == ".toml": + loaded = tomllib.loads(path.read_text()) + else: + raise ValueError(f"unsupported config format {suffix!r} (use .yaml/.yml/.toml)") + except (yaml.YAMLError, tomllib.TOMLDecodeError) as exc: + raise ValueError(f"could not parse config file {path}: {exc}") from exc + if not isinstance(loaded, dict): + raise ValueError(f"config file {path} must contain a mapping at the top level") + return loaded + + +def _deep_merge(base: dict[str, Any], override: dict[str, Any]) -> dict[str, Any]: + """Recursively merge ``override`` into ``base`` (override wins), non-mutating.""" + merged = dict(base) + for key, value in override.items(): + existing = merged.get(key) + if isinstance(existing, dict) and isinstance(value, dict): + merged[key] = _deep_merge(existing, value) + else: + merged[key] = value + return merged diff --git a/md_workflows/core/exceptions.py b/md_workflows/core/exceptions.py new file mode 100644 index 0000000..4a21af4 --- /dev/null +++ b/md_workflows/core/exceptions.py @@ -0,0 +1,72 @@ +"""Structured exceptions raised by core functions. + +Core code never calls ``sys.exit`` or terminates the process — it raises these so any +caller (CLI, SDK, or a future orchestrator) can catch and handle them. The CLI layer +is the only place that maps these to process exit codes. +""" + +from __future__ import annotations + +from pathlib import Path + + +class MDWorkflowError(Exception): + """Base class for all md-workflows domain errors.""" + + +class MissingInputError(MDWorkflowError): + """A step was asked to run but one or more required input files are absent. + + Raised by the input guard (``check_inputs``) before any external tool is invoked. + """ + + def __init__(self, step: str, missing: list[Path]): + self.step = step + self.missing = list(missing) + listed = "\n ".join(str(p) for p in self.missing) + super().__init__(f"{step}: missing required input file(s):\n {listed}") + + +class StepToolError(MDWorkflowError): + """An external tool (gmx, tleap, ChimeraX, AmberTools, ...) exited non-zero.""" + + def __init__( + self, + tool: str, + returncode: int, + *, + cmd: list[str] | None = None, + stderr: str | None = None, + log_path: Path | None = None, + ): + self.tool = tool + self.returncode = returncode + self.cmd = list(cmd) if cmd else None + self.stderr = stderr + self.log_path = log_path + parts = [f"{tool} exited with code {returncode}"] + if cmd: + parts.append("command: " + " ".join(str(c) for c in cmd)) + if log_path: + parts.append(f"log: {log_path}") + if stderr: + tail = "\n".join(stderr.strip().splitlines()[-20:]) + parts.append(f"stderr (tail):\n{tail}") + super().__init__("\n".join(parts)) + + +class GmxOutputParseError(MDWorkflowError): + """A value expected in GROMACS stdout/log (e.g. a molecule count) was not found. + + These parses are tied to gmx output strings that shift across versions, so failures + are surfaced explicitly rather than as a bare ``RuntimeError``. + """ + + def __init__(self, what: str, source: Path | str, *, gmx_version: str | None = None): + self.what = what + self.source = source + self.gmx_version = gmx_version + msg = f"could not parse {what} from {source}" + if gmx_version: + msg += f" (gmx {gmx_version})" + super().__init__(msg) diff --git a/md_workflows/core/gmx.py b/md_workflows/core/gmx.py new file mode 100644 index 0000000..2a63abc --- /dev/null +++ b/md_workflows/core/gmx.py @@ -0,0 +1,347 @@ +"""GROMACS-level primitives shared across steps. + +These wrap one external operation each (grompp+mdrun, solvate, energy, topology edit) +so the higher-level steps and the resolvation flow reuse a single implementation +instead of duplicating the min/equil/solvate motif. All failures raise structured +exceptions; nothing here calls ``sys.exit``. +""" + +from __future__ import annotations + +import hashlib +import subprocess +from pathlib import Path + +from pydantic import BaseModel + +from .config import GromacsRunProfile +from .exceptions import GmxOutputParseError, StepToolError + +_WATER_RESNAMES = frozenset({"WAT", "SOL", "HOH", "H2O"}) +_WATER_OXYGENS = frozenset({"OW", "O", "OW1"}) + + +# --------------------------------------------------------------------------- # +# Process / hashing helpers +# --------------------------------------------------------------------------- # +def run_tool( + cmd: list[str], + *, + tool: str, + cwd: str | Path | None = None, + input: str | None = None, + log_path: str | Path | None = None, + env: dict[str, str] | None = None, +) -> subprocess.CompletedProcess[str]: + """Run an external tool, capturing output; raise :class:`StepToolError` on failure. + + When ``log_path`` is given, the combined stdout+stderr is written there (matching + the shell scripts' ``>& tool.log`` redirects). + """ + proc = subprocess.run( + [str(c) for c in cmd], + cwd=str(cwd) if cwd is not None else None, + input=input, + text=True, + capture_output=True, + env=env, + ) + if log_path is not None: + Path(log_path).write_text((proc.stdout or "") + (proc.stderr or "")) + if proc.returncode != 0: + raise StepToolError( + tool, + proc.returncode, + cmd=[str(c) for c in cmd], + stderr=proc.stderr, + log_path=Path(log_path) if log_path is not None else None, + ) + return proc + + +def sha256_file(path: str | Path) -> str: + """Return the hex SHA-256 of a file (streamed).""" + h = hashlib.sha256() + with open(path, "rb") as fh: + for chunk in iter(lambda: fh.read(1 << 20), b""): + h.update(chunk) + return h.hexdigest() + + +def checksums(paths: dict[str, Path]) -> dict[str, str]: + """SHA-256 each existing path in a name->path mapping; skip missing.""" + return {name: sha256_file(p) for name, p in paths.items() if Path(p).exists()} + + +def gmx_version(gmx_bin: str = "gmx") -> str: + """Best-effort GROMACS version string; returns 'unknown' if it can't be read.""" + try: + proc = subprocess.run([gmx_bin, "--version"], text=True, capture_output=True, timeout=30) + except (OSError, subprocess.SubprocessError): + return "unknown" + for line in (proc.stdout or "").splitlines(): + if "version" in line.lower(): + return line.split(":", 1)[-1].strip() or line.strip() + return "unknown" + + +# --------------------------------------------------------------------------- # +# Structure helpers +# --------------------------------------------------------------------------- # +def count_waters(structure: str | Path) -> int: + """Count water molecules in a ``.gro`` or ``.pdb`` file (one oxygen per molecule). + + Unifies the two heuristics the scripts used (WAT-atoms/3 in make_waterbox, OW-count + in pressure_interpolate): counting water oxygens gives the exact molecule count. + """ + path = Path(structure) + is_gro = path.suffix.lower() == ".gro" + count = 0 + with open(path) as fh: + lines = fh.readlines() + if is_gro: + # GRO fixed columns: resname [5:10], atom name [10:15]. Skip title + count + + # the trailing box-vector line. + for line in lines[2:-1]: + resname = line[5:10].strip() + atom = line[10:15].strip() + if resname in _WATER_RESNAMES and atom in _WATER_OXYGENS: + count += 1 + else: + for line in lines: + if not line.startswith(("ATOM", "HETATM")): + continue + resname = line[17:20].strip() + atom = line[12:16].strip() + if resname in _WATER_RESNAMES and atom in _WATER_OXYGENS: + count += 1 + return count + + +# --------------------------------------------------------------------------- # +# mdrun primitives +# --------------------------------------------------------------------------- # +class GmxRunResult(BaseModel): + """Files produced by a grompp+mdrun invocation (paths; may not all exist).""" + + deffnm: str + tpr: Path + gro: Path + edr: Path + log: Path + + +def _grompp( + profile: GromacsRunProfile, + *, + mdp: Path, + structure: Path, + top: Path, + tpr: Path, + ref: Path | None, + maxwarn: int, + workdir: Path, + grompp_log: Path, +) -> None: + cmd = [ + profile.gmx_bin, + "grompp", + "-f", + str(mdp), + "-c", + str(structure), + "-o", + str(tpr), + "-p", + str(top), + ] + if ref is not None: + cmd += ["-r", str(ref)] + if maxwarn: + cmd += ["-maxwarn", str(maxwarn)] + run_tool(cmd, tool="gmx grompp", cwd=workdir, log_path=grompp_log) + + +def _mdrun(profile: GromacsRunProfile, *, tpr: Path, deffnm: str, workdir: Path) -> None: + # -deffnm is relative, so mdrun must run with cwd=workdir; outputs land there. + cmd = [ + profile.gmx_bin, + "mdrun", + "-s", + str(tpr), + *profile.to_mdrun_flags(), + "-deffnm", + deffnm, + "-v", + ] + run_tool(cmd, tool="gmx mdrun", cwd=workdir, log_path=workdir / f"mdrun_{deffnm}.log") + + +def run_min( + structure: Path, + top: Path, + mdp: Path, + deffnm: str, + profile: GromacsRunProfile, + workdir: Path, +) -> GmxRunResult: + """grompp + mdrun for an energy minimization; outputs ``.{tpr,gro,edr,log}``.""" + workdir = Path(workdir) + tpr = workdir / f"{deffnm}.tpr" + _grompp( + profile, + mdp=mdp, + structure=structure, + top=top, + tpr=tpr, + ref=None, + maxwarn=0, + workdir=workdir, + grompp_log=workdir / f"grompp_{deffnm}.log", + ) + _mdrun(profile, tpr=tpr, deffnm=deffnm, workdir=workdir) + return GmxRunResult( + deffnm=deffnm, + tpr=tpr, + gro=workdir / f"{deffnm}.gro", + edr=workdir / f"{deffnm}.edr", + log=workdir / f"{deffnm}.log", + ) + + +def run_equil( + structure: Path, + top: Path, + mdp: Path, + ref: Path, + deffnm: str, + profile: GromacsRunProfile, + workdir: Path, + *, + maxwarn: int = 0, +) -> GmxRunResult: + """grompp (with restraint ref ``-r``) + mdrun for an equilibration.""" + workdir = Path(workdir) + tpr = workdir / f"{deffnm}.tpr" + _grompp( + profile, + mdp=mdp, + structure=structure, + top=top, + tpr=tpr, + ref=ref, + maxwarn=maxwarn, + workdir=workdir, + grompp_log=workdir / f"grompp_{deffnm}.log", + ) + _mdrun(profile, tpr=tpr, deffnm=deffnm, workdir=workdir) + return GmxRunResult( + deffnm=deffnm, + tpr=tpr, + gro=workdir / f"{deffnm}.gro", + edr=workdir / f"{deffnm}.edr", + log=workdir / f"{deffnm}.log", + ) + + +# --------------------------------------------------------------------------- # +# solvate / energy primitives +# --------------------------------------------------------------------------- # +class AddWaterResult(BaseModel): + out_pdb: Path + nwat_added: int + log: Path + + +def add_water( + input_gro: Path, + reservoir_gro: Path, + maxsol: int | None, + out_pdb: Path, + workdir: Path, + *, + gmx_bin: str = "gmx", + log_path: Path | None = None, +) -> AddWaterResult: + """``gmx solvate`` to add water from a reservoir; parse the count actually added. + + ``maxsol=None`` runs an unbounded trial solvation (used to size the fill). + """ + workdir = Path(workdir) + log_path = Path(log_path) if log_path is not None else workdir / "gmx_solvate.log" + cmd = [gmx_bin, "solvate", "-cp", str(input_gro), "-cs", str(reservoir_gro), "-o", str(out_pdb)] + if maxsol is not None: + cmd += ["-maxsol", str(maxsol)] + run_tool(cmd, tool="gmx solvate", cwd=workdir, log_path=log_path) + added = _parse_solvent_count(log_path) + return AddWaterResult(out_pdb=out_pdb, nwat_added=added, log=log_path) + + +def _parse_solvent_count(log_path: Path) -> int: + for line in Path(log_path).read_text().splitlines(): + if "Number of solvent molecules:" in line: + for token in reversed(line.split()): + if token.isdigit(): + return int(token) + raise GmxOutputParseError("number of solvent molecules", log_path) + + +def read_mean_pressure(edr: Path, workdir: Path, *, gmx_bin: str = "gmx") -> float: + """Mean pressure (bar) from an ``.edr`` via ``gmx energy``.""" + workdir = Path(workdir) + tmp = workdir / f".pressure_{Path(edr).stem}.xvg" + try: + proc = run_tool( + [gmx_bin, "energy", "-f", str(edr), "-o", str(tmp)], + tool="gmx energy", + cwd=workdir, + input="Pressure\n0\n", + ) + finally: + tmp.unlink(missing_ok=True) + output = (proc.stdout or "") + (proc.stderr or "") + for line in output.splitlines(): + if line.strip().startswith("Pressure"): + parts = line.split() + if len(parts) >= 2: + try: + return float(parts[1]) + except ValueError: + continue + raise GmxOutputParseError("mean Pressure", edr, gmx_version=gmx_version(gmx_bin)) + + +# --------------------------------------------------------------------------- # +# Topology WAT-block management +# --------------------------------------------------------------------------- # +def manage_wat_topology(top: Path, count: int, *, mode: str = "append") -> None: + """Maintain the ordered ``WAT`` blocks in a topology's ``[ molecules ]`` section. + + Waters are added at several stages and sit in non-contiguous coordinate blocks + (initial waters -> Cl-/Na+ -> each resolvation stage's waters), so the topology + must carry a matching *sequence* of ``WAT`` lines, not one summed line. + + - ``mode="append"``: append ``WAT `` (idempotent — skips if the file already + ends with that exact block, so re-running a stage does not double-add). + - ``mode="drop_last_then_append"``: drop the trailing ``WAT`` block (the discarded + run-2 probe) then append ``WAT `` — the final resolvation stage. + """ + top = Path(top) + lines = top.read_text().splitlines() + if "[ molecules ]" not in "\n".join(lines): + raise ValueError(f"{top}: no '[ molecules ]' section to manage WAT blocks in") + + while lines and not lines[-1].strip(): + lines.pop() + + if mode == "drop_last_then_append": + if lines and lines[-1].split()[:1] == ["WAT"]: + lines.pop() + elif mode != "append": + raise ValueError(f"unknown manage_wat_topology mode {mode!r}") + + target = f"WAT {count}" + if not (lines and lines[-1].strip() == target): + lines.append(target) + + top.write_text("\n".join(lines) + "\n") diff --git a/md_workflows/core/paths.py b/md_workflows/core/paths.py new file mode 100644 index 0000000..b520bba --- /dev/null +++ b/md_workflows/core/paths.py @@ -0,0 +1,16 @@ +"""Path-resolution helpers shared by step resolvers.""" + +from __future__ import annotations + +from pathlib import Path + + +def under(base: Path, path: str | Path) -> Path: + """Resolve ``path`` against ``base`` unless it is already absolute. + + Used so a relative config value (e.g. ``mdp_dir="artifacts"``) is located inside the + run's ``workdir`` while an absolute path is honored as-is — keeping steps + cwd-independent. + """ + p = Path(path) + return p if p.is_absolute() else base / p diff --git a/md_workflows/core/pipelines/__init__.py b/md_workflows/core/pipelines/__init__.py new file mode 100644 index 0000000..58084e2 --- /dev/null +++ b/md_workflows/core/pipelines/__init__.py @@ -0,0 +1,5 @@ +"""Hardened multi-step pipelines composed from the individual steps. + +Pipelines are plain function composition over ``core.steps`` — no orchestration +framework. The same composition can later be wrapped as a Prefect ``@flow`` unchanged. +""" diff --git a/md_workflows/core/pipelines/standard_md.py b/md_workflows/core/pipelines/standard_md.py new file mode 100644 index 0000000..c9c8efb --- /dev/null +++ b/md_workflows/core/pipelines/standard_md.py @@ -0,0 +1,90 @@ +"""The standard crystalline-MD prep pipeline, as plain function composition. + +``param_prot -> make_crystal -> make_waterbox -> solvate -> minimize -> equilibrate -> +resolvate`` (the last stage runs the adaptive pressure loop). Every step reads/writes +the shared ``workdir`` under fixed conventional names, so each stage's inputs are +resolved from that workdir; with ``resume=True`` a stage whose durable outputs already +exist is skipped. + +This function is the literal thing a Prefect ``@flow`` would wrap later. +""" + +from __future__ import annotations + +from pathlib import Path + +from pydantic import BaseModel + +from ..config import SystemConfig +from ..results import StepResult +from ..steps import ( + equilibrate as equilibrate_step, +) +from ..steps import ( + make_crystal as make_crystal_step, +) +from ..steps import ( + make_waterbox as make_waterbox_step, +) +from ..steps import ( + minimize as minimize_step, +) +from ..steps import ( + param_prot as param_prot_step, +) +from ..steps import ( + resolvate as resolvate_step, +) +from ..steps import ( + solvate as solvate_step, +) + + +class PipelineResult(BaseModel): + workdir: Path + steps: list[StepResult] + + @property + def final(self) -> StepResult: + return self.steps[-1] + + +def standard_md_pipeline( + workdir: Path, + cfg: SystemConfig, + *, + resume: bool = False, +) -> PipelineResult: + """Run the full prep pipeline in ``workdir`` under ``cfg`` and return every result.""" + workdir = Path(workdir) + + pp = param_prot_step.param_prot(param_prot_step.resolve_inputs(workdir, cfg), resume=resume) + mc = make_crystal_step.make_crystal( + make_crystal_step.resolve_inputs(workdir, cfg), cfg.crystal, resume=resume + ) + wb = make_waterbox_step.make_waterbox( + make_waterbox_step.resolve_inputs(workdir, cfg), + cfg.waterbox, + cfg.profile("waterbox_min"), + cfg.profile("waterbox_equil"), + resume=resume, + ) + sv = solvate_step.solvate(solvate_step.resolve_inputs(workdir, cfg), cfg.solvate, resume=resume) + mn = minimize_step.minimize( + minimize_step.resolve_inputs(workdir, cfg), cfg.profile("min"), resume=resume + ) + eq = equilibrate_step.equilibrate( + equilibrate_step.resolve_inputs(workdir, cfg), + cfg.equilibrate, + cfg.profile("equil"), + resume=resume, + ) + rv = resolvate_step.resolvate( + resolvate_step.resolve_inputs(workdir, cfg), + cfg.resolvate, + cfg.profile("resolv_min"), + cfg.profile("resolv_equil"), + resume=resume, + ) + + return PipelineResult(workdir=workdir, steps=[pp, mc, wb, sv, mn, eq, rv]) diff --git a/md_workflows/core/results.py b/md_workflows/core/results.py new file mode 100644 index 0000000..114db8a --- /dev/null +++ b/md_workflows/core/results.py @@ -0,0 +1,71 @@ +"""Typed inputs and results shared by every core step. + +These are Pydantic models so results are serializable (JSON round-trip) and carry +enough metadata — parameters used, input checksums, tool versions — to later prove a +CLI run and an orchestrated run are scientifically equivalent. Per-step ``*Inputs`` and +``*Result`` subclasses live next to each step in ``core/steps``. +""" + +from __future__ import annotations + +from enum import Enum +from pathlib import Path +from typing import Any + +from pydantic import BaseModel, ConfigDict, Field + +from .config import GromacsRunProfile +from .exceptions import MissingInputError + + +class StepStatus(str, Enum): + COMPLETED = "completed" + SKIPPED = "skipped" # outputs already present and inputs unchanged (--resume) + FAILED = "failed" + + +class StepInputs(BaseModel): + """Explicit, typed inputs for one step. Subclasses add the concrete file fields. + + Core steps take one of these instead of assuming filenames in ``cwd``. + """ + + model_config = ConfigDict(extra="forbid") + + workdir: Path + + def consumed_paths(self) -> list[Path]: + """Files that must exist before the step runs. Overridden per step.""" + return [] + + def check_exists(self, step: str) -> None: + """Raise :class:`MissingInputError` if any required input is absent.""" + missing = [p for p in self.consumed_paths() if not Path(p).exists()] + if missing: + raise MissingInputError(step, missing) + + +class StepResult(BaseModel): + """Typed result of a step: what it produced plus reproducibility metadata.""" + + model_config = ConfigDict(extra="forbid") + + step: str + status: StepStatus = StepStatus.COMPLETED + workdir: Path + outputs: dict[str, Path] = Field(default_factory=dict) + params: dict[str, Any] = Field(default_factory=dict) + run_profiles_used: dict[str, GromacsRunProfile] = Field(default_factory=dict) + input_checksums: dict[str, str] = Field(default_factory=dict) + tool_versions: dict[str, str] = Field(default_factory=dict) + metrics: dict[str, float | int] = Field(default_factory=dict) + log_paths: list[Path] = Field(default_factory=list) + + def output(self, name: str) -> Path: + """Return a named output path, or raise KeyError with a helpful message.""" + try: + return self.outputs[name] + except KeyError as exc: + raise KeyError( + f"{self.step}: no output named {name!r}; have {sorted(self.outputs)}" + ) from exc diff --git a/md_workflows/core/steps/__init__.py b/md_workflows/core/steps/__init__.py new file mode 100644 index 0000000..d282cb2 --- /dev/null +++ b/md_workflows/core/steps/__init__.py @@ -0,0 +1,11 @@ +"""Individual MD prep steps. + +Each module exposes the same contract: + +* a ``*Inputs`` model (explicit, typed file paths + ``workdir``), +* ``resolve_inputs(workdir, cfg)`` — map the conventional filenames in a working + directory to explicit input paths, +* ``check_inputs(inputs)`` — guard that raises ``MissingInputError`` before any tool + runs, +* the step function itself, which takes ``*Inputs`` and returns a ``StepResult``. +""" diff --git a/md_workflows/core/steps/equilibrate.py b/md_workflows/core/steps/equilibrate.py new file mode 100644 index 0000000..a982352 --- /dev/null +++ b/md_workflows/core/steps/equilibrate.py @@ -0,0 +1,282 @@ +"""Set up position restraints and run the NPT equilibration. + +Corresponds to ``equilibrate.sh``: isolate one asymmetric-unit copy, split it into +chains, generate per-chain position restraints, weave ``#ifdef POSRES_*`` blocks into +the topology, then equilibrate. + +Ligand handling is deliberately minimal here (full support is a later feature): + +* ``chain_split_mode="ignore_ligand"`` (default) restrains protein chains only (taylor + behavior). +* ``chain_split_mode="split_intelligently"`` isolates copy 1 up to the ligand before + splitting (original behavior, generalized to any ``ligand_resname``). +* ``handle_ligand_restraint`` (default ``False``) additionally restrains the ligand via + a ``make_ndx`` group — provisional and off by default. +""" + +from __future__ import annotations + +import shutil +from pathlib import Path + +from ...pdb_file_processing import find_ligands_in_legacy_pdb_file +from ..config import EquilibrateParams, GromacsRunProfile, SystemConfig +from ..exceptions import MDWorkflowError +from ..gmx import checksums, gmx_version, run_equil, run_tool +from ..paths import under +from ..results import StepInputs, StepResult, StepStatus + +STEP = "equilibrate" +DEFFNM = "md_equil" + + +class EquilibrateInputs(StepInputs): + pdb_clean: Path + model_top: Path + model_pdb: Path + min_gro: Path + equil_mdp: Path + + def consumed_paths(self) -> list[Path]: + return [self.pdb_clean, self.model_top, self.model_pdb, self.min_gro, self.equil_mdp] + + +class EquilibrateResult(StepResult): + @property + def gro(self) -> Path: + return self.output("gro") + + @property + def edr(self) -> Path: + return self.output("edr") + + @property + def posre_top(self) -> Path: + return self.output("posre_top") + + +def resolve_inputs(workdir: Path, cfg: SystemConfig) -> EquilibrateInputs: + workdir = Path(workdir) + mdp_dir = under(workdir, cfg.mdp_dir) + return EquilibrateInputs( + workdir=workdir, + pdb_clean=workdir / "pdb_clean.pdb", + model_top=workdir / "md_model.top", + model_pdb=workdir / "md_model.pdb", + min_gro=workdir / "md_min.gro", + equil_mdp=mdp_dir / cfg.equilibrate.equil_mdp, + ) + + +def check_inputs(inputs: EquilibrateInputs) -> None: + inputs.check_exists(STEP) + + +def equilibrate( + inputs: EquilibrateInputs, + params: EquilibrateParams, + profile: GromacsRunProfile, + *, + resume: bool = False, +) -> EquilibrateResult: + check_inputs(inputs) + wd = inputs.workdir + + posre_top = wd / "md_model_posre.top" + tpr = wd / f"{DEFFNM}.tpr" + gro = wd / f"{DEFFNM}.gro" + edr = wd / f"{DEFFNM}.edr" + outputs = {"posre_top": posre_top, "tpr": tpr, "gro": gro, "edr": edr} + consumed = { + "pdb_clean": inputs.pdb_clean, + "model_top": inputs.model_top, + "model_pdb": inputs.model_pdb, + "min_gro": inputs.min_gro, + "equil_mdp": inputs.equil_mdp, + } + + if resume and gro.exists() and posre_top.exists(): + return EquilibrateResult( + step=STEP, + status=StepStatus.SKIPPED, + workdir=wd, + outputs=outputs, + params=params.model_dump(), + input_checksums=checksums(consumed), + ) + + _clean_stale(wd) + _build_first_copy(wd, inputs.pdb_clean, params) + chain_files = _split_chains(wd) + _generate_restraints(wd, chain_files, params) + + shutil.copy(inputs.model_top, posre_top) + _weave_posre_includes(posre_top, chain_files) + + if params.handle_ligand_restraint: + _add_ligand_restraint(wd, inputs.pdb_clean, posre_top, params, len(chain_files)) + + run = run_equil( + inputs.min_gro, posre_top, inputs.equil_mdp, inputs.model_pdb, DEFFNM, profile, wd + ) + + return EquilibrateResult( + step=STEP, + status=StepStatus.COMPLETED, + workdir=wd, + outputs={"posre_top": posre_top, "tpr": run.tpr, "gro": run.gro, "edr": run.edr}, + params=params.model_dump(), + input_checksums=checksums(consumed), + run_profiles_used={"equil": profile}, + tool_versions={"gmx": gmx_version(profile.gmx_bin)}, + metrics={"n_chains": len(chain_files)}, + log_paths=[run.log], + ) + + +def _clean_stale(workdir: Path) -> None: + """Remove leftover chain fragments and restraint files from a prior run.""" + for pattern in ("part??", "part??_amber.pdb", "posre_part??.itp"): + for f in workdir.glob(pattern): + f.unlink() + + +def _build_first_copy(workdir: Path, pdb_clean: Path, params: EquilibrateParams) -> None: + """Write ``first_copy_prot.pdb`` — the protein atoms to restrain.""" + lines = [ + ln for ln in pdb_clean.read_text().splitlines(keepends=True) if not ln.startswith("JRNL") + ] + + if params.chain_split_mode == "ignore_ligand": + prot = [ln for ln in lines if ln.startswith(("ATOM", "TER"))] + (workdir / "first_copy_prot.pdb").write_text("".join(prot)) + return + + # split_intelligently: isolate copy 1 up to the first ligand residue, then drop it. + resn = _resolve_ligand_resname(params, pdb_clean) + kept: list[str] = [] + found = False + for line in lines: + if resn in line: + found = True + kept.append(line) + elif found: + break + else: + kept.append(line) + prot = [ln for ln in kept if ln.startswith(("ATOM", "HETATM", "TER")) and resn not in ln] + (workdir / "first_copy_prot.pdb").write_text("".join(prot)) + + +def _split_chains(workdir: Path) -> list[str]: + """Split ``first_copy_prot.pdb`` after each TER into part00, part01, ... files.""" + lines = (workdir / "first_copy_prot.pdb").read_text().splitlines(keepends=True) + chunks: list[str] = [] + buf: list[str] = [] + for line in lines: + buf.append(line) + if line.startswith("TER"): + chunks.append("".join(buf)) + buf = [] + if buf: + chunks.append("".join(buf)) + + files: list[str] = [] + for i, chunk in enumerate(c for c in chunks if c.strip()): + # Skip fragments with no atoms (empty-part guard, taylor). + if not any(ln.startswith("ATOM") for ln in chunk.splitlines()): + continue + name = f"part{i:02d}" + (workdir / name).write_text(chunk) + files.append(name) + return files + + +def _generate_restraints(workdir: Path, chain_files: list[str], params: EquilibrateParams) -> None: + fc = str(params.restraint_fc) + for f in chain_files: + run_tool( + ["pdb4amber", "-i", f, "-o", f"{f}_amber.pdb"], + tool="pdb4amber", + cwd=workdir, + log_path=workdir / f"pdb4amber_{f}.log", + ) + run_tool( + ["gmx", "genrestr", "-fc", fc, fc, fc, "-f", f"{f}_amber.pdb", "-o", f"posre_{f}.itp"], + tool="gmx genrestr", + cwd=workdir, + input=f"{params.restraint_group}\nq\n", + log_path=workdir / f"genrestr_{f}.log", + ) + + +def _weave_posre_includes(posre_top: Path, chain_files: list[str]) -> None: + """Insert ``#ifdef POSRES_partXX`` include blocks before each chain's moleculetype.""" + for offset, f in enumerate(chain_files): + target = offset + 2 # 1-indexed, skipping the first moleculetype + lines = posre_top.read_text().splitlines(keepends=True) + out: list[str] = [] + count = 0 + for line in lines: + if "moleculetype" in line.lower(): + count += 1 + if count == target: + out.append(f'#ifdef POSRES_{f}\n#include "posre_{f}.itp"\n#endif\n\n') + out.append(line) + posre_top.write_text("".join(out)) + + +def _add_ligand_restraint( + workdir: Path, pdb_clean: Path, posre_top: Path, params: EquilibrateParams, n_parts: int +) -> None: + """Provisional ligand restraint via a make_ndx group (off by default).""" + resn = _resolve_ligand_resname(params, pdb_clean) + f = f"part{n_parts:02d}" + lig_lines = [ + ln + for ln in pdb_clean.read_text().splitlines(keepends=True) + if ln.startswith(("ATOM", "HETATM")) and resn in ln + ] + if not lig_lines: + return + (workdir / f).write_text("".join(lig_lines)) + run_tool(["pdb4amber", "-i", f, "-o", f"{f}_amber.pdb"], tool="pdb4amber", cwd=workdir) + run_tool( + ["gmx", "make_ndx", "-f", f"{f}_amber.pdb", "-o", f"{f}_amber.ndx"], + tool="gmx make_ndx", + cwd=workdir, + input=f"! a H*\nname 3 {resn}-H\nq\n", + ) + fc = str(params.restraint_fc) + run_tool( + [ + "gmx", + "genrestr", + "-fc", + fc, + fc, + fc, + "-f", + f"{f}_amber.pdb", + "-o", + f"posre_{f}.itp", + "-n", + f"{f}_amber.ndx", + ], + tool="gmx genrestr", + cwd=workdir, + input=f"{resn}-H\nq\n", + ) + + +def _resolve_ligand_resname(params: EquilibrateParams, pdb_clean: Path) -> str: + if params.ligand_resname: + return params.ligand_resname + hits = find_ligands_in_legacy_pdb_file(pdb_clean) + unique = sorted({h["resname"] for h in hits}) + if len(unique) == 1: + return unique[0] + raise MDWorkflowError( + "equilibrate: ligand_resname is required for this mode but could not be " + f"auto-detected (found {unique or 'none'}); set equilibrate.ligand_resname." + ) diff --git a/md_workflows/core/steps/make_crystal.py b/md_workflows/core/steps/make_crystal.py new file mode 100644 index 0000000..d50c885 --- /dev/null +++ b/md_workflows/core/steps/make_crystal.py @@ -0,0 +1,178 @@ +"""Build a crystal supercell from the parameterized protein. + +Corresponds to ``make_crystal.sh``: dry the protein, restore the CRYST1 cell, expand +the asymmetric unit to the full unit cell with ChimeraX, set the spacegroup to P1, then +replicate into a supercell with PropPDB. +""" + +from __future__ import annotations + +import shutil +from pathlib import Path + +from ..config import CrystalParams, SystemConfig +from ..gmx import checksums +from ..results import StepInputs, StepResult, StepStatus + +STEP = "make_crystal" + + +class MakeCrystalInputs(StepInputs): + prot_pdb: Path + pdb_clean: Path + + def consumed_paths(self) -> list[Path]: + return [self.prot_pdb, self.pdb_clean] + + +class MakeCrystalResult(StepResult): + @property + def xtal(self) -> Path: + return self.output("xtal") + + +def resolve_inputs(workdir: Path, cfg: SystemConfig) -> MakeCrystalInputs: + workdir = Path(workdir) + return MakeCrystalInputs( + workdir=workdir, + prot_pdb=workdir / "prot.pdb", + pdb_clean=workdir / "pdb_clean.pdb", + ) + + +def check_inputs(inputs: MakeCrystalInputs) -> None: + inputs.check_exists(STEP) + + +def make_crystal( + inputs: MakeCrystalInputs, + params: CrystalParams, + *, + resume: bool = False, +) -> MakeCrystalResult: + from ..gmx import run_tool # local import keeps module import light for tests + + check_inputs(inputs) + wd = inputs.workdir + ix = params.ix + iy = params.iy if params.iy is not None else ix + iz = params.iz if params.iz is not None else ix + + prot_dry = wd / "prot_dry.pdb" + prot_dry_cell = wd / "prot_dry_cell.pdb" + xtal = wd / "xtal.pdb" + + outputs = {"prot_dry": prot_dry, "prot_dry_cell": prot_dry_cell, "xtal": xtal} + consumed = {"prot_pdb": inputs.prot_pdb, "pdb_clean": inputs.pdb_clean} + + if resume and xtal.exists(): + return MakeCrystalResult( + step=STEP, + status=StepStatus.SKIPPED, + workdir=wd, + outputs=outputs, + params=params.model_dump(), + input_checksums=checksums(consumed), + ) + + # 1. Dry the protein (strip waters/ions) -> prot_dry.pdb + run_tool( + ["pdb4amber", "-i", str(inputs.prot_pdb), "-o", str(prot_dry), "--dry"], + tool="pdb4amber", + cwd=wd, + log_path=wd / "pdb4amber_dry.log", + ) + + # 2. Restore the CRYST1 cell from pdb_clean; drop stray ions / any wrong CRYST1. + _prepend_cryst1(inputs.pdb_clean, prot_dry) + + # 3. Expand the asymmetric unit to the full unit cell with ChimeraX. + _expand_unit_cell(wd, prot_dry, prot_dry_cell, params.chimerax_exec) + + # 4. Rewrite the CRYST1 spacegroup to P1 and prepend to the cell PDB. + _set_p1_spacegroup(prot_dry, prot_dry_cell) + + # 5. Replicate the P1 cell into the requested supercell. + if ix > 0 or iy > 0 or iz > 0: + run_tool( + [ + "PropPDB", + "-p", + str(prot_dry_cell), + "-o", + str(xtal), + "-ix", + str(ix), + "-iy", + str(iy), + "-iz", + str(iz), + ], + tool="PropPDB", + cwd=wd, + log_path=wd / "proppdb.log", + ) + else: + shutil.copy(prot_dry_cell, xtal) + + return MakeCrystalResult( + step=STEP, + status=StepStatus.COMPLETED, + workdir=wd, + outputs=outputs, + params=params.model_dump(), + input_checksums=checksums(consumed), + metrics={"ix": ix, "iy": iy, "iz": iz}, + ) + + +def _prepend_cryst1(source_pdb: Path, target_pdb: Path) -> None: + """Copy CRYST1 from source to the top of target, dropping Na+/Cl-/old CRYST1.""" + cryst1 = "" + with open(source_pdb) as fh: + for line in fh: + if line.startswith("CRYST1"): + cryst1 = line + break + + with open(target_pdb) as fh: + lines = fh.readlines() + filtered = [ + line + for line in lines + if "Na+" not in line and "Cl-" not in line and not line.startswith("CRYST1") + ] + with open(target_pdb, "w") as fh: + fh.write(cryst1) + fh.writelines(filtered) + + +def _expand_unit_cell(workdir: Path, dry_pdb: Path, cell_pdb: Path, chimerax_exec: str) -> None: + from ..gmx import run_tool + + cxc = workdir / "expand.cxc" + cxc.write_text( + f"open {dry_pdb}\nchangechains #1 A\nunitcell #1\ncombine #2\nsave {cell_pdb} #3\nquit\n" + ) + run_tool( + [chimerax_exec, "--offscreen", "--nogui", str(cxc)], + tool="ChimeraX", + cwd=workdir, + log_path=workdir / "chimerax_expand.log", + ) + + +def _set_p1_spacegroup(dry_pdb: Path, cell_pdb: Path) -> None: + """Rewrite the CRYST1 line with a P 1 spacegroup and prepend to the cell PDB.""" + cryst1_p1 = "" + with open(dry_pdb) as fh: + for line in fh: + if line.startswith("CRYST1"): + cryst1_p1 = line[:55] + "P 1\n" + break + + with open(cell_pdb) as fh: + cell_content = fh.read() + with open(cell_pdb, "w") as fh: + fh.write(cryst1_p1) + fh.write(cell_content) diff --git a/md_workflows/core/steps/make_waterbox.py b/md_workflows/core/steps/make_waterbox.py new file mode 100644 index 0000000..26e6f01 --- /dev/null +++ b/md_workflows/core/steps/make_waterbox.py @@ -0,0 +1,208 @@ +"""Build an equilibrated bulk-water reservoir matching the crystal unit cell. + +Corresponds to the canonical (taylor) ``make_waterbox.sh``: subdivide the cell by +``nc_scale``, fill the sub-cell with water at ``conc`` mol/L, tile it back up +``nc_scale``x, restore the full crystal CRYST1, then minimize + equilibrate the box. + +The water count written to the topology is the exact molecule count of the expanded box +(``count_waters``), superseding the scripts' log-derived estimates. +""" + +from __future__ import annotations + +from pathlib import Path + +from ..config import GromacsRunProfile, SystemConfig, WaterboxParams +from ..gmx import checksums, count_waters, run_min, run_tool +from ..paths import under +from ..results import StepInputs, StepResult, StepStatus + +STEP = "make_waterbox" + + +class MakeWaterboxInputs(StepInputs): + xtal: Path + wat_pdb: Path + prot_top: Path + min_mdp: Path + equil_mdp: Path + + def consumed_paths(self) -> list[Path]: + return [self.xtal, self.wat_pdb, self.prot_top, self.min_mdp, self.equil_mdp] + + +class MakeWaterboxResult(StepResult): + @property + def water_equil_gro(self) -> Path: + return self.output("water_equil_gro") + + +def resolve_inputs(workdir: Path, cfg: SystemConfig) -> MakeWaterboxInputs: + workdir = Path(workdir) + mdp_dir = under(workdir, cfg.mdp_dir) + return MakeWaterboxInputs( + workdir=workdir, + xtal=workdir / "xtal.pdb", + wat_pdb=workdir / "WAT.pdb", + prot_top=workdir / "prot.top", + min_mdp=mdp_dir / cfg.waterbox.min_mdp, + equil_mdp=mdp_dir / cfg.waterbox.equil_mdp, + ) + + +def check_inputs(inputs: MakeWaterboxInputs) -> None: + inputs.check_exists(STEP) + + +def make_waterbox( + inputs: MakeWaterboxInputs, + params: WaterboxParams, + min_profile: GromacsRunProfile, + equil_profile: GromacsRunProfile, + *, + resume: bool = False, +) -> MakeWaterboxResult: + check_inputs(inputs) + wd = inputs.workdir + wb = wd / "waterbox" + wb.mkdir(parents=True, exist_ok=True) + nc = params.nc_scale + + box_solv_expand = wb / "box_solv_expand.pdb" + waterbox_top = wb / "waterbox.top" + water_min_gro = wb / "water_min.gro" + water_equil_gro = wb / "water_equil.gro" + outputs = { + "box_solv_expand": box_solv_expand, + "waterbox_top": waterbox_top, + "water_min_gro": water_min_gro, + "water_equil_gro": water_equil_gro, + } + consumed = { + "xtal": inputs.xtal, + "wat_pdb": inputs.wat_pdb, + "prot_top": inputs.prot_top, + "min_mdp": inputs.min_mdp, + "equil_mdp": inputs.equil_mdp, + } + + if resume and water_equil_gro.exists(): + return MakeWaterboxResult( + step=STEP, + status=StepStatus.SKIPPED, + workdir=wd, + outputs=outputs, + params=params.model_dump(), + input_checksums=checksums(consumed), + ) + + cryst1_xtal = _extract_cryst1(inputs.xtal) + (wb / "cryst1_xtal.pdb").write_text(cryst1_xtal) + _create_box_pdb(inputs.xtal, wb / "box.pdb", nc) + + # Fill the sub-cell with water, then tile it up nc x nc x nc. + run_tool( + [ + "gmx", + "insert-molecules", + "-f", + str(wb / "box.pdb"), + "-ci", + str(inputs.wat_pdb), + "-conc", + str(params.conc), + "-o", + str(wb / "box_solv.pdb"), + ], + tool="gmx insert-molecules", + cwd=wb, + log_path=wb / "insert-molecules.log", + ) + run_tool( + [ + "PropPDB", + "-p", + str(wb / "box_solv.pdb"), + "-o", + str(box_solv_expand), + "-ix", + str(nc), + "-iy", + str(nc), + "-iz", + str(nc), + ], + tool="PropPDB", + cwd=wb, + log_path=wb / "proppdb.log", + ) + _restore_cryst1(box_solv_expand, cryst1_xtal) + + nwat = count_waters(box_solv_expand) + _write_topology(inputs.prot_top, waterbox_top, nwat) + + # Water box has no restraint reference, so equilibration is a plain grompp+mdrun. + min_run = run_min(box_solv_expand, waterbox_top, inputs.min_mdp, "water_min", min_profile, wb) + equil_run = run_min( + min_run.gro, waterbox_top, inputs.equil_mdp, "water_equil", equil_profile, wb + ) + + return MakeWaterboxResult( + step=STEP, + status=StepStatus.COMPLETED, + workdir=wd, + outputs=outputs, + params=params.model_dump(), + input_checksums=checksums(consumed), + run_profiles_used={"waterbox_min": min_profile, "waterbox_equil": equil_profile}, + metrics={"nwat": nwat, "nc_scale": nc}, + log_paths=[min_run.log, equil_run.log], + ) + + +def _extract_cryst1(xtal: Path) -> str: + with open(xtal) as fh: + for line in fh: + if line.startswith("CRYST1"): + return line + raise ValueError(f"{xtal}: no CRYST1 record found") + + +def _create_box_pdb(xtal: Path, box_pdb: Path, nc: int) -> None: + """Write a CRYST1-only box.pdb with cell dimensions divided by ``nc``.""" + with open(xtal) as fh: + for line in fh: + if line.startswith("CRYST1"): + a = float(line[6:15]) / nc + b = float(line[15:24]) / nc + c = float(line[24:33]) / nc + alpha = float(line[33:40]) + beta = float(line[40:47]) + gamma = float(line[47:54]) + box_pdb.write_text( + f"CRYST1{a:9.3f}{b:9.3f}{c:9.3f}{alpha:7.2f}{beta:7.2f}{gamma:7.2f}\n" + ) + return + raise ValueError(f"{xtal}: no CRYST1 record found") + + +def _restore_cryst1(expanded_pdb: Path, cryst1: str) -> None: + with open(expanded_pdb) as fh: + body = [line for line in fh if not line.startswith(("CRYST1", "HEADER"))] + with open(expanded_pdb, "w") as fh: + fh.write(cryst1) + fh.writelines(body) + + +def _write_topology(prot_top: Path, waterbox_top: Path, nwat: int) -> None: + header: list[str] = [] + with open(prot_top) as fh: + for line in fh: + if "molecules" in line.lower(): + break + header.append(line) + with open(waterbox_top, "w") as fh: + fh.writelines(header) + fh.write("[ molecules ]\n") + fh.write("; Compound #mols\n") + fh.write(f"WAT {nwat}\n") diff --git a/md_workflows/core/steps/minimize.py b/md_workflows/core/steps/minimize.py new file mode 100644 index 0000000..c16079f --- /dev/null +++ b/md_workflows/core/steps/minimize.py @@ -0,0 +1,96 @@ +"""Energy-minimize the solvated MD model. + +Reference implementation of the step contract (Inputs / resolve / guard / typed +Result) that the other steps follow. Corresponds to ``minimize.sh``: +``gmx grompp -f min.mdp -c md_model.pdb -o md_min.tpr -p md_model.top`` then +``gmx mdrun -deffnm md_min``. +""" + +from __future__ import annotations + +from pathlib import Path + +from ..config import GromacsRunProfile, SystemConfig +from ..gmx import checksums, gmx_version, run_min +from ..paths import under +from ..results import StepInputs, StepResult, StepStatus + +STEP = "minimize" +DEFFNM = "md_min" + + +class MinimizeInputs(StepInputs): + model_pdb: Path + model_top: Path + min_mdp: Path + + def consumed_paths(self) -> list[Path]: + return [self.model_pdb, self.model_top, self.min_mdp] + + +class MinimizeResult(StepResult): + @property + def gro(self) -> Path: + return self.output("gro") + + @property + def tpr(self) -> Path: + return self.output("tpr") + + +def resolve_inputs(workdir: Path, cfg: SystemConfig) -> MinimizeInputs: + """Map the conventional filenames in ``workdir`` to explicit input paths.""" + workdir = Path(workdir) + mdp_dir = under(workdir, cfg.mdp_dir) + return MinimizeInputs( + workdir=workdir, + model_pdb=workdir / "md_model.pdb", + model_top=workdir / "md_model.top", + min_mdp=mdp_dir / cfg.minimize.min_mdp, + ) + + +def check_inputs(inputs: MinimizeInputs) -> None: + inputs.check_exists(STEP) + + +def minimize( + inputs: MinimizeInputs, + profile: GromacsRunProfile, + *, + resume: bool = False, +) -> MinimizeResult: + """Run energy minimization and return a typed, metadata-carrying result.""" + check_inputs(inputs) + + consumed = { + "model_pdb": inputs.model_pdb, + "model_top": inputs.model_top, + "min_mdp": inputs.min_mdp, + } + tpr = inputs.workdir / f"{DEFFNM}.tpr" + gro = inputs.workdir / f"{DEFFNM}.gro" + + if resume and gro.exists(): + return MinimizeResult( + step=STEP, + status=StepStatus.SKIPPED, + workdir=inputs.workdir, + outputs={"tpr": tpr, "gro": gro}, + run_profiles_used={"min": profile}, + input_checksums=checksums(consumed), + ) + + run = run_min( + inputs.model_pdb, inputs.model_top, inputs.min_mdp, DEFFNM, profile, inputs.workdir + ) + return MinimizeResult( + step=STEP, + status=StepStatus.COMPLETED, + workdir=inputs.workdir, + outputs={"tpr": run.tpr, "gro": run.gro}, + run_profiles_used={"min": profile}, + input_checksums=checksums(consumed), + tool_versions={"gmx": gmx_version(profile.gmx_bin)}, + log_paths=[run.log], + ) diff --git a/md_workflows/core/steps/param_prot.py b/md_workflows/core/steps/param_prot.py new file mode 100644 index 0000000..7706468 --- /dev/null +++ b/md_workflows/core/steps/param_prot.py @@ -0,0 +1,196 @@ +"""Parameterize the protein: clean the PDB, build the Amber topology, convert to GROMACS. + +Corresponds to ``param_prot.sh``: strip the entry PDB, run pdb4amber, build a minimal +solvation with tleap to obtain solvent templates, convert Amber->GROMACS with ParmEd, +extract single Na+/Cl-/WAT templates, then rebuild the final protein topology. + +This is the first step; it downloads ``.pdb`` into the workdir if absent, so it +has no required pre-existing input files. +""" + +from __future__ import annotations + +import sys +import textwrap +from pathlib import Path + +from ...pdb_file_processing import ensure_entry_pdb_file +from ..config import SystemConfig +from ..gmx import checksums, run_tool +from ..results import StepInputs, StepResult, StepStatus + +STEP = "param_prot" + +_TLEAP_FF = """\ +source leaprc.protein.ff19SB +source leaprc.DNA.OL15 +source leaprc.RNA.OL3 +source leaprc.water.spceb +source leaprc.gaff2 +""" + + +class ParamProtInputs(StepInputs): + pdb_id: str + + def consumed_paths(self) -> list[Path]: + return [] # entry PDB is downloaded if missing + + +class ParamProtResult(StepResult): + @property + def prot_top(self) -> Path: + return self.output("prot_top") + + @property + def prot_pdb(self) -> Path: + return self.output("prot_pdb") + + +def resolve_inputs(workdir: Path, cfg: SystemConfig) -> ParamProtInputs: + return ParamProtInputs(workdir=Path(workdir), pdb_id=cfg.pdb_id) + + +def check_inputs(inputs: ParamProtInputs) -> None: + inputs.check_exists(STEP) + + +def param_prot(inputs: ParamProtInputs, *, resume: bool = False) -> ParamProtResult: + check_inputs(inputs) + wd = inputs.workdir + wd.mkdir(parents=True, exist_ok=True) + + outputs = { + "pdb_clean": wd / "pdb_clean.pdb", + "pdb_clean_amber": wd / "pdb_clean_amber.pdb", + "prot_parm7": wd / "prot.parm7", + "prot_rst7": wd / "prot.rst7", + "prot_top": wd / "prot.top", + "prot_pdb": wd / "prot.pdb", + "wat_pdb": wd / "WAT.pdb", + "na_pdb": wd / "Na+.pdb", + "cl_pdb": wd / "Cl-.pdb", + } + + if resume and outputs["prot_top"].exists() and outputs["prot_pdb"].exists(): + return ParamProtResult( + step=STEP, + status=StepStatus.SKIPPED, + workdir=wd, + outputs=outputs, + params={"pdb_id": inputs.pdb_id}, + ) + + entry_pdb = ensure_entry_pdb_file(inputs.pdb_id, wd) + + _clean_pdb(wd, entry_pdb) + _tleap(wd, "tleap_temp.in", _initial_solvation_input()) + _amber_to_gromacs(wd) + _extract_solvent_pdbs(wd) + _tleap(wd, "tleap_prot.in", _final_topology_input()) + _amber_to_gromacs(wd) + + return ParamProtResult( + step=STEP, + status=StepStatus.COMPLETED, + workdir=wd, + outputs=outputs, + params={"pdb_id": inputs.pdb_id}, + input_checksums=checksums({"entry_pdb": entry_pdb}), + ) + + +def _clean_pdb(workdir: Path, entry_pdb: Path) -> None: + """Strip REMARK/KEYWDS, keep structural records, then run pdb4amber (--prot).""" + kept: list[str] = [] + with open(entry_pdb) as fh: + for line in fh: + if line.startswith(("REMARK", "KEYWDS")): + continue + if line.startswith(("CRYST1", "ATOM", "HETATM", "TER", "END")): + kept.append(line) + (workdir / "pdb_clean.pdb").write_text("".join(kept)) + run_tool( + ["pdb4amber", "-i", "pdb_clean.pdb", "--prot", "-o", "pdb_clean_amber.pdb"], + tool="pdb4amber", + cwd=workdir, + log_path=workdir / "pdb4amber_clean.log", + ) + + +def _initial_solvation_input() -> str: + return _TLEAP_FF + textwrap.dedent("""\ + p = loadpdb pdb_clean_amber.pdb + x = combine{p} + addions2 x Cl- 0 + addions2 x Na+ 0 + addions2 x Na+ 1 + addions2 x Cl- 1 + solvateBox x SPCBOX 1. + set default PBradii mbondi3 + set default nocenter on + saveAmberParm x prot.parm7 prot.rst7 + quit + """) + + +def _final_topology_input() -> str: + return _TLEAP_FF + textwrap.dedent("""\ + p = loadpdb pdb_clean_amber.pdb + w = loadpdb WAT.pdb + x = combine{p w} + addions2 x Na+ 0 + addions2 x Cl- 0 + addions2 x Na+ 1 + addions2 x Cl- 1 + set default PBradii mbondi3 + set default nocenter on + saveAmberParm x prot.parm7 prot.rst7 + quit + """) + + +def _tleap(workdir: Path, script_name: str, tleap_input: str) -> None: + (workdir / script_name).write_text(tleap_input) + run_tool( + ["tleap", "-f", script_name], + tool="tleap", + cwd=workdir, + log_path=workdir / f"{script_name}.log", + ) + # prot.top/prot.pdb are (re)generated by ParmEd; remove any stale copies first. + for f in ("prot.top", "prot.pdb"): + (workdir / f).unlink(missing_ok=True) + + +def _amber_to_gromacs(workdir: Path) -> None: + """Convert Amber parm7/rst7 to GROMACS topology + PDB with ParmEd.""" + script = textwrap.dedent("""\ + import parmed as pmd + parm = pmd.load_file("prot.parm7", "prot.rst7") + parm.save("prot.top") + parm.save("prot.pdb") + """) + (workdir / "amber_to_gromacs.py").write_text(script) + run_tool( + [sys.executable, "amber_to_gromacs.py"], + tool="parmed (amber->gromacs)", + cwd=workdir, + log_path=workdir / "amber_to_gromacs.log", + ) + + +def _extract_solvent_pdbs(workdir: Path) -> None: + """Pull single Na+, Cl-, and WAT coordinate templates from the solvated PDB.""" + with open(workdir / "prot.pdb") as fh: + hetatm = [line for line in fh if line.startswith("HETATM")] + + na = [line for line in hetatm if "Na+" in line] + cl = [line for line in hetatm if "Cl-" in line] + wat = [line for line in hetatm if "WAT" in line] + if na: + (workdir / "Na+.pdb").write_text(na[0]) + if cl: + (workdir / "Cl-.pdb").write_text(cl[0]) + if wat: + (workdir / "WAT.pdb").write_text("".join(wat[:3])) diff --git a/md_workflows/core/steps/pressure_interpolate.py b/md_workflows/core/steps/pressure_interpolate.py new file mode 100644 index 0000000..d09e102 --- /dev/null +++ b/md_workflows/core/steps/pressure_interpolate.py @@ -0,0 +1,154 @@ +"""Two-state linear interpolation for the resolvation water correction. + +Port of the taylor ``pressure_interpolate.py`` (Wych & Wall MDPreparation) as a pure +function. Given two ``(water count, mean pressure)`` states it solves for the water +delta that brings the system to the target pressure: + + dNw = round((target - P1) * (Nw2 - Nw1) / (P2 - P1)) + +The script's ``argparse``/``print`` command-suggestion UX is dropped — that belongs to +the CLI/flow layer — and errors are raised rather than ``sys.exit``. + +The numeric solve (:func:`interpolate`) is separated from the gmx I/O +(:func:`pressure_interpolate`) so it is unit-testable without GROMACS. + +Two robustness guards beyond the original script (both use the same rule: if the system +is already within ``pressure_tol`` of target, report convergence with ``dNw=0``; +otherwise raise, because the linear model cannot be trusted to reach target): + +* **Zero / near-zero slope** (``P2 ≈ P1``): adding water did not move the pressure, so + the slope is undefined — the original ``P1 == P2`` guard, generalized to an epsilon. +* **Absurd ``dNw``**: a tiny (but nonzero) slope can produce an impossible water count; + reject any ``|dNw|`` beyond ``max_dNw_factor * Nw1``. +""" + +from __future__ import annotations + +from pathlib import Path + +from pydantic import BaseModel + +from ..exceptions import MDWorkflowError +from ..gmx import count_waters, read_mean_pressure + +_SLOPE_EPS = 1e-9 # bar; below this |P2-P1| is treated as zero slope + + +class InterpolationResult(BaseModel): + dNw: int + Nw_target: int + ref_gro: Path + P1: float + P2: float + Nw1: int + Nw2: int + converged: bool = False # True when already within pressure_tol of target + warning: str | None = None + + +def interpolate( + p1: float, + p2: float, + nw1: int, + nw2: int, + ref_gro: Path, + *, + target: float = 1.0, + pressure_tol: float = 100.0, + max_dNw_factor: float = 1.0, +) -> InterpolationResult: + """Pure numeric solve for the water delta (no gmx I/O).""" + ref = Path(ref_gro) + near_target = abs(p1 - target) <= pressure_tol + + def _converged(reason: str) -> InterpolationResult: + return _result(0, nw1, ref, p1, p2, nw1, nw2, converged=True, warning=reason) + + # Guard 1: zero / near-zero slope — linear model is undefined. + if abs(p2 - p1) < _SLOPE_EPS: + if near_target: + return _converged( + f"pressure insensitive to water (P1={p1:.1f}≈P2={p2:.1f} bar) but within " + f"{pressure_tol:.0f} bar of target — keeping current solvation." + ) + raise MDWorkflowError( + f"pressure_interpolate: zero slope (P1={p1:.1f}≈P2={p2:.1f} bar) and pressure " + f"is {abs(p1 - target):.0f} bar from target ({target:.1f}); cannot interpolate." + ) + + dNw = int(round((target - p1) * (nw2 - nw1) / (p2 - p1))) + + # Guard 2: sanity-bound the correction against tiny-slope blow-ups. + if abs(dNw) > max_dNw_factor * nw1: + if near_target: + return _converged( + f"computed dNw={dNw} exceeds sanity bound ({max_dNw_factor}*{nw1}) but " + f"pressure is within {pressure_tol:.0f} bar of target — keeping current." + ) + raise MDWorkflowError( + f"pressure_interpolate: computed dNw={dNw} exceeds sanity bound " + f"({max_dNw_factor}*{nw1} waters); the two-state model is unreliable here." + ) + + warning = None + if (p1 - target) * (p2 - target) > 0: + warning = ( + f"extrapolation: P1={p1:.1f} and P2={p2:.1f} bar are on the same side of " + f"target={target:.1f} bar; result may be inaccurate." + ) + return _result(dNw, nw1 + dNw, ref, p1, p2, nw1, nw2, converged=False, warning=warning) + + +def pressure_interpolate( + gro1: Path, + edr1: Path, + gro2: Path, + edr2: Path, + workdir: Path, + *, + target: float = 1.0, + pressure_tol: float = 100.0, + max_dNw_factor: float = 1.0, + ref_gro: Path | None = None, + gmx_bin: str = "gmx", +) -> InterpolationResult: + """Read the two states from gmx and interpolate the water delta to ``target``.""" + p1 = read_mean_pressure(edr1, workdir, gmx_bin=gmx_bin) + p2 = read_mean_pressure(edr2, workdir, gmx_bin=gmx_bin) + nw1 = count_waters(gro1) + nw2 = count_waters(gro2) + return interpolate( + p1, + p2, + nw1, + nw2, + Path(ref_gro) if ref_gro is not None else Path(gro1), + target=target, + pressure_tol=pressure_tol, + max_dNw_factor=max_dNw_factor, + ) + + +def _result( + dNw: int, + nw_target: int, + ref: Path, + p1: float, + p2: float, + nw1: int, + nw2: int, + *, + converged: bool, + warning: str | None, +) -> InterpolationResult: + return InterpolationResult( + dNw=dNw, + Nw_target=nw_target, + ref_gro=ref, + P1=p1, + P2=p2, + Nw1=nw1, + Nw2=nw2, + converged=converged, + warning=warning, + ) diff --git a/md_workflows/core/steps/resolvate.py b/md_workflows/core/steps/resolvate.py new file mode 100644 index 0000000..7d7049f --- /dev/null +++ b/md_workflows/core/steps/resolvate.py @@ -0,0 +1,299 @@ +"""Adaptive resolvation to bring the equilibrated model to the target pressure. + +Corresponds to the taylor ``resolvate_1.sh`` -> ``resolvate_2.sh`` -> ``resolvate_final.sh`` +sequence. The reusable :func:`resolvation_cycle` (solvate -> manage WAT topology -> +minimize -> equilibrate) is invoked up to three times; the branching between them is the +dynamic control flow that would become a Prefect ``@flow`` body unchanged: + +1. run 1: add ``trial_fraction`` of a trial solvation, minimize + equilibrate. +2. read run-1 pressure. Case A (under target): add more water (run-2 probe) and + interpolate between run 1 and run 2. Case B (overshot): interpolate between the + pre-resolvation equilibration and run 1. +3. final: apply the interpolated ``dNw`` — positive => add to the chosen reference, + negative => restart from ``md_equil`` with fewer waters, zero/converged => keep run 1. + +WAT topology is managed as ordered blocks (``manage_wat_topology``): each stage appends +its block; the final stage drops the discarded run-2 probe block first. +""" + +from __future__ import annotations + +import shutil +from pathlib import Path + +from pydantic import BaseModel + +from ..config import GromacsRunProfile, ResolvateParams, SystemConfig +from ..gmx import ( + add_water, + checksums, + gmx_version, + manage_wat_topology, + read_mean_pressure, + run_equil, + run_min, +) +from ..paths import under +from ..results import StepInputs, StepResult, StepStatus +from .pressure_interpolate import pressure_interpolate + +STEP = "resolvate" +FINAL_PREFIX = "md_resolv_final" + + +class ResolvateInputs(StepInputs): + md_equil_gro: Path + md_equil_edr: Path + water_equil: Path + posre_top: Path + model_pdb: Path + min_mdp: Path + equil_mdp: Path + + def consumed_paths(self) -> list[Path]: + return [ + self.md_equil_gro, + self.md_equil_edr, + self.water_equil, + self.posre_top, + self.model_pdb, + self.min_mdp, + self.equil_mdp, + ] + + +class ResolvationResult(BaseModel): + """Result of one resolvation cycle (internal, not a StepResult).""" + + prefix: str + pdb: Path + min_gro: Path + equil_gro: Path + equil_edr: Path + nwat_added: int + + +class ResolvateResult(StepResult): + @property + def final_equil_gro(self) -> Path: + return self.output("final_equil_gro") + + @property + def final_equil_edr(self) -> Path: + return self.output("final_equil_edr") + + +def resolve_inputs(workdir: Path, cfg: SystemConfig) -> ResolvateInputs: + workdir = Path(workdir) + mdp_dir = under(workdir, cfg.mdp_dir) + return ResolvateInputs( + workdir=workdir, + md_equil_gro=workdir / "md_equil.gro", + md_equil_edr=workdir / "md_equil.edr", + water_equil=workdir / "waterbox" / "water_equil.gro", + posre_top=workdir / "md_model_posre.top", + model_pdb=workdir / "md_model.pdb", + min_mdp=mdp_dir / cfg.minimize.min_mdp, + equil_mdp=mdp_dir / cfg.equilibrate.equil_mdp, + ) + + +def check_inputs(inputs: ResolvateInputs) -> None: + inputs.check_exists(STEP) + + +def resolvation_cycle( + *, + input_gro: Path, + maxsol: int, + prefix: str, + ref_pdb: Path, + reservoir: Path, + top: Path, + min_mdp: Path, + equil_mdp: Path, + min_profile: GromacsRunProfile, + equil_profile: GromacsRunProfile, + workdir: Path, + wat_mode: str = "append", +) -> ResolvationResult: + """One resolvation cycle: solvate -> WAT topology -> minimize -> equilibrate.""" + workdir = Path(workdir) + aw = add_water( + input_gro, + reservoir, + maxsol, + workdir / f"{prefix}.pdb", + workdir, + gmx_bin=min_profile.gmx_bin, + log_path=workdir / f"gmx_{prefix}.log", + ) + manage_wat_topology(top, aw.nwat_added, mode=wat_mode) + mn = run_min(aw.out_pdb, top, min_mdp, f"{prefix}_min", min_profile, workdir) + eq = run_equil( + mn.gro, top, equil_mdp, ref_pdb, f"{prefix}_equil", equil_profile, workdir, maxwarn=2 + ) + return ResolvationResult( + prefix=prefix, + pdb=aw.out_pdb, + min_gro=mn.gro, + equil_gro=eq.gro, + equil_edr=eq.edr, + nwat_added=aw.nwat_added, + ) + + +def resolvate( + inputs: ResolvateInputs, + params: ResolvateParams, + min_profile: GromacsRunProfile, + equil_profile: GromacsRunProfile, + *, + resume: bool = False, +) -> ResolvateResult: + check_inputs(inputs) + wd = inputs.workdir + gmx_bin = min_profile.gmx_bin + + final_equil_gro = wd / f"{FINAL_PREFIX}_equil.gro" + final_equil_edr = wd / f"{FINAL_PREFIX}_equil.edr" + outputs = { + "posre_top": inputs.posre_top, + "final_equil_gro": final_equil_gro, + "final_equil_edr": final_equil_edr, + } + consumed = { + "md_equil_gro": inputs.md_equil_gro, + "md_equil_edr": inputs.md_equil_edr, + "water_equil": inputs.water_equil, + "posre_top": inputs.posre_top, + "model_pdb": inputs.model_pdb, + "min_mdp": inputs.min_mdp, + "equil_mdp": inputs.equil_mdp, + } + + if resume and final_equil_gro.exists(): + return ResolvateResult( + step=STEP, + status=StepStatus.SKIPPED, + workdir=wd, + outputs=outputs, + params=params.model_dump(), + input_checksums=checksums(consumed), + ) + + def cycle(input_gro: Path, maxsol: int, prefix: str, ref_pdb: Path, wat_mode: str): + return resolvation_cycle( + input_gro=input_gro, + maxsol=maxsol, + prefix=prefix, + ref_pdb=ref_pdb, + reservoir=inputs.water_equil, + top=inputs.posre_top, + min_mdp=inputs.min_mdp, + equil_mdp=inputs.equil_mdp, + min_profile=min_profile, + equil_profile=equil_profile, + workdir=wd, + wat_mode=wat_mode, + ) + + # Trial solvation just sizes run 1 (no topology change). + trial = add_water( + inputs.md_equil_gro, + inputs.water_equil, + None, + wd / "tmp_resolv_trial.pdb", + wd, + gmx_bin=gmx_bin, + log_path=wd / "gmx_resolv_trial.log", + ) + maxsol1 = int(trial.nwat_added * params.trial_fraction) + + r1 = cycle(inputs.md_equil_gro, maxsol1, "md_resolv1", inputs.model_pdb, "append") + p1 = read_mean_pressure(r1.equil_edr, wd, gmx_bin=gmx_bin) + + if p1 < params.target_bar: # Case A: still under-solvated -> probe + bracket r1<->r2 + maxsol2 = round(r1.nwat_added * params.scale_add) + r2 = cycle(r1.equil_gro, maxsol2, "md_resolv2", inputs.model_pdb, "append") + itp = pressure_interpolate( + r1.pdb, + r1.equil_edr, + r2.pdb, + r2.equil_edr, + wd, + target=params.target_bar, + pressure_tol=params.pressure_tol, + max_dNw_factor=params.max_dNw_factor, + ref_gro=r1.equil_gro, + gmx_bin=gmx_bin, + ) + case = "A" + else: # Case B: overshot on run 1 -> bracket md_equil<->r1 + itp = pressure_interpolate( + inputs.md_equil_gro, + inputs.md_equil_edr, + r1.pdb, + r1.equil_edr, + wd, + target=params.target_bar, + pressure_tol=params.pressure_tol, + max_dNw_factor=params.max_dNw_factor, + ref_gro=inputs.md_equil_gro, + gmx_bin=gmx_bin, + ) + case = "B" + + _write_interpolation_report(wd / "interpolation.txt", case, p1, itp) + + if itp.converged or itp.dNw == 0: + # Already at target: run 1 is the final state; publish it under the final names. + shutil.copy(r1.equil_gro, final_equil_gro) + shutil.copy(r1.equil_edr, final_equil_edr) + final_added = 0 + elif itp.dNw > 0: + cycle(itp.ref_gro, itp.dNw, FINAL_PREFIX, inputs.model_pdb, "drop_last_then_append") + final_added = itp.dNw + else: # dNw < 0: cannot remove water by solvating -> restart from md_equil with fewer + cycle( + inputs.md_equil_gro, + max(0, itp.Nw_target), + FINAL_PREFIX, + inputs.model_pdb, + "drop_last_then_append", + ) + final_added = itp.Nw_target + + return ResolvateResult( + step=STEP, + status=StepStatus.COMPLETED, + workdir=wd, + outputs=outputs, + params=params.model_dump(), + input_checksums=checksums(consumed), + run_profiles_used={"resolv_min": min_profile, "resolv_equil": equil_profile}, + tool_versions={"gmx": gmx_version(gmx_bin)}, + metrics={ + "trial_nwat": trial.nwat_added, + "maxsol1": maxsol1, + "p1_bar": p1, + "dNw": itp.dNw, + "Nw_target": itp.Nw_target, + "final_added": final_added, + "converged": int(itp.converged), + }, + log_paths=[wd / "interpolation.txt"], + ) + + +def _write_interpolation_report(path: Path, case: str, p1: float, itp) -> None: + lines = [ + f"case: {case}", + f"run1 mean pressure: {p1:.1f} bar", + f"P1={itp.P1:.1f} bar Nw1={itp.Nw1}", + f"P2={itp.P2:.1f} bar Nw2={itp.Nw2}", + f"dNw={itp.dNw} Nw_target={itp.Nw_target} converged={itp.converged}", + ] + if itp.warning: + lines.append(f"WARNING: {itp.warning}") + path.write_text("\n".join(lines) + "\n") diff --git a/md_workflows/core/steps/run_params_gaussian.py b/md_workflows/core/steps/run_params_gaussian.py new file mode 100644 index 0000000..0abc1fa --- /dev/null +++ b/md_workflows/core/steps/run_params_gaussian.py @@ -0,0 +1,317 @@ +"""Ligand parameterization with Gaussian + AmberTools (standalone utility). + +Corresponds to ``run_params_gaussian.sh`` (run under ``ligand/``). This is kept off the +main prep chain — wiring the resulting ligand parameters into the protein topology is +part of the deferred full ligand-handling feature. + +Ported to be cwd-independent: instead of ``os.chdir(ligand)`` + mutating ``os.environ``, +every external tool runs with ``cwd=lig_dir`` and an explicit ``env`` carrying +``g16root``/``OMP_NUM_THREADS``. + +Requires ``.pdb`` (the ligand coordinates) to already exist in ``ligand/``; the +residue name is auto-detected from the RCSB legacy PDB. +""" + +from __future__ import annotations + +import os +from pathlib import Path + +from ...pdb_file_processing import prepare_pdb_and_resn_files +from ..config import GaussianParams, SystemConfig +from ..exceptions import MDWorkflowError +from ..gmx import checksums, run_tool +from ..results import StepInputs, StepResult, StepStatus + +STEP = "run_params_gaussian" + + +class RunParamsGaussianInputs(StepInputs): + pdb_id: str + + def consumed_paths(self) -> list[Path]: + return [] # legacy PDB is downloaded; .pdb is checked after detection + + +class RunParamsGaussianResult(StepResult): + @property + def frcmod(self) -> Path: + return self.output("frcmod") + + @property + def mol2(self) -> Path: + return self.output("mol2") + + +def resolve_inputs(workdir: Path, cfg: SystemConfig) -> RunParamsGaussianInputs: + return RunParamsGaussianInputs(workdir=Path(workdir), pdb_id=cfg.pdb_id) + + +def check_inputs(inputs: RunParamsGaussianInputs) -> None: + inputs.check_exists(STEP) + + +def run_params_gaussian( + inputs: RunParamsGaussianInputs, + params: GaussianParams, + *, + resume: bool = False, +) -> RunParamsGaussianResult: + if not params.g16root: + raise MDWorkflowError("run_params_gaussian: gaussian.g16root must be set") + check_inputs(inputs) + + lig_dir = inputs.workdir / "ligand" + lig_dir.mkdir(parents=True, exist_ok=True) + resn, _ = prepare_pdb_and_resn_files(lig_dir=lig_dir, pdb_id=inputs.pdb_id) + + ligand_pdb = lig_dir / f"{resn}.pdb" + if not ligand_pdb.exists(): + raise MDWorkflowError( + f"run_params_gaussian: expected ligand coordinates {ligand_pdb} (residue " + f"{resn}); provide it before running." + ) + + outputs = { + "mol2": lig_dir / f"{resn}_resp.mol2", + "frcmod": lig_dir / f"{resn}_resp.frcmod", + "parm7": lig_dir / f"{resn}_resp.parm7", + "rst7": lig_dir / f"{resn}_resp.rst7", + "pdb": lig_dir / f"{resn}_resp.pdb", + } + if resume and outputs["frcmod"].exists(): + return RunParamsGaussianResult( + step=STEP, + status=StepStatus.SKIPPED, + workdir=inputs.workdir, + outputs=outputs, + params=params.model_dump(), + ) + + env = _g16_env(params.g16root, params.nproc) + + run_tool( + [ + "antechamber", + "-fi", + "pdb", + "-fo", + "gcrt", + "-i", + f"{resn}.pdb", + "-o", + f"{resn}.gau", + "-nc", + str(params.net_charge), + "-m", + "1", + ], + tool="antechamber", + cwd=lig_dir, + env=env, + log_path=lig_dir / "antechamber_gcrt.log", + ) + _patch_gaussian_input(lig_dir, resn, params.nproc, params.method) + _run_gaussian_opt(lig_dir, resn, env) + _run_gaussian_esp(lig_dir, resn, params.method, env) + _process_resp_charges(lig_dir, resn, env) + _build_amber_lib(lig_dir, resn, env) + + return RunParamsGaussianResult( + step=STEP, + status=StepStatus.COMPLETED, + workdir=inputs.workdir, + outputs=outputs, + params=params.model_dump(), + input_checksums=checksums({"ligand_pdb": ligand_pdb}), + metrics={"net_charge": params.net_charge}, + ) + + +def _g16_env(g16root: str, nproc: int) -> dict[str, str]: + env = dict(os.environ) + env["g16root"] = g16root + env["OMP_NUM_THREADS"] = str(nproc) + return env + + +def _patch_gaussian_input(lig_dir: Path, resn: str, nproc: int, method: str) -> None: + """Insert NprocShared and switch the HF header to the requested method.""" + lines = (lig_dir / f"{resn}.gau").read_text().splitlines(keepends=True) + patched: list[str] = [] + for line in lines: + patched.append(line) + if "Link" in line: + patched.append(f"%NprocShared={nproc}\n") + text = "".join(patched).replace("#HF", f"#{method}") + (lig_dir / "tmp").write_text(text) + + +def _run_gaussian_opt(lig_dir: Path, resn: str, env: dict[str, str]) -> None: + """Run the Gaussian geometry optimization if the input changed or is new.""" + gau = lig_dir / f"{resn}.gau" + log = lig_dir / f"{resn}.log" + tmp = lig_dir / "tmp" + if log.exists() and gau.read_text() == tmp.read_text(): + tmp.unlink(missing_ok=True) + return + os.replace(tmp, gau) + run_tool(["g16", f"{resn}.gau"], tool="g16 (opt)", cwd=lig_dir, env=env) + + +def _run_gaussian_esp(lig_dir: Path, resn: str, method: str, env: dict[str, str]) -> None: + """Run the Gaussian ESP / CHELPG calculation if the input changed or is new.""" + text = (lig_dir / f"{resn}.gau").read_text() + text = text.replace("opt", "pop(chelpg,regular)").replace("molecule", "grid") + resp_gau = lig_dir / f"{resn}_resp.gau" + resp_log = lig_dir / f"{resn}_resp.log" + if resp_log.exists() and resp_gau.exists() and resp_gau.read_text() == text: + return + resp_gau.write_text(text) + run_tool(["g16", f"{resn}_resp.gau"], tool="g16 (esp)", cwd=lig_dir, env=env) + + +def _process_resp_charges(lig_dir: Path, resn: str, env: dict[str, str]) -> None: + """Derive RESP charges, graft original coordinates, and build GAFF parameters.""" + run_tool( + [ + "antechamber", + "-fi", + "gout", + "-i", + f"{resn}_resp.log", + "-cf", + f"{resn}_resp.crg", + "-c", + "resp", + "-o", + f"{resn}_gauss.ac", + "-fo", + "ac", + "-rn", + resn, + ], + tool="antechamber (resp)", + cwd=lig_dir, + env=env, + ) + run_tool( + [ + "antechamber", + "-fi", + "gout", + "-i", + f"{resn}_resp.log", + "-o", + f"{resn}_gauss.pdb", + "-fo", + "pdb", + "-rn", + resn, + ], + tool="antechamber (pdb)", + cwd=lig_dir, + env=env, + ) + + coords = _extract_coords_from_pdb(lig_dir / f"{resn}.pdb") + _graft_coords_to_ac(lig_dir / f"{resn}_gauss.ac", coords, lig_dir / f"{resn}_resp.ac") + _correct_charge(lig_dir / f"{resn}_resp.ac") + + run_tool( + [ + "antechamber", + "-fi", + "ac", + "-i", + f"{resn}_resp.ac", + "-fo", + "mol2", + "-o", + f"{resn}_resp.mol2", + "-rn", + resn, + ], + tool="antechamber (mol2)", + cwd=lig_dir, + env=env, + ) + run_tool( + ["atomtype", "-i", f"{resn}_resp.ac", "-o", f"{resn}_resp_gaff.ac", "-p", "gaff"], + tool="atomtype", + cwd=lig_dir, + env=env, + ) + run_tool( + ["prepgen", "-i", f"{resn}_resp_gaff.ac", "-o", f"{resn}_resp_gaff.prepc", "-f", "car"], + tool="prepgen", + cwd=lig_dir, + env=env, + ) + run_tool( + ["parmchk2", "-i", f"{resn}_resp_gaff.prepc", "-o", f"{resn}_resp.frcmod", "-f", "prepc"], + tool="parmchk2", + cwd=lig_dir, + env=env, + ) + + +def _build_amber_lib(lig_dir: Path, resn: str, env: dict[str, str]) -> None: + import textwrap + + tleap_input = textwrap.dedent(f"""\ + source leaprc.protein.ff19SB + source leaprc.gaff + loadamberparams {resn}_resp.frcmod + loadamberprep {resn}_resp_gaff.prepc + lig = loadmol2 {resn}_resp.mol2 + savepdb lig {resn}_resp.pdb + saveamberparm lig {resn}_resp.parm7 {resn}_resp.rst7 + quit + """) + (lig_dir / "tleap_lig.in").write_text(tleap_input) + run_tool(["tleap", "-f", "tleap_lig.in"], tool="tleap (ligand)", cwd=lig_dir, env=env) + + +def _extract_coords_from_pdb(pdb_file: Path) -> list[str]: + coords: list[str] = [] + with open(pdb_file) as fh: + for line in fh: + if line.startswith(("ATOM", "HETATM")): + coords.append(line[30:53]) + return coords + + +def _graft_coords_to_ac(ac_file: Path, coords: list[str], out_file: Path) -> None: + anum = 0 + with open(ac_file) as fh, open(out_file, "w") as out: + for line in fh: + if line.startswith(("ATOM", "HETATM")): + out.write(line[:30] + coords[anum] + line[53:]) + anum += 1 + else: + out.write(line) + + +def _correct_charge(ac_file: Path) -> None: + """Adjust the largest-magnitude charge so the total is an exact integer.""" + lines = ac_file.read_text().splitlines(keepends=True) + charges = [float(line[54:63]) for line in lines if line.startswith("ATOM")] + if not charges: + return + correction = round(sum(charges)) - sum(charges) + max_idx = max(range(len(charges)), key=lambda i: abs(charges[i])) + new_charge = charges[max_idx] + correction + + out: list[str] = [] + anum = 0 + for line in lines: + if line.startswith("ATOM"): + if anum == max_idx: + out.append(f"{line[:54]}{new_charge:9.6f}{line[63:]}") + else: + out.append(line) + anum += 1 + else: + out.append(line) + ac_file.write_text("".join(out)) diff --git a/md_workflows/core/steps/solvate.py b/md_workflows/core/steps/solvate.py new file mode 100644 index 0000000..6b7ba1c --- /dev/null +++ b/md_workflows/core/steps/solvate.py @@ -0,0 +1,291 @@ +"""Solvate the crystal with water and add neutralizing/ionic-strength ions. + +Corresponds to ``solvate.sh``: count symmetry copies, fill the crystal voids from the +equilibrated water reservoir, compute Na+/Cl- counts for the target ionic strength plus +charge neutralization, insert the ions, and finalize the model topology. + +The target ionic strength is a tunable parameter (``solvate.ionic_strength``, default +0.1 M) rather than a hard-coded constant, so runs can request a different salt +concentration without editing code. +""" + +from __future__ import annotations + +import re +import shutil +from pathlib import Path + +from ..config import SolvateParams, SystemConfig +from ..exceptions import GmxOutputParseError +from ..gmx import checksums, count_waters, run_tool +from ..results import StepInputs, StepResult, StepStatus + +STEP = "solvate" + + +class SolvateInputs(StepInputs): + prot_dry: Path + xtal: Path + water_equil: Path + prot_pdb: Path + prot_top: Path + cl_pdb: Path + na_pdb: Path + + def consumed_paths(self) -> list[Path]: + return [ + self.prot_dry, + self.xtal, + self.water_equil, + self.prot_pdb, + self.prot_top, + self.cl_pdb, + self.na_pdb, + ] + + +class SolvateResult(StepResult): + @property + def md_model_pdb(self) -> Path: + return self.output("md_model_pdb") + + @property + def md_model_top(self) -> Path: + return self.output("md_model_top") + + +def resolve_inputs(workdir: Path, cfg: SystemConfig) -> SolvateInputs: + workdir = Path(workdir) + return SolvateInputs( + workdir=workdir, + prot_dry=workdir / "prot_dry.pdb", + xtal=workdir / "xtal.pdb", + water_equil=workdir / "waterbox" / "water_equil.gro", + prot_pdb=workdir / "prot.pdb", + prot_top=workdir / "prot.top", + cl_pdb=workdir / "Cl-.pdb", + na_pdb=workdir / "Na+.pdb", + ) + + +def check_inputs(inputs: SolvateInputs) -> None: + inputs.check_exists(STEP) + + +def solvate( + inputs: SolvateInputs, + params: SolvateParams, + *, + resume: bool = False, +) -> SolvateResult: + check_inputs(inputs) + wd = inputs.workdir + + xtal_solv = wd / "xtal_solv.pdb" + xtal_solv_cl_na = wd / "xtal_solv_cl_na.pdb" + md_model_pdb = wd / "md_model.pdb" + md_model_top = wd / "md_model.top" + outputs = { + "xtal_solv": xtal_solv, + "xtal_solv_cl_na": xtal_solv_cl_na, + "md_model_pdb": md_model_pdb, + "md_model_top": md_model_top, + } + consumed = { + "prot_dry": inputs.prot_dry, + "xtal": inputs.xtal, + "water_equil": inputs.water_equil, + "prot_pdb": inputs.prot_pdb, + "prot_top": inputs.prot_top, + "cl_pdb": inputs.cl_pdb, + "na_pdb": inputs.na_pdb, + } + + if resume and md_model_pdb.exists() and md_model_top.exists(): + return SolvateResult( + step=STEP, + status=StepStatus.SKIPPED, + workdir=wd, + outputs=outputs, + params=params.model_dump(), + input_checksums=checksums(consumed), + ) + + ncopies = _count_copies(inputs.prot_dry, inputs.xtal) + + # Fill crystal voids with the equilibrated water reservoir. + solvate_log = wd / "gmx_solvate.log" + run_tool( + [ + "gmx", + "solvate", + "-cp", + str(inputs.xtal), + "-cs", + str(inputs.water_equil), + "-o", + str(xtal_solv), + ], + tool="gmx solvate", + cwd=wd, + log_path=solvate_log, + ) + nwat_initial = _read_solvate_nwat(solvate_log) + + _write_topology_header(inputs.prot_top, md_model_top) + + ions_pos, ions_neg = _count_ions(inputs.prot_pdb) + net_ion_charge = (ions_pos - ions_neg) * ncopies + nna, ncl = _compute_ion_counts( + nwat_initial, net_ion_charge, params.ionic_strength, params.water_molarity + ) + + _insert_ions(wd, xtal_solv, inputs.cl_pdb, inputs.na_pdb, xtal_solv_cl_na, ncl, nna) + + nwat = count_waters(xtal_solv_cl_na) + _finalize_topology(md_model_top, ncopies, nwat, ncl, nna) + shutil.copy(xtal_solv_cl_na, md_model_pdb) + + return SolvateResult( + step=STEP, + status=StepStatus.COMPLETED, + workdir=wd, + outputs=outputs, + params=params.model_dump(), + input_checksums=checksums(consumed), + metrics={ + "ncopies": ncopies, + "nwat_initial": nwat_initial, + "ions_pos": ions_pos, + "ions_neg": ions_neg, + "net_ion_charge": net_ion_charge, + "nna": nna, + "ncl": ncl, + "nwat": nwat, + }, + log_paths=[solvate_log], + ) + + +def _count_copies(prot_dry: Path, xtal: Path) -> int: + def _natoms(path: Path) -> int: + n = 0 + with open(path) as fh: + for line in fh: + if line.startswith(("ATOM", "HETATM")): + n += 1 + return n + + nats_one = _natoms(prot_dry) + if nats_one == 0: + raise GmxOutputParseError("asymmetric-unit atom count", prot_dry) + return _natoms(xtal) // nats_one + + +def _read_solvate_nwat(log_path: Path) -> int: + for line in Path(log_path).read_text().splitlines(): + m = re.search(r"Output configuration contains\s+(\d+)", line) + if m: + return int(m.group(1)) + raise GmxOutputParseError("output water count", log_path) + + +def _write_topology_header(prot_top: Path, md_model_top: Path) -> None: + header: list[str] = [] + with open(prot_top) as fh: + for line in fh: + if "molecules" in line.lower(): + break + header.append(line) + with open(md_model_top, "w") as fh: + fh.writelines(header) + + +def _count_ions(prot_pdb: Path) -> tuple[int, int]: + pos = neg = 0 + with open(prot_pdb) as fh: + for line in fh: + if line.startswith("HETATM"): + if "Na+" in line: + pos += 1 + elif "Cl-" in line: + neg += 1 + return pos, neg + + +def _compute_ion_counts( + nwat: int, net_ion_charge: int, ionic_strength: float, water_molarity: float +) -> tuple[int, int]: + """Na+/Cl- counts for the requested ionic strength plus charge neutralization. + + ``base`` is the salt-pair count implied by ``ionic_strength`` (mol/L) relative to + water's molarity; the net charge is then neutralized by adding to the counter-ion. + """ + base = nwat * ionic_strength // water_molarity + if net_ion_charge >= 0: + ncl = base + nna = ncl + net_ion_charge + else: + nna = base + ncl = nna - net_ion_charge + return int(nna), int(ncl) + + +def _insert_ions( + workdir: Path, + solv_pdb: Path, + cl_pdb: Path, + na_pdb: Path, + out_pdb: Path, + ncl: int, + nna: int, +) -> None: + cl_out = workdir / "xtal_solv_cl.pdb" + run_tool( + [ + "gmx", + "insert-molecules", + "-f", + str(solv_pdb), + "-ci", + str(cl_pdb), + "-o", + str(cl_out), + "-replace", + "SOL", + "-nmol", + str(ncl), + ], + tool="gmx insert-molecules", + cwd=workdir, + log_path=workdir / "insert_cl.log", + ) + run_tool( + [ + "gmx", + "insert-molecules", + "-f", + str(cl_out), + "-ci", + str(na_pdb), + "-o", + str(out_pdb), + "-replace", + "SOL", + "-nmol", + str(nna), + ], + tool="gmx insert-molecules", + cwd=workdir, + log_path=workdir / "insert_na.log", + ) + + +def _finalize_topology(md_model_top: Path, ncopies: int, nwat: int, ncl: int, nna: int) -> None: + with open(md_model_top, "a") as fh: + fh.write("[ molecules ]\n") + fh.write("; Compound #mols\n") + fh.write("system1 1\n" * ncopies) + fh.write(f"WAT {nwat}\n") + fh.write(f"Cl- {ncl}\n") + fh.write(f"Na+ {nna}\n") diff --git a/md_workflows/equilibrate.py b/md_workflows/equilibrate.py deleted file mode 100644 index 791ab98..0000000 --- a/md_workflows/equilibrate.py +++ /dev/null @@ -1,166 +0,0 @@ -"""Set up positional restraints and run NPT equilibration. - -Corresponds to run_all.sh line: - bash scripts/equilibrate.sh -""" - -import glob -import os -import subprocess -from pathlib import Path - - -def run(ntomp: int = 26): - artifacts_dir = Path("artifacts") - _extract_first_copy() - chain_files = _split_chains() - _generate_restraints(chain_files) - _build_restrained_topology(chain_files) - - subprocess.run( - [ - "gmx", - "grompp", - "-f", - str(artifacts_dir / "equil.mdp"), - "-c", - "md_min.gro", - "-o", - "md_equil.tpr", - "-p", - "md_model_posre.top", - "-r", - "md_model.pdb", - ], - capture_output=True, - text=True, - check=True, - ) - - subprocess.run( - [ - "gmx", - "mdrun", - "-ntmpi", - "1", - "-ntomp", - str(ntomp), - "-deffnm", - "md_equil", - "-v", - ], - check=True, - ) - - -def _extract_first_copy(): - """Extract the first copy of the asymmetric unit (up to first GOL residue).""" - with open("pdb_clean.pdb") as fh: - lines = fh.readlines() - - lines = [line for line in lines if not line.startswith("JRNL")] - - kept = [] - found_gol = False - gol_next_idx = None - for i, line in enumerate(lines): - if "GOL" in line: - found_gol = True - kept.append(line) - elif found_gol and gol_next_idx is None: - gol_next_idx = i - break - else: - kept.append(line) - - with open("first_copy.pdb", "w") as fh: - fh.writelines(kept) - - prot_lines = [ - line for line in kept if (line.startswith(("ATOM", "HETATM", "TER")) and "GOL" not in line) - ] - with open("first_copy_prot.pdb", "w") as fh: - fh.writelines(prot_lines) - - -def _split_chains() -> list[str]: - """Split first_copy_prot.pdb at TER records into part00, part01, etc.""" - for f in glob.glob("part??"): - os.remove(f) - - with open("first_copy_prot.pdb") as fh: - lines = fh.readlines() - - # Split after each TER line (PDB chain end). Python re forbids variable-width lookbehinds - # like (?<=TER.*\n). - chunks: list[str] = [] - buf: list[str] = [] - for line in lines: - buf.append(line) - if line.startswith("TER"): - chunks.append("".join(buf)) - buf = [] - if buf: - chunks.append("".join(buf)) - chunks = [c for c in chunks if c.strip()] - - files = [] - for i, chunk in enumerate(chunks): - fname = f"part{i:02d}" - with open(fname, "w") as fh: - fh.write(chunk) - files.append(fname) - return files - - -def _generate_restraints(chain_files: list[str]): - """Run pdb4amber + gmx genrestr for each chain fragment.""" - for f in chain_files: - subprocess.run(["pdb4amber", "-i", f, "-o", f"{f}_amber.pdb"], check=True) - subprocess.run( - [ - "gmx", - "genrestr", - "-fc", - "209.2", - "209.2", - "209.2", - "-f", - f"{f}_amber.pdb", - "-o", - f"posre_{f}.itp", - ], - input="Protein-H\nq\n", - text=True, - check=True, - ) - - -def _build_restrained_topology(chain_files: list[str]): - """Insert #ifdef POSRES_partXX blocks into md_model_posre.top.""" - import shutil - - shutil.copy("md_model.top", "md_model_posre.top") - - for molnum_offset, f in enumerate(chain_files): - target_moltype_count = molnum_offset + 2 # 1-indexed, skip first - - with open("md_model_posre.top") as fh: - lines = fh.readlines() - - new_lines = [] - cnt = 0 - for line in lines: - if "moleculetype" in line.lower(): - cnt += 1 - if cnt == target_moltype_count: - posre_block = f'#ifdef POSRES_{f}\n#include "posre_{f}.itp"\n#endif\n\n' - new_lines.append(posre_block) - new_lines.append(line) - - with open("md_model_posre.top", "w") as fh: - fh.writelines(new_lines) - - -if __name__ == "__main__": - run() diff --git a/md_workflows/make_crystal.py b/md_workflows/make_crystal.py deleted file mode 100644 index 00a9fa1..0000000 --- a/md_workflows/make_crystal.py +++ /dev/null @@ -1,124 +0,0 @@ -"""Build a crystal supercell from the protein structure. - -Corresponds to run_all.sh line: - bash scripts/make_crystal.sh 1 -""" - -import subprocess -from pathlib import Path - - -def run( - ix: int = 1, - iy: int | None = None, - iz: int | None = None, - chimerax_exec: str = "/usr/bin/chimerax-daily", -): - """ - Args: - ix: Number of unit-cell replications in x (also used for y, z if iy/iz - are not given). - iy: Optional separate y replication count. - iz: Optional separate z replication count. - chimerax_exec: Path to the ChimeraX executable. - """ - if iy is None: - iy = ix - if iz is None: - iz = ix - - workdir = Path.cwd() - - subprocess.run(["pdb4amber", "-i", "prot.pdb", "-o", "prot_dry.pdb", "--dry"], check=True) - - _prepend_cryst1("pdb_clean.pdb", "prot_dry.pdb") - _expand_unit_cell(workdir, chimerax_exec) - _set_p1_spacegroup("prot_dry.pdb", "prot_dry_cell.pdb") - _propagate_crystal(ix, iy, iz) - - -def _prepend_cryst1(source_pdb: str, target_pdb: str): - """Copy the CRYST1 line from source_pdb to the top of target_pdb, - removing any Na+/Cl-/existing CRYST1 lines from target_pdb.""" - cryst1 = "" - with open(source_pdb) as fh: - for line in fh: - if line.startswith("CRYST1"): - cryst1 = line - break - - with open(target_pdb) as fh: - lines = fh.readlines() - - filtered = [ - line - for line in lines - if "Na+" not in line and "Cl-" not in line and not line.startswith("CRYST1") - ] - - with open(target_pdb, "w") as fh: - fh.write(cryst1) - fh.writelines(filtered) - - -def _expand_unit_cell(workdir: Path, chimerax_exec: str): - """Use ChimeraX to expand the unit cell.""" - cxc_script = f"""\ -open {workdir}/prot_dry.pdb -changechains #1 A -unitcell #1 -combine #2 -save {workdir}/prot_dry_cell.pdb #3 -quit -""" - with open("expand.cxc", "w") as fh: - fh.write(cxc_script) - subprocess.run([chimerax_exec, "--offscreen", "--nogui", "expand.cxc"], check=True) - - -def _set_p1_spacegroup(dry_pdb: str, cell_pdb: str): - """Rewrite the CRYST1 line with P 1 spacegroup and prepend to cell PDB.""" - with open(dry_pdb) as fh: - for line in fh: - if line.startswith("CRYST1"): - cryst1_p1 = line[:55] + "P 1\n" - break - - with open("cryst1_p1.pdb", "w") as fh: - fh.write(cryst1_p1) - - with open(cell_pdb) as fh: - cell_content = fh.read() - - with open(cell_pdb, "w") as fh: - fh.write(cryst1_p1) - fh.write(cell_content) - - -def _propagate_crystal(ix: int, iy: int, iz: int): - """Use PropPDB to replicate the unit cell, or just copy if 0.""" - if ix > 0 or iy > 0 or iz > 0: - subprocess.run( - [ - "PropPDB", - "-p", - "prot_dry_cell.pdb", - "-o", - "xtal.pdb", - "-ix", - str(ix), - "-iy", - str(iy), - "-iz", - str(iz), - ], - check=True, - ) - else: - import shutil - - shutil.copy("prot_dry_cell.pdb", "xtal.pdb") - - -if __name__ == "__main__": - run() diff --git a/md_workflows/make_waterbox.py b/md_workflows/make_waterbox.py deleted file mode 100644 index 222dc22..0000000 --- a/md_workflows/make_waterbox.py +++ /dev/null @@ -1,202 +0,0 @@ -"""Build an equilibrated bulk-water box matching the crystal unit cell. - -Corresponds to run_all.sh line: - bash scripts/make_waterbox.sh -""" - -import subprocess -from pathlib import Path - - -def run(ntomp: int = 26): - workdir = Path.cwd() - artifacts_dir = workdir / "artifacts" - wb_dir = workdir / "waterbox" - wb_dir.mkdir(parents=True, exist_ok=True) - - _extract_cryst1(workdir, wb_dir) - _create_box_pdb(workdir, wb_dir) - _insert_water(workdir, wb_dir) - _expand_waterbox(workdir, wb_dir) - # Topology must match the coordinate file passed to grompp (expanded box). - nwat = _count_wat_molecules(wb_dir / "box_solv_expand.pdb") - _write_topology(workdir, wb_dir, nwat) - _minimize_waterbox(artifacts_dir, wb_dir, ntomp) - _equilibrate_waterbox(artifacts_dir, wb_dir, ntomp) - - -def _extract_cryst1(workdir: Path, wb_dir: Path): - with open(workdir / "xtal.pdb") as fh: - for line in fh: - if line.startswith("CRYST1"): - with open(wb_dir / "cryst1_xtal.pdb", "w") as out: - out.write(line) - return - - -def _create_box_pdb(workdir: Path, wb_dir: Path): - """Create a box.pdb with CRYST1 dimensions scaled from Angstroms to nm.""" - with open(workdir / "xtal.pdb") as fh: - for line in fh: - if line.startswith("CRYST1"): - a = float(line[6:15]) / 10.0 - b = float(line[15:24]) / 10.0 - c = float(line[24:33]) / 10.0 - alpha = float(line[33:40]) - beta = float(line[40:47]) - gamma = float(line[47:54]) - cryst1 = f"CRYST1{a:9.3f}{b:9.3f}{c:9.3f}{alpha:7.2f}{beta:7.2f}{gamma:7.2f}\n" - with open(wb_dir / "box.pdb", "w") as out: - out.write(cryst1) - return - - -def _insert_water(workdir: Path, wb_dir: Path): - subprocess.run( - [ - "gmx", - "insert-molecules", - "-f", - str(wb_dir / "box.pdb"), - "-ci", - str(workdir / "WAT.pdb"), - "-conc", - "58.0", - "-o", - str(wb_dir / "box_solv.pdb"), - ], - capture_output=True, - text=True, - cwd=str(wb_dir), - check=True, - ) - - -def _expand_waterbox(workdir: Path, wb_dir: Path): - subprocess.run( - [ - "PropPDB", - "-p", - str(wb_dir / "box_solv.pdb"), - "-o", - str(wb_dir / "box_solv_expand.pdb"), - "-ix", - "10", - "-iy", - "10", - "-iz", - "10", - ], - check=True, - ) - - with open(wb_dir / "cryst1_xtal.pdb") as fh: - cryst1 = fh.read() - - with open(wb_dir / "box_solv_expand.pdb") as fh: - lines = [line for line in fh if not line.startswith(("CRYST1", "HEADER"))] - - with open(wb_dir / "box_solv_expand.pdb", "w") as fh: - fh.write(cryst1) - fh.writelines(lines) - - -def _count_wat_molecules(pdb_path: Path) -> int: - """Count WAT residues assuming 3-site water (matches GROMACS WAT.pdb).""" - wat_atoms = 0 - with open(pdb_path) as fh: - for line in fh: - if line.startswith(("ATOM", "HETATM")) and " WAT " in line: - wat_atoms += 1 - if wat_atoms % 3 != 0: - raise ValueError(f"{pdb_path}: expected multiple of 3 WAT atoms, got {wat_atoms}") - return wat_atoms // 3 - - -def _write_topology(workdir: Path, wb_dir: Path, nwat: int): - with open(workdir / "prot.top") as fh: - header_lines = [] - for line in fh: - if "molecules" in line.lower(): - break - header_lines.append(line) - - with open(wb_dir / "waterbox.top", "w") as fh: - fh.writelines(header_lines) - fh.write("[ molecules ]\n") - fh.write("; Compound #mols\n") - fh.write(f"WAT {nwat}\n") - - -def _minimize_waterbox(artifacts_dir: Path, wb_dir: Path, ntomp: int): - subprocess.run( - [ - "gmx", - "grompp", - "-f", - str(artifacts_dir / "min_water.mdp"), - "-c", - str(wb_dir / "box_solv_expand.pdb"), - "-o", - str(wb_dir / "water_min.tpr"), - "-p", - str(wb_dir / "waterbox.top"), - ], - cwd=str(wb_dir), - check=True, - ) - - subprocess.run( - [ - "gmx", - "mdrun", - "-ntmpi", - "1", - "-ntomp", - str(ntomp), - "-deffnm", - "water_min", - "-v", - ], - cwd=str(wb_dir), - check=True, - ) - - -def _equilibrate_waterbox(artifacts_dir: Path, wb_dir: Path, ntomp: int): - subprocess.run( - [ - "gmx", - "grompp", - "-f", - str(artifacts_dir / "equil_water.mdp"), - "-c", - str(wb_dir / "water_min.gro"), - "-o", - str(wb_dir / "water_equil.tpr"), - "-p", - str(wb_dir / "waterbox.top"), - ], - cwd=str(wb_dir), - check=True, - ) - - subprocess.run( - [ - "gmx", - "mdrun", - "-ntmpi", - "1", - "-ntomp", - str(ntomp), - "-deffnm", - "water_equil", - "-v", - ], - cwd=str(wb_dir), - check=True, - ) - - -if __name__ == "__main__": - run() diff --git a/md_workflows/minimize.py b/md_workflows/minimize.py deleted file mode 100644 index 71ca540..0000000 --- a/md_workflows/minimize.py +++ /dev/null @@ -1,48 +0,0 @@ -"""Energy-minimize the solvated MD model. - -Corresponds to run_all.sh line: - bash scripts/minimize.sh -""" - -import subprocess -from pathlib import Path - - -def run(ntomp: int = 26): - artifacts_dir = Path("artifacts") - subprocess.run( - [ - "gmx", - "grompp", - "-f", - str(artifacts_dir / "min.mdp"), - "-c", - "md_model.pdb", - "-o", - "md_min.tpr", - "-p", - "md_model.top", - ], - capture_output=True, - text=True, - check=True, - ) - - subprocess.run( - [ - "gmx", - "mdrun", - "-ntmpi", - "1", - "-ntomp", - str(ntomp), - "-deffnm", - "md_min", - "-v", - ], - check=True, - ) - - -if __name__ == "__main__": - run() diff --git a/md_workflows/param_prot.py b/md_workflows/param_prot.py deleted file mode 100644 index 06f0fe3..0000000 --- a/md_workflows/param_prot.py +++ /dev/null @@ -1,144 +0,0 @@ -"""Parameterize the protein: clean PDB, build Amber topology, convert to GROMACS. - -Corresponds to run_all.sh line: - bash scripts/param_prot.sh 6B8X -""" - -import subprocess -import textwrap -from pathlib import Path - -from .pdb_file_processing import ensure_entry_pdb_file - - -def run(pdb_id: str = "6B8X"): - _clean_pdb(pdb_id) - _initial_solvation() - _amber_to_gromacs() - _extract_solvent_pdbs() - _final_protein_topology() - _amber_to_gromacs() - - -def _clean_pdb(pdb_id: str): - """Strip REMARK/KEYWDS, keep CRYST1/ATOM/HETATM/TER/END lines.""" - pdb_file = str(ensure_entry_pdb_file(pdb_id, Path.cwd())) - with open(pdb_file) as fh: - lines = fh.readlines() - - kept = [] - for line in lines: - if line.startswith(("REMARK", "KEYWDS")): - continue - if line.startswith(("CRYST1", "ATOM", "HETATM", "TER", "END")): - kept.append(line) - - with open("pdb_clean.pdb", "w") as fh: - fh.writelines(kept) - - subprocess.run( - [ - "pdb4amber", - "-i", - "pdb_clean.pdb", - "--prot", - "-o", - "pdb_clean_amber.pdb", - ], - check=True, - ) - - -def _initial_solvation(): - """Run tleap with minimal solvation to obtain solvent PDB templates.""" - tleap_input = textwrap.dedent("""\ - source leaprc.protein.ff19SB - source leaprc.DNA.OL15 - source leaprc.RNA.OL3 - source leaprc.water.spceb - source leaprc.gaff2 - p = loadpdb pdb_clean_amber.pdb - x = combine{p} - addions2 x Cl- 0 - addions2 x Na+ 0 - addions2 x Na+ 1 - addions2 x Cl- 1 - solvateBox x SPCBOX 1. - set default PBradii mbondi3 - set default nocenter on - saveAmberParm x prot.parm7 prot.rst7 - quit - """) - with open("tleap_temp.in", "w") as fh: - fh.write(tleap_input) - subprocess.run(["tleap", "-f", "tleap_temp.in"], check=True) - - for f in ["prot.top", "prot.pdb"]: - Path(f).unlink(missing_ok=True) - - -def _amber_to_gromacs(): - """Convert Amber parm7/rst7 to GROMACS topology and PDB.""" - script = textwrap.dedent("""\ - import parmed as pmd - parm = pmd.load_file("prot.parm7", "prot.rst7") - parm.save("prot.top") - parm.save("prot.pdb") - """) - with open("amber_to_gromacs.py", "w") as fh: - fh.write(script) - subprocess.run(["python", "amber_to_gromacs.py"], check=True) - - -def _extract_solvent_pdbs(): - """Pull single Na+, Cl-, and WAT coordinate templates from the solvated PDB.""" - with open("prot.pdb") as fh: - lines = fh.readlines() - - hetatm_lines = [line for line in lines if line.startswith("HETATM")] - - na_lines = [line for line in hetatm_lines if "Na+" in line] - cl_lines = [line for line in hetatm_lines if "Cl-" in line] - wat_lines = [line for line in hetatm_lines if "WAT" in line] - - if na_lines: - with open("Na+.pdb", "w") as fh: - fh.write(na_lines[0]) - if cl_lines: - with open("Cl-.pdb", "w") as fh: - fh.write(cl_lines[0]) - if wat_lines: - with open("WAT.pdb", "w") as fh: - fh.writelines(wat_lines[:3]) - - -def _final_protein_topology(): - """Build the final protein topology with one water + ions.""" - tleap_input = textwrap.dedent("""\ - source leaprc.protein.ff19SB - source leaprc.DNA.OL15 - source leaprc.RNA.OL3 - source leaprc.water.spceb - source leaprc.gaff2 - p = loadpdb pdb_clean_amber.pdb - w = loadpdb WAT.pdb - x = combine{p w} - addions2 x Na+ 0 - addions2 x Cl- 0 - addions2 x Na+ 1 - addions2 x Cl- 1 - set default PBradii mbondi3 - set default nocenter on - saveAmberParm x prot.parm7 prot.rst7 - quit - """) - with open("tleap_prot.in", "w") as fh: - fh.write(tleap_input) - subprocess.run(["tleap", "-f", "tleap_prot.in"], check=True) - - for f in ["prot.top", "prot.pdb"]: - Path(f).unlink(missing_ok=True) - - -if __name__ == "__main__": - run() diff --git a/md_workflows/resolvate.py b/md_workflows/resolvate.py deleted file mode 100644 index 2654e17..0000000 --- a/md_workflows/resolvate.py +++ /dev/null @@ -1,151 +0,0 @@ -"""Re-solvate after equilibration and run a second round of min+equil. - -Corresponds to run_all.sh line: - bash scripts/resolvate.sh -""" - -import re -import subprocess -from pathlib import Path - - -def run(ntmpi: int = 8, ntomp: int = 1): - maxsol = _compute_maxsol() - _resolvate(maxsol) - _update_topology(maxsol) - _minimize(ntmpi, ntomp) - _equilibrate(ntmpi, ntomp) - - -def _compute_maxsol() -> int: - """Do a trial solvation to figure out 25 % fill.""" - result = subprocess.run( - [ - "gmx", - "solvate", - "-cp", - "md_equil.gro", - "-cs", - "waterbox/water_equil.gro", - "-o", - "tmp.pdb", - ], - capture_output=True, - text=True, - ) - log_text = result.stdout + result.stderr - with open("gmx_solvate.log", "w") as fh: - fh.write(log_text) - - for line in log_text.splitlines(): - m = re.search(r"Number of solvent molecules:\s+(\d+)", line) - if m: - return int(int(m.group(1)) * 0.25) - raise RuntimeError("Could not parse solvent molecule count from gmx solvate output") - - -def _resolvate(maxsol: int): - result = subprocess.run( - [ - "gmx", - "solvate", - "-cp", - "md_equil.gro", - "-cs", - "waterbox/water_equil.gro", - "-o", - "md_resolv.pdb", - "-maxsol", - str(maxsol), - ], - capture_output=True, - text=True, - check=True, - ) - with open("gmx_resolvate.log", "w") as fh: - fh.write(result.stdout + result.stderr) - - -def _update_topology(maxsol: int): - with open("md_model_posre.top", "a") as fh: - fh.write(f"WAT {maxsol}\n") - - -def _minimize(ntmpi: int, ntomp: int): - artifacts_dir = Path("artifacts") - subprocess.run( - [ - "gmx", - "grompp", - "-f", - str(artifacts_dir / "min.mdp"), - "-c", - "md_resolv.pdb", - "-o", - "md_resolv_min.tpr", - "-p", - "md_model_posre.top", - ], - capture_output=True, - text=True, - check=True, - ) - - subprocess.run( - [ - "gmx", - "mdrun", - "-ntmpi", - str(ntmpi), - "-ntomp", - str(ntomp), - "-deffnm", - "md_resolv_min", - "-v", - ], - check=True, - ) - - -def _equilibrate(ntmpi: int, ntomp: int): - artifacts_dir = Path("artifacts") - subprocess.run( - [ - "gmx", - "grompp", - "-f", - str(artifacts_dir / "equil.mdp"), - "-c", - "md_resolv_min.gro", - "-o", - "md_resolv_equil.tpr", - "-p", - "md_model_posre.top", - "-r", - "md_model.pdb", - "-maxwarn", - "2", - ], - capture_output=True, - text=True, - check=True, - ) - - subprocess.run( - [ - "gmx", - "mdrun", - "-ntmpi", - str(ntmpi), - "-ntomp", - str(ntomp), - "-deffnm", - "md_resolv_equil", - "-v", - ], - check=True, - ) - - -if __name__ == "__main__": - run() diff --git a/md_workflows/run_params_gaussian.py b/md_workflows/run_params_gaussian.py deleted file mode 100644 index ac72d31..0000000 --- a/md_workflows/run_params_gaussian.py +++ /dev/null @@ -1,316 +0,0 @@ -"""Ligand parameterization using Gaussian and AmberTools. - -Corresponds to run_all.sh lines: - mkdir -p ligand - cd ligand - bash ../run_params_gaussian.sh - cd - -""" - -from __future__ import annotations - -import os -import subprocess -import textwrap -from pathlib import Path - -from .pdb_file_processing import prepare_pdb_and_resn_files - - -def _patch_gaussian_input(resn: str, nproc: int): - """Insert NprocShared and switch HF -> B3LYP/6-31+G(d,p).""" - with open(f"{resn}.gau") as fh: - lines = fh.readlines() - - patched = [] - for line in lines: - if "Link" in line: - patched.append(line) - patched.append(f"%NprocShared={nproc}\n") - else: - patched.append(line) - - text = "".join(patched) - text = text.replace("#HF", "#B3LYP/6-31+G(d,p)") - - with open("tmp", "w") as fh: - fh.write(text) - - -def _run_gaussian_opt(resn: str, nproc: int): - """Run Gaussian geometry optimization if needed.""" - gau_file = f"{resn}.gau" - log_file = f"{resn}.log" - env = {**os.environ, "OMP_NUM_THREADS": str(nproc)} - - should_run = False - if not os.path.exists(log_file): - print("Running Gaussian optimization") - should_run = True - else: - with open(gau_file) as a, open("tmp") as b: - if a.read() != b.read(): - print("Gaussian input file is different. Running Gaussian optimization.") - should_run = True - else: - print("Gaussian input file is the same. Skipping run.") - - if should_run: - os.replace("tmp", gau_file) - subprocess.run(["g16", gau_file], env=env, check=True) - - -def _run_gaussian_esp(resn: str, nproc: int): - """Run Gaussian ESP / CHELPG calculation if needed.""" - with open(f"{resn}.gau") as fh: - text = fh.read() - - text = text.replace("opt", "pop(chelpg,regular)") - text = text.replace("molecule", "grid") - - resp_gau = f"{resn}_resp.gau" - resp_log = f"{resn}_resp.log" - env = {**os.environ, "OMP_NUM_THREADS": str(nproc)} - - should_run = False - if not os.path.exists(resp_log): - print("Running Gaussian ESP calculation.") - should_run = True - else: - existing = "" - if os.path.exists(resp_gau): - with open(resp_gau) as fh: - existing = fh.read() - if existing != text: - print("Gaussian input file is different. Running Gaussian ESP calculation.") - should_run = True - else: - print("Gaussian input file is the same. Skipping run.") - - if should_run: - with open(resp_gau, "w") as fh: - fh.write(text) - subprocess.run(["g16", resp_gau], env=env, check=True) - - -def _process_resp_charges(resn: str): - """Derive RESP charges from Gaussian output and correct net charge.""" - subprocess.run( - [ - "antechamber", - "-fi", - "gout", - "-i", - f"{resn}_resp.log", - "-cf", - f"{resn}_resp.crg", - "-c", - "resp", - "-o", - f"{resn}_gauss.ac", - "-fo", - "ac", - "-rn", - resn, - ], - check=True, - ) - - subprocess.run( - [ - "antechamber", - "-fi", - "gout", - "-i", - f"{resn}_resp.log", - "-o", - f"{resn}_gauss.pdb", - "-fo", - "pdb", - "-rn", - resn, - ], - check=True, - ) - - orig_coords = _extract_coords_from_pdb(f"{resn}.pdb") - _graft_coords_to_ac(f"{resn}_gauss.ac", orig_coords, f"{resn}_resp.ac") - _correct_charge(f"{resn}_resp.ac") - - subprocess.run( - [ - "antechamber", - "-fi", - "ac", - "-i", - f"{resn}_resp.ac", - "-fo", - "mol2", - "-o", - f"{resn}_resp.mol2", - "-rn", - resn, - ], - check=True, - ) - - subprocess.run( - [ - "atomtype", - "-i", - f"{resn}_resp.ac", - "-o", - f"{resn}_resp_gaff.ac", - "-p", - "gaff", - ], - check=True, - ) - - subprocess.run( - [ - "prepgen", - "-i", - f"{resn}_resp_gaff.ac", - "-o", - f"{resn}_resp_gaff.prepc", - "-f", - "car", - ], - check=True, - ) - - subprocess.run( - [ - "parmchk2", - "-i", - f"{resn}_resp_gaff.prepc", - "-o", - f"{resn}_resp.frcmod", - "-f", - "prepc", - ], - check=True, - ) - - -def _extract_coords_from_pdb(pdb_file: str) -> list[str]: - coords = [] - with open(pdb_file) as fh: - for line in fh: - if line.startswith(("ATOM", "HETATM")): - coords.append(line[30:53]) - return coords - - -def _graft_coords_to_ac(ac_file: str, coords: list[str], out_file: str): - anum = 0 - with open(ac_file) as fh, open(out_file, "w") as out: - for line in fh: - if line.startswith(("ATOM", "HETATM")): - out.write(line[:30] + coords[anum] + line[53:]) - anum += 1 - else: - out.write(line) - - -def _correct_charge(ac_file: str): - """Adjust the largest-magnitude charge to make the total an exact integer.""" - charges = [] - with open(ac_file) as fh: - lines = fh.readlines() - for line in lines: - if line.startswith("ATOM"): - charges.append(float(line[54:63])) - - total = sum(charges) - nearest_int = round(total) - correction = nearest_int - total - - max_idx = max(range(len(charges)), key=lambda i: abs(charges[i])) - new_charge = charges[max_idx] + correction - - anum = 0 - new_lines = [] - for line in lines: - if line.startswith("ATOM"): - if anum == max_idx: - new_lines.append(f"{line[:54]}{new_charge:9.6f}{line[63:]}") - else: - new_lines.append(line) - anum += 1 - else: - new_lines.append(line) - - with open(ac_file, "w") as fh: - fh.writelines(new_lines) - - -def _build_amber_lib(resn: str): - """Run tleap to create Amber parameter/topology files.""" - tleap_input = textwrap.dedent(f"""\ - source leaprc.protein.ff19SB - source leaprc.gaff - loadamberparams {resn}_resp.frcmod - loadamberprep {resn}_resp_gaff.prepc - lig = loadmol2 {resn}_resp.mol2 - savepdb lig {resn}_resp.pdb - saveamberparm lig {resn}_resp.parm7 {resn}_resp.rst7 - quit - """) - with open("tleap_lig.in", "w") as fh: - fh.write(tleap_input) - subprocess.run(["tleap", "-f", "tleap_lig.in"], check=True) - - -def _run_parameterization(resn: str, g16root: str, nproc: int): - os.environ["g16root"] = g16root - g16_profile = os.path.join(g16root, "g16", "bsd", "g16.profile") - if os.path.exists(g16_profile): - subprocess.run(["bash", "-c", f"source {g16_profile}"], check=True) - - subprocess.run( - [ - "antechamber", - "-fi", - "pdb", - "-fo", - "gcrt", - "-i", - f"{resn}.pdb", - "-o", - f"{resn}.gau", - "-nc", - "-2", - "-m", - "1", - ], - check=True, - ) - - _patch_gaussian_input(resn, nproc) - _run_gaussian_opt(resn, nproc) - _run_gaussian_esp(resn, nproc) - _process_resp_charges(resn) - _build_amber_lib(resn) - - -def run( - g16root: str = "/Users/mewall/packages", - nproc: int = 8, -): - base_dir = Path.cwd() - lig_dir = base_dir / "ligand" - lig_dir.mkdir(parents=True, exist_ok=True) - resn, _ = prepare_pdb_and_resn_files(lig_dir=lig_dir) - - os.chdir(lig_dir) - - try: - _run_parameterization(resn, g16root, nproc) - finally: - os.chdir(base_dir) - - -if __name__ == "__main__": - run() diff --git a/md_workflows/sdk/__init__.py b/md_workflows/sdk/__init__.py new file mode 100644 index 0000000..a3cc33d --- /dev/null +++ b/md_workflows/sdk/__init__.py @@ -0,0 +1,74 @@ +"""Public Python API for md-workflows. + +Thin, curated re-exports of the pure ``core`` functions plus a couple of ergonomic +wrappers. Import from here for a stable surface; ``core`` internals may move. + + from md_workflows.sdk import run_standard_md, standard_md_pipeline, SystemConfig +""" + +from __future__ import annotations + +from ..core.config import ( + CrystalParams, + EquilibrateParams, + GaussianParams, + GromacsRunProfile, + MinimizeParams, + ResolvateParams, + SolvateParams, + SystemConfig, + WaterboxParams, +) +from ..core.exceptions import ( + GmxOutputParseError, + MDWorkflowError, + MissingInputError, + StepToolError, +) +from ..core.pipelines.standard_md import PipelineResult, standard_md_pipeline +from ..core.results import StepResult, StepStatus +from ..core.steps.equilibrate import equilibrate +from ..core.steps.make_crystal import make_crystal +from ..core.steps.make_waterbox import make_waterbox +from ..core.steps.minimize import minimize +from ..core.steps.param_prot import param_prot +from ..core.steps.pressure_interpolate import pressure_interpolate +from ..core.steps.resolvate import resolvate, resolvation_cycle +from ..core.steps.run_params_gaussian import run_params_gaussian +from ..core.steps.solvate import solvate +from .convenience import run_standard_md + +__all__ = [ + # pipeline + convenience + "standard_md_pipeline", + "PipelineResult", + "run_standard_md", + # steps + "param_prot", + "make_crystal", + "make_waterbox", + "solvate", + "minimize", + "equilibrate", + "resolvate", + "resolvation_cycle", + "pressure_interpolate", + "run_params_gaussian", + # config + results + "SystemConfig", + "GromacsRunProfile", + "CrystalParams", + "WaterboxParams", + "SolvateParams", + "MinimizeParams", + "EquilibrateParams", + "ResolvateParams", + "GaussianParams", + "StepResult", + "StepStatus", + # exceptions + "MDWorkflowError", + "MissingInputError", + "StepToolError", + "GmxOutputParseError", +] diff --git a/md_workflows/sdk/convenience.py b/md_workflows/sdk/convenience.py new file mode 100644 index 0000000..6bb71a2 --- /dev/null +++ b/md_workflows/sdk/convenience.py @@ -0,0 +1,28 @@ +"""Ergonomic wrappers over the core pipeline for SDK users.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +from ..core.config import SystemConfig +from ..core.pipelines.standard_md import PipelineResult, standard_md_pipeline + + +def run_standard_md( + pdb_id: str, + workdir: str | Path, + *, + config: str | Path | None = None, + resume: bool = False, + **overrides: Any, +) -> PipelineResult: + """Run the standard MD prep pipeline for ``pdb_id`` in ``workdir``. + + ``config`` is an optional YAML/TOML file; keyword ``overrides`` (deep-merged, so + nested groups like ``crystal={"ix": 5}`` work) win over it. Example:: + + run_standard_md("4LZT", "runs/4lzt", crystal={"ix": 5}, resume=True) + """ + cfg = SystemConfig.load(config, overrides={"pdb_id": pdb_id, **overrides}) + return standard_md_pipeline(Path(workdir), cfg, resume=resume) diff --git a/md_workflows/solvate.py b/md_workflows/solvate.py deleted file mode 100644 index 7fd1d65..0000000 --- a/md_workflows/solvate.py +++ /dev/null @@ -1,174 +0,0 @@ -"""Solvate the crystal with water and add ions. - -Corresponds to run_all.sh line: - bash scripts/solvate.sh -""" - -import re -import subprocess - - -def run(): - ncopies = _count_copies() - _solvate_crystal() - nwat_initial = _read_solvate_nwat() - _write_topology_header(ncopies) - ions_pos, ions_neg = _count_ions() - net_ion_charge = (ions_pos - ions_neg) * ncopies - print(f"Positive ions, negative ions, net charge: {ions_pos} {ions_neg} {net_ion_charge}") - - nna, ncl = _compute_ion_counts(nwat_initial, net_ion_charge) - _insert_ions(ncl, nna) - nwat = _count_final_water() - _finalize_topology(ncopies, nwat, ncl, nna) - - import shutil - - shutil.copy("xtal_solv_cl_na.pdb", "md_model.pdb") - - -def _count_copies() -> int: - nats_one = 0 - with open("prot_dry.pdb") as fh: - for line in fh: - if line.startswith(("ATOM", "HETATM")): - nats_one += 1 - - nats_cell = 0 - with open("xtal.pdb") as fh: - for line in fh: - if line.startswith(("ATOM", "HETATM")): - nats_cell += 1 - - ncopies = nats_cell // nats_one - print(f"There are {ncopies} copies of the asymmetric unit") - return ncopies - - -def _solvate_crystal(): - # Match scripts/solvate.sh: >& gmx_solvate.log so _read_solvate_nwat() can parse it. - with open("gmx_solvate.log", "w") as log_fh: - subprocess.run( - [ - "gmx", - "solvate", - "-cp", - "xtal.pdb", - "-cs", - "waterbox/water_equil.gro", - "-o", - "xtal_solv.pdb", - ], - stdout=log_fh, - stderr=subprocess.STDOUT, - check=True, - ) - - -def _read_solvate_nwat() -> int: - with open("gmx_solvate.log") as fh: - for line in fh: - m = re.search(r"Output configuration contains\s+(\d+)", line) - if m: - return int(m.group(1)) - raise RuntimeError("Could not parse water count from gmx_solvate.log") - - -def _write_topology_header(ncopies: int): - with open("prot.top") as fh: - header_lines = [] - for line in fh: - if "molecules" in line.lower(): - break - header_lines.append(line) - - with open("md_model.top", "w") as fh: - fh.writelines(header_lines) - - -def _count_ions() -> tuple[int, int]: - ions_pos = 0 - ions_neg = 0 - with open("prot.pdb") as fh: - for line in fh: - if line.startswith("HETATM"): - if "Na+" in line: - ions_pos += 1 - elif "Cl-" in line: - ions_neg += 1 - return ions_pos, ions_neg - - -def _compute_ion_counts(nwat: int, net_ion_charge: int) -> tuple[int, int]: - """Compute Na+ and Cl- counts for ~0.1 M ionic strength plus neutralization.""" - if net_ion_charge >= 0: - ncl = nwat * 0.1 // 55 - nna = ncl + net_ion_charge - else: - nna = nwat * 0.1 // 55 - ncl = nna - net_ion_charge - return int(nna), int(ncl) - - -def _insert_ions(ncl: int, nna: int): - subprocess.run( - [ - "gmx", - "insert-molecules", - "-f", - "xtal_solv.pdb", - "-ci", - "Cl-.pdb", - "-o", - "xtal_solv_cl.pdb", - "-replace", - "SOL", - "-nmol", - str(ncl), - ], - check=True, - ) - - subprocess.run( - [ - "gmx", - "insert-molecules", - "-f", - "xtal_solv_cl.pdb", - "-ci", - "Na+.pdb", - "-o", - "xtal_solv_cl_na.pdb", - "-replace", - "SOL", - "-nmol", - str(nna), - ], - check=True, - ) - - -def _count_final_water() -> int: - wat_atoms = 0 - with open("xtal_solv_cl_na.pdb") as fh: - for line in fh: - if "WAT" in line: - wat_atoms += 1 - return wat_atoms // 3 - - -def _finalize_topology(ncopies: int, nwat: int, ncl: int, nna: int): - topcopy = "system1 1\n" - topall = topcopy * ncopies - - with open("md_model.top", "a") as fh: - fh.write("[ molecules ]\n") - fh.write("; Compound #mols\n") - fh.write(topall) - fh.write(f"WAT {nwat}\n") - fh.write(f"Cl- {ncl}\n") - fh.write(f"Na+ {nna}\n") - - -if __name__ == "__main__": - run() diff --git a/md_workflows/workflows/__init__.py b/md_workflows/workflows/__init__.py deleted file mode 100644 index 7d31bda..0000000 --- a/md_workflows/workflows/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Orchestration helpers (e.g. :mod:`md_workflows.workflows.mdmx`).""" diff --git a/md_workflows/workflows/mdmx.py b/md_workflows/workflows/mdmx.py deleted file mode 100644 index 0c162b9..0000000 --- a/md_workflows/workflows/mdmx.py +++ /dev/null @@ -1,100 +0,0 @@ -"""Run the MD pipeline in the order defined by ``scripts/run_all.sh``. - -Matches the current shell script: - -1. ``run_params_gaussian`` (under ``ligand/``, i.e. ``cd ligand && bash - ../run_params_gaussian.sh``) -2. ``param_prot`` -3. ``make_crystal`` -4. ``make_waterbox`` -5. ``solvate`` -6. ``minimize`` -7. ``equilibrate`` -8. ``resolvate`` - -Adjust defaults via CLI flags below. -""" - -from __future__ import annotations - -import argparse - -from ..equilibrate import run as run_equilibrate -from ..make_crystal import run as run_make_crystal -from ..make_waterbox import run as run_make_waterbox -from ..minimize import run as run_minimize -from ..param_prot import run as run_param_prot -from ..resolvate import run as run_resolvate -from ..solvate import run as run_solvate - - -def main( - *, - ntomp: int = 26, - param_pdb_id: str = "6B8X", - crystal_ix: int = 1, - crystal_iy: int | None = None, - crystal_iz: int | None = None, - chimerax_exec: str = "/usr/bin/chimerax-daily", - resolv_ntmpi: int = 8, - resolv_ntomp: int = 1, -) -> None: - """Execute workflow stages in ``run_all.sh`` order.""" - run_param_prot(pdb_id=param_pdb_id) - run_make_crystal(ix=crystal_ix, iy=crystal_iy, iz=crystal_iz, chimerax_exec=chimerax_exec) - run_make_waterbox(ntomp=ntomp) - run_solvate() - run_minimize(ntomp=ntomp) - run_equilibrate(ntomp=ntomp) - run_resolvate(ntmpi=resolv_ntmpi, ntomp=resolv_ntomp) - - -def _cli() -> None: - parser = argparse.ArgumentParser(description="Run full pipeline per scripts/run_all.sh") - parser.add_argument( - "--ntomp", - type=int, - default=26, - help="OpenMP threads for GROMACS steps (waterbox, minimize, equilibrate)", - ) - parser.add_argument( - "--param-pdb-id", - default="6B8X", - help="PDB ID passed to param_prot (Coordinates file should be .pdb in cwd)", - ) - parser.add_argument( - "--ix", - type=int, - default=1, - help="make_crystal supercell replication (x; also y/z if omitted)", - ) - parser.add_argument( - "--iy", type=int, default=None, help="make_crystal y replication (optional)" - ) - parser.add_argument( - "--iz", type=int, default=None, help="make_crystal z replication (optional)" - ) - parser.add_argument( - "--chimerax-exec", - default="/usr/bin/chimerax-daily", - help="ChimeraX executable for make_crystal", - ) - - parser.add_argument("--resolv-ntmpi", type=int, default=8, help="resolvate gmx mdrun -ntmpi") - parser.add_argument("--resolv-ntomp", type=int, default=1, help="resolvate gmx mdrun -ntomp") - - args = parser.parse_args() - main( - ntomp=args.ntomp, - param_pdb_id=args.param_pdb_id, - crystal_ix=args.ix, - crystal_iy=args.iy, - crystal_iz=args.iz, - chimerax_exec=args.chimerax_exec, - resolv_ntmpi=args.resolv_ntmpi, - resolv_ntomp=args.resolv_ntomp, - ) - - -if __name__ == "__main__": - _cli() diff --git a/pyproject.toml b/pyproject.toml index 0e42fd4..d9ce1e8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,27 +7,26 @@ name = "md-workflows" version = "0.1.0" description = "Molecular dynamics workflow orchestration scripts" requires-python = ">=3.10" +dependencies = [ + "pydantic>=2", + "pyyaml>=6", + "typer>=0.12", +] [project.optional-dependencies] # Dev/CI tooling only — kept out of [project.dependencies] so it is not baked into the # runtime image. CI installs it via `pip install .[dev]` (or plain `pip install ruff`). dev = [ "ruff>=0.15.19", + "pytest>=8", ] [project.scripts] -"md_workflows.param_prot" = "md_workflows.cli:param_prot_cli" -"md_workflows.make_crystal" = "md_workflows.cli:make_crystal_cli" -"md_workflows.make_waterbox" = "md_workflows.cli:make_waterbox_cli" -"md_workflows.solvate" = "md_workflows.cli:solvate_cli" -"md_workflows.minimize" = "md_workflows.cli:minimize_cli" -"md_workflows.equilibrate" = "md_workflows.cli:equilibrate_cli" -"md_workflows.resolvate" = "md_workflows.cli:resolvate_cli" -"md_workflows.run_params_gaussian" = "md_workflows.cli:run_params_gaussian_cli" -"md_workflows.mdmx" = "md_workflows.workflows.mdmx:_cli" - -[tool.setuptools] -packages = ["md_workflows", "md_workflows.workflows"] +# Single Typer app; subcommands replace the old per-step dotted entry points. +md-workflows = "md_workflows.cli:app" + +[tool.setuptools.packages.find] +include = ["md_workflows*"] [tool.ruff] # Pin the knobs so local runs and CI agree regardless of ruff's shifting defaults. @@ -37,3 +36,12 @@ target-version = "py310" [tool.ruff.lint] # pycodestyle errors/warnings, pyflakes, isort, pyupgrade, bugbear. select = ["E", "F", "W", "I", "UP", "B"] + +[tool.ruff.lint.per-file-ignores] +# Typer's public API puts Option()/Argument() calls in parameter defaults by design. +"md_workflows/cli.py" = ["B008"] + +[dependency-groups] +dev = [ + "pytest>=9.1.1", +] diff --git a/tests/test_cli.py b/tests/test_cli.py new file mode 100644 index 0000000..0b2bcf1 --- /dev/null +++ b/tests/test_cli.py @@ -0,0 +1,50 @@ +from pathlib import Path + +from typer.testing import CliRunner + +from md_workflows.cli import app + +runner = CliRunner() + + +def test_help_lists_all_commands(): + result = runner.invoke(app, ["--help"]) + assert result.exit_code == 0 + for cmd in [ + "param-prot", + "make-crystal", + "make-waterbox", + "solvate", + "minimize", + "equilibrate", + "resolvate", + "run-params-gaussian", + "run-pipeline", + ]: + assert cmd in result.output + + +def test_subcommand_help_exit_zero(): + assert runner.invoke(app, ["run-pipeline", "--help"]).exit_code == 0 + + +def test_missing_input_exits_2(tmp_path: Path): + result = runner.invoke(app, ["--workdir", str(tmp_path), "minimize"]) + assert result.exit_code == 2 + assert "missing required input" in result.output + + +def test_bad_config_exits_1(tmp_path: Path): + bad = tmp_path / "bad.yaml" + bad.write_text("a: b: c: [\n") + result = runner.invoke(app, ["--workdir", str(tmp_path), "--config", str(bad), "minimize"]) + assert result.exit_code == 1 + assert "config error" in result.output + + +def test_valid_config_parses_then_reports_missing_inputs(tmp_path: Path): + cfg = tmp_path / "ok.yaml" + cfg.write_text("pdb_id: 4LZT\ncrystal:\n ix: 5\n") + result = runner.invoke(app, ["--workdir", str(tmp_path), "--config", str(cfg), "minimize"]) + # config parsed fine; fails only on missing inputs (exit 2) + assert result.exit_code == 2 diff --git a/tests/test_config.py b/tests/test_config.py new file mode 100644 index 0000000..5844a04 --- /dev/null +++ b/tests/test_config.py @@ -0,0 +1,84 @@ +from pathlib import Path + +import pytest + +from md_workflows.core.config import RUN_PROFILE_KEYS, GromacsRunProfile, SystemConfig + + +def test_default_profiles_render_taylor_flags(): + cfg = SystemConfig() + assert cfg.profile("min").to_mdrun_flags() == [ + "-ntmpi", + "1", + "-ntomp", + "16", + "-nb", + "gpu", + "-pme", + "cpu", + "-bonded", + "cpu", + "-notunepme", + ] + assert cfg.profile("equil").to_mdrun_flags() == [ + "-ntmpi", + "1", + "-ntomp", + "16", + "-nb", + "gpu", + "-pme", + "gpu", + "-bonded", + "gpu", + "-notunepme", + ] + + +def test_to_mdrun_flags_omits_none_and_handles_tunepme(): + p = GromacsRunProfile(ntmpi=None, ntomp=8, nb="gpu", pme=None, bonded=None, tunepme=True) + assert p.to_mdrun_flags() == ["-ntomp", "8", "-nb", "gpu"] + + +def test_overrides_win_and_merge_into_defaults(): + cfg = SystemConfig.load( + overrides={"pdb_id": "4LZT", "crystal": {"ix": 5}, "run_profiles": {"min": {"ntomp": 26}}} + ) + assert cfg.pdb_id == "4LZT" + assert cfg.crystal.ix == 5 + # overriding one field of one profile keeps that profile's other defaults ... + assert "-pme" in cfg.profile("min").to_mdrun_flags() + assert cfg.profile("min").pme == "cpu" + assert cfg.profile("min").ntomp == 26 + # ... and leaves the other profiles present + assert set(cfg.run_profiles) == set(RUN_PROFILE_KEYS) + assert cfg.profile("equil").pme == "gpu" + + +def test_yaml_round_trip(tmp_path: Path): + cfg_file = tmp_path / "c.yaml" + cfg_file.write_text("pdb_id: 6LYZ\nwaterbox:\n nc_scale: 3\n") + cfg = SystemConfig.load(cfg_file) + assert cfg.pdb_id == "6LYZ" + assert cfg.waterbox.nc_scale == 3 + + +def test_flag_overrides_beat_file(tmp_path: Path): + cfg_file = tmp_path / "c.yaml" + cfg_file.write_text("pdb_id: 6LYZ\n") + cfg = SystemConfig.load(cfg_file, overrides={"pdb_id": "4LZT"}) + assert cfg.pdb_id == "4LZT" + + +def test_bad_yaml_raises_valueerror(tmp_path: Path): + bad = tmp_path / "bad.yaml" + bad.write_text("a: b: c: [\n") + with pytest.raises(ValueError): + SystemConfig.load(bad) + + +def test_unknown_format_raises(tmp_path: Path): + weird = tmp_path / "c.ini" + weird.write_text("x=1\n") + with pytest.raises(ValueError): + SystemConfig.load(weird) diff --git a/tests/test_gmx.py b/tests/test_gmx.py new file mode 100644 index 0000000..b4eea8b --- /dev/null +++ b/tests/test_gmx.py @@ -0,0 +1,72 @@ +from pathlib import Path + +import pytest + +from md_workflows.core.gmx import checksums, count_waters, manage_wat_topology, sha256_file + + +def test_count_waters_pdb(tmp_path: Path): + pdb = tmp_path / "x.pdb" + pdb.write_text( + "ATOM 1 N ALA A 1 0.0 0.0 0.0\n" + "HETATM 2 O WAT A 2 1.0 0.0 0.0\n" + "HETATM 3 H1 WAT A 2 1.1 0.0 0.0\n" + "HETATM 4 H2 WAT A 2 1.2 0.0 0.0\n" + "HETATM 5 O WAT A 3 2.0 0.0 0.0\n" + "HETATM 6 H1 WAT A 3 2.1 0.0 0.0\n" + "HETATM 7 H2 WAT A 3 2.2 0.0 0.0\n" + ) + assert count_waters(pdb) == 2 + + +def test_count_waters_gro(tmp_path: Path): + gro = tmp_path / "y.gro" + gro.write_text( + "title\n 4\n" + " 1ALA N 1 0.000 0.000 0.000\n" + " 2SOL OW 2 1.000 0.000 0.000\n" + " 2SOL HW1 3 1.100 0.000 0.000\n" + " 2SOL HW2 4 1.200 0.000 0.000\n" + " 2.0 2.0 2.0\n" + ) + assert count_waters(gro) == 1 + + +def _top(tmp_path: Path) -> Path: + top = tmp_path / "m.top" + top.write_text("[ molecules ]\n; Compound #mols\nsystem1 1\nWAT 100\nCl- 5\nNa+ 7\n") + return top + + +def test_manage_wat_topology_append_is_idempotent(tmp_path: Path): + top = _top(tmp_path) + manage_wat_topology(top, 863, mode="append") + manage_wat_topology(top, 863, mode="append") # re-run must not double-add + assert top.read_text().count("WAT 863") == 1 + assert top.read_text().rstrip().endswith("WAT 863") + + +def test_manage_wat_topology_multi_block_sequence(tmp_path: Path): + top = _top(tmp_path) + manage_wat_topology(top, 863, mode="append") # run 1 + manage_wat_topology(top, 1294, mode="append") # run 2 probe + manage_wat_topology(top, 412, mode="drop_last_then_append") # final drops probe + wat_lines = [ln for ln in top.read_text().splitlines() if ln.startswith("WAT")] + # initial 100 kept, run-1 863 kept, probe 1294 dropped, final 412 appended + assert wat_lines == ["WAT 100", "WAT 863", "WAT 412"] + + +def test_manage_wat_topology_requires_molecules_section(tmp_path: Path): + top = tmp_path / "bad.top" + top.write_text("[ atoms ]\n") + with pytest.raises(ValueError): + manage_wat_topology(top, 10) + + +def test_checksums_skips_missing(tmp_path: Path): + f = tmp_path / "a.txt" + f.write_text("hello") + result = checksums({"a": f, "missing": tmp_path / "nope"}) + assert set(result) == {"a"} + assert result["a"] == sha256_file(f) + assert len(result["a"]) == 64 diff --git a/tests/test_pressure_interpolate.py b/tests/test_pressure_interpolate.py new file mode 100644 index 0000000..86ad429 --- /dev/null +++ b/tests/test_pressure_interpolate.py @@ -0,0 +1,47 @@ +from pathlib import Path + +import pytest + +from md_workflows.core.exceptions import MDWorkflowError +from md_workflows.core.steps.pressure_interpolate import interpolate + +REF = Path("ref.gro") + + +def test_normal_interpolation(): + # dNw = round((1-(-500))*(1200-1000)/(300-(-500))) = round(125.25) = 125 + r = interpolate(-500.0, 300.0, 1000, 1200, REF, target=1.0) + assert r.dNw == 125 + assert r.Nw_target == 1125 + assert not r.converged + assert r.warning is None + + +def test_extrapolation_same_side_warns_and_can_be_negative(): + r = interpolate(400.0, 900.0, 1000, 1200, REF, target=1.0) + assert r.dNw < 0 + assert r.warning is not None + + +def test_zero_slope_near_target_converges(): + r = interpolate(50.0, 50.0, 1000, 1000, REF, target=1.0, pressure_tol=100.0) + assert r.dNw == 0 + assert r.converged + + +def test_zero_slope_far_from_target_raises(): + with pytest.raises(MDWorkflowError): + interpolate(-800.0, -800.0, 1000, 1050, REF, target=1.0, pressure_tol=100.0) + + +def test_absurd_dNw_far_from_target_raises(): + with pytest.raises(MDWorkflowError): + interpolate(-5000.0, -4999.999, 1000, 1001, REF, target=1.0, max_dNw_factor=1.0) + + +def test_absurd_dNw_near_target_converges(): + r = interpolate( + 50.0, 50.001, 1000, 1001, REF, target=1.0, pressure_tol=100.0, max_dNw_factor=1.0 + ) + assert r.dNw == 0 + assert r.converged diff --git a/tests/test_steps_contract.py b/tests/test_steps_contract.py new file mode 100644 index 0000000..f243427 --- /dev/null +++ b/tests/test_steps_contract.py @@ -0,0 +1,83 @@ +from pathlib import Path + +import pytest + +from md_workflows.core.config import SystemConfig +from md_workflows.core.exceptions import MissingInputError +from md_workflows.core.results import StepStatus +from md_workflows.core.steps import ( + equilibrate, + make_crystal, + make_waterbox, + minimize, + param_prot, + resolvate, + solvate, +) + +# Steps with required pre-existing inputs (param_prot / gaussian download or detect theirs). +STEPS_WITH_REQUIRED_INPUTS = [ + make_crystal, + make_waterbox, + solvate, + minimize, + equilibrate, + resolvate, +] + + +@pytest.mark.parametrize("mod", STEPS_WITH_REQUIRED_INPUTS) +def test_resolve_inputs_are_under_workdir(mod, tmp_path: Path): + cfg = SystemConfig() + inputs = mod.resolve_inputs(tmp_path, cfg) + assert inputs.workdir == tmp_path + consumed = inputs.consumed_paths() + assert consumed, f"{mod.STEP} should declare required inputs" + for p in consumed: + assert tmp_path in p.parents, f"{p} not under workdir" + + +@pytest.mark.parametrize("mod", STEPS_WITH_REQUIRED_INPUTS) +def test_check_inputs_raises_listing_all_missing(mod, tmp_path: Path): + cfg = SystemConfig() + inputs = mod.resolve_inputs(tmp_path, cfg) + with pytest.raises(MissingInputError) as exc: + mod.check_inputs(inputs) + assert exc.value.step == mod.STEP + assert len(exc.value.missing) == len(inputs.consumed_paths()) + + +def test_param_prot_has_no_required_files(tmp_path: Path): + cfg = SystemConfig() + inputs = param_prot.resolve_inputs(tmp_path, cfg) + assert inputs.consumed_paths() == [] + # guard does not raise when nothing is required + assert param_prot.check_inputs(inputs) is None + + +def test_minimize_resume_skips_without_running_gmx(tmp_path: Path): + (tmp_path / "artifacts").mkdir() + for name in ("md_model.pdb", "md_model.top"): + (tmp_path / name).write_text("x") + (tmp_path / "artifacts" / "min.mdp").write_text("integrator=steep\n") + (tmp_path / "md_min.gro").write_text("fake existing output") + + cfg = SystemConfig() + inputs = minimize.resolve_inputs(tmp_path, cfg) + result = minimize.minimize(inputs, cfg.profile("min"), resume=True) + + assert result.status == StepStatus.SKIPPED + assert result.outputs["gro"] == tmp_path / "md_min.gro" + assert set(result.input_checksums) == {"model_pdb", "model_top", "min_mdp"} + + +def test_mdp_dir_resolves_under_workdir(tmp_path: Path): + cfg = SystemConfig() # mdp_dir default "artifacts" (relative) + inputs = minimize.resolve_inputs(tmp_path, cfg) + assert inputs.min_mdp == tmp_path / "artifacts" / "min.mdp" + + +def test_absolute_mdp_dir_is_honored(tmp_path: Path): + cfg = SystemConfig.load(overrides={"mdp_dir": "/opt/mdps"}) + inputs = minimize.resolve_inputs(tmp_path, cfg) + assert inputs.min_mdp == Path("/opt/mdps/min.mdp")