Remove SamplingContext.view and helper View class

Author: p-zubieta

pysages/backends/ase.py

@@ -16,7 +16,7 @@
     build_data_querier,
 )
 from pysages.backends.utils import view
-from pysages.typing import Callable, NamedTuple
+from pysages.typing import Callable
 from pysages.utils import ToCPU, copy
 
 
@@ -29,7 +29,7 @@ class Sampler(Calculator):
     """
 
     def __init__(self, context, method_bundle, callback: Callable):
-        initial_snapshot, initialize, mehod_update = method_bundle
+        initial_snapshot, initialize, method_update = method_bundle
 
         atoms = context.atoms
         self.implemented_properties = atoms.calc.implemented_properties
@@ -41,7 +41,7 @@ def __init__(self, context, method_bundle, callback: Callable):
         self.callback = callback
         self.snapshot = initial_snapshot
         self.state = initialize()
-        self.update = mehod_update
+        self.update = method_update
 
         sig = signature(atoms.calc.calculate).parameters
         self._calculator = atoms.calc
@@ -151,10 +151,6 @@ def dimensionality():
     return helpers
 
 
-class View(NamedTuple):
-    synchronize: Callable
-
-
 def bind(sampling_context: SamplingContext, callback: Callable, **kwargs):
     """
     Entry point for the backend code, it gets called when the simulation
@@ -166,6 +162,5 @@ def bind(sampling_context: SamplingContext, callback: Callable, **kwargs):
     helpers = build_helpers(sampling_context, sampling_method)
     method_bundle = sampling_method.build(snapshot, helpers)
     sampler = Sampler(context, method_bundle, callback)
-    sampling_context.view = View((lambda: None))
     sampling_context.run = context.run
     return sampler

pysages/backends/core.py

