Merge 5dcd421 into 56f52ad

openpathsampling/__init__.py

@@ -172,6 +172,7 @@
 from .pathmovers.move_schemes import *
 
 import openpathsampling.numerics as numerics
+import openpathsampling.beta
 
 from openpathsampling.engines import Trajectory, BaseSnapshot
 

openpathsampling/beta/hooks.py

@@ -0,0 +1,105 @@
+import openpathsampling as paths
+from openpathsampling.pathsimulators.path_simulator import PathSimulator
+from openpathsampling.netcdfplus import StorableNamedObject
+
+"""
+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.
+"""
+
+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.
+    """
+    implemented_for = ['before_simulation', 'after_step',
+                       'after_simulation']
+    def __init__(self, storage=None, frequency=None):
+        self.storage = storage
+        self.frequency = frequency
+
+    def before_simulation(self, sim):
+        if self.storage is None:
+            self.storage = sim.storage
+        if self.frequency is None:
+            self.frequency = sim.save_frequency
+
+    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:
+                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.
+
+    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
+
+    def before_simulation(self, sim):
+        if self.output_stream is None:
+            self.output_stream = sim.output_stream
+        if self.allow_refresh is None:
+            self.allow_refresh  = sim.allow_refresh
+
+    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
+        )
+
+
+# TODO: here are some other hooks to implement in the future
+# class LiveVisualizerHook(PathSimulatorHook):
+    # pass
+
+
+# class PathSamplingOutputHook(PathSimulatorHook):
+    # pass

openpathsampling/pathsimulators/path_simulator.py

@@ -65,11 +65,11 @@ class PathSimulator(with_metaclass(abc.ABCMeta, StorableNamedObject)):
         This is likely to be overridden when a pathsimulator is wrapped in
         another simulation.
     """
-    #__metaclass__ = abc.ABCMeta
-
     calc_name = "PathSimulator"
     _excluded_attr = ['sample_set', 'step', 'save_frequency',
                       'output_stream']
+    hook_names = ['before_simulation', 'before_step', 'after_step',
+                  'after_simulation']
 
     def __init__(self, storage):
         super(PathSimulator, self).__init__()
@@ -84,6 +84,8 @@ def __init__(self, storage):
         self.sample_set = None
         self.output_stream = sys.stdout  # user can change to file handler
         self.allow_refresh = True
+        self.hooks = self.empty_hooks()
+        self.attach_default_hooks()
 
     def sync_storage(self):
         """
@@ -92,6 +94,50 @@ def sync_storage(self):
         if self.storage is not None:
             self.storage.sync_all()
 
+    def attach_default_hooks(self):
+        pass
+
+    def empty_hooks(self):
+        """Return a hook dictionary with no hooks."""
+        return {k: [] for k in self.hook_names}
+
+    def attach_hook(self, hook, hook_for=None):
+        """Attach a hook class or method to this simulation.
+
+        Parameters
+        ----------
+        hook : :class:`.PathSimulatorHook` or method
+            Hook to add
+        hook_for : str or None
+            If None (default) then the ``hook`` must be a class with methods
+            named to match the hook names in this simulator. If ``hook`` is
+            a method, then ``hook_for`` must be the name of the hook it
+            represents
+        """
+        def add_hook_method(hook_method, hook_name):
+            try:
+                self.hooks[hook_name].append(hook_method)
+            except KeyError:
+                raise TypeError("No hook '" + hook_name + "' in " +
+                                str(self.__class__.__name__))
+
+        if hook_for is None:
+            for hook_name in hook.implemented_for:
+                hook_method = getattr(hook, hook_name)
+                add_hook_method(hook_method, hook_name)
+        else:
+            add_hook_method(hook, hook_for)
+
+    def run_hooks(self, hook_name, **kwargs):
+        """Run the hooks for the given ``hook_name``"""
+        hook_name_state = {}
+        for hook in self.hooks[hook_name]:
+            result = hook(**kwargs)
+            if result is not None:
+                hook_name_state.update({hook: result})
+        if hook_name_state != {}:
+            return hook_name_state
+
     @abc.abstractmethod
     def run(self, n_steps):
         """
@@ -102,7 +148,7 @@ def run(self, n_steps):
         n_steps : int
             number of step to be run
         """
-        pass
+        raise NotImplementedError()
 
     def save_initial_step(self):
         """

