Merge pull request #1343 from glotzerlab/fix_msd_nanobind

Fix msd

Author: DomFijan

ChangeLog.md

@@ -12,6 +12,7 @@ The format is based on [Keep a Changelog] and this project adheres to
 
 ### Fixed
 * Type hinting in `UnitCell` no longer infers incorrect type.
+* Fixed `MSD` compute to correctly preprocess the data for msd calculation.
 
 ## 3.3.1 -- 2025-04-28
 

freud/msd.py

@@ -205,7 +205,11 @@ def compute(self, positions, images=None, reset=True):
         # Make sure we aren't modifying the provided array
         if self._box is not None and images is not None:
             unwrapped_positions = positions.copy()
-            positions = self._box.wrap(unwrapped_positions, images)
+            for i in range(positions.shape[0]):
+                unwrapped_positions[i, :, :] = self._box.unwrap(
+                    unwrapped_positions[i, :, :], images[i, :, :]
+                )
+            positions = unwrapped_positions
 
         if self.mode == "window":
             # First compute the first term r^2(k+m) - r^2(k)

tests/test_msd_msd.py

@@ -81,6 +81,26 @@ def simple_msd(positions):
             npt.assert_allclose(solution, simple, atol=1e-6)
             npt.assert_allclose(solution_particle, simple_particle, atol=1e-5)
 
+    def test_MSD_compute(self):
+        n_frames = 5
+        n_particles = 10
+
+        positions = np.zeros((n_frames, n_particles, 3), dtype=np.float32)
+        for t in range(n_frames):
+            # at frame t, x = t, y = z = 0
+            positions[t, :, 0] = t
+
+        images = np.zeros((n_frames, n_particles, 3), dtype=np.int32)
+        box = freud.box.Box.cube(30)
+
+        msd = freud.msd.MSD(box=box, mode="window")
+        msd.compute(positions=positions, images=images, reset=False)
+
+        expected_msd = np.arange(n_frames) ** 2
+        expected_particle = np.tile(expected_msd, (n_particles, 1)).T
+        assert np.allclose(msd.msd, expected_msd)
+        assert np.allclose(msd.particle_msd, expected_particle)
+
     def test_repr(self):
         msd = freud.msd.MSD()
         assert str(msd) == str(eval(repr(msd)))