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'])
 
     # -------------------------------------------------------------------------