Merge pull request #785 from nkeim/fix-msd-multi-index

BUG: Let imsd() use output of subtract_drift()

Author: nkeim

trackpy/motion.py

@@ -212,7 +212,7 @@ def imsd(traj, mpp, fps, max_lagtime=100, statistic='msd', pos_columns=None):
     msds = []
     # Note: Index is set by msd, so we don't need to worry
     # about conformity here.
-    for pid, ptraj in traj.groupby('particle'):
+    for pid, ptraj in traj.reset_index(drop=True).groupby('particle'):
         msds.append(msd(ptraj, mpp, fps, max_lagtime, False, pos_columns))
         ids.append(pid)
     results = pandas_concat(msds, keys=ids)

trackpy/tests/test_motion.py

@@ -195,6 +195,18 @@ def test_linear_emsd(self):
         actual.index = expected.index
         assert_series_equal(np.round(actual), expected)
 
+    def test_linear_msd_drift(self):
+        '''Smoke test that imsd() and emsd() can use output of subtract_drift()'''
+        N = 10
+        EARLY = 7
+        drift = DataFrame(np.zeros((N - 1, 2)),
+                          np.arange(1, N, dtype=int)).astype('float64')
+        drift.columns = ['x', 'y']
+        drift.index.name = 'frame'
+        subtracted = tp.subtract_drift(self.many_walks, drift)
+        actual_imsd = tp.imsd(subtracted, 1, 1, max_lagtime=EARLY)
+        actual_emsd = tp.emsd(subtracted, 1, 1, max_lagtime=EARLY)
+
     def test_linear_emsd_gaps(self):
         A = 1
         EARLY = 4  # only early lag times have good stats