Skip to content

Commit 4a41b8e

Browse files
committed
OpenMMEngine: Persist platform
This allows OpenMM engines to persist their platform (if given at initialization). It also allow override of `openmm_properties` (and `platform`) in the `initialization()` method.
1 parent 8503204 commit 4a41b8e

2 files changed

Lines changed: 158 additions & 37 deletions

File tree

‎openpathsampling/engines/openmm/engine.py‎

Lines changed: 79 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -79,7 +79,8 @@ def __init__(
7979
system,
8080
integrator,
8181
openmm_properties=None,
82-
options=None):
82+
options=None,
83+
platform=None):
8384
"""
8485
Parameters
8586
----------
@@ -95,8 +96,6 @@ def __init__(
9596
keys include GPU floating point precision.
9697
Note that by default the engine selects the fastest currently
9798
available OpenMM platform.
98-
If you want to specify the platform you will have to call
99-
`engine.initialize(platform)` after creating the engine.
10099
options : dict
101100
a dictionary that provides additional settings for the OPS engine.
102101
Allowed are
@@ -106,6 +105,11 @@ def __init__(
106105
'n_frames_max' : int, default: 5000,
107106
the maximal number of frames allowed for a returned
108107
trajectory object
108+
platform : str or :class:`openmm.Platform` or None
109+
optional default platform for creating the OpenMM simulation.
110+
This value is used by ``initialize`` and serialization. The
111+
``.platform`` property reflects the initialized simulation context
112+
and remains ``None`` until initialization occurs.
109113
110114
Notes
111115
-----
@@ -138,6 +142,8 @@ def __init__(
138142
openmm_properties = {}
139143

140144
self.openmm_properties = openmm_properties
145+
self._normalize_platform(platform) # validate serializability
146+
self._platform = platform
141147

142148
# set no cached snapshot
143149
self._current_snapshot = None
@@ -202,7 +208,8 @@ def from_new_options(
202208
self.system,
203209
integrator,
204210
openmm_properties=openmm_properties,
205-
options=new_options)
211+
options=new_options,
212+
platform=self._platform)
206213

207214
if self._simulation is not None and \
208215
integrator is self.integrator and \
@@ -220,14 +227,35 @@ def from_new_options(
220227
@property
221228
def platform(self):
222229
"""
223-
str : Return the name of the currently used platform
230+
str or None : Return the name of the currently used platform.
231+
232+
This value is populated from the simulation context and is ``None``
233+
until the simulation has been initialized. Configured platform values
234+
may exist before initialization for serialization, but are not exposed
235+
through this property.
224236
225237
"""
226238
if self._simulation is not None:
227239
return self._simulation.context.getPlatform().getName()
228240
else:
229241
return None
230242

243+
@staticmethod
244+
def _normalize_platform(platform):
245+
if platform is None:
246+
return None
247+
248+
if isinstance(platform, str):
249+
return platform
250+
251+
try:
252+
return platform.getName()
253+
except (AttributeError, TypeError) as err:
254+
raise TypeError(
255+
"platform must be None, a platform name string, or an "
256+
"openmm.Platform instance"
257+
) from err
258+
231259
@property
232260
def simulation(self):
233261
if self._simulation is None:
@@ -261,7 +289,7 @@ def unload_context(self):
261289
def _default_trajectory_writer(self):
262290
return TRRTrajectoryWriter()
263291

264-
def initialize(self, platform=None):
292+
def initialize(self, platform=None, openmm_properties=None):
265293
"""
266294
Create the final OpenMMEngine
267295
@@ -270,6 +298,9 @@ def initialize(self, platform=None):
270298
platform : str or :class:`openmm.Platform` or None
271299
either a string with a name of the platform or a platform object
272300
if None it will default to the fastest currently available platform
301+
openmm_properties : dict or None
302+
optional platform properties for this initialize call. If ``None``
303+
the engine-stored ``openmm_properties`` are used.
273304
274305
Notes
275306
-----
@@ -281,37 +312,45 @@ def initialize(self, platform=None):
281312
"""
282313

283314
if self._simulation is None:
284-
if type(platform) is str:
285-
self._simulation = openmm.app.Simulation(
286-
topology=self.topology.mdtraj.to_openmm(),
287-
system=self.system,
288-
integrator=self.integrator,
289-
platform=openmm.Platform.getPlatformByName(platform),
290-
platformProperties=self.openmm_properties
291-
)
292-
elif platform is None:
293-
# as of OpenMM 8.1, we can't give an empty props dict when
294-
# platform is None. This will still raise the internal
295-
# OpenMM error is platform is None and properties are
296-
# provided.
297-
openmm_props = self.openmm_properties
298-
if openmm_props == {}:
299-
openmm_props = None
315+
if platform is None:
316+
effective_platform = self._platform
317+
else:
318+
self._normalize_platform(platform) # validate serializability
319+
effective_platform = platform
300320

301-
self._simulation = openmm.app.Simulation(
302-
topology=self.topology.mdtraj.to_openmm(),
303-
system=self.system,
304-
integrator=self.integrator,
305-
platformProperties=openmm_props,
306-
)
321+
if openmm_properties is None:
322+
effective_properties = self.openmm_properties
307323
else:
308-
self._simulation = openmm.app.Simulation(
309-
topology=self.topology.mdtraj.to_openmm(),
310-
system=self.system,
311-
integrator=self.integrator,
312-
platform=platform,
313-
platformProperties=self.openmm_properties
314-
)
324+
effective_properties = openmm_properties
325+
326+
simulation_kwargs = {
327+
'topology': self.topology.mdtraj.to_openmm(),
328+
'system': self.system,
329+
'integrator': self.integrator,
330+
}
331+
332+
if effective_platform is None:
333+
if not effective_properties:
334+
openmm_props = None
335+
else:
336+
openmm_props = effective_properties
337+
raise ValueError(
338+
"OpenMM platform-specific properties were provided, "
339+
"but no platform was specified."
340+
)
341+
simulation_kwargs['platformProperties'] = openmm_props
342+
else:
343+
if isinstance(effective_platform, str):
344+
resolved_platform = openmm.Platform.getPlatformByName(
345+
effective_platform
346+
)
347+
else:
348+
resolved_platform = effective_platform
349+
350+
simulation_kwargs['platform'] = resolved_platform
351+
simulation_kwargs['platformProperties'] = effective_properties
352+
353+
self._simulation = openmm.app.Simulation(**simulation_kwargs)
315354

316355
logger.info(
317356
'Initialized OpenMM engine using platform `%s`' %
@@ -333,7 +372,8 @@ def to_dict(self):
333372
'integrator_xml': integrator_xml,
334373
'topology': self.topology,
335374
'options': self.options,
336-
'properties': self.openmm_properties
375+
'properties': self.openmm_properties,
376+
'platform': self._normalize_platform(self._platform)
337377
}
338378

339379
@classmethod
@@ -343,6 +383,7 @@ def from_dict(cls, dct):
343383
topology = dct['topology']
344384
options = dct['options']
345385
properties = dct['properties']
386+
platform = dct.get('platform', None)
346387

347388
# we need to have str as keys
348389
properties = {str(key): str(value)
@@ -355,7 +396,8 @@ def from_dict(cls, dct):
355396
system=openmm.XmlSerializer.deserialize(system_xml),
356397
integrator=integrator,
357398
options=options,
358-
openmm_properties=properties
399+
openmm_properties=properties,
400+
platform=platform
359401
)
360402

361403
@property

‎openpathsampling/tests/test_openmm_engine.py‎

Lines changed: 79 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -83,6 +83,13 @@ def setup_method(self):
8383
)
8484
integrator.setConstraintTolerance(0.00001)
8585

86+
uninitialized_integrator = mm.LangevinIntegrator(
87+
300*u.kelvin,
88+
old_div(1.0, u.picoseconds),
89+
2.0*u.femtoseconds
90+
)
91+
uninitialized_integrator.setConstraintTolerance(0.00001)
92+
8693
# Engine options
8794
options = {
8895
'n_steps_per_frame': 2,
@@ -96,6 +103,13 @@ def setup_method(self):
96103
options=options
97104
)
98105

106+
self.uninitialized_engine = peng.Engine(
107+
template.topology,
108+
system,
109+
uninitialized_integrator,
110+
options=options
111+
)
112+
99113
self.engine.initialize('CPU')
100114
context = self.engine.simulation.context
101115
zero_array = np.zeros((template.topology.n_atoms, 3))
@@ -296,3 +310,68 @@ def test_export_trajectory(self, tmp_path):
296310
assert len(reloaded) == len(traj)
297311
npt.assert_allclose(reloaded.xyz, traj.xyz)
298312

313+
def test_serialization_cycle(self):
314+
integrator = mm.LangevinIntegrator(
315+
300*u.kelvin,
316+
old_div(1.0, u.picoseconds),
317+
2.0*u.femtoseconds
318+
)
319+
integrator.setConstraintTolerance(0.00001)
320+
engine = peng.Engine(
321+
template.topology,
322+
system,
323+
integrator,
324+
options={'n_steps_per_frame': 2, 'n_frames_max': 5},
325+
platform='CPU'
326+
)
327+
328+
dct = engine.to_dict()
329+
rebuilt = peng.Engine.from_dict(dct)
330+
dct2 = rebuilt.to_dict()
331+
assert dct == dct2
332+
assert dct2['platform'] == 'CPU'
333+
334+
def test_initialize_uses_exact_platform_object(self):
335+
platform_obj = mm.Platform.getPlatformByName('CPU')
336+
self.uninitialized_engine.initialize(platform=platform_obj)
337+
338+
assert self.uninitialized_engine.platform == 'CPU'
339+
assert self.uninitialized_engine._platform is None
340+
assert self.uninitialized_engine.to_dict()['platform'] is None
341+
342+
def test_from_dict_without_platform_key_is_backward_compatible(self):
343+
serialized = self.engine.to_dict()
344+
del serialized['platform']
345+
346+
restored = peng.Engine.from_dict(serialized)
347+
assert restored is not None
348+
assert restored.platform is None
349+
350+
def test_initialize_uses_stored_properties_when_override_is_none(
351+
self):
352+
self.uninitialized_engine.openmm_properties = {'__ops_bad__': '1'}
353+
with pytest.raises(Exception, match='Illegal property name'):
354+
self.uninitialized_engine.initialize(
355+
platform='CPU',
356+
openmm_properties=None
357+
)
358+
359+
def test_initialize_explicit_properties_override_stored(self):
360+
self.uninitialized_engine.openmm_properties = {'__ops_bad__': '1'}
361+
self.uninitialized_engine.initialize(
362+
platform='CPU',
363+
openmm_properties={}
364+
)
365+
assert isinstance(self.uninitialized_engine.simulation, mm.app.Simulation)
366+
assert self.uninitialized_engine.platform == 'CPU'
367+
368+
def test_initialize_nonempty_properties_without_platform_raises(self):
369+
self.uninitialized_engine.openmm_properties = {'Threads': '1'}
370+
with pytest.raises(ValueError, match='no platform was specified'):
371+
self.uninitialized_engine.initialize()
372+
373+
def test_initialize_empty_properties_without_platform_is_valid(self):
374+
self.uninitialized_engine.openmm_properties = {}
375+
self.uninitialized_engine.initialize()
376+
assert isinstance(self.uninitialized_engine.simulation, mm.app.Simulation)
377+
assert self.uninitialized_engine.platform is not None

0 commit comments

Comments
 (0)