From bc56ee675f62d42da0a73564c55d81b24a2c5030 Mon Sep 17 00:00:00 2001
From: hausmanns <37653619+hausmanns@users.noreply.github.com>
Date: Wed, 27 Apr 2022 11:51:14 +0200
Subject: [PATCH] Edit to umap functions to allow for less than 3 components,
 docs (#11)

* minor edits to umap fncs

* added pcutoff for umap, set x,y to 0, all keyp. op
---
 README.md                  |  4 ++--
 dlc2kinematics/mainfxns.py | 23 ++++++++++++++++-------
 dlc2kinematics/plotting.py | 22 ++++++++++++++++------
 3 files changed, 34 insertions(+), 15 deletions(-)

diff --git a/README.md b/README.md
index a17b388..25656a6 100644
--- a/README.md
+++ b/README.md
@@ -110,9 +110,9 @@ dlc2kinematics.plot_3d_pca_reconstruction(df_vel, n_components=10, framenumber=5
 
 ### UMAP Embeddings
 ``` python
-data = dlc2kinematics.compute_umap(df, key=['LeftForelimb', 'RightForelimb'], chunk_length=30, fit_transform=True, n_neighbors=30, n_components=3,metric="euclidean")
+embedding, transformed_data = dlc2kinematics.compute_umap(df, key=['LeftForelimb', 'RightForelimb'], chunk_length=30, fit_transform=True, n_neighbors=30, n_components=3,metric="euclidean")
 
-dlc2kinematics.plot_umap(data, size=5, alpha=1, color="indigo", figsize=(10, 6))
+dlc2kinematics.plot_umap(transformed_data, size=5, alpha=1, color="indigo", figsize=(10, 6))
 ```
 
 ## Contributing
diff --git a/dlc2kinematics/mainfxns.py b/dlc2kinematics/mainfxns.py
index 1fe8ead..4d1e863 100644
--- a/dlc2kinematics/mainfxns.py
+++ b/dlc2kinematics/mainfxns.py
@@ -193,7 +193,8 @@ def extract_kinematic_synergies(
 
 def compute_umap(
     df,
-    key,
+    keypoints=None,
+    pcutoff=0.6,
     chunk_length=30,
     fit_transform=True,
     n_neighbors=30,
@@ -234,7 +235,9 @@ def compute_umap(
     ----------
     df: original dataframe df, _, _ = dlc2kinematics.load_data(DLC_2D_file)
 
-    key: list of limbs of interests ex: key = ['LeftForelimb', 'RightForelimb']
+    keypoints: list of limbs of interests ex: key = ['LeftForelimb', 'RightForelimb'], if None, all bodyparts are taken into account
+
+    pcutoff: likelihood at which the keypoints are kept for analysis, if under pcutoff, set coord to 0
 
     chunk_length: Number of frames per segment. #TODO: add the posibility of a sliding window?
 
@@ -299,11 +302,17 @@ def compute_umap(
 
     # pandas dataframe managing
     df_clean = df.copy()
-    df_clean.columns = df_clean.columns.droplevel(
-        0
-    )  # Drop scorer row to make it easier to navigate in the df table
-    key2 = ["x", "y"]
-    df_limbs = df_clean.loc[:, (key, key2)]  # Drop likelihood column
+
+    if keypoints is None: # If no keypoints specified, use all
+        keypoints = df_clean.columns.get_level_values('bodyparts').unique().to_list()
+    
+    df_limbs = df_clean.loc[:, pd.IndexSlice[:, keypoints]]
+
+    temp = df_limbs.stack(level=['scorer', 'bodyparts']) # Stack with likelihood, x, y
+    temp.loc[temp['likelihood'] < pcutoff, ['x','y']] = 0.0 # Set values under pcutoff to 0.0 to exclude
+    unstacked_temp = temp.unstack(level=['scorer', 'bodyparts']) # Unstack again
+    unstacked_temp.reorder_levels(['scorer','bodyparts','coords'], axis=1).reindex_like(df_limbs) # Re-index like original df
+
     n_frames, n_bodyparts = df_limbs.shape
     n_chunks = n_frames // chunk_length
 
diff --git a/dlc2kinematics/plotting.py b/dlc2kinematics/plotting.py
index f17a6c0..a996583 100644
--- a/dlc2kinematics/plotting.py
+++ b/dlc2kinematics/plotting.py
@@ -273,9 +273,19 @@ def plot_umap(Y, size=5, alpha=1, color="indigo", figsize=(10, 6)):
     >>> c = np.random.rand(Y.shape[0],3)
     >>> dlc2kinematics.plot_umap(Y, 5, 1, c)
     """
-
-    fig = plt.figure(figsize=figsize)
-    ax = Axes3D(fig)
-    ax.view_init(0, 20)
-    ax.dist = 8
-    ax.scatter(Y[:, 0], Y[:, 2], Y[:, 1], alpha=alpha, s=size, c=color)
+    
+    if Y.shape[1] < 2:
+        print("Please pick at least 2 components for plotting.")
+    else:
+        fig, ax = plt.subplots(figsize=figsize)
+
+        if Y.shape[1] > 2:
+            ax = Axes3D(fig)
+            ax.view_init(0, 20)
+            ax.dist = 8
+            ax.scatter(Y[:, 0], Y[:, 1], Y[:, 2], alpha=alpha, s=size, c=color)
+        else:
+            ax.scatter(Y[:, 0], Y[:, 1], alpha=alpha, s=size, c=color)
+    
+
+