ThermodynamicState support for more barostats

docs/releasehistory.rst

@@ -1,6 +1,11 @@
 Release History
 ***************
 
+Enhancements
+------------
+- Added support for anisotropic and membrane barostats in `ThermodynamicState`
+- Added support for platform properties in ContextCache (e.g. for mixed and double precision CUDA in multistate sampler)
+
 0.18.3 - Storage enhancements and bugfixes
 ==========================================
 

openmmtools/cache.py

@@ -229,6 +229,9 @@ class ContextCache(object):
     platform : simtk.openmm.Platform, optional
         The OpenMM platform to use to create Contexts. If None, OpenMM
         tries to select the fastest one available (default is None).
+    platform_properties : dict, optional
+        A dictionary of platform properties for the OpenMM platform.
+        Only valid if the platform is not None (default is None).
     **kwargs
         Parameters to pass to the underlying LRUCache constructor such
         as capacity and time_to_live.
@@ -303,8 +306,11 @@ class ContextCache(object):
 
     """
 
-    def __init__(self, platform=None, **kwargs):
+    def __init__(self, platform=None, platform_properties=None, **kwargs):
         self._platform = platform
+        self._platform_properties = platform_properties
+        if platform_properties is not None and platform is None:
+            raise ValueError("To set platform_properties, you need to also specify the platform.")
         self._lru = LRUCache(**kwargs)
 
     def __len__(self):
@@ -429,7 +435,7 @@ def get_context(self, thermodynamic_state, integrator=None):
             try:
                 context = self._lru[context_id]
             except KeyError:
-                context = thermodynamic_state.create_context(integrator, self._platform)
+                context = thermodynamic_state.create_context(integrator, self._platform, self._platform_properties)
             self._lru[context_id] = context
         context_integrator = context.getIntegrator()
 
@@ -447,13 +453,17 @@ def __getstate__(self):
         else:
             platform_serialization = None
         return dict(platform=platform_serialization, capacity=self.capacity,
-                    time_to_live=self.time_to_live)
+                    time_to_live=self.time_to_live, platform_properties=self._platform_properties)
 
     def __setstate__(self, serialization):
         if serialization['platform'] is None:
             self._platform = None
         else:
             self._platform = openmm.Platform.getPlatformByName(serialization['platform'])
+        if not 'platform_properties' in serialization:
+            self._platform_properties = None
+        else:
+            self._platform_properties = serialization["platform_properties"]
         self._lru = LRUCache(serialization['capacity'], serialization['time_to_live'])
 
     # -------------------------------------------------------------------------

openmmtools/states.py

@@ -242,6 +242,23 @@ def _box_vectors_volume(box_vectors):
     return np.linalg.det(box_matrix) * a.unit**3
 
 
+def _box_vectors_area(box_vectors):
+    """Return the xy-area of the box vectors.
+
+    Parameters
+    ----------
+    box_vectors : simtk.unit.Quantity
+        Vectors defining the box.
+
+    Returns
+    -------
+
+    area : simtk.unit.Quantity
+        The box area in units of length^2.
+    """
+    return box_vectors[0][0] * box_vectors[1][1]
+
+
 # =============================================================================
 # CUSTOM EXCEPTIONS
 # =============================================================================
@@ -282,7 +299,7 @@ class ThermodynamicsError(Exception):
      NO_BAROSTAT,
      UNSUPPORTED_BAROSTAT,
      UNSUPPORTED_ANISOTROPIC_BAROSTAT,
-     UNSUPPORTED_MEMBRANE_BAROSTAT,
+     SURFACE_TENSION_NOT_SUPPORTED,
      INCONSISTENT_BAROSTAT,
      BAROSTATED_NONPERIODIC,
      INCONSISTENT_INTEGRATOR,
@@ -298,8 +315,8 @@ class ThermodynamicsError(Exception):
         UNSUPPORTED_BAROSTAT: "Found unsupported barostat {} in system.",
         UNSUPPORTED_ANISOTROPIC_BAROSTAT:
             "MonteCarloAnisotropicBarostat is only supported if the pressure along all scaled axes is the same.",
-        UNSUPPORTED_MEMBRANE_BAROSTAT:
-            "MonteCarloAnisotropicBarostat is only supported if the surface tension is zero.",
+        SURFACE_TENSION_NOT_SUPPORTED:
+            "Surface tension can only be set for states that have a system with a MonteCarloBarostat.",
         NO_BAROSTAT: "System does not have a barostat specifying the pressure.",
         INCONSISTENT_BAROSTAT: "System barostat is inconsistent with thermodynamic state.",
         BAROSTATED_NONPERIODIC: "Non-periodic systems cannot have a barostat.",
@@ -368,11 +385,19 @@ class ThermodynamicState(object):
     state. It can be used to create new OpenMM Contexts, or to convert
     an existing Context to this particular thermodynamic state.
 
-    Only NVT and NPT ensembles are supported. The temperature must
+    NVT, NPT and NPgammaT ensembles are supported. The temperature must
     be specified in the constructor, either implicitly via a thermostat
     force in the system, or explicitly through the temperature
     parameter, which overrides an eventual thermostat indication.
 
+    To set a ThermodynamicState up in the NPgammaT ensemble, the system
+    passed to the constructor has to have a MonteCarloMembraneBarostat.
+
+    To set a ThermodynamicState up with anisotropic pressure control,
+    the system passed to the constructor has to have a MonteCarloAnisotropicBarostat.
+    Currently the MonteCarloAnisotropicBarostat is only supported if
+    the pressure is equal for all axes that are under pressure control.
+
     Parameters
     ----------
     system : simtk.openmm.System
@@ -394,6 +419,7 @@ class ThermodynamicState(object):
     system
     temperature
     pressure
+    surface_tension
     volume
     n_particles
 
@@ -756,6 +782,26 @@ def is_periodic(self):
         """True if the system is in a periodic box (read-only)."""
         return self._standard_system.usesPeriodicBoundaryConditions()
 
+    @property
+    def surface_tension(self):
+        """Surface tension; None if the barostat is not a MonteCarloMembraneBarostat"""
+        barostat = self._find_barostat(self._standard_system)
+        if isinstance(barostat, openmm.MonteCarloMembraneBarostat):
+            return barostat.getDefaultSurfaceTension()
+        else:
+            return None
+
+    @surface_tension.setter
+    def surface_tension(self, gamma):
+        # working around a bug in the unit conversion https://github.com/openmm/openmm/issues/2406
+        if isinstance(gamma, unit.Quantity):
+            gamma = gamma.value_in_unit(unit.bar * unit.nanometer)
+        barostat = self._find_barostat(self._standard_system)
+        if isinstance(barostat, openmm.MonteCarloMembraneBarostat):
+            barostat.setDefaultSurfaceTension(gamma)
+        else:
+            raise ThermodynamicsError(ThermodynamicsError.SURFACE_TENSION_NOT_SUPPORTED)
+
     def reduced_potential(self, context_state):
         """Reduced potential in this thermodynamic state.
 