@@ -45,14 +45,12 @@ def __init__(
 
         self.context = context
         self.method = sampling_method
-        self.view = None
         self.run = None
 
         backend = import_module("." + self._backend_name, package="pysages.backends")
         self.sampler = backend.bind(self, callback, **kwargs)
 
-        # `self.view` and `self.run` *must* be set by the backend bind function.
-        assert self.view is not None
+        # `self.run` *must* be set by the backend bind function.
         assert self.run is not None
 
     @property

pysages/backends/hoomd.py

@@ -7,21 +7,12 @@
 
 import hoomd
 from hoomd import md
-from hoomd.dlext import (
-    AccessLocation,
-    AccessMode,
-    DLExtSampler,
-    SystemView,
-    images,
-    net_forces,
-    positions_types,
-    rtags,
-    velocities_masses,
-)
+from hoomd.dlext import AccessLocation, AccessMode, DLExtSampler, SystemView
 from jax import jit
 from jax import numpy as np
 from jax.dlpack import from_dlpack
 
+from pysages.backends import snapshot as pbs
 from pysages.backends.core import SamplingContext
 from pysages.backends.snapshot import (
     Box,
@@ -30,17 +21,17 @@
     SnapshotMethods,
     build_data_querier,
 )
-from pysages.backends.snapshot import restore as _restore
-from pysages.typing import Callable
-from pysages.utils import check_device_array, copy
+from pysages.typing import Callable, Optional
+from pysages.utils import copy
+
+SamplerBase = DLExtSampler
 
 # TODO: Figure out a way to automatically tie the lifetime of Sampler
 # objects to the contexts they bind to
 CONTEXTS_SAMPLERS = {}
 
 
 if getattr(hoomd, "__version__", "").startswith("2."):
-    SamplerBase = DLExtSampler
 
     def is_on_gpu(context):
         return context.on_gpu()
@@ -61,10 +52,7 @@ def remove_half_step_hook(context):
         context.integrator.cpp_integrator.removeHalfStepHook()
 
 else:
-    if hasattr(hoomd.dlext, "__version__"):
-        SamplerBase = DLExtSampler
-
-    else:
+    if not hasattr(hoomd.dlext, "__version__"):
 
         class SamplerBase(DLExtSampler, md.HalfStepHook):
             def __init__(self, sysview, update, location, mode):
@@ -91,41 +79,71 @@ def remove_half_step_hook(context):
         context.operations.integrator.half_step_hook = None
 
 
+if hasattr(AccessLocation, "OnDevice"):
+
+    def default_location():
+        return AccessLocation.OnDevice
+
+else:
+
+    def default_location():
+        return AccessLocation.OnHost
+
+
 class Sampler(SamplerBase):
-    def __init__(self, sysview, method_bundle, bias, callback: Callable, restore):
-        initial_snapshot, initialize, method_update = method_bundle
+    """
+    HOOMD-blue HalfStepHook that connects PySAGES sampling methods to HOOMD-blue simulations.
 
-        def update(positions, vel_mass, rtags, images, forces, timestep):
-            snapshot = self._pack_snapshot(positions, vel_mass, forces, rtags, images)
-            self.state = method_update(snapshot, self.state)
-            self.bias(snapshot, self.state)
-            if self.callback:
-                self.callback(snapshot, self.state, timestep)
+    Parameters
+    ----------
 
-        super().__init__(sysview, update, default_location(), AccessMode.Read)
-        self.state = initialize()
-        self.bias = bias
-        self.box = initial_snapshot.box
+    context: hoomd.Simulation
+        The HOOMD-blue simulation instance to which the PySAGES sampling machinery will be hooked.
+
+    sampling_method: pysages.methods.SamplingMethod
+        The sampling method used.
+
+    callbacks: Optional[Callback]
+        A callback. Some methods define one for logging, but it can also be user-defined.
+
+    location: hoomd.dlext.AccessLocation
+        Device where the simulation data will be retrieved.
+    """
+
+    def __init__(
+        self, context, sampling_method, callback: Optional[Callable], location=default_location()
+    ):
+        self.context = context
         self.callback = callback
-        self.dt = initial_snapshot.dt
-        self._restore = restore
+        self.location = location
+        self.view = SystemView(get_system(context))
 
-    def restore(self, prev_snapshot):
-        def restore_callback(positions, vel_mass, rtags, images, forces, n):
-            snapshot = self._pack_snapshot(positions, vel_mass, forces, rtags, images)
-            self._restore(snapshot, prev_snapshot)
+        self.box = Box(*get_global_box(self.view))
+        self.dt = get_timestep(self.context)
+        self.update_box = lambda: self.box  # NOTE: extend for NPT simulations
 
-        self.forward_data(restore_callback, default_location(), AccessMode.Overwrite, 0)
+        super().__init__(self.view, self._update_callback, location, AccessMode.Read)
 
-    def take_snapshot(self):
-        container = []
+        # Create initial snapshot and setup sampling method
+        snapshot = self.take_snapshot()  # sets `self.snapshot`
+        helpers, restore, bias = build_helpers(context, sampling_method)
+        _, initialize, method_update = sampling_method.build(snapshot, helpers)
 
-        def snapshot_callback(positions, vel_mass, rtags, images, forces, n):
-            snapshot = self._pack_snapshot(positions, vel_mass, forces, rtags, images)
-            container.append(copy(snapshot))
+        # Initialize state and store method references
+        self.state = initialize()
+        self._restore = restore
+        self._method_update = method_update
+        self._bias = bias
+
+    def restore(self, prev_snapshot):
+        """Restore the simulation state from a previous snapshot."""
+        self.snapshot = prev_snapshot
+        self.forward_data(self._restore_callback, self.location, AccessMode.Overwrite, 0)
 
-        self.forward_data(snapshot_callback, default_location(), AccessMode.Read, 0)
-        return container[0]
+    def take_snapshot(self):
+        """Take a snapshot of the current simulation state."""
+        self.forward_data(self._snapshot_callback, self.location, AccessMode.Read, 0)
+        return self.snapshot
 
     def _pack_snapshot(self, positions, vel_mass, forces, rtags, images):
         return Snapshot(
@@ -134,44 +152,25 @@ def _pack_snapshot(self, positions, vel_mass, forces, rtags, images):
             from_dlpack(forces),
             from_dlpack(rtags),
             from_dlpack(images),
-            self.box,
+            self.update_box(),
             self.dt,
         )
 
+    # NOTE: The order of the callbacks arguments do not match that of the `Snapshot` attributes
+    def _restore_callback(self, positions, vel_mass, rtags, images, forces, _):
+        snapshot = self._pack_snapshot(positions, vel_mass, forces, rtags, images)
+        self._restore(snapshot, self.snapshot)
 
-if hasattr(AccessLocation, "OnDevice"):
+    def _snapshot_callback(self, positions, vel_mass, rtags, images, forces, _):
+        snapshot = self._pack_snapshot(positions, vel_mass, forces, rtags, images)
+        self.snapshot = copy(snapshot)
 
-    def default_location():
-        return AccessLocation.OnDevice
-
-else:
-
-    def default_location():
-        return AccessLocation.OnHost
-
-
-def take_snapshot(sampling_context, location=default_location()):
-    context = sampling_context.context
-    sysview = sampling_context.view
-    positions = copy(from_dlpack(positions_types(sysview, location, AccessMode.Read)))
-    vel_mass = copy(from_dlpack(velocities_masses(sysview, location, AccessMode.Read)))
-    forces = copy(from_dlpack(net_forces(sysview, location, AccessMode.ReadWrite)))
-    ids = copy(from_dlpack(rtags(sysview, location, AccessMode.Read)))
-    imgs = copy(from_dlpack(images(sysview, location, AccessMode.Read)))
-
-    check_device_array(positions)  # currently, we only support `DeviceArray`s
-
-    box = sysview.particle_data.getGlobalBox()
-    L = box.getL()
-    xy = box.getTiltFactorXY()
-    xz = box.getTiltFactorXZ()
-    yz = box.getTiltFactorYZ()
-    lo = box.getLo()
-    H = ((L.x, xy * L.y, xz * L.z), (0.0, L.y, yz * L.z), (0.0, 0.0, L.z))
-    origin = (lo.x, lo.y, lo.z)
-    dt = get_integrator(context).dt
-
-    return Snapshot(positions, vel_mass, forces, ids, imgs, Box(H, origin), dt)
+    def _update_callback(self, positions, vel_mass, rtags, images, forces, timestep):
+        snapshot = self._pack_snapshot(positions, vel_mass, forces, rtags, images)
+        self.state = self._method_update(snapshot, self.state)
+        self._bias(snapshot, self.state, self.view.synchronize)
+        if self.callback:
+            self.callback(snapshot, self.state, timestep)
 
 
 def build_snapshot_methods(sampling_method):
@@ -235,30 +234,40 @@ def dimensionality():
 
     snapshot_methods = build_snapshot_methods(sampling_method)
     flags = sampling_method.snapshot_flags
-    restore = partial(_restore, view)
+    restore = partial(pbs.restore, view)
     helpers = HelperMethods(build_data_querier(snapshot_methods, flags), dimensionality)
 
     return helpers, restore, bias
 
 
+def get_global_box(system_view):
+    """Get the box and origin of a HOOMD-blue simulation."""
+    box = system_view.particle_data.getGlobalBox()
+    L = box.getL()
+    xy = box.getTiltFactorXY()
+    xz = box.getTiltFactorXZ()
+    yz = box.getTiltFactorYZ()
+    lo = box.getLo()
+    H = ((L.x, xy * L.y, xz * L.z), (0.0, L.y, yz * L.z), (0.0, 0.0, L.z))
+    origin = (lo.x, lo.y, lo.z)
+    return H, origin
+
+
+def get_timestep(context):
+    """Get the timestep of a HOOMD-blue simulation."""
+    return get_integrator(context).dt
+
+
 def bind(sampling_context: SamplingContext, callback: Callable, **kwargs):
     context = sampling_context.context
     sampling_method = sampling_context.method
-    sysview = SystemView(get_system(context))
-    sampling_context.view = sysview
-    sampling_context.run = get_run_method(context)
-    helpers, restore, bias = build_helpers(context, sampling_method)
 
-    with sysview:
-        snapshot = take_snapshot(sampling_context)
-
-    method_bundle = sampling_method.build(snapshot, helpers)
-    sync_and_bias = partial(bias, sync_backend=sysview.synchronize)
-    sampler = Sampler(sysview, method_bundle, sync_and_bias, callback, restore)
+    sampler = Sampler(context, sampling_method, callback)
     set_half_step_hook(context, sampler)
-
     CONTEXTS_SAMPLERS[context] = sampler
 
+    sampling_context.run = get_run_method(context)
+
     return sampler
 
 

pysages/backends/jax-md.py

@@ -13,7 +13,7 @@
     SnapshotMethods,
     build_data_querier,
 )
-from pysages.typing import Callable, NamedTuple
+from pysages.typing import Callable
 from pysages.utils import check_device_array, copy
 
 
@@ -117,10 +117,6 @@ def run(timesteps):
     return run
 
 
-class View(NamedTuple):
-    synchronize: Callable
-
-
 def bind(sampling_context: SamplingContext, callback: Callable, **kwargs):
     context = sampling_context.context
     sampling_method = sampling_context.method
@@ -129,7 +125,6 @@ def bind(sampling_context: SamplingContext, callback: Callable, **kwargs):
     helpers = build_helpers(context, sampling_method)
     method_bundle = sampling_method.build(snapshot, helpers)
     sampler = Sampler(method_bundle, context_state, callback)
-    sampling_context.view = View((lambda: None))
     sampling_context.run = build_runner(
         context, sampler, jit_compile=kwargs.get("jit_compile", True)
     )