Experimental HybridTopProtocol (no-protein)

devtools/gen_serialized_results.py

@@ -14,6 +14,8 @@
 import gufe
 import openfe
 from gufe.tokenization import JSON_HANDLER
+from kartograf import KartografAtomMapper
+from kartograf.atom_aligner import align_mol_shape
 from openff.toolkit import AmberToolsToolkitWrapper, Molecule, RDKitToolkitWrapper
 from openff.toolkit.utils.toolkit_registry import (
     ToolkitRegistry,
@@ -22,11 +24,13 @@
 from openff.units import unit
 
 from pontibus.components import ExtendedSolventComponent
+from pontibus.protocols.relative import HybridTopProtocol
 from pontibus.protocols.solvation import ASFEProtocol
 
 logger = logging.getLogger(__name__)
 
 LIGA = "[H]C([H])([H])C([H])([H])C(=O)C([H])([H])C([H])([H])[H]"
+LIGB = "[H]C([H])([H])C(=O)C([H])([H])C([H])([H])C([H])([H])[H]"
 
 amber_rdkit = ToolkitRegistry([RDKitToolkitWrapper(), AmberToolsToolkitWrapper()])
 
@@ -111,7 +115,62 @@ def generate_asfe_json_octanol(smc):
     execute_and_serialize(dag, protocol, "ASFEProtocol_octanol")
 
 
+def generate_rfe_settings():
+    settings = HybridTopProtocol.default_settings()
+    settings.simulation_settings.equilibration_length = 10 * unit.picosecond
+    settings.simulation_settings.production_length = 250 * unit.picosecond
+    return settings
+
+
+def generate_rfe_inputs(molA, molB):
+    a_molB = align_mol_shape(molB, ref_mol=molA)
+    mapper = KartografAtomMapper(atom_map_hydrogens=True)
+    mapping = next(mapper.suggest_mappings(molA, a_molB))
+
+    return mapping, molA, a_molB
+
+
+def generate_hybridtop_rfe_vacuum_json(molA, molB):
+    settings = generate_rfe_settings()
+    settings.forcefield_settings.nonbonded_method = "nocutoff"
+    protocol = HybridTopProtocol(settings=settings)
+
+    mapping, molA, molB = generate_rfe_inputs(molA, molB)
+
+    systemA = openfe.ChemicalSystem({"ligand": molA})
+    systemB = openfe.ChemicalSystem({"ligand": molB})
+
+    dag = protocol.create(
+        stateA=systemA,
+        stateB=systemB,
+        mapping=mapping,
+    )
+
+    execute_and_serialize(dag, protocol, "HybridTopProtocol_vacuum")
+
+
+def generate_hybridtop_rfe_solvent_json(molA, molB):
+    settings = generate_rfe_settings()
+    protocol = HybridTopProtocol(settings=settings)
+
+    mapping, molA, molB = generate_rfe_inputs(molA, molB)
+
+    systemA = openfe.ChemicalSystem({"ligand": molA, "solvent": openfe.SolventComponent()})
+    systemB = openfe.ChemicalSystem({"ligand": molB, "solvent": openfe.SolventComponent()})
+
+    dag = protocol.create(
+        stateA=systemA,
+        stateB=systemB,
+        mapping=mapping,
+    )
+
+    execute_and_serialize(dag, protocol, "HybridTopProtocol_solvent")
+
+
 if __name__ == "__main__":
     molA = get_molecule(LIGA, "ligandA")
-    # generate_asfe_json_water(molA)
+    molB = get_molecule(LIGB, "ligandB")
+    generate_asfe_json_water(molA)
     generate_asfe_json_octanol(molA)
+    generate_hybridtop_rfe_vacuum_json(molA, molB)
+    generate_hybridtop_rfe_solvent_json(molA, molB)

environment.yml

@@ -3,7 +3,7 @@ channels:
   - conda-forge
 dependencies:
   - pip
-  - gufe >=1.5
+  - gufe >=1.6
   - openff-toolkit >0.16.0
   - openff-interchange >=0.4
   - openff-nagl-base >=0.3.3
@@ -20,12 +20,13 @@ dependencies:
   - pydantic >=2.0
   - pyyaml
   - coverage
-  - cinnabar ~=0.4.0
+  - cinnabar ~=0.5.0
   - click
   - typing-extensions
   - openmm >=8.0.0,!=8.1.0,<8.4.0
   - openmmtools >=0.25.0
   - openmmforcefields
+  - openfe-analysis >=0.3.1
   - plugcli
   - tqdm
   # Testing deps
@@ -36,4 +37,4 @@ dependencies:
   - pip:
     - git+https://github.com/OpenFreeEnergy/openfe@main
     - git+https://github.com/OpenFreeEnergy/ofe-sphinx-theme@a45f3edd5bc3e973c1a01b577c71efa1b62a65d6
-    - git+https://github.com/openforcefield/openff-interchange@sltcap
+    - git+https://github.com/openforcefield/openff-interchange@main

src/pontibus/protocols/relative/__init__.py

@@ -0,0 +1,19 @@
+# This code is part of OpenFE and is licensed under the MIT license.
+# For details, see https://github.com/OpenFreeEnergy/openfe
+"""
+Run relative free energy calculations using OpenMM and OpenMMTools.
+"""
+
+from .hybridtop_protocol import (
+    HybridTopProtocol,
+    HybridTopProtocolResult,
+)
+from .hybridtop_units import HybridTopProtocolUnit
+from .settings import HybridTopProtocolSettings
+
+__all__ = [
+    "HybridTopProtocol",
+    "HybridTopProtocolSettings",
+    "HybridTopProtocolResult",
+    "HybridTopProtocolUnit",
+]

src/pontibus/protocols/relative/hybridtop_protocol.py

@@ -0,0 +1,145 @@
+# This code is part of OpenFE and is licensed under the MIT license.
+# For details, see https://github.com/OpenFreeEnergy/openfe
+"""Equilibrium Relative Free Energy methods using OpenMM and OpenMMTools in a
+Perses-like manner.
+This module implements the necessary methodology toolking to run calculate a
+ligand relative free energy transformation using OpenMM tools and one of the
+following methods:
+    - Hamiltonian Replica Exchange
+    - Self-adjusted mixture sampling
+    - Independent window sampling
+
+Acknowledgements
+----------------
+This Protocol is a subclass of the OpenFE RelativeHybridTopologyProtocol.
+This Protocol is based on, and leverages components originating from
+the Perses toolkit (https://github.com/choderalab/perses).
+"""
+
+import uuid
+
+from gufe import ChemicalSystem, ComponentMapping, ProtocolDAGResult, ProtocolUnit
+from gufe.settings import ThermoSettings
+from openfe.protocols.openmm_rfe.equil_rfe_methods import (
+    RelativeHybridTopologyProtocol,
+    RelativeHybridTopologyProtocolResult,
+    _validate_alchemical_components,
+)
+from openfe.protocols.openmm_rfe.equil_rfe_settings import (
+    AlchemicalSettings,
+    LambdaSettings,
+)
+from openfe.protocols.openmm_utils import system_validation
+from openfe.protocols.openmm_utils.omm_settings import (
+    IntegratorSettings,
+    MultiStateOutputSettings,
+    MultiStateSimulationSettings,
+    OpenFFPartialChargeSettings,
+    OpenMMEngineSettings,
+)
+from openff.units import unit
+
+from pontibus.protocols.relative.hybridtop_units import HybridTopProtocolUnit
+from pontibus.protocols.relative.settings import HybridTopProtocolSettings
+from pontibus.utils.settings import InterchangeFFSettings, PackmolSolvationSettings
+
+
+class HybridTopProtocolResult(RelativeHybridTopologyProtocolResult):
+    """
+    Results class for the HybridTopologyProtocol class.
+    Inherits from
+    :class:`openfe.protocols.openmm_rfe.RelativeHybridTopologyProtocolResult`.
+    """
+
+
+class HybridTopProtocol(RelativeHybridTopologyProtocol):
+    """
+    Relative Free Energy calculations using OpenMM and OpenMMTools.
+
+    Based on `Perses <https://github.com/choderalab/perses>`_
+
+    See Also
+    --------
+    :mod:`openfe.protocols`
+    :class:`openfe.protocols.openmm_rfe.RelativeHybridTopologyProtocol`
+    :class:`pontibus.protocols.relative.HybridTopSettings`
+    :class:`pontibus.protocols.relative.HybridTopResult`
+    :class:`pontibus.protocols.relative.HybridTopProtocolUnit`
+    """
+
+    result_cls = HybridTopProtocolResult
+    _settings_cls = HybridTopProtocolSettings
+    _settings: HybridTopProtocolSettings
+
+    @classmethod
+    def _default_settings(cls):
+        """A dictionary of initial settings for this creating this Protocol
+        These settings are intended as a suitable starting point for creating
+        an instance of this protocol.  It is recommended, however that care is
+        taken to inspect and customize these before performing a Protocol.
+        Returns
+        -------
+        Settings
+          a set of default settings
+        """
+        return HybridTopProtocolSettings(
+            protocol_repeats=3,
+            forcefield_settings=InterchangeFFSettings(),
+            thermo_settings=ThermoSettings(
+                temperature=298.15 * unit.kelvin,
+                pressure=1 * unit.bar,
+            ),
+            partial_charge_settings=OpenFFPartialChargeSettings(),
+            solvation_settings=PackmolSolvationSettings(),
+            alchemical_settings=AlchemicalSettings(softcore_LJ="gapsys"),
+            lambda_settings=LambdaSettings(),
+            simulation_settings=MultiStateSimulationSettings(
+                equilibration_length=1.0 * unit.nanosecond,
+                production_length=5.0 * unit.nanosecond,
+            ),
+            engine_settings=OpenMMEngineSettings(),
+            integrator_settings=IntegratorSettings(),
+            output_settings=MultiStateOutputSettings(),
+        )
+
+    def _create(
+        self,
+        stateA: ChemicalSystem,
+        stateB: ChemicalSystem,
+        mapping: ComponentMapping | list[ComponentMapping] | None,
+        extends: ProtocolDAGResult | None,
+    ) -> list[ProtocolUnit]:
+        if extends:  # pragma: no cover
+            raise NotImplementedError("Can't extend simulations yet")
+
+        # Get alchemical components & validate them + mapping
+        alchem_comps = system_validation.get_alchemical_components(stateA, stateB)
+        _validate_alchemical_components(alchem_comps, mapping)
+        ligandmapping = mapping[0] if isinstance(mapping, list) else mapping
+
+        # Validate solvent component
+        nonbond = self.settings.forcefield_settings.nonbonded_method
+        system_validation.validate_solvent(stateA, nonbond)
+
+        # Validate protein component
+        system_validation.validate_protein(stateA)
+
+        # actually create and return Units
+        Anames = ",".join(c.name for c in alchem_comps["stateA"])
+        Bnames = ",".join(c.name for c in alchem_comps["stateB"])
+        # our DAG has no dependencies, so just list units
+        n_repeats = self.settings.protocol_repeats
+        units = [
+            HybridTopProtocolUnit(
+                protocol=self,
+                stateA=stateA,
+                stateB=stateB,
+                ligandmapping=ligandmapping,
+                generation=0,
+                repeat_id=int(uuid.uuid4()),
+                name=f"{Anames} to {Bnames} repeat {i} generation 0",
+            )
+            for i in range(n_repeats)
+        ]
+
+        return units