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

src/openfe/__init__.py

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

src/openfe/protocols/openmm_afe/abfe_units.py

@@ -15,7 +15,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
@@ -69,6 +69,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
 
 
@@ -105,7 +112,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
@@ -474,7 +481,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,6 +32,7 @@
 
 import gufe
 from gufe import (
+    BaseSolventComponent,
     ChemicalSystem,
     ProteinComponent,
     SmallMoleculeComponent,
@@ -171,7 +172,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,
@@ -400,6 +402,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 +416,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,19 @@ 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 +370,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 +404,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 +592,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}
@@ -638,6 +645,11 @@ def run(
         # f. Save pdb of entire system
         if output_settings.preminimized_structure:
             with open(shared_basepath / output_settings.preminimized_structure, "w") as f:
+                # roundtrip box vectors to remove vec3 issues
+                box = to_openmm(from_openmm(
+                    stateA_system.getDefaultPeriodicBoxVectors()
+                ))
+                stateA_topology.setPeriodicBoxVectors(box)
                 openmm.app.PDBFile.writeFile(
                     stateA_topology, stateA_positions, file=f, keepIds=True
                 )

src/openfe/protocols/openmm_rfe/_rfe_utils/relative.py

@@ -470,9 +470,12 @@ def _check_and_store_system_forces(self):
 
         def _check_unknown_forces(forces, system_name):
             # TODO: double check that CMMotionRemover is ok being here
-            known_forces = {'HarmonicBondForce', 'HarmonicAngleForce',
-                            'PeriodicTorsionForce', 'NonbondedForce',
-                            'MonteCarloBarostat', 'CMMotionRemover', 'CMAPTorsionForce'}
+            known_forces = {
+                'HarmonicBondForce', 'HarmonicAngleForce',
+                'PeriodicTorsionForce', 'NonbondedForce',
+                'MonteCarloBarostat', 'CMMotionRemover',
+                'CMAPTorsionForce', 'MonteCarloMembraneBarostat',
+            }
 
             force_names = forces.keys()
             unknown_forces = set(force_names) - set(known_forces)
@@ -548,10 +551,17 @@ def _handle_box(self):
         """
         # Check that if there is a barostat in the old system,
         # it is added to the hybrid system
-        if "MonteCarloBarostat" in self._old_system_forces.keys():
+        present_barostat = [
+            i for i in self._old_system_forces.keys()
+            if i in ["MonteCarloBarostat", "MonteCarloMembraneBarostat"]
+        ]
+        if len(present_barostat) == 1:
             barostat = copy.deepcopy(
-                self._old_system_forces["MonteCarloBarostat"])
+                self._old_system_forces[present_barostat[0]])
             self._hybrid_system.addForce(barostat)
+        if len(present_barostat) > 1:
+            errmsg = "More than 1 barostat are present which is not supported"
+            raise ValueError(errmsg)
 
         # Copy over the box vectors from the old system
         box_vectors = self._old_system.getDefaultPeriodicBoxVectors()

src/openfe/protocols/openmm_rfe/hybridtop_protocols.py

@@ -25,6 +25,7 @@
     ComponentMapping,
     LigandAtomMapping,
     ProteinComponent,
+    ProteinMembraneComponent,
     SmallMoleculeComponent,
     SolventComponent,
     settings,
@@ -196,9 +197,20 @@ def _adaptive_settings(
             protocol_settings.lambda_settings.lambda_windows = 22
 
         # adapt the solvation padding based on the system components
-        if stateA.contains(ProteinComponent) and stateB.contains(ProteinComponent):
+        if stateA.contains(ProteinComponent):
             protocol_settings.solvation_settings.solvent_padding = 1 * offunit.nanometer
 
+        # adapt the barostat based on the system components
+        if stateA.contains(ProteinMembraneComponent):
+            protocol_settings.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
@@ -539,6 +551,11 @@ def _validate(
         # 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 charge difference
         # Note: validation depends on the mapping & solvent component checks
         if stateA.contains(SolventComponent):

src/openfe/protocols/openmm_septop/base.py

@@ -1374,4 +1374,11 @@ def run(
                 **unit_result_dict,
             }
         else:
-            return {"debug": {"sampler": sampler}}
+            return {
+                "debug": {
+                    "sampler": sampler,
+                    "alchem_system": system,
+                    "selection_indices": self.selection_indices,
+                    "positions": equil_positions,
+                }
+            }