forked from openpathsampling/openpathsampling
-
Notifications
You must be signed in to change notification settings - Fork 0
/
hooks.py
367 lines (305 loc) · 13.4 KB
/
hooks.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
"""
Hooks to change class:`.PathSimulator` behavior.
These hooks group several methods together for use as part of a
:class:`.PathSimulator` ``run`` method. They allow for additional
calculations or output at several points in the simulation.
"""
import time
import openpathsampling as paths
from openpathsampling.netcdfplus import StorableNamedObject
class SimulationNotFoundError(RuntimeError):
"""
Raised when a hook tries to access its parent simulation before knowing it.
"""
pass
class PathSimulatorHook(StorableNamedObject):
"""Superclass for PathSimulator hooks.
This implementation is a do-nothing hook. Subclasses should subclass the
relevant method in order to add hooks PathSimulator objects.
"""
implemented_for = ['before_simulation', 'before_step', 'after_step',
'after_simulation']
def before_simulation(self, sim):
pass # pragma: no-cover
def before_step(self, sim, step_number, step_info, state):
pass # pragma: no-cover
def after_step(self, sim, step_number, step_info, state, results,
hook_state):
pass # pragma: no-cover
def after_simulation(self, sim):
pass # pragma: no-cover
class StorageHook(PathSimulatorHook):
"""
Standard hook for storage.
NOTE: Arguments passed to init take precedence over the corresponding
parameters of the PathSimulator this hook is attached to. They can
only be accessed through this hook, e.g. as hook.live_visualizer.
Parameters
----------
storage : :class:`.Storage
where to save to; default ``None`` uses the simulation's
``storage``
frequency : int
save frequency measured in steps; default ``None`` uses the
simulation's value for ``save_frequency``
"""
implemented_for = ['before_simulation', 'after_step',
'after_simulation']
def __init__(self, storage=None, frequency=None):
self.storage = storage
self.frequency = frequency
self._simulation = None
@property
def frequency(self):
if self._frequency is not None:
return self._frequency
elif self._simulation is not None:
return self._simulation.save_frequency
else:
raise SimulationNotFoundError("'frequency' has not "
+ "been set and no hosting "
+ "simulation known to get a value."
)
@frequency.setter
def frequency(self, val):
self._frequency = val
@property
def storage(self):
if self._storage is not None:
return self._storage
elif self._simulation is not None:
return self._simulation.storage
else:
raise SimulationNotFoundError("'storage' has not "
+ "been set and no hosting "
+ "simulation known to get a value."
)
@storage.setter
def storage(self, val):
self._storage = val
def before_simulation(self, sim):
self._simulation = sim
def after_step(self, sim, step_number, step_info, state, results,
hook_state):
if self.storage is not None:
self.storage.save(results)
if step_number % self.frequency == 0:
if sim.sample_set is not None:
# some PathSimulators never set their sample_set
# but PathSimulator.__init__ sets it to None
sim.sample_set.sanity_check()
self.storage.sync_all()
def after_simulation(self, sim):
if self.storage is not None:
sim.storage.sync_all()
class ShootFromSnapshotsOutputHook(PathSimulatorHook):
"""Default (serial) output for ShootFromSnapshotsSimulation objects.
Updates every time a new snapshot is shot from.
NOTE: Arguments passed to init take precedence over the corresponding
parameters of the PathSimulator this hook is attached to. They can
only be accessed through this hook, e.g. as hook.live_visualizer.
Parameters
----------
output_stream : stream
where to write the results; default ``None`` uses the simulation's
``output_stream``
allow_refresh : bool
whether to allow refresh (see :meth:`.refresh_output`); default
``None`` uses the simulation's value
"""
implemented_for = ['before_simulation', 'before_step']
def __init__(self, output_stream=None, allow_refresh=None):
self.output_stream = output_stream
self.allow_refresh = allow_refresh
self._simulation = None
@property
def output_stream(self):
if self._output_stream is not None:
return self._output_stream
elif self._simulation is not None:
return self._simulation.output_stream
else:
raise SimulationNotFoundError("'output_stream' has not "
+ "been set and no hosting "
+ "simulation known to get a value."
)
@output_stream.setter
def output_stream(self, val):
self._output_stream = val
@property
def allow_refresh(self):
if self._allow_refresh is not None:
return self._allow_refresh
elif self._simulation is not None:
return self._simulation.allow_refresh
else:
raise SimulationNotFoundError("'allow_refresh' has not "
+ "been set and no hosting "
+ "simulation known to get a value."
)
@allow_refresh.setter
def allow_refresh(self, val):
self._allow_refresh = val
def before_simulation(self, sim):
self._simulation = sim
def before_step(self, sim, step_number, step_info, state):
snap_num, n_snapshots, step, n_per_snapshot = step_info
paths.tools.refresh_output(
"Working on snapshot %d / %d; shot %d / %d" % (
snap_num+1, n_snapshots, step+1, n_per_snapshot
),
output_stream=self.output_stream,
refresh=self.allow_refresh
)
class LiveVisualizerHook(PathSimulatorHook):
"""
LiveVisualization using the :class:`openpathsampling.StepVisualizer2D`.
Updates every `simulation.status_update_frequency` MCSteps, where
simulation is the `PathSimulator` this hook is attached to.
NOTE: You will have to set PathSimulator.allow_refresh = False
Otherwise the LiveVisualization will get refreshed away
(i.e. deleted) right after creation.
NOTE: Arguments passed to init take precedence over the corresponding
parameters of the PathSimulator this hook is attached to. They can
only be accessed through this hook, e.g. as hook.live_visualizer.
Parameters
----------
live_visualizer : :class:`openpathsampling.StepVisualizer2D`
default `None` uses the simulations live_visualizer
status_update_frequency : int
number of steps between two refreshs of the visualization;
default `None` uses the simulations value (PathSampling default=1)
"""
# NOTE: we visualize after step, because otherwise the 'next' MCstep
# would depend on the 'previous' one just for viualization
# this deviates from the previous implementation but avoids
# having to pass the previous MCstep to before_step hooks
implemented_for = ['before_simulation', 'after_step']
def __init__(self, live_visualizer=None, status_update_frequency=None):
self.live_visualizer = live_visualizer
self.status_update_frequency = status_update_frequency
self._simulation = None
@property
def live_visualizer(self):
if self._live_visualizer is not None:
return self._live_visualizer
elif self._simulation is not None:
return self._simulation.live_visualizer
else:
# hejung: I think here we should return None, because that is the
# default and leaving the default value unchanged should not lead
# to an error :)
return None
@live_visualizer.setter
def live_visualizer(self, val):
self._live_visualizer = val
@property
def status_update_frequency(self):
if self._status_update_frequency is not None:
return self._status_update_frequency
elif self._simulation is not None:
return self._simulation.status_update_frequency
else:
raise SimulationNotFoundError("'status_update_frequency' has not "
+ "been set and no hosting "
+ "simulation known to get a value."
)
@status_update_frequency.setter
def status_update_frequency(self, val):
self._status_update_frequency = val
def before_simulation(self, sim):
self._simulation = sim
def after_step(self, sim, step_number, step_info, state, results,
hook_state):
if step_number % self.status_update_frequency == 0:
# do we visualize this step?
if self.live_visualizer is not None and results is not None:
# do we visualize at all?
self.live_visualizer.draw_ipynb(results)
class PathSamplingOutputHook(PathSimulatorHook):
"""
Default (serial) output for PathSamplingSimulation objects.
Updates every `PathSampling.status_update_frequency` MCSteps.
NOTE: Arguments passed to init take precedence over the corresponding
parameters of the PathSimulator this hook is attached to. They can
only be accessed through this hook, e.g. as hook.output_stream.
Parameters
----------
output_stream : stream
where to write the results; default ``None`` uses the simulation's
``output_stream``
allow_refresh : bool
whether to allow refresh (see :meth:`.refresh_output`); default
``None`` uses the simulation's value
status_update_frequency : int
number of steps between two refreshs of the visualization;
default `None` uses the simulations value (PathSampling default=1)
"""
implemented_for = ['before_simulation', 'before_step', 'after_simulation']
def __init__(self, output_stream=None, allow_refresh=None,
status_update_frequency=None):
self.output_stream = output_stream
self.allow_refresh = allow_refresh
self.status_update_frequency = status_update_frequency
self._simulation = None
@property
def output_stream(self):
if self._output_stream is not None:
return self._output_stream
elif self._simulation is not None:
return self._simulation.output_stream
else:
raise SimulationNotFoundError("'output_stream' has not "
+ "been set and no hosting "
+ "simulation known to get a value."
)
@output_stream.setter
def output_stream(self, val):
self._output_stream = val
@property
def allow_refresh(self):
if self._allow_refresh is not None:
return self._allow_refresh
elif self._simulation is not None:
return self._simulation.allow_refresh
else:
raise SimulationNotFoundError("'allow_refresh' has not "
+ "been set and no hosting "
+ "simulation known to get a value."
)
@allow_refresh.setter
def allow_refresh(self, val):
self._allow_refresh = val
@property
def status_update_frequency(self):
if self._status_update_frequency is not None:
return self._status_update_frequency
elif self._simulation is not None:
return self._simulation.status_update_frequency
else:
raise SimulationNotFoundError("'status_update_frequency' has not "
+ "been set and no hosting "
+ "simulation known to get a value."
)
@status_update_frequency.setter
def status_update_frequency(self, val):
self._status_update_frequency = val
def before_simulation(self, sim):
self._simulation = sim
self._initial_time = time.time()
def before_step(self, sim, step_number, step_info, state):
if step_number % self.status_update_frequency == 0:
nn, n_steps = step_info
elapsed = time.time() - self._initial_time
paths.tools.refresh_output(
"Working on Monte Carlo cycle number " + str(step_number)
+ "\n" + paths.tools.progress_string(nn, n_steps, elapsed),
refresh=self.allow_refresh,
output_stream=self.output_stream
)
def after_simulation(self, sim):
paths.tools.refresh_output(
"DONE! Completed " + str(sim.step) + " Monte Carlo cycles.\n",
refresh=False,
output_stream=self.output_stream
)