Added tests for preprocess.py

deepof/preprocess.py

@@ -248,7 +248,7 @@ class project:
             dats = []
             for clique in cliques:
                 dat = pd.DataFrame(
-                    angle_trio(np.array(tab[clique]).reshape(3, tab.shape[0], 2))
+                    angle_trio(np.array(tab[clique]).reshape([3, tab.shape[0], 2]))
                 ).T
 
                 orders = [[0, 1, 2], [0, 2, 1], [1, 0, 2]]
@@ -383,15 +383,9 @@ class coordinates:
                 ]
 
         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, deriv=order + 1)
-                    vel.columns = cols
-                    tabs[key] = vel
+            for key, tab in tabs.items():
+                vel = rolling_speed(tab, deriv=speed + 1, center=center)
+                tabs[key] = vel
 
         if length:
             for key, tab in tabs.items():
@@ -429,15 +423,9 @@ class coordinates:
         if self.distances is not None:
 
             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, deriv=order + 1)
-                        vel.columns = cols
-                        tabs[key] = vel
+                for key, tab in tabs.items():
+                    vel = rolling_speed(tab, deriv=speed + 1, typ="dists")
+                    tabs[key] = vel
 
             if length:
                 for key, tab in tabs.items():
@@ -460,15 +448,9 @@ class coordinates:
                 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, deriv=order + 1)
-                        vel.columns = cols
-                        tabs[key] = vel
+                for key, tab in tabs.items():
+                    vel = rolling_speed(tab, deriv=speed + 1, typ="angles")
+                    tabs[key] = vel
 
             if length:
                 for key, tab in tabs.items():

deepof/utils.py

@@ -484,7 +484,12 @@ def climb_wall(
 
 
 def rolling_speed(
-    dframe: pd.DatetimeIndex, window: int = 5, rounds: int = 10, deriv: int = 1
+    dframe: pd.DatetimeIndex,
+    window: int = 5,
+    rounds: int = 10,
+    deriv: int = 1,
+    center: str = None,
+    typ: str = "coords",
 ) -> pd.DataFrame:
     """Returns the average speed over n frames in pixels per frame
 
@@ -494,29 +499,44 @@ def rolling_speed(
             - rounds (int): float rounding decimals
             - deriv (int): position derivative order; 1 for speed,
             2 for acceleration, 3 for jerk, etc
+            - center (str): for internal usage only; solves an issue
+            with pandas.MultiIndex that arises when centering frames
+            to a specific body part
 
         Returns:
             - speeds (pd.DataFrame): containing 2D speeds for each body part
             in the original data or their consequent derivatives"""
 
     original_shape = dframe.shape
-    body_parts = dframe.columns.levels[0]
+    if center:
+        body_parts = [bp for bp in dframe.columns.levels[0] if bp != center]
+    else:
+        try:
+            body_parts = dframe.columns.levels[0]
+        except AttributeError:
+            body_parts = dframe.columns
+
     speeds = pd.DataFrame
 
     for der in range(deriv):
+
+        features = 2 if der == 0 and typ == "coords" else 1
+
         distances = np.concatenate(
             [
-                np.array(dframe).reshape([-1, (2 if der == 0 else 1)], order="F"),
-                np.array(dframe.shift()).reshape(
-                    [-1, (2 if der == 0 else 1)], order="F"
-                ),
+                np.array(dframe).reshape([-1, features], order="F"),
+                np.array(dframe.shift()).reshape([-1, features], order="F"),
             ],
             axis=1,
         )
 
         distances = np.array(compute_dist(distances))
         distances = distances.reshape(
-            [original_shape[0], original_shape[1] // 2], order="F"
+            [
+                original_shape[0],
+                (original_shape[1] // 2 if typ == "coords" else original_shape[1]),
+            ],
+            order="F",
         )
         distances = pd.DataFrame(distances, index=dframe.index)
         speeds = np.round(distances.rolling(window).mean(), rounds)

tests/test_preprocess.py

 # @author lucasmiranda42
 
 from hypothesis import given
-from hypothesis import HealthCheck
 from hypothesis import settings
 from hypothesis import strategies as st
-from hypothesis.extra.numpy import arrays
-from hypothesis.extra.pandas import range_indexes, columns, data_frames
-from scipy.spatial import distance
+from collections import defaultdict
 from deepof.utils import *
 import deepof.preprocess
 import pytest
@@ -126,3 +123,66 @@ def test_run(nodes, ego):
     ).run(verbose=True)
 
     assert type(prun) == deepof.preprocess.coordinates
+
+
+@settings(deadline=None)
+@given(
+    nodes=st.integers(min_value=0, max_value=1),
+    ego=st.integers(min_value=0, max_value=2),
+    sampler=st.data(),
+)
+def test_get_table_dicts(nodes, ego, sampler):
+
+    nodes = ["All", ["Center", "Nose", "Tail_base"]][nodes]
+    ego = [False, "Center", "Nose"][ego]
+
+    prun = deepof.preprocess.project(
+        path=os.path.join(".", "tests", "test_examples"),
+        arena="circular",
+        arena_dims=[380],
+        video_format=".mp4",
+        table_format=".h5",
+        distances=nodes,
+        ego=ego,
+    ).run(verbose=True)
+
+    coords = prun.get_coords(
+        center=sampler.draw(st.one_of(st.just("arena"), st.just("Center"))),
+        polar=sampler.draw(st.booleans()),
+        length=sampler.draw(st.one_of(st.just(False), st.just("00:10:00"))),
+        align=sampler.draw(st.one_of(st.just(False), st.just("Nose"))),
+    )
+    speeds = prun.get_coords(
+        center=sampler.draw(st.one_of(st.just("arena"), st.just("Center"))),
+        polar=sampler.draw(st.booleans()),
+        length=sampler.draw(st.one_of(st.just(False), st.just("00:10:00"))),
+        speed=sampler.draw(st.integers(min_value=0, max_value=5)),
+    )
+    distances = prun.get_distances(
+        length=sampler.draw(st.one_of(st.just(False), st.just("00:10:00"))),
+        speed=sampler.draw(st.integers(min_value=0, max_value=5)),
+    )
+    angles = prun.get_angles(
+        degrees=sampler.draw(st.booleans()),
+        length=sampler.draw(st.one_of(st.just(False), st.just("00:10:00"))),
+        speed=sampler.draw(st.integers(min_value=0, max_value=5)),
+    )
+
+    # deepof.coordinates testing
+
+    assert type(coords) == deepof.preprocess.table_dict
+    assert type(speeds) == deepof.preprocess.table_dict
+    assert type(distances) == deepof.preprocess.table_dict
+    assert type(angles) == deepof.preprocess.table_dict
+    assert type(prun.get_videos()) == list
+    assert prun.get_exp_conditions is None
+    assert type(prun.get_quality()) == defaultdict
+    assert type(prun.get_arenas) == tuple
+
+    # deepof.table_dict testing
+
+    table = sampler.draw(
+        st.one_of(st.just(coords), st.just(speeds), st.just(distances), st.just(angles))
+    )
+
+    #table.filter()
\ No newline at end of file