Add semantic latent trajectory views

app.py

@@ -67,6 +67,7 @@ scalability = load_csv("scalability_summary.csv")
 neuron_shap = load_csv("neuron_shap_summary.csv")
 trial_shapley = load_csv("trial_shapley_summary.csv")
 latent_samples = load_csv("latent_samples.csv")
+latent_trajectories = load_csv("latent_trajectories.csv")
 
 
 def present_rows(df: pd.DataFrame) -> pd.DataFrame:
@@ -104,6 +105,56 @@ DATASETS = ordered_unique(prediction.get("dataset", pd.Series(dtype=str)))
 MODELS = sorted(set(prediction.get("model", pd.Series(dtype=str)).dropna().astype(str)) | set(MODEL_INFO))
 DEFAULT_SELECTED_MODELS = [m for m in DEFAULT_MODELS if m in MODELS] or MODELS[:10]
 
+CONDITION_LABELS = {
+    "monkey": {
+        "0": "Up",
+        "1": "Up-right",
+        "2": "Right",
+        "3": "Down-right",
+        "4": "Down",
+        "5": "Down-left",
+        "6": "Left",
+        "7": "Up-left",
+    },
+    "allen_neuropixels": {
+        str(i): f"{angle} deg"
+        for i, angle in enumerate([0, 45, 90, 135, 180, 225, 270, 315])
+    },
+    "speech": {
+        "0": "Do nothing",
+        "1": "ban",
+        "2": "choice",
+        "3": "day",
+        "4": "feel",
+        "5": "kite",
+        "6": "though",
+        "7": "were",
+    },
+}
+
+
+def condition_sort_key(value: object) -> tuple[int, float | str]:
+    try:
+        return (0, float(value))
+    except (TypeError, ValueError):
+        return (1, str(value))
+
+
+def condition_label(dataset: str, condition: object) -> str:
+    if pd.isna(condition):
+        return "Unknown"
+    text = str(condition)
+    mapped = CONDITION_LABELS.get(dataset, {}).get(text)
+    if mapped is not None:
+        return mapped
+    if dataset == "ratinabox":
+        try:
+            idx = int(float(text))
+            return f"x{idx % 10}, y{idx // 10}"
+        except ValueError:
+            return text
+    return text
+
 
 def model_metadata_frame() -> pd.DataFrame:
     rows = []
@@ -256,10 +307,26 @@ def latent_space_figure(dataset: str, models: list[str] | None) -> go.Figure:
     for col in ["x", "y", "z"]:
         plot_df[col] = pd.to_numeric(plot_df[col], errors="coerce")
     plot_df["condition_num"] = pd.to_numeric(plot_df["condition"], errors="coerce")
+    plot_df["condition_label"] = plot_df["condition"].map(
+        lambda value: condition_label(dataset, value)
+    )
     plot_df = plot_df.dropna(subset=["x", "y", "z"])
     if plot_df.empty:
         return empty_figure("Latent-space samples are empty after filtering.")
 
+    trajectory_df = latent_trajectories[
+        (latent_trajectories["dataset"].astype(str) == str(dataset))
+        & (latent_trajectories["model"].astype(str) == str(model))
+    ].copy()
+    for col in ["x", "y", "z"]:
+        if col in trajectory_df:
+            trajectory_df[col] = pd.to_numeric(trajectory_df[col], errors="coerce")
+    if not trajectory_df.empty:
+        trajectory_df["condition_label"] = trajectory_df["condition"].map(
+            lambda value: condition_label(dataset, value)
+        )
+        trajectory_df = trajectory_df.dropna(subset=["x", "y", "z"])
+
     sessions = ordered_unique(plot_df["session_label"])
     n_cols = 2
     n_rows = int(np.ceil(len(sessions) / n_cols))
@@ -272,7 +339,10 @@ def latent_space_figure(dataset: str, models: list[str] | None) -> go.Figure:
         vertical_spacing=0.105,
     )
 
-    condition_values = sorted(plot_df["condition"].astype(str).unique())
+    condition_values = sorted(
+        plot_df["condition"].astype(str).unique(),
+        key=condition_sort_key,
+    )
     use_categorical_conditions = len(condition_values) <= 12
     direction_palette = [
         "#B23AEE",
@@ -292,7 +362,16 @@ def latent_space_figure(dataset: str, models: list[str] | None) -> go.Figure:
         condition: direction_palette[idx % len(direction_palette)]
         for idx, condition in enumerate(condition_values)
     }
