[lammps] Fix handling types when restoring (#344)

For LAMMPS, only unique masses are stored, instead the mass "type" of each particle is stored, and the mass of the `i`th particle can be retrieved as `masses[types][i]`. For this reason we store `(masses, types)` in place of a `masses`, this requires special handling when restoring which we missed initially. Fixing it here.

Author: p-zubieta

pysages/backends/lammps.py

@@ -154,6 +154,16 @@ def add_bias(forces, biases):
         def add_bias(forces, biases):
             forces[:, :3] += factor * biases
 
+    def restore_vm(view, snapshot, prev_snapshot):
+        velocities = view(snapshot.vel_mass[0])
+        masses_types = snapshot.vel_mass[1]
+        masses = view(masses_types[0])
+        types = view(masses_types[1])
+        prev_masses_types = prev_snapshot.vel_mass[1]
+        velocities[:] = view(prev_snapshot.vel_mass[0])
+        masses[:] = view(prev_masses_types[0])
+        types[:] = view(prev_masses_types[1])
+
     # TODO: check if this can be sped up.  # pylint: disable=W0511
     def bias(snapshot, state):
         """Adds the computed bias to the forces."""
@@ -166,7 +176,7 @@ def bias(snapshot, state):
 
     snapshot_methods = build_snapshot_methods(sampling_method, on_gpu)
     flags = sampling_method.snapshot_flags
-    restore = partial(restore_fn, view)
+    restore = partial(restore_fn, view, restore_vm=restore_vm)
     helpers = HelperMethods(build_data_querier(snapshot_methods, flags), lambda: dim)
 
     return helpers, restore, bias