Move images to the new Snapshot.extras attribute

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

@@ -41,7 +41,6 @@ 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
@@ -91,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):

tests/test_snapshots.py

@@ -19,7 +19,7 @@ def test_copying():
     box = Box(H, origin)
     dt = 0.1
 
-    old = Snapshot(positions, vel_mass, forces, ids, None, box, dt)
+    old = Snapshot(positions, vel_mass, forces, ids, box, dt)
     new = copy(old)
 
     old_ptr = old.positions.unsafe_buffer_pointer()