Update Figure 2B

src/ExtendedDataFigure3.py

@@ -15,28 +15,24 @@
 
 # # Extended Data Figure 3: CEBRA produces consistent, highly decodable embeddings
 
-# - Additional rat data shown for all algorithms we benchmarked (see Methods). CEBRA was trained with output latent 3D (the minimum) and all other methods were obtained with a 2D latent
+# - Additional rat data shown for all algorithms we benchmarked (see Methods). CEBRA was trained with output latent on the 2-sphere (the minimum) and all other methods were obtained with a 2D latent in Euclidean space.
 
 # +
 import pandas as pd
 import matplotlib.pyplot as plt
 import seaborn as sns
 
+df = pd.concat([
+  pd.read_hdf("../data/EDFigure3.h5", key="data"),
+  pd.read_hdf("../data/EDFigure3_addition.h5", key="data")
+], axis = 0, ignore_index = True)
 
 def scatter(data, index, ax, s=0.01, alpha=0.5):
     mask = index[:, 1] > 0
     ax.scatter(*data[mask].T, c=index[mask, 0], s=s, cmap="viridis", alpha=alpha)
     ax.scatter(*data[~mask].T, c=index[~mask, 0], s=s, cmap="cool", alpha=alpha)
 
 
-df = pd.concat(
-    [
-        pd.read_hdf("../data/EDFigure3.h5", key="data"),
-        pd.read_hdf("../data/EDFigure3_extension.h5", key="data"),
-    ],
-    axis=0,
-).reset_index(drop=True)
-
 fig = plt.figure(figsize=(4 * 3, 7 * 3), dpi=600)
 for i in df.index:
     ax = fig.add_subplot(7, 4, i + 1)
@@ -51,4 +47,45 @@ def scatter(data, index, ax, s=0.01, alpha=0.5):
         ax.set_title(f"Rat {df.loc[i, 'animal']}", fontsize=18)
     # first column labels
     if i % 4 == 0:
-        ax.set_ylabel(df.loc[i, "method"], fontsize=18)
+        ax.set_ylabel(df.loc[i, "method"])
+
+
+# -
+
+# For a higher resolution plot, we export each row as a separate file:
+
+# +
+def scatter(data, index, ax, s=0.01, alpha=0.5):
+    mask = index[:, 1] > 0
+    ax.scatter(*data[mask].T, c=index[mask, 0], s=s, cmap="viridis", alpha=alpha)
+    ax.scatter(*data[~mask].T, c=index[~mask, 0], s=s, cmap="cool", alpha=alpha)
+
+def export_highres():
+  for method in df.method.unique():
+      print(method)
+      fig = plt.figure(figsize=(4 * 3, 1 * 3), dpi=600)    
+      entry = df[df.method == method].set_index("animal")
+      for i, animal in enumerate(sorted(entry.index)):
+          ax = fig.add_subplot(1, 4, i + 1)
+          scatter(
+            entry.loc[animal, "emission"][:, :2],
+            entry.loc[animal, "labels"],
+            ax=ax, s=0.5, alpha=0.7
+          )
+          ax.set_yticklabels([])
+          ax.set_xticklabels([])
+          ax.set_xticks([])
+          ax.set_yticks([])
+          ax.set_aspect("equal")
+          sns.despine(bottom=True, left=True, ax=ax)
+          # first row labels
+          #if i // 4 == 0:
+          #    ax.set_title(f"Rat {df.loc[i, 'animal']}")
+          # first column labels
+          if i % 4 == 0:
+              ax.set_ylabel(method)
+      method = method.replace('/', '-')
+      plt.savefig(f'edf3_{method}.png', bbox_inches = "tight", transparent = True)
+      plt.show()
+
+export_highres()

src/Figure2.py

@@ -30,7 +30,10 @@
 
 # -
 
-data = pd.read_hdf("../data/Figure2.h5", key="data")
+data = pd.concat([
+  pd.read_hdf("../data/Figure2.h5", key="data"),
+  pd.read_hdf("../data/Figure2_addition.h5", key="data")
+], axis = 0)
 
 # ## Figure 2b
 #
@@ -40,21 +43,36 @@
 method_viz = data["visualization"]
 
 fig = plt.figure(figsize=(20, 5))
-for i, model in enumerate(["hypothesis", "shuffled", "discovery", "hybrid"]):
-    ax = fig.add_subplot(1, 4, i + 1, projection="3d")
+for i, (model, title) in enumerate(zip(
+  ["behavior_mse", "discovery", "hypothesis", "hybrid", "shuffled"],
+  ["Behavioral data:\nposition", "Discovery:\n time only", "Hypothesis:\nposition", "Hybrid:\ntime+behavior", "Shuffled Labels"]
+)):
     emb = method_viz[model]["embedding"]
     label = method_viz[model]["label"]
     r = label[:, 1] == 1
     l = label[:, 2] == 1
     idx1, idx2, idx3 = (0, 1, 2)
-    if i == 3:
-        idx1, idx2, idx3 = (1, 2, 0)
-    ax.scatter(
-        emb[l, idx1], emb[l, idx2], emb[l, idx3], c=label[l, 0], cmap="cool", s=0.1
-    )
-    ax.scatter(emb[r, idx1], emb[r, idx2], emb[r, idx3], c=label[r, 0], s=0.1)
+    if i == 0:
+        ax = fig.add_subplot(1, 5, i + 1)
+        ax.scatter(
+            emb[l, 1], emb[l, 0], c=label[l, 0], cmap="cool", s=0.1
+        )
+        ax.scatter(emb[r, 1], emb[r, 0], c=label[r, 0], s=0.1)
+        #ax.set_xlim()
+    else:
+        ax = fig.add_subplot(1, 5, i + 1, projection="3d")
+        if i == 4:
+            idx1, idx2, idx3 = (1, 2, 0)
+        ax.scatter(
+            emb[l, idx1], emb[l, idx2], emb[l, idx3], c=label[l, 0], cmap="cool", s=0.1
+        )
+        ax.scatter(emb[r, idx1], emb[r, idx2], emb[r, idx3], c=label[r, 0], s=0.1)
+        lim = -.6, .6
+        ax.set_xlim(lim)
+        ax.set_ylim(lim)
+        ax.set_zlim(lim)
     ax.axis("off")
-    ax.set_title(f"{model}", fontsize=20)
+    ax.set_title(title, fontsize=20)
 # -
 
 # ## Figure 2c