Smoothing 25x faster

source/classes.py

@@ -31,9 +31,9 @@ class project:
         smooth_alpha=0.1,
         arena_dims=[1],
         center_coords=True,
-        distances=False,
+        distances='All',
         ego=False,
-        angles=False,
+        angles=True,
         connectivity=None,
     ):
 
@@ -282,7 +282,7 @@ class coordinates:
         else:
             return "DLC analysis of {} videos".format(len(self._videos))
 
-    def get_coords(self, center=True, polar=False):
+    def get_coords(self, center=True, polar=False, speed=0):
         tabs = deepcopy(self._tables)
 
         if center:
@@ -300,6 +300,17 @@ class coordinates:
             for key, tab in tabs.items():
                 tabs[key] = tab2polar(tab)
 
+        if speed:
+            for order in range(speed):
+                for key, tab in tabs.items():
+                    try:
+                        cols = tab.columns.levels[0]
+                    except AttributeError:
+                        cols = tab.columns
+                    vel = rolling_speed(tab, typ="coords", order=order + 1)
+                    vel.columns = cols
+                    tabs[key] = vel
+
         return table_dict(
             tabs,
             "coords",
@@ -309,23 +320,50 @@ class coordinates:
             polar=polar,
         )
 
-    def get_distances(self):
+    def get_distances(self, speed=0):
+
+        tabs = deepcopy(self.distances)
+
         if self.distances is not None:
-            return table_dict(self.distances, typ="dists")
+
+            if speed:
+                for order in range(speed):
+                    for key, tab in tabs.items():
+                        try:
+                            cols = tab.columns.levels[0]
+                        except AttributeError:
+                            cols = tab.columns
+                        vel = rolling_speed(tab, typ="dists", order=order + 1)
+                        vel.columns = cols
+                        tabs[key] = vel
+
+            return table_dict(tabs, typ="dists")
+
         raise ValueError(
             "Distances not computed. Read the documentation for more details"
         )
 
-    def get_angles(self, degrees=False):
+    def get_angles(self, degrees=False, speed=0):
+
+        tabs = deepcopy(self.angles)
+
         if self.angles is not None:
-            if degrees == True:
-                return table_dict(
-                    {key: np.degrees(tab) for key, tab in self.angles.items()},
-                    typ="angles",
-                )
+            if degrees:
+                tabs = {key: np.degrees(tab) for key, tab in tabs.items()}
+
+            if speed:
+                for order in range(speed):
+                    for key, tab in tabs.items():
+                        try:
+                            cols = tab.columns.levels[0]
+                        except AttributeError:
+                            cols = tab.columns
+                        vel = rolling_speed(tab, typ="dists", order=order + 1)
+                        vel.columns = cols
+                        tabs[key] = vel
+
+            return table_dict(tabs, typ="angles")
 
-            else:
-                return table_dict(self.angles, typ="angles")
         raise ValueError("Angles not computed. Read the documentation for more details")
 
     def get_videos(self, play=False):

source/utils.py

@@ -244,11 +244,26 @@ def climb_wall(arena, pos_dict, fnum, tol, mouse):
     return np.linalg.norm(nose - center) > arena[2] + tol
 
 
-def rolling_speed(dframe, pause=10, rounds=5):
+def rolling_speed(dframe, typ, pause=10, rounds=5, order=1):
     """Returns the average speed over 10 frames in pixels per frame"""
 
-    distances = np.linalg.norm(np.array(dframe) - np.array(dframe.shift()), axis=1)
-    distances = pd.Series(distances, index=dframe.index)
+    s = dframe.shape[0]
+
+    if typ == "coords":
+        bp = dframe.shape[1] / 2 if order == 1 else dframe.shape[1]
+        d = 2 if order == 1 else 1
+
+    else:
+        bp = dframe.shape[1]
+        d = 1
+
+    distances = np.linalg.norm(
+        np.array(dframe).reshape(s, int(bp), d)
+        - np.array(dframe.shift()).reshape(s, int(bp), d),
+        axis=2,
+    )
+
+    distances = pd.DataFrame(distances, index=dframe.index)
     speeds = np.round(distances.rolling(pause).mean(), rounds)
     speeds[np.isnan(speeds)] = 0.0