@@ -772,15 +818,17 @@ def reduced_potential(self, context_state):
 
         Notes
         -----
-        The reduced potential is defined as in Ref. [1]
+        The reduced potential is defined as in Ref. [1],
+        with a additional term for the surface tension
 
-        u = \beta [U(x) + p V(x) + \mu N(x)]
+        u = \beta [U(x) + p V(x) + \mu N(x) - \gamma A]
 
         where the thermodynamic parameters are
 
         \beta = 1/(kB T) is the inverse temperature
         p is the pressure
         \mu is the chemical potential
+        \gamma is the surface tension
 
         and the configurational properties are
 
@@ -789,6 +837,7 @@ def reduced_potential(self, context_state):
         V(x) is the instantaneous box volume
         N(x) the numbers of various particle species (e.g. protons of
              titratible groups)
+        A(x) is the xy-area of the box.
 
         References
         ----------
@@ -833,17 +882,18 @@ def reduced_potential(self, context_state):
             openmm_state = context_state.getState(getEnergy=True)
             potential_energy = openmm_state.getPotentialEnergy()
             volume = openmm_state.getPeriodicBoxVolume()
+            area = _box_vectors_area(openmm_state.getPeriodicBoxVectors())
         else:
             n_particles = context_state.n_particles
             potential_energy = context_state.potential_energy
             volume = context_state.volume
+            area = context_state.area
 
         # Check compatibility.
         if n_particles != self.n_particles:
             raise ThermodynamicsError(ThermodynamicsError.INCOMPATIBLE_SAMPLER_STATE)
