Add support for optional Snapshot attributes (#365)

Add support for optional Snapshot attributes

Author: p-zubieta

pysages/backends/ase.py

@@ -112,8 +112,7 @@ def take_snapshot(simulation, forces=None):
     origin = (0.0, 0.0, 0.0)
     dt = simulation.dt
 
-    # ASE doesn't use images explicitely
-    return Snapshot(positions, vel_mass, forces, ids, None, Box(H, origin), dt)
+    return Snapshot(positions, vel_mass, forces, ids, Box(H, origin), dt)
 
 
 def _calculator_defaults(sig, arg, default=[]):

pysages/backends/hoomd.py

@@ -151,9 +151,9 @@ def _pack_snapshot(self, positions, vel_mass, forces, rtags, images):
             from_dlpack(vel_mass),
             from_dlpack(forces),
             from_dlpack(rtags),
-            from_dlpack(images),
             self.update_box(),
             self.dt,
+            dict(images=from_dlpack(images)),  # extras
         )
 
     # NOTE: The order of the callbacks arguments do not match that of the `Snapshot` attributes
@@ -178,7 +178,7 @@ def build_snapshot_methods(sampling_method):
 
         def positions(snapshot):
             L = np.diag(snapshot.box.H)
-            return snapshot.positions[:, :3] + L * snapshot.images
+            return snapshot.positions[:, :3] + L * snapshot.extras["images"]
 
     else:
 

pysages/backends/lammps.py

@@ -92,14 +92,15 @@ def _partial_snapshot(self, include_masses: bool = False):
         velocities = from_dlpack(dlext.velocities(self.view, self.location))
         forces = from_dlpack(dlext.forces(self.view, self.location))
         tags_map = from_dlpack(dlext.tags_map(self.view, self.location))
-        imgs = from_dlpack(dlext.images(self.view, self.location))
+        images = from_dlpack(dlext.images(self.view, self.location))
 
         masses = None
         if include_masses:
             masses = from_dlpack(dlext.masses(self.view, self.location))
         vel_mass = (velocities, (masses, types))
+        extras = dict(images=images)
 
-        return Snapshot(positions, vel_mass, forces, tags_map, imgs, None, None)
+        return Snapshot(positions, vel_mass, forces, tags_map, None, None, extras)
 
     def _update_snapshot(self):
         s = self._partial_snapshot()
@@ -109,7 +110,7 @@ def _update_snapshot(self):
         box = self._update_box()
         dt = self.snapshot.dt
 
-        return Snapshot(s.positions, vel_mass, s.forces, s.ids[1:], s.images, box, dt)
+        return Snapshot(s.positions, vel_mass, s.forces, s.ids[1:], box, dt, s.extras)
 
     def restore(self, prev_snapshot):
         """Replaces this sampler's snapshot with `prev_snapshot`."""
@@ -122,7 +123,7 @@ def take_snapshot(self):
         dt = get_timestep(self.context)
 
         return Snapshot(
-            copy(s.positions), copy(s.vel_mass), copy(s.forces), s.ids[1:], copy(s.images), box, dt
+            copy(s.positions), copy(s.vel_mass), copy(s.forces), s.ids[1:], box, dt, copy(s.extras)
         )
 
 
@@ -198,7 +199,7 @@ def unpack(image):
 
         def positions(snapshot):
             L = np.diag(snapshot.box.H)
-            return snapshot.positions[:, :3] + L * vmap(unpack)(snapshot.images)
+            return snapshot.positions[:, :3] + L * vmap(unpack)(snapshot.extras["images"])
 
     else:
 

pysages/backends/openmm.py

@@ -86,8 +86,7 @@ def _take_snapshot(self):
         origin = (0.0, 0.0, 0.0)
         dt = context.getIntegrator().getStepSize() / unit.picosecond
 
-        # OpenMM doesn't have images
-        return Snapshot(positions, vel_mass, forces, ids, None, Box(H, origin), dt)
+        return Snapshot(positions, vel_mass, forces, ids, Box(H, origin), dt)
 
 
 def is_on_gpu(view: ContextView):

pysages/backends/snapshot.py

@@ -4,7 +4,16 @@
 from jax import jit
 from jax import numpy as np
 
-from pysages.typing import Callable, JaxArray, NamedTuple, Optional, Tuple, Union
+from pysages.typing import (
+    Any,
+    Callable,
+    Dict,
+    JaxArray,
+    NamedTuple,
+    Optional,
+    Tuple,
+    Union,
+)
 from pysages.utils import copy, dispatch, identity
 
 AbstractBox = NamedTuple("AbstractBox", [("H", JaxArray), ("origin", JaxArray)])
