Membrane support

environment.yml

@@ -19,8 +19,8 @@ dependencies:
   - openff-toolkit-base >=0.16.2
   - openff-units==0.3.1  # https://github.com/OpenFreeEnergy/openfe/pull/1374
   - openmm ~=8.2.0  # omit 8.3.0 and 8.3.1 due to https://github.com/openmm/openmm/pull/5069, unpin once we've qualified 8.3.2
-  - openmmforcefields >=0.15.0  # min needed for https://github.com/OpenFreeEnergy/openfe/pull/1695
-  - openmmtools >=0.25.3  # fix to support numpy >=2.3: https://github.com/choderalab/openmmtools/pull/793
+  - openmmforcefields >=0.15.1  # min needed for https://github.com/OpenFreeEnergy/openfe/pull/414
+  - openmmtools >=0.26  # fix to support membrane barostat: https://github.com/choderalab/openmmtools/pull/798
   - packaging
   - pandas
   - parmed >=4.3.1  # fix to support numpy >=2.3: https://github.com/ParmEd/ParmEd/pull/1387

news/membrane.rst

@@ -0,0 +1,23 @@
+**Added:**
+
+* This PR adds membrane support to the protocols PlainMDProtocol, RelativeHybridTopologyProtocol, SepTopProtocol, AbsoluteBindingProtocol.
+
+**Changed:**
+
+* <news item>
+
+**Deprecated:**
+
+* <news item>
+
+**Removed:**
+
+* <news item>
+
+**Fixed:**
+
+* <news item>
+
+**Security:**
+
+* <news item>

src/openfe/__init__.py

@@ -54,6 +54,7 @@
     LigandAtomMapping,
     NonTransformation,
     ProteinComponent,
+    ProteinMembraneComponent,
     SmallMoleculeComponent,
     SolventComponent,
     Transformation,

src/openfe/protocols/openmm_afe/abfe_units.py

@@ -16,7 +16,7 @@
 from gufe import (
     SolventComponent,
 )
-from gufe.components import Component
+from gufe.components import Component, SolvatedPDBComponent
 from openff.units import Quantity
 from openff.units.openmm import to_openmm
 from openmm import System
@@ -70,6 +70,13 @@ def _get_components(self):
         # an error will have been raised when calling `validate_solvent`
         # in the Protocol's `_create`.
         # Similarly we don't need to check prot_comp
+
+        # If there is an SolvatedPDBComponent, we set the solv_comp
+        # in the complex to the SolvatedPDBComponent, as the SolventComponent
+        # is only used in the solvent leg
+        if isinstance(prot_comp, SolvatedPDBComponent):
+            solv_comp = prot_comp
+
         return alchem_comps, solv_comp, prot_comp, off_comps
 
 
@@ -106,7 +113,7 @@ def _get_settings(self) -> dict[str, SettingsBaseModel]:
         settings["alchemical_settings"] = prot_settings.alchemical_settings
         settings["lambda_settings"] = prot_settings.complex_lambda_settings
         settings["engine_settings"] = prot_settings.engine_settings
-        settings["integrator_settings"] = prot_settings.integrator_settings
+        settings["integrator_settings"] = prot_settings.complex_integrator_settings
         settings["equil_simulation_settings"] = prot_settings.complex_equil_simulation_settings
         settings["equil_output_settings"] = prot_settings.complex_equil_output_settings
         settings["simulation_settings"] = prot_settings.complex_simulation_settings
@@ -475,7 +482,7 @@ def _get_settings(self) -> dict[str, SettingsBaseModel]:
         settings["alchemical_settings"] = prot_settings.alchemical_settings
         settings["lambda_settings"] = prot_settings.solvent_lambda_settings
         settings["engine_settings"] = prot_settings.engine_settings
-        settings["integrator_settings"] = prot_settings.integrator_settings
+        settings["integrator_settings"] = prot_settings.solvent_integrator_settings
         settings["equil_simulation_settings"] = prot_settings.solvent_equil_simulation_settings
         settings["equil_output_settings"] = prot_settings.solvent_equil_output_settings
         settings["simulation_settings"] = prot_settings.solvent_simulation_settings