-
         return self._compute_reduced_potential(potential_energy, self.temperature,
-                                               volume, self.pressure)
+                                               volume, self.pressure, area, self.surface_tension)
 
     @classmethod
     def reduced_potential_at_states(cls, context, thermodynamic_states):
@@ -1024,7 +1074,7 @@ def is_context_compatible(self, context):
         is_compatible = self._standard_system_hash == context_system_hash
         return is_compatible
 
-    def create_context(self, integrator, platform=None):
+    def create_context(self, integrator, platform=None, platform_properties=None):
         """Create a context in this ThermodynamicState.
 
         The context contains a copy of the system. If the integrator
@@ -1048,6 +1098,9 @@ def create_context(self, integrator, platform=None):
         platform : simtk.openmm.Platform, optional
            Platform to use. If None, OpenMM tries to select the fastest
            available platform. Default is None.
+        platform_properties : dict, optional
+           A dictionary of platform properties. Requires platform to be
+           specified.
 
         Returns
         -------
@@ -1059,6 +1112,8 @@ def create_context(self, integrator, platform=None):
         ThermodynamicsError
             If the integrator has a temperature different from this
             ThermodynamicState.
+        ValueError
+            If platform_properties is specified, but platform is None
 
         Examples
         --------
@@ -1099,9 +1154,13 @@ def create_context(self, integrator, platform=None):
 
         # Create context.
         if platform is None:
+            if platform_properties is not None:
+                raise ValueError("To set platform_properties, you need to also specify the platform.")
             return openmm.Context(system, integrator)
-        else:
+        elif platform_properties is None:
             return openmm.Context(system, integrator, platform)
+        else:
+            return openmm.Context(system, integrator, platform, platform_properties)
 
     def apply_to_context(self, context):
         """Apply this ThermodynamicState to the context.
@@ -1147,7 +1206,7 @@ def apply_to_context(self, context):
         310.0
 
         """
-        self._set_context_barostat(context, update_pressure=True, update_temperature=True)
+        self._set_context_barostat(context, update_pressure=True, update_temperature=True, update_surface_tension=True)
         self._set_context_thermostat(context)
 
     # -------------------------------------------------------------------------
@@ -1416,7 +1475,7 @@ def _update_standard_system(self, standard_system):
     # Internal-usage: context handling
     # -------------------------------------------------------------------------
 
-    def _set_context_barostat(self, context, update_pressure, update_temperature):
+    def _set_context_barostat(self, context, update_pressure, update_temperature, update_surface_tension):
         """Set the barostat parameters in the Context."""
         barostat = self._find_barostat(context.getSystem())
 
@@ -1432,6 +1491,9 @@ def _set_context_barostat(self, context, update_pressure, update_temperature):
         if update_pressure:
             self._set_barostat_pressure(barostat, self.pressure)
             self._set_barostat_pressure_in_context(barostat, self.pressure, context)
+        if self.surface_tension is not None and update_surface_tension:
+            self._set_barostat_surface_tension_in_context(barostat, self.surface_tension, context)
+
         if update_temperature:
             self._set_barostat_temperature(barostat, self.temperature)
             # TODO remove try except when drop openmm7.0 support
@@ -1470,9 +1532,10 @@ def _apply_to_context_in_state(self, context, thermodynamic_state):
         """
         update_pressure = self.pressure != thermodynamic_state.pressure
         update_temperature = self.temperature != thermodynamic_state.temperature
+        update_surface_tension = self.surface_tension != thermodynamic_state.surface_tension
 
         if update_pressure or update_temperature:
-            self._set_context_barostat(context, update_pressure, update_temperature)
+            self._set_context_barostat(context, update_pressure, update_temperature, update_surface_tension)
         if update_temperature:
             self._set_context_thermostat(context)
 
@@ -1588,13 +1651,9 @@ def _find_barostat(cls, system, get_index=False):
                 scaled = [barostat.getScaleX(), barostat.getScaleY(), barostat.getScaleY()]
                 if sum(scaled) == 0:
                     raise ThermodynamicsError(ThermodynamicsError.UNSUPPORTED_ANISOTROPIC_BAROSTAT)
-                active_pressures = [pressure for pressure, active in zip(pressures, scaled)]
+                active_pressures = [pressure for pressure, active in zip(pressures, scaled) if active]
                 if any(abs(pressure - active_pressures[0]) > 0 for pressure in active_pressures):
                     raise ThermodynamicsError(ThermodynamicsError.UNSUPPORTED_ANISOTROPIC_BAROSTAT)