@@ -32,9 +41,9 @@ class Snapshot(NamedTuple):
     vel_mass: Union[JaxArray, Tuple[JaxArray, JaxArray]]
     forces: JaxArray
     ids: JaxArray
-    images: Optional[JaxArray]
     box: Box
     dt: Union[JaxArray, float]
+    extras: Optional[Dict[str, Any]] = None
 
     def __repr__(self):
         return "PySAGES " + type(self).__name__
@@ -81,9 +90,9 @@ def restore(view, snapshot, prev_snapshot, restore_vm=restore_vm):
     # Special handling for velocities and masses
     restore_vm(view, snapshot, prev_snapshot)
     # Overwrite images if the backend uses them
-    if snapshot.images is not None:
-        images = view(snapshot.images)
-        images[:] = view(prev_snapshot.images)
+    if hasattr(snapshot.extras, "images"):
+        images = view(snapshot.extras["images"])
+        images[:] = view(prev_snapshot.extras["images"])
 
 
 def build_data_querier(snapshot_methods, flags):

pysages/methods/abf.py

@@ -69,7 +69,7 @@ class ABFState(NamedTuple):
     force: JaxArray
     Wp: JaxArray
     Wp_: JaxArray
-    ncalls: int
+    ncalls: int = 0
 
     def __repr__(self):
         return repr("PySAGES " + type(self).__name__)
@@ -187,7 +187,7 @@ def initialize():
         force = np.zeros(dims)
         Wp = np.zeros(dims)
         Wp_ = np.zeros(dims)
-        return ABFState(xi, bias, hist, Fsum, force, Wp, Wp_, 0)
+        return ABFState(xi, bias, hist, Fsum, force, Wp, Wp_)
 
     def update(state, data):
         """

pysages/methods/ann.py

@@ -70,7 +70,7 @@ class ANNState(NamedTuple):
     phi: JaxArray
     prob: JaxArray
     nn: NNData
-    ncalls: int
+    ncalls: int = 0
 
     def __repr__(self):
         return repr("PySAGES " + type(self).__name__)
@@ -148,7 +148,7 @@ def initialize():
         phi = np.zeros(shape)
         prob = np.ones(shape)
         nn = NNData(ps, np.array(0.0), np.array(1.0))
-        return ANNState(xi, bias, hist, phi, prob, nn, 0)
+        return ANNState(xi, bias, hist, phi, prob, nn)
 
     def update(state, data):
         ncalls = state.ncalls + 1

pysages/methods/cff.py

@@ -93,7 +93,7 @@ class CFFState(NamedTuple):
     Wp_: JaxArray
     nn: NNData
     fnn: NNData
-    ncalls: int
+    ncalls: int = 0
 
     def __repr__(self):
         return repr("PySAGES " + type(self).__name__)
@@ -218,7 +218,7 @@ def initialize():
         nn = NNData(ps, np.array(0.0), np.array(1.0))
         fnn = NNData(fps, np.zeros(dims), np.array(1.0))
 
-        return CFFState(xi, bias, hist, histp, prob, fe, Fsum, force, Wp, Wp_, nn, fnn, 0)
+        return CFFState(xi, bias, hist, histp, prob, fe, Fsum, force, Wp, Wp_, nn, fnn)
 
     def update(state, data):
         # During the intial stage, when there are not enough collected samples, use ABF

pysages/methods/ffs.py

@@ -24,7 +24,7 @@
 class FFSState(NamedTuple):
     xi: JaxArray
     bias: Optional[JaxArray]
-    ncalls: int
+    ncalls: int = 0
 
     def __repr__(self):
         return repr("PySAGES " + type(self).__name__)
@@ -211,7 +211,7 @@ def _ffs(method, snapshot, helpers):
     # initialize method
     def initialize():
         xi = cv(helpers.query(snapshot))
-        return FFSState(xi, None, 0)
+        return FFSState(xi, None)
 
     def update(state, data):
         xi = cv(data)

pysages/methods/funn.py

@@ -78,7 +78,7 @@ class FUNNState(NamedTuple):
     Wp: JaxArray
     Wp_: JaxArray
     nn: NNData
-    ncalls: int
+    ncalls: int = 0
 
     def __repr__(self):
         return repr("PySAGES " + type(self).__name__)
@@ -182,7 +182,7 @@ def initialize():
         Wp = np.zeros(dims)
         Wp_ = np.zeros(dims)
         nn = NNData(ps, F, F)
-        return FUNNState(xi, bias, hist, Fsum, F, Wp, Wp_, nn, 0)
+        return FUNNState(xi, bias, hist, Fsum, F, Wp, Wp_, nn)
 
     def update(state, data):
         # During the intial stage, when there are not enough collected samples, use ABF