[ WIP ] Membrane support prototype

docs/environment.yaml

@@ -26,7 +26,7 @@ dependencies:
 - threadpoolctl
 - tqdm
 - pip:
-  - git+https://github.com/OpenFreeEnergy/gufe@main
+  - git+https://github.com/OpenFreeEnergy/gufe@membrane_prototype
   - git+https://github.com/OpenFreeEnergy/ofe-sphinx-theme@v0.3.1
   # pip install these so that we can make sure docs build on main while these packages' docs are under development
   - git+https://github.com/OpenFreeEnergy/kartograf@main

environment.yml

@@ -19,8 +19,6 @@ 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.0
   - packaging
   - pandas
   - perses>=0.10.3
@@ -52,4 +50,6 @@ dependencies:
   # Control blas/openmp threads
   - threadpoolctl
   - pip:
-    - git+https://github.com/OpenFreeEnergy/gufe@main
+    - git+https://github.com/OpenFreeEnergy/gufe@membrane_prototype
+    - git+https://github.com/hannahbaumann/openmmtools@update_disabling_barostat
+    - git+https://github.com/hannahbaumann/openmmforcefields@systemgenerator_bastostat

openfe/__init__.py

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

openfe/protocols/openmm_afe/equil_binding_afe_method.py

@@ -39,6 +39,7 @@
 from gufe import (
     ChemicalSystem,
     ProteinComponent,
+    ProteinMembraneComponent,
     SmallMoleculeComponent,
     SolventComponent,
     settings,
@@ -655,13 +656,19 @@ def _validate_endstates(
           If stateB contains any unique Components.
           If the alchemical species is charged.
         """
-        if not (stateA.contains(ProteinComponent) and stateB.contains(ProteinComponent)):
-            errmsg = "No ProteinComponent found"
-            raise ValueError(errmsg)
+        components = (ProteinComponent, ProteinMembraneComponent)
+        system_validation.require_components(
+            systems=[stateA, stateB],
+            component_types=components,
+            msg="No ProteinComponent or ProteinMembraneComponent found",
+        )
 
-        if not (stateA.contains(SolventComponent) and stateB.contains(SolventComponent)):
-            errmsg = "No SolventComponent found"
-            raise ValueError(errmsg)
+        components = (SolventComponent, ProteinMembraneComponent)
+        system_validation.require_components(
+            systems=[stateA, stateB],
+            component_types=components,
+            msg="No SolventComponent or ProteinMembraneComponent found",
+        )
 
         # Needs gufe 1.3
         diff = stateA.component_diff(stateB)
@@ -944,6 +951,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 a ProteinMembraneComponent, we set the SolventComponent
+        # in the complex to None, as it is only used in the solvent leg
+        if isinstance(prot_comp, ProteinMembraneComponent):
+            solv_comp = None
+
         return alchem_comps, solv_comp, prot_comp, off_comps
 
     def _handle_settings(self) -> dict[str, SettingsBaseModel]:

openfe/protocols/openmm_md/plain_md_methods.py

@@ -24,6 +24,7 @@
 import openmm.unit as omm_unit
 from gufe import (
     ChemicalSystem,
+    ProteinMembraneComponent,
     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,
@@ -187,7 +189,11 @@ def _create(
         # 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 solvent_comp is not None or isinstance(protein_comp, ProteinMembraneComponent)
+            else "Vacuum MD"
+        )
 
         for comp in [protein_comp] + small_mols:
             if comp is not None:
@@ -363,9 +369,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 +403,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 +591,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 +640,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:

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

openfe/protocols/openmm_rfe/equil_rfe_methods.py

@@ -45,6 +45,7 @@
     ComponentMapping,
     LigandAtomMapping,
     ProteinComponent,
+    ProteinMembraneComponent,
     SmallMoleculeComponent,
     SolventComponent,
     settings,
@@ -1182,20 +1183,26 @@ def run(
 
             sampler.energy_context_cache = energy_context_cache
             sampler.sampler_context_cache = sampler_context_cache
-
             if not dry:  # pragma: no-cover
                 # minimize
+                print(
+                    "box1", sampler._thermodynamic_states[0].system.getDefaultPeriodicBoxVectors()
+                )
                 if verbose:
                     self.logger.info("Running minimization")
 
                 sampler.minimize(max_iterations=sampler_settings.minimization_steps)
-
+                print(
+                    "box2", sampler._thermodynamic_states[0].system.getDefaultPeriodicBoxVectors()
+                )
                 # equilibrate
                 if verbose:
                     self.logger.info("Running equilibration phase")
 
                 sampler.equilibrate(int(equil_steps / steps_per_iteration))
-
+                print(
+                    "box3", sampler._thermodynamic_states[0].system.getDefaultPeriodicBoxVectors()
+                )
                 # production
                 if verbose:
                     self.logger.info("Running production phase")

openfe/protocols/openmm_septop/base.py

@@ -168,7 +168,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
@@ -1373,4 +1377,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,
+                }
+            }