bioai-platform/backend/app/tools/md_config.py

"""MD force field / solvent configuration with startup verification.

The exposed menu of force fields and implicit-solvent models is not a static
list: every (force field x solvent) combination is verified at import time by
running a real OpenMM ``createSystem()`` on an embedded alanine-dipeptide
structure (the canonical test system), using the exact same construction path
as the simulation engine (PDBFile -> Modeller -> addHydrogens -> createSystem
with a non-periodic 2.0 nm cutoff). Only combinations that build successfully
are exposed through ``GET /api/md/forcefields`` and accepted by the run
endpoints. An unknown or unverified combination raises an explicit ValueError
instead of silently falling back to AMBER14/OBC2.
"""

from __future__ import annotations

import io
import logging

logger = logging.getLogger(__name__)

# User-facing force field keys -> bundled OpenMM XML force field files.
FF_XML: dict[str, str] = {
    "amber14": "amber14-all.xml",                # legacy all-residue AMBER14
    "ff14sb": "amber14/protein.ff14SB.xml",
    "ff15ipq": "amber14/protein.ff15ipq.xml",
    "ff19sb": "amber19/protein.ff19SB.xml",
    "amberfb15": "amberfb15.xml",
    "charmm36": "charmm36.xml",
}

FF_LABELS: dict[str, str] = {
    "amber14": "AMBER14 (ff14SB)",
    "ff14sb": "AMBER ff14SB",
    "ff15ipq": "AMBER ff15ipq",
    "ff19sb": "AMBER ff19SB",
    "amberfb15": "AMBER ff15 (amberfb15)",
    "charmm36": "CHARMM36",
}

# Implicit-solvent (Generalized Born) models -> bundled XML files.
SOLVENT_XML: dict[str, str] = {
    "obc1": "implicit/obc1.xml",
    "obc2": "implicit/obc2.xml",
    "gbn2": "implicit/gbn2.xml",
}

DEFAULT_FORCEFIELD = "amber14"
DEFAULT_SOLVENT = "obc2"

# Heavy-atom alanine dipeptide (ALA-ALA with C-terminal OXT) used as the
# startup verification probe. It exercises the same Modeller/addHydrogens
# construction path the real simulation uses, is fully self-contained, and
# needs no network access. Its geometry is approximate; the probe only checks
# that createSystem() builds a valid System, so this is acceptable.
_DIPEPTIDE_PDB = """\
REMARK   1 CREATED WITH OPENMM 8.5.2
CRYST1   40.960   18.650   22.520  90.00  90.77  90.00 P 1           1
ATOM      1  N   ALA A   1      17.047  14.099   3.625  1.00  0.00           N
ATOM      2  CA  ALA A   1      16.967  12.784   4.338  1.00  0.00           C
ATOM      3  C   ALA A   1      15.685  12.755   5.133  1.00  0.00           C
ATOM      4  O   ALA A   1      15.268  13.825   5.594  1.00  0.00           O
ATOM      5  CB  ALA A   1      18.170  12.703   5.337  1.00  0.00           C
ATOM      6  N   ALA A   2      15.115  11.555   5.265  1.00  0.00           N
ATOM      7  CA  ALA A   2      13.856  11.469   6.066  1.00  0.00           C
ATOM      8  C   ALA A   2      14.164  10.785   7.379  1.00  0.00           C
ATOM      9  O   ALA A   2      14.993   9.862   7.443  1.00  0.00           O
ATOM     10  CB  ALA A   2      12.732  10.711   5.261  1.00  0.00           C
ATOM     11  OXT ALA A   2      13.343  11.699   7.316  1.00  0.00           O
TER      11      ALA A   2
END
"""

_VERIFIED: dict[str, tuple[str, ...]] | None = None
_VERIFY_ERRORS: dict[tuple[str, str], str] | None = None