-    condition_prefix = "Reach direction" if dataset == "monkey" else "Condition"
+    if dataset == "monkey":
+        condition_prefix = "Reach direction"
+    elif dataset == "ratinabox":
+        condition_prefix = "Position bin"
+    elif dataset == "allen_neuropixels":
+        condition_prefix = "Orientation"
+    elif dataset == "speech":
+        condition_prefix = "Cue"
+    else:
+        condition_prefix = "Condition"
 
     for session_idx, session in enumerate(sessions):
         session_df = plot_df[plot_df["session_label"].astype(str) == session]
@@ -306,7 +385,11 @@ def latent_space_figure(dataset: str, models: list[str] | None) -> go.Figure:
                 cond_df = session_df[session_df["condition"].astype(str) == condition]
                 if cond_df.empty:
                     continue
-                trace_name = f"{condition_prefix} {condition}"
+                trace_name = condition_label(dataset, condition)
+                session_traj = trajectory_df[
+                    (trajectory_df["session_label"].astype(str) == session)
+                    & (trajectory_df["condition"].astype(str) == condition)
+                ].sort_values("time_index")
                 fig.add_trace(
                     go.Scatter3d(
                         x=cond_df["x"],
@@ -317,13 +400,14 @@ def latent_space_figure(dataset: str, models: list[str] | None) -> go.Figure:
                         legendgroup=condition,
                         showlegend=session_idx == 0,
                         marker=dict(
-                            size=2.6,
-                            opacity=0.72,
+                            size=1.9 if not session_traj.empty else 2.6,
+                            opacity=0.28 if not session_traj.empty else 0.72,
                             color=condition_colors[condition],
                         ),
                         customdata=np.stack(
                             [
                                 cond_df["session_label"].astype(str),
+                                cond_df["condition_label"].astype(str),
                                 cond_df["trial_index"].astype(str),
                                 cond_df["time_index"].astype(str),
                             ],
@@ -331,14 +415,34 @@ def latent_space_figure(dataset: str, models: list[str] | None) -> go.Figure:
                         ),
                         hovertemplate=(
                             "session=%{customdata[0]}<br>"
-                            f"{condition_prefix.lower()}={condition}<br>"
-                            "trial=%{customdata[1]} time=%{customdata[2]}"
+                            f"{condition_prefix.lower()}=%{{customdata[1]}}<br>"
+                            "trial=%{customdata[2]} time=%{customdata[3]}"
                             "<extra></extra>"
                         ),
                     ),
                     row=row,
                     col=col,
                 )
+                if not session_traj.empty:
+                    fig.add_trace(
+                        go.Scatter3d(
+                            x=session_traj["x"],
+                            y=session_traj["y"],
+                            z=session_traj["z"],
+                            mode="lines",
+                            name=trace_name,
+                            legendgroup=condition,
+                            showlegend=False,
+                            line=dict(color=condition_colors[condition], width=5),
+                            hovertemplate=(
+                                f"{condition_prefix.lower()}={trace_name}<br>"
+                                "time=%{customdata}<extra></extra>"
+                            ),
+                            customdata=session_traj["time_index"],
+                        ),
+                        row=row,
+                        col=col,
+                    )
         else:
             fig.add_trace(
                 go.Scatter3d(
@@ -359,6 +463,7 @@ def latent_space_figure(dataset: str, models: list[str] | None) -> go.Figure:
                     customdata=np.stack(
                         [
                             session_df["condition"].astype(str),
+                            session_df["condition"].map(lambda value: condition_label(dataset, value)).astype(str),
                             session_df["trial_index"].astype(str),
                             session_df["time_index"].astype(str),
                         ],
@@ -366,8 +471,8 @@ def latent_space_figure(dataset: str, models: list[str] | None) -> go.Figure:
                     ),
                     hovertemplate=(
                         "session=%{fullData.name}<br>"
-                        f"{condition_prefix.lower()}=%{{customdata[0]}}<br>"
-                        "trial=%{customdata[1]} time=%{customdata[2]}"
+                        f"{condition_prefix.lower()}=%{{customdata[1]}}<br>"
+                        "trial=%{customdata[2]} time=%{customdata[3]}"
                         "<extra></extra>"
                     ),
                 ),