Membrane support

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,12 @@ 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 None, as it is only used in the solvent leg
+        if isinstance(prot_comp, SolvatedPDBComponent):
+            solv_comp = None
+
         return alchem_comps, solv_comp, prot_comp, off_comps
 
 
@@ -105,7 +111,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 +480,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/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_protein_barostat(
+            stateA, self.settings.complex_integrator_settings.barostat
+        )
+
         # Validate solvation settings
         settings_validation.validate_openmm_solvation_settings(
             self.settings.solvent_solvation_settings
@@ -411,7 +418,11 @@ def _validate(
         # Validate integrator things
         settings_validation.validate_timestep(
             self.settings.forcefield_settings.hydrogen_mass,
-            self.settings.integrator_settings.timestep,
+            self.settings.complex_integrator_settings.timestep,
+        )
+        settings_validation.validate_timestep(
+            self.settings.forcefield_settings.hydrogen_mass,
+            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_protein_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}
@@ -634,6 +641,11 @@ def run(
                 stateA_topology,
                 molecules=[s.to_openff() for s in small_mols],
             )
+        box = [
+            openmm.Vec3(*v.value_in_unit(omm_unit.nanometer))
+            for v in stateA_system.getDefaultPeriodicBoxVectors()
+        ] * omm_unit.nanometer
+        stateA_topology.setPeriodicBoxVectors(box)
 
         # f. Save pdb of entire system
         if output_settings.preminimized_structure:

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,9 +551,11 @@ 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():
+        known_barostats = ["MonteCarloBarostat", "MonteCarloMembraneBarostat"]
+        present_barostat = [i for i in self._old_system_forces.keys() if i in known_barostats]
+        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)
 
         # Copy over the box vectors from the old system

src/openfe/protocols/openmm_rfe/hybridtop_protocols.py

@@ -25,6 +25,7 @@
     ComponentMapping,
     LigandAtomMapping,
     ProteinComponent,
+    ProteinMembraneComponent,
     SmallMoleculeComponent,
     SolventComponent,
     settings,
@@ -199,6 +200,17 @@ def _adaptive_settings(
         if stateA.contains(ProteinComponent) and stateB.contains(ProteinComponent):
             protocol_settings.solvation_settings.solvent_padding = 1 * offunit.nanometer
 
+        # adapt the barostat based on the system components
+        if stateA.contains(ProteinMembraneComponent) and stateB.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_protein_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

@@ -169,7 +169,11 @@ def _pre_equilibrate(
 
     # Don't do anything if we're doing a dry run
     if dry:
-        return positions, system.getDefaultPeriodicBoxVectors()
+        box = [
+            openmm.Vec3(*v.value_in_unit(omm_unit.nanometer))
+            for v in system.getDefaultPeriodicBoxVectors()
+        ] * omm_unit.nanometer
+        return positions, box
 
     # TODO: Refactor this part to live outside the method call
     # We have to modify the output settings to have different output
@@ -1374,4 +1378,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,
+                }
+            }