def _build_dipeptide_system(forcefield: str, solvent: str) -> None:
    """Run the real construction path on the dipeptide (raises on failure)."""
    from openmm.app import PDBFile, ForceField, Modeller, CutoffNonPeriodic
    from openmm import unit

    pdb = PDBFile(io.StringIO(_DIPEPTIDE_PDB))
    forcefield_xml = ForceField(FF_XML[forcefield], SOLVENT_XML[solvent])
    modeller = Modeller(pdb.topology, pdb.positions)
    modeller.addHydrogens(forcefield_xml)
    forcefield_xml.createSystem(
        modeller.topology,
        nonbondedMethod=CutoffNonPeriodic,
        nonbondedCutoff=2.0 * unit.nanometer,
    )


def verify_ff_solvent_combos() -> dict[str, tuple[str, ...]]:
    """Verify every (force field x solvent) combo with a real createSystem().

    Returns a dict mapping each passing force field to the tuple of solvents
    it supports. Results are cached after the first call. When OpenMM is not
    installed, returns an empty mapping (MD degrades to structural analysis).
    """
    global _VERIFIED, _VERIFY_ERRORS
    if _VERIFIED is not None:
        return _VERIFIED

    verified: dict[str, list[str]] = {ff: [] for ff in FF_XML}
    errors: dict[tuple[str, str], str] = {}
    try:
        import openmm  # noqa: F401
    except Exception as e:
        logger.warning("OpenMM unavailable — MD force field verification skipped: %s", e)
        _VERIFIED = {}
        _VERIFY_ERRORS = {}
        return _VERIFIED

    for ff in FF_XML:
        for solvent in SOLVENT_XML:
            try:
                _build_dipeptide_system(ff, solvent)
                verified[ff].append(solvent)
            except Exception as exc:
                errors[(ff, solvent)] = f"{type(exc).__name__}: {exc}"
                logger.warning(
                    "MD combo %s x %s FAILED verification: %s: %s",
                    ff, solvent, type(exc).__name__, exc,
                )
    _VERIFIED = {ff: tuple(ff_solvents) for ff, ff_solvents in verified.items()}
    _VERIFY_ERRORS = errors
    logger.info(
        "MD force field verification complete: %d force fields, combos=%s",
        len(_VERIFIED), {ff: list(s) for ff, s in _VERIFIED.items()},
    )
    return _VERIFIED


def get_forcefields_menu() -> dict:
    """Menu payload for GET /api/md/forcefields (verified combos only)."""
    verified = verify_ff_solvent_combos()
    return {
        "forcefields": [
            {"value": ff, "label": FF_LABELS.get(ff, ff)}
            for ff in FF_XML
            if verified.get(ff)
        ],
        "solvents": [{"value": s, "label": f"Implicit · {s.upper()}"} for s in SOLVENT_XML],
        "combos": {ff: list(solvents) for ff, solvents in verified.items()},
        "defaults": {"forcefield": DEFAULT_FORCEFIELD, "solvent": DEFAULT_SOLVENT},
        "probe": {
            "system": "alanine dipeptide (ALA-ALA)",
            "forcefields_tested": len(FF_XML),
            "solvents_tested": len(SOLVENT_XML),
        },
    }


def resolve_combo(forcefield: str | None, solvent: str | None) -> tuple[str, str]:
    """Normalize and validate a requested (forcefield, solvent) pair.

    Raises ValueError with a clear message when the force field or solvent is
    unknown, or when the specific combination did not pass startup
    verification. No silent fallback: an invalid request is an explicit error.
    """
    ff_key = (forcefield or DEFAULT_FORCEFIELD).strip().lower()
    sol_key = (solvent or DEFAULT_SOLVENT).strip().lower()

    if ff_key not in FF_XML:
        raise ValueError(
            f"Unsupported force field: {forcefield!r}. Choose from: "
            + ", ".join(FF_XML)
        )
    if sol_key not in SOLVENT_XML:
        raise ValueError(
            f"Unsupported solvent model: {solvent!r}. Choose from: "
            + ", ".join(SOLVENT_XML)
        )

    verified = verify_ff_solvent_combos()
    if verified and sol_key not in verified.get(ff_key, ()):
        raise ValueError(
            f"Force field '{ff_key}' is not available with solvent '{sol_key}' "
            f"on this deployment. Supported solvents for '{ff_key}': "
            + (", ".join(verified.get(ff_key, ())) if verified.get(ff_key) else "none")
        )
    return ff_key, sol_key