Implemented rotation alignment in preprocess.py

model_training.py

@@ -232,7 +232,7 @@ DLC_social_1_coords = DLC_social_1.run(verbose=True)
 DLC_social_2_coords = DLC_social_2.run(verbose=True)
 
 # Coordinates for training data
-coords1 = DLC_social_1_coords.get_coords(center="B_Center", polar=False)
+coords1 = DLC_social_1_coords.get_coords(center="B_Center", align="B_Nose")
 distances1 = DLC_social_1_coords.get_distances()
 angles1 = DLC_social_1_coords.get_angles()
 coords_distances1 = merge_tables(coords1, distances1)
@@ -241,7 +241,7 @@ dists_angles1 = merge_tables(distances1, angles1)
 coords_dist_angles1 = merge_tables(coords1, distances1, angles1)
 
 # Coordinates for validation data
-coords2 = DLC_social_2_coords.get_coords(center="B_Center", polar=False)
+coords2 = DLC_social_2_coords.get_coords(center="B_Center", align="B_Nose")
 distances2 = DLC_social_2_coords.get_distances()
 angles2 = DLC_social_2_coords.get_angles()
 coords_distances2 = merge_tables(coords2, distances2)
@@ -258,6 +258,7 @@ input_dict_train = {
         random_state=42,
         filter="gaussian",
         sigma=55,
+        align=True,
     ),
     "dists": distances1.preprocess(
         window_size=11,
@@ -266,6 +267,7 @@ input_dict_train = {
         random_state=42,
         filter="gaussian",
         sigma=55,
+        align=True,
     ),
     "angles": angles1.preprocess(
         window_size=11,
@@ -274,6 +276,7 @@ input_dict_train = {
         random_state=42,
         filter="gaussian",
         sigma=55,
+        align=True,
     ),
     "coords+dist": coords_distances1.preprocess(
         window_size=11,
@@ -282,6 +285,7 @@ input_dict_train = {
         random_state=42,
         filter="gaussian",
         sigma=55,
+        align=True,
     ),
     "coords+angle": coords_angles1.preprocess(
         window_size=11,
@@ -290,6 +294,7 @@ input_dict_train = {
         random_state=42,
         filter="gaussian",
         sigma=55,
+        align=True,
     ),
     "dists+angle": dists_angles1.preprocess(
         window_size=11,
@@ -298,6 +303,7 @@ input_dict_train = {
         random_state=42,
         filter="gaussian",
         sigma=55,
+        align=True,
     ),
     "coords+dist+angle": coords_dist_angles1.preprocess(
         window_size=11,
@@ -306,54 +312,60 @@ input_dict_train = {
         random_state=42,
         filter="gaussian",
         sigma=55,
+        align=True,
     ),
 }
 
 input_dict_val = {
     "coords": coords2.preprocess(
         window_size=11,
-        window_step=1,
+        window_step=10,
         scale=True,
         random_state=42,
         filter="gaussian",
         sigma=55,
         shuffle=True,
+        align=True,
     ),
     "dists": distances2.preprocess(
         window_size=11,
-        window_step=1,
+        window_step=10,
         scale=True,
         random_state=42,
         filter="gaussian",
         sigma=55,
         shuffle=True,
+        align=True,
     ),
     "angles": angles2.preprocess(
         window_size=11,
-        window_step=1,
+        window_step=10,
         scale=True,
         random_state=42,
         filter="gaussian",
         sigma=55,
         shuffle=True,
+        align=True,
     ),
     "coords+dist": coords_distances2.preprocess(
         window_size=11,
-        window_step=1,
+        window_step=10,
         scale=True,
         random_state=42,
         filter="gaussian",
         sigma=55,
         shuffle=True,
+        align=True,
     ),
     "coords+angle": coords_angles2.preprocess(
         window_size=11,
-        window_step=1,
+        window_step=10,
         scale=True,
         random_state=42,
         filter="gaussian",
         sigma=55,
         shuffle=True,
+        align=True,
     ),
     "dists+angle": dists_angles2.preprocess(
         window_size=11,
@@ -363,15 +375,17 @@ input_dict_val = {
         filter="gaussian",
         sigma=55,
         shuffle=True,
+        align=True,
     ),
     "coords+dist+angle": coords_dist_angles2.preprocess(
         window_size=11,
-        window_step=1,
+        window_step=10,
         scale=True,
         random_state=42,
         filter="gaussian",
         sigma=55,
         shuffle=True,
+        align=True,
     ),
 }
 

source/preprocess.py

@@ -309,7 +309,9 @@ class coordinates:
         else:
             return "DLC analysis of {} videos".format(len(self._videos))
 
-    def get_coords(self, center="arena", polar=False, speed=0, length=None):
+    def get_coords(
+        self, center="arena", polar=False, speed=0, length=None, align=False
+    ):
         tabs = deepcopy(self._tables)
 
         if polar:
@@ -384,6 +386,21 @@ class coordinates:
                     "00:00:00", length, periods=tab.shape[0] + 1, closed="left"
                 )
 
+        if align:
+            assert (
+                align in list(tabs.values())[0].columns.levels[0]
+            ), "align must be set to the name of a bodypart"
+
+            for key, tab in tabs.items():
+
+                # Bring forward the column to align
+                columns = [i for i in tab.columns if align not in i]
+                columns = [
+                    (align, ("phi" if polar else "x")),
+                    (align, ("rho" if polar else "y")),
+                ] + columns
+                tabs[key] = tab[columns]
+
         return table_dict(
             tabs,
             "coords",
@@ -524,17 +541,17 @@ class table_dict(dict):
         self,
         window_size=1,
         window_step=1,
-        scale=True,
+        scale="standard",
         test_proportion=0,
         random_state=None,
         verbose=False,
         filter=None,
         sigma=None,
         shift=0,
-        standard_scaler=True,
         shuffle=False,
+        align=False,
     ):
-        """Builds a sliding window. If desired, splits train and test and
+        """Builds a sliding window. If specified, splits train and test and
            Z-scores the data using sklearn's standard scaler"""
 
         X_train = self.get_training_set()
@@ -550,16 +567,20 @@ class table_dict(dict):
             if verbose:
                 print("Scaling data...")
 
-            if standard_scaler:
+            if scale == "standard":
                 scaler = StandardScaler()
-            else:
+            elif scale == "minmax":
                 scaler = MinMaxScaler()
+            else:
+                raise ValueError(
+                    "Invalid scaler. Select one of standard, minmax or None"
+                )
 
             X_train = scaler.fit_transform(
                 X_train.reshape(-1, X_train.shape[-1])
             ).reshape(X_train.shape)
 
-            if standard_scaler:
+            if scale == "standard":
                 assert np.allclose(np.mean(X_train), 0)
                 assert np.allclose(np.std(X_train), 1)
 
@@ -573,6 +594,9 @@ class table_dict(dict):
 
         X_train = rolling_window(X_train, window_size, window_step)
 
+        if align:
+            X_train = align_trajectories(X_train)
+
         if filter == "gaussian":
             r = range(-int(window_size / 2), int(window_size / 2) + 1)
             r = [i - shift for i in r]

source/utils.py

@@ -11,6 +11,7 @@ import pims
 import re
 import scipy
 import seaborn as sns
+from copy import deepcopy
 from itertools import cycle, combinations, product
 from joblib import Parallel, delayed
 from numba import jit
@@ -87,6 +88,32 @@ def angle_trio(array, degrees=False):
     return np.array([angle(a, b, c), angle(a, c, b), angle(b, a, c),])
 
 
+def rotate(p, angles, origin=np.array([0, 0])):
+    R = np.array([[np.cos(angles), -np.sin(angles)], [np.sin(angles), np.cos(angles)]])
+
+    o = np.atleast_2d(origin)
+    p = np.atleast_2d(p)
+
+    return np.squeeze((R @ (p.T - o.T) + o.T).T)
+
+
+def align_trajectories(data):
+
+    data = deepcopy(data)
+
+    center_time = (data.shape[1] - 1) // 2
+    angles = np.arctan2(data[:, center_time, 0], data[:, center_time, 1])
+
+    aligned_trajs = np.zeros(data.shape)
+
+    for frame in range(data.shape[0]):
+        aligned_trajs[frame] = rotate(
+            data[frame].reshape([data.shape[1] * data.shape[2] // 2, 2]), angles[frame],
+        ).reshape(data.shape[1:])
+
+    return aligned_trajs
+
+
 def smooth_boolean_array(a):
     """Returns a boolean array in which isolated appearances of a feature are smoothened"""
     for i in range(1, len(a) - 1):