-            elif isinstance(barostat, openmm.MonteCarloMembraneBarostat):
-                surface_tension = barostat.getDefaultSurfaceTension()
-                if abs(surface_tension.value_in_unit(unit.bar*unit.nanometer)) > 0:
-                    raise ThermodynamicsError(ThermodynamicsError.UNSUPPORTED_MEMBRANE_BAROSTAT)
         if get_index:
             return force_idx, barostat
         return barostat
@@ -1698,6 +1757,18 @@ def _set_barostat_temperature(barostat, temperature):
         """Set barostat temperature."""
         barostat.setDefaultTemperature(temperature)
 
+    @staticmethod
+    def _set_barostat_surface_tension_in_context(barostat, surface_tension, context):
+        """Set barostat surface tension."""
+        # work around a unit conversion issue in openmm
+        if isinstance(surface_tension, unit.Quantity):
+            surface_tension = surface_tension.value_in_unit(unit.nanometer*unit.bar)
+        try:
+            context.getParameter(barostat.SurfaceTension())
+        except Exception as e:
+            raise ThermodynamicsError(ThermodynamicsError.INCOMPATIBLE_ENSEMBLE)
+        context.setParameter(barostat.SurfaceTension(), surface_tension)
+
     # -------------------------------------------------------------------------
     # Internal-usage: thermostat handling
     # -------------------------------------------------------------------------
@@ -1762,12 +1833,14 @@ def _set_system_temperature(cls, system, temperature):
     # -------------------------------------------------------------------------
 
     @staticmethod
-    def _compute_reduced_potential(potential_energy, temperature, volume, pressure):
+    def _compute_reduced_potential(potential_energy, temperature, volume, pressure, area=None, surface_tension=None):
         """Convert potential energy into reduced potential."""
         beta = 1.0 / (unit.BOLTZMANN_CONSTANT_kB * temperature)
         reduced_potential = potential_energy / unit.AVOGADRO_CONSTANT_NA
         if pressure is not None:
             reduced_potential += pressure * volume
+        if area is not None and surface_tension is not None:
+            reduced_potential -= surface_tension * area
         return beta * reduced_potential
 
     def _find_force_groups_to_update(self, context, thermodynamic_state, memo):
@@ -2032,6 +2105,11 @@ def volume(self):
         """The volume of the box (read-only)"""
         return _box_vectors_volume(self.box_vectors)
 
+    @property
+    def area(self):
+        """The xy-area of the box (read-only)"""
+        return _box_vectors_area(self.box_vectors)
+
     @property
     def n_particles(self):
         """Number of particles (read-only)."""

openmmtools/tests/test_cache.py

@@ -357,3 +357,33 @@ def test_platform_property(self):
         cache.get_context(self.compatible_states[0], integrator)
         with nose.tools.assert_raises(RuntimeError):
             cache.platform = platforms[0]
+
+    def test_platform_properties(self):
+        # Failure test
+        platform_properties = {"CpuThreads": "2"}
+        with nose.tools.assert_raises(ValueError) as cm:
+            ContextCache(platform=None, platform_properties=platform_properties)
+        assert str(cm.exception) == "To set platform_properties, you need to also specify the platform."
+
+        # test
+        platform = openmm.Platform.getPlatformByName("CPU")
+        cache = ContextCache(
+                platform=platform,
+                platform_properties=platform_properties
+        )
+        thermodynamic_state = copy.deepcopy(self.water_300k)
+        integrator = integrators.LangevinIntegrator(temperature=300 * unit.kelvin, measure_heat=True,
+                                                    measure_shadow_work=True)
+        context, _ = cache.get_context(thermodynamic_state, integrator)
+        assert context.getPlatform().getPropertyValue(context, "CpuThreads") == "2"
+
+        # test serialization
+        cache.__setstate__(cache.__getstate__())
+        assert cache._platform_properties == {"CpuThreads": "2"}
+
+        # test serialization no 2
+        state = cache.__getstate__()
+        state["platform_properties"] = {"CpuThreads": "3"}
+        cache.__setstate__(state)
+        assert cache._platform_properties == {"CpuThreads": "3"}
+        del context