src/openfe/protocols/openmm_afe/base_afe_units.py

@@ -29,6 +29,7 @@
 import openmm
 import openmmtools
 from gufe import (
+    BaseSolventComponent,
     ProteinComponent,
     SmallMoleculeComponent,
     SolventComponent,
@@ -356,7 +357,7 @@ def _assign_partial_charges(
     @staticmethod
     def _get_system_generator(
         settings: dict[str, SettingsBaseModel],
-        solvent_component: SolventComponent | None,
+        solvent_component: BaseSolventComponent | None,
         openff_molecules: list[OFFMolecule],
         ffcache: pathlib.Path | None,
     ) -> SystemGenerator:
@@ -368,8 +369,8 @@ def _get_system_generator(
         ----------
         settings : dict[str, SettingsBaseModel]
           A dictionary of settings object for the unit.
-        solvent_comp : SolventComponent | None
-          The solvent component of this system, if there is one.
+        solvent_comp : BaseSolventComponent | None
+          The BaseSolventComponent of this system, if there is one.
         openff_molecules : list[openff.toolkit.Molecule] | None
           A list of OpenFF Molecules to generate templates for, if any.
         ffcache : pathlib.Path | None
@@ -401,7 +402,7 @@ def _get_system_generator(
     @staticmethod
     def _get_modeller(
         protein_component: ProteinComponent | None,
-        solvent_component: SolventComponent | None,
+        solvent_component: BaseSolventComponent | None,
         small_mols: dict[SmallMoleculeComponent, OFFMolecule],
         system_generator: SystemGenerator,
         solvation_settings: BaseSolvationSettings,
@@ -414,8 +415,8 @@ def _get_modeller(
         ----------
         protein_component : ProteinComponent | None
           Protein Component, if it exists.
-        solvent_component : SolventComponent | None
-          Solvent Component, if it exists.
+        solvent_component : BaseSolventComponent | None
+          Base Solvent Component, if it exists.
         small_mols : dict[SmallMoleculeComponent, openff.toolkit.Molecule]
           Dictionary of OpenFF Molecules to add, keyed by
           SmallMoleculeComponent.
@@ -447,7 +448,7 @@ def _get_omm_objects(
         self,
         settings: dict[str, SettingsBaseModel],
         protein_component: ProteinComponent | None,
-        solvent_component: SolventComponent | None,
+        solvent_component: BaseSolventComponent | None,
         small_mols: dict[SmallMoleculeComponent, OFFMolecule],
     ) -> tuple[
         app.Topology,
@@ -465,8 +466,8 @@ def _get_omm_objects(
           Protocol settings
         protein_component : ProteinComponent | None
           Protein component for the system.
-        solvent_component : SolventComponent | None
-          Solvent component for the system.
+        solvent_component : BaseSolventComponent | None
+          BaseSolventComponent for the system.
         small_mols : dict[str, openff.toolkit.Molecule]
           Dictionary of SmallMoleculeComponents and OpenFF Molecules
           defining the ligands to be added to the system
@@ -847,7 +848,7 @@ def _get_states(
         box_vectors: openmm.unit.Quantity,
         thermodynamic_settings: ThermoSettings,
         lambdas: dict[str, list[float]],
-        solvent_component: SolventComponent | None,
+        solvent_component: BaseSolventComponent | None,
         alchemically_restrained: bool,
     ) -> tuple[list[SamplerState], list[ThermodynamicState]]:
         """
@@ -866,8 +867,8 @@ def _get_states(
           Settings controlling the thermodynamic parameters.
         lambdas : dict[str, list[float]]
           A dictionary of lambda scales.
-        solvent_component : SolventComponent | None
-          The solvent component of the system, if there is one.
+        solvent_component : BaseSolventComponent | None
+          The base solvent component of the system, if there is one.
         alchemically_restrained : bool
           Whether or not the system requires a control parameter
           for any alchemical restraints.

src/openfe/protocols/openmm_afe/equil_afe_settings.py

@@ -387,10 +387,15 @@ def must_be_positive(cls, v):
     """
 
     # Sampling State defining things
-    integrator_settings: IntegratorSettings
+    solvent_integrator_settings: IntegratorSettings
     """
     Settings for controlling the integrator, such as the timestep and
-    barostat settings.
+    barostat settings in the solvent.
+    """
+    complex_integrator_settings: IntegratorSettings
+    """
+    Settings for controlling the integrator, such as the timestep and
+    barostat settings in the complex.
     """
 
     # Simulation run settings

src/openfe/protocols/openmm_afe/equil_binding_afe_method.py

@@ -32,8 +32,10 @@
 
 import gufe
 from gufe import (
+    BaseSolventComponent,
     ChemicalSystem,
     ProteinComponent,
+    ProteinMembraneComponent,
     SmallMoleculeComponent,
     SolventComponent,
     settings,
@@ -171,7 +173,8 @@ def _default_settings(cls):
             ),
             solvent_solvation_settings=OpenMMSolvationSettings(),
             engine_settings=OpenMMEngineSettings(),
-            integrator_settings=IntegratorSettings(),
+            solvent_integrator_settings=IntegratorSettings(),
+            complex_integrator_settings=IntegratorSettings(),
             restraint_settings=BoreschRestraintSettings(),
             solvent_equil_simulation_settings=MDSimulationSettings(
                 equilibration_length_nvt=0.1 * offunit.nanosecond,
@@ -208,6 +211,53 @@ def _default_settings(cls):
         )
         # fmt: on
 
+    @classmethod
+    def _adaptive_settings(
+        cls,
+        stateA: ChemicalSystem,
+        stateB: ChemicalSystem,
+        initial_settings: None | AbsoluteBindingSettings = None,
+    ) -> AbsoluteBindingSettings:
+        """
+        Get the recommended OpenFE settings for this protocol based on the input states involved in the
+        transformation.
+
+        These are intended as a suitable starting point for creating an instance of this protocol, which can be further
+        customized before performing a Protocol.
+
+        Parameters
+        ----------
+        stateA : ChemicalSystem
+            The initial state of the transformation.
+        stateB : ChemicalSystem
+            The final state of the transformation.
+        initial_settings : None | AbsoluteBindingSettings, optional
+            Initial settings to base the adaptive settings on. If None, default settings are used.
+
+        Returns
+        -------
+        AbsoluteBindingSettings
+            The recommended settings for this protocol based on the input states.
+        """
+        # use initial settings or default settings
+        if initial_settings is not None:
+            protocol_settings = initial_settings.model_copy(deep=True)
+        else:
+            protocol_settings = cls.default_settings()
+
+        # adapt the barostat and lipid forcefield based on the ProteinComponent
+        if stateA.contains(ProteinMembraneComponent):
+            protocol_settings.complex_integrator_settings.barostat = "MonteCarloMembraneBarostat"
+            protocol_settings.forcefield_settings.forcefields = [
+                "amber/ff14SB.xml",
+                "amber/tip3p_standard.xml",
+                "amber/tip3p_HFE_multivalent.xml",
+                "amber/lipid17_merged.xml",
+                "amber/phosaa10.xml",
+            ]
+
+        return protocol_settings
+
     @staticmethod
     def _validate_endstates(
         stateA: ChemicalSystem,
@@ -400,6 +450,11 @@ def _validate(
         # Use the more complete system validation solvent checks
         system_validation.validate_solvent(stateA, nonbonded_method)
 
+        # Validate the barostat used in combination with the protein component
+        system_validation.validate_barostat(
+            stateA, self.settings.complex_integrator_settings.barostat
+        )
+
         # Validate solvation settings
         settings_validation.validate_openmm_solvation_settings(
             self.settings.solvent_solvation_settings
@@ -409,9 +464,15 @@ def _validate(
         )
 
         # Validate integrator things
+        # We validate the timstep for both the complex & solvent settings
+        settings_validation.validate_timestep(
+            self.settings.forcefield_settings.hydrogen_mass,
+            self.settings.complex_integrator_settings.timestep,
+        )
+
         settings_validation.validate_timestep(
             self.settings.forcefield_settings.hydrogen_mass,
-            self.settings.integrator_settings.timestep,
+            self.settings.solvent_integrator_settings.timestep,
         )
 
     def _create(

src/openfe/protocols/openmm_md/plain_md_methods.py

@@ -23,6 +23,7 @@
 import openmm
 import openmm.unit as omm_unit
 from gufe import (
+    BaseSolventComponent,
     ChemicalSystem,
     SmallMoleculeComponent,
     settings,
@@ -32,6 +33,7 @@
 from openff.toolkit.topology import Molecule as OFFMolecule
 from openff.units import Quantity, unit
 from openff.units.openmm import from_openmm, to_openmm
+from openmm import MonteCarloBarostat, MonteCarloMembraneBarostat
 
 from openfe.protocols.openmm_md.plain_md_settings import (
     IntegratorSettings,
@@ -180,14 +182,17 @@ def _create(
         # Validate protein component
         system_validation.validate_protein(stateA)
 
+        # Validate the barostat used in combination with the protein component
+        system_validation.validate_barostat(stateA, self.settings.integrator_settings.barostat)
+
         # Validate solvation settings
         settings_validation.validate_openmm_solvation_settings(self.settings.solvation_settings)
 
         # actually create and return Units
         # TODO: Deal with multiple ProteinComponents
         solvent_comp, protein_comp, small_mols = system_validation.get_components(stateA)
 
-        system_name = "Solvent MD" if solvent_comp is not None else "Vacuum MD"
+        system_name = "Solvent MD" if stateA.contains(BaseSolventComponent) else "Vacuum MD"
 
         for comp in [protein_comp] + small_mols:
             if comp is not None:
@@ -363,9 +368,9 @@ def _run_MD(
                 logger.info("Running NVT equilibration")
 
             # Set barostat frequency to zero for NVT
-            for x in simulation.context.getSystem().getForces():
-                if x.getName() == "MonteCarloBarostat":
-                    x.setFrequency(0)
+            for force in simulation.context.getSystem().getForces():
+                if isinstance(force, (MonteCarloBarostat, MonteCarloMembraneBarostat)):
+                    force.setFrequency(0)
 
             simulation.context.setVelocitiesToTemperature(to_openmm(temperature))
 
@@ -397,9 +402,9 @@ def _run_MD(
         simulation.context.setVelocitiesToTemperature(to_openmm(temperature))
 
         # Enable the barostat for NPT
-        for x in simulation.context.getSystem().getForces():
-            if x.getName() == "MonteCarloBarostat":
-                x.setFrequency(barostat_frequency.m)
+        for force in simulation.context.getSystem().getForces():
+            if isinstance(force, (MonteCarloBarostat, MonteCarloMembraneBarostat)):
+                force.setFrequency(barostat_frequency.m)
 
         t0 = time.time()
         simulation.step(equil_steps_npt)
@@ -585,7 +590,7 @@ def run(
         solvent_comp, protein_comp, small_mols = system_validation.get_components(stateA)
 
         # 1. Create stateA system
-        # Create a dictionary of OFFMol for each SMC for bookeeping
+        # Create a dictionary of OFFMol for each SMC for bookkeeping
         smc_components: dict[SmallMoleculeComponent, OFFMolecule]
 
         smc_components = {i: i.to_openff() for i in small_mols}
@@ -637,6 +642,9 @@ def run(
 
         # f. Save pdb of entire system
         if output_settings.preminimized_structure:
+            # roundtrip box vectors to remove vec3 issues
+            box = to_openmm(from_openmm(stateA_system.getDefaultPeriodicBoxVectors()))
+            stateA_topology.setPeriodicBoxVectors(box)
             with open(shared_basepath / output_settings.preminimized_structure, "w") as f:
                 openmm.app.PDBFile.writeFile(
                     stateA_topology, stateA_positions, file=f, keepIds=True