openpathsampling/pathsimulators/shoot_snapshots.py

@@ -3,6 +3,7 @@
 
 logger = logging.getLogger(__name__)
 from .path_simulator import PathSimulator, MCStep
+from openpathsampling.beta import hooks
 
 class ShootFromSnapshotsSimulation(PathSimulator):
     """
@@ -17,6 +18,17 @@ class ShootFromSnapshotsSimulation(PathSimulator):
     user can create a simulation of this sort on-the-fly for some weird
     ensembles.
 
+    .. note::
+
+        **Hooks**: This can interact with the :class:`.PathSimulatorHook`
+        framework. In these hooks, ``step_info`` is the 4-tuple
+        ``(snap, n_snaps, step, n_steps)`` where ``snap`` is the snapshot
+        number, ``step`` is the step number within that snapshot, and
+        ``n_snaps`` and ``n_steps`` are the total number of each,
+        respectively. The ``state`` will provide the initial snapshot we are
+        shooting from, and ``results`` is the :class:`.MCStep` that comes
+        out of each step.
+
     Parameters
     ----------
     storage : :class:`.Storage`
@@ -62,7 +74,6 @@ def __init__(self, storage, engine, starting_volume, forward_ensemble,
             ensemble=self.starting_ensemble,
             target_ensemble=self.backward_ensemble
         )
-
         # subclasses will often override this
         self.mover = paths.RandomChoiceMover([self.forward_mover,
                                               self.backward_mover])
@@ -93,6 +104,9 @@ def from_dict(cls, dct):
         obj.mover = dct['mover']
         return obj
 
+    def attach_default_hooks(self):
+        self.attach_hook(hooks.StorageHook())
+        self.attach_hook(hooks.ShootFromSnapshotsOutputHook())
 
     def run(self, n_per_snapshot, as_chain=False):
         """Run the simulation.
@@ -110,19 +124,19 @@ def run(self, n_per_snapshot, as_chain=False):
         """
         self.step = 0
         snap_num = 0
-        self.output_stream.write("\n")
+        n_snapshots = len(self.initial_snapshots)
+        hook_state = None
+        self.run_hooks('before_simulation', sim=self)
         for snapshot in self.initial_snapshots:
+            # before_snapshot
             start_snap = snapshot
             # do what we need to get the snapshot set up
             for step in range(n_per_snapshot):
-                paths.tools.refresh_output(
-                    "Working on snapshot %d / %d; shot %d / %d\n" % (
-                        snap_num+1, len(self.initial_snapshots),
-                        step+1, n_per_snapshot
-                    ),
-                    output_stream=self.output_stream,
-                    refresh=self.allow_refresh
-                )
+                step_number = self.step
+                step_info = (snap_num, n_snapshots, step, n_per_snapshot)
+                self.run_hooks('before_step', sim=self,
+                               step_number=step_number, step_info=step_info,
+                               state=start_snap)
 
                 if as_chain:
                     start_snap = self.randomizer(start_snap)
@@ -135,6 +149,8 @@ def run(self, n_per_snapshot, as_chain=False):
                                  ensemble=self.starting_ensemble)
                 ])
                 sample_set.sanity_check()
+
+                # shoot_snapshot_task (start)
                 new_pmc = self.mover.move(sample_set)
                 samples = new_pmc.results
                 new_sample_set = sample_set.apply_samples(samples)
@@ -146,15 +162,18 @@ def run(self, n_per_snapshot, as_chain=False):
                     active=new_sample_set,
                     change=new_pmc
                 )
+                # shoot_snapshot_task (end)
 
-                if self.storage is not None:
-                    self.storage.steps.save(mcstep)
-                    if self.step % self.save_frequency == 0:
-                        self.sync_storage()
+                hook_state = self.run_hooks(
+                    'after_step', sim=self, step_number=step_number,
+                    step_info=step_info, state=start_snap, results=mcstep,
+                    hook_state=hook_state
+                )
 
                 self.step += 1
+            # after_snapshot
             snap_num += 1
-
+        self.run_hooks('after_simulation', sim=self)
 
 
 class CommittorSimulation(ShootFromSnapshotsSimulation):