Allow for any barostat in SystemGenerator

openmmforcefields/generators/system_generators.py

@@ -303,25 +303,15 @@ def _modify_forces(self, system):
         """
         Add barostat and modify forces if requested.
         """
-
         # Add barostat if requested and the system uses periodic boundary conditions
         if (self.barostat is not None) and system.usesPeriodicBoundaryConditions():
             import numpy as np
-            import openmm
+            import copy
 
             MAXINT = np.iinfo(np.int32).max
 
-            # Determine pressure, temperature, and frequency
-            pressure = self.barostat.getDefaultPressure()
-            if hasattr(self.barostat, "getDefaultTemperature"):
-                temperature = self.barostat.getDefaultTemperature()
-            else:
-                temperature = self.barostat.getTemperature()
-            frequency = self.barostat.getFrequency()
-
-            # Create the barostat
-            # TODO: Make sure we can support other kinds of barostats?
-            barostat = openmm.MonteCarloBarostat(pressure, temperature, frequency)
+            # Get the barostat
+            barostat = copy.deepcopy(self.barostat)
             seed = np.random.randint(MAXINT)
             barostat.setRandomNumberSeed(seed)
             system.addForce(barostat)

openmmforcefields/tests/test_system_generator.py

@@ -7,6 +7,7 @@
 import pytest
 from openff.toolkit.topology import Molecule
 from openmm.app import LJPME, PME, CutoffNonPeriodic, Modeller, PDBFile
+from openmm import unit, MonteCarloBarostat, MonteCarloMembraneBarostat
 
 from openmmforcefields.generators import SystemGenerator
 from openmmforcefields.utils import Timer, get_data_filename
@@ -139,21 +140,52 @@ def test_create(self):
         # Create an empty system generator
         SystemGenerator()
 
-    def test_barostat(self):
-        """Test that barostat addition works correctly"""
+    @pytest.mark.parametrize(
+        "barostat_class, args",
+        [
+            # MonteCarloBarostat
+            (
+                MonteCarloBarostat,
+                [0.95 * unit.atmospheres, 301.0 * unit.kelvin, 23],
+            ),
+            # MonteCarloMembraneBarostat
+            (
+                MonteCarloMembraneBarostat,
+                [
+                    1.0 * unit.atmospheres,
+                    10.0 * unit.millinewton / unit.meter,
+                    301.0 * unit.kelvin,
+                    MonteCarloMembraneBarostat.XYIsotropic,
+                    MonteCarloMembraneBarostat.ZFree,
+                    23,
+                ],
+            ),
+        ],
+    )
+    def test_barostat(self, barostat_class, args):
+        """Test that different barostats are correctly applied to the system"""
         # Create a protein SystemGenerator
         generator = SystemGenerator(forcefields=self.amber_forcefields)
 
-        # Create a template barostat
-        from openmm import MonteCarloBarostat, unit
-
-        pressure = 0.95 * unit.atmospheres
-        temperature = 301.0 * unit.kelvin
-        frequency = 23
-        generator.barostat = MonteCarloBarostat(pressure, temperature, frequency)
+        # Create the barostat
+        generator.barostat = barostat_class(*args)
+
+        # Derive expected values based on barostat type
+        if barostat_class is MonteCarloBarostat:
+            expected = {
+                "pressure": args[0],
+                "temperature": args[1],
+                "frequency": args[2],
+            }
+        else:  # MonteCarloMembraneBarostat
+            expected = {
+                "pressure": args[0],
+                "surface_tension": args[1],
+                "temperature": args[2],
+                "frequency": args[-1],
+            }
 
         # Load a PDB file
-
         pdb_filename = get_data_filename(os.path.join("perses_jacs_systems", "mcl1", "MCL1_protein.pdb"))
         pdbfile = PDBFile(pdb_filename)
 
@@ -176,13 +208,18 @@ def test_barostat(self):
 
         # Check barostat is present
         forces = {force.__class__.__name__: force for force in system.getForces()}
-        assert "MonteCarloBarostat" in forces.keys()
+        name = barostat_class.__name__
+        assert name in forces, f"{name} not found in system forces"
 
         # Check barostat parameters
-        force = forces["MonteCarloBarostat"]
-        assert force.getDefaultPressure() == pressure
-        assert force.getDefaultTemperature() == temperature
-        assert force.getFrequency() == frequency
+        force = forces[name]
+        assert force.getDefaultTemperature() == expected["temperature"]
+        assert force.getDefaultPressure() == expected["pressure"]
+        assert force.getFrequency() == expected["frequency"]
+
+        # Conditional check
+        if hasattr(force, "getDefaultSurfaceTension"):
+            assert force.getDefaultSurfaceTension() == expected["surface_tension"]
 
     @pytest.mark.parametrize(
         "small_molecule_forcefield",