Skip to content

Commit 10a5729

Browse files
dwhswensondwhswenson
authored
Merge pull request #1180 from dwhswenson/openmm-platform-serialization
OpenMMEngine: Persist `platform`; allow `openmm_properties` override in `initialize`
2 parents 45f4f40 + 637e386 commit 10a5729

2 files changed

Lines changed: 177 additions & 38 deletions

File tree

‎openpathsampling/engines/openmm/engine.py‎

Lines changed: 77 additions & 38 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+
match platform:
246+
case None:
247+
return None
248+
case str() as name:
249+
return name
250+
case candidate if isinstance(candidate, openmm.Platform):
251+
return candidate.getName()
252+
case _:
253+
raise TypeError(
254+
"platform must be None, a platform name string, or an "
255+
"openmm.Platform instance. "
256+
f"Got {type(platform)} instead"
257+
)
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,42 @@ 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
300-
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-
)
315+
if platform is None:
316+
effective_platform = self._platform
307317
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-
)
318+
self._normalize_platform(platform) # validate serializability
319+
effective_platform = platform
320+
321+
if openmm_properties is None:
322+
effective_properties = self.openmm_properties
323+
else:
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 effective_properties:
334+
raise ValueError(
335+
"OpenMM platform-specific properties were provided, "
336+
"but no platform was specified."
337+
)
338+
simulation_kwargs['platformProperties'] = None
339+
else:
340+
if isinstance(effective_platform, str):
341+
resolved_platform = openmm.Platform.getPlatformByName(
342+
effective_platform
343+
)
344+
else:
345+
resolved_platform = effective_platform
346+
347+
simulation_kwargs['platform'] = resolved_platform
348+
simulation_kwargs['platformProperties'] = effective_properties
349+
350+
self._simulation = openmm.app.Simulation(**simulation_kwargs)
315351

316352
logger.info(
317353
'Initialized OpenMM engine using platform `%s`' %
@@ -333,7 +369,8 @@ def to_dict(self):
333369
'integrator_xml': integrator_xml,
334370
'topology': self.topology,
335371
'options': self.options,
336-
'properties': self.openmm_properties
372+
'properties': self.openmm_properties,
373+
'platform': self._normalize_platform(self._platform)
337374
}
338375

339376
@classmethod
@@ -343,6 +380,7 @@ def from_dict(cls, dct):
343380
topology = dct['topology']
344381
options = dct['options']
345382
properties = dct['properties']
383+
platform = dct.get('platform', None)
346384

347385
# we need to have str as keys
348386
properties = {str(key): str(value)
@@ -355,7 +393,8 @@ def from_dict(cls, dct):
355393
system=openmm.XmlSerializer.deserialize(system_xml),
356394
integrator=integrator,
357395
options=options,
358-
openmm_properties=properties
396+
openmm_properties=properties,
397+
platform=platform
359398
)
360399

361400
@property

‎openpathsampling/tests/test_openmm_engine.py‎

Lines changed: 100 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,89 @@ 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_duck_typed_platform_like_object_raises_typeerror(self):
343+
class PlatformLikeObject(object):
344+
def getName(self):
345+
return 'CPU'
346+
347+
integrator = mm.LangevinIntegrator(
348+
300*u.kelvin,
349+
old_div(1.0, u.picoseconds),
350+
2.0*u.femtoseconds
351+
)
352+
integrator.setConstraintTolerance(0.00001)
353+
354+
with pytest.raises(TypeError, match=r'platform must be None.*PlatformLikeObject'):
355+
peng.Engine(
356+
template.topology,
357+
system,
358+
integrator,
359+
options={'n_steps_per_frame': 2, 'n_frames_max': 5},
360+
platform=PlatformLikeObject()
361+
)
362+
363+
def test_from_dict_without_platform_key_is_backward_compatible(self):
364+
serialized = self.engine.to_dict()
365+
del serialized['platform']
366+
367+
restored = peng.Engine.from_dict(serialized)
368+
assert restored is not None
369+
assert restored.platform is None
370+
371+
def test_initialize_uses_stored_properties_when_override_is_none(
372+
self):
373+
self.uninitialized_engine.openmm_properties = {'__ops_bad__': '1'}
374+
with pytest.raises(Exception, match='Illegal property name'):
375+
self.uninitialized_engine.initialize(
376+
platform='CPU',
377+
openmm_properties=None
378+
)
379+
380+
def test_initialize_explicit_properties_override_stored(self):
381+
self.uninitialized_engine.openmm_properties = {'__ops_bad__': '1'}
382+
self.uninitialized_engine.initialize(
383+
platform='CPU',
384+
openmm_properties={}
385+
)
386+
assert isinstance(self.uninitialized_engine.simulation, mm.app.Simulation)
387+
assert self.uninitialized_engine.platform == 'CPU'
388+
389+
def test_initialize_nonempty_properties_without_platform_raises(self):
390+
self.uninitialized_engine.openmm_properties = {'Threads': '1'}
391+
with pytest.raises(ValueError, match='no platform was specified'):
392+
self.uninitialized_engine.initialize()
393+
394+
def test_initialize_empty_properties_without_platform_is_valid(self):
395+
self.uninitialized_engine.openmm_properties = {}
396+
self.uninitialized_engine.initialize()
397+
assert isinstance(self.uninitialized_engine.simulation, mm.app.Simulation)
398+
assert self.uninitialized_engine.platform is not None

0 commit comments

Comments
 (0)