Fix latent consistency visualization

app.py

@@ -189,6 +189,24 @@ def empty_figure(message: str) -> go.Figure:
     return fig_layout(fig)
 
 
+def latent_fig_layout(fig: go.Figure, *, height: int) -> go.Figure:
+    fig.update_layout(
+        height=height,
+        paper_bgcolor="white",
+        plot_bgcolor="white",
+        margin=dict(l=8, r=8, t=76, b=84),
+        font=dict(family="Inter, Arial, sans-serif", size=13, color="#1f2933"),
+        legend=dict(
+            orientation="h",
+            yanchor="top",
+            y=-0.08,
+            xanchor="left",
+            x=0,
+        ),
+    )
+    return fig
+
+
 def parse_float_list(value: object) -> list[float]:
     if pd.isna(value):
         return []
@@ -226,53 +244,102 @@ def latent_space_figure(dataset: str, models: list[str] | None) -> go.Figure:
     ].copy()
     if not score_df.empty:
         score_df["mean_r2"] = pd.to_numeric(score_df["mean_r2"], errors="coerce")
-        plot_models = (
+        model = (
             score_df.sort_values("mean_r2", ascending=False)["model"]
             .astype(str)
-            .head(4)
-            .tolist()
+            .iloc[0]
         )
     else:
-        plot_models = sorted(df["model"].astype(str).unique())[:4]
+        model = sorted(df["model"].astype(str).unique())[0]
 
-    plot_df = df[df["model"].astype(str).isin(plot_models)].copy()
+    plot_df = df[df["model"].astype(str) == model].copy()
     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 = plot_df.dropna(subset=["x", "y", "z"])
     if plot_df.empty:
         return empty_figure("Latent-space samples are empty after filtering.")
 
-    score_lookup = (
-        score_df.set_index("model")["mean_r2"].to_dict() if not score_df.empty else {}
-    )
-    titles = []
-    for model in plot_models:
-        score = score_lookup.get(model)
-        suffix = "" if pd.isna(score) else f"  R2={score:.2f}"
-        titles.append(f"{model}{suffix}")
-
+    sessions = ordered_unique(plot_df["session_label"])
     n_cols = 2
-    n_rows = int(np.ceil(len(plot_models) / n_cols))
+    n_rows = int(np.ceil(len(sessions) / n_cols))
     fig = make_subplots(
         rows=n_rows,
         cols=n_cols,
         specs=[[{"type": "scene"} for _ in range(n_cols)] for _ in range(n_rows)],
-        subplot_titles=titles,
-        horizontal_spacing=0.03,
-        vertical_spacing=0.08,
+        subplot_titles=sessions,
+        horizontal_spacing=0.035,
+        vertical_spacing=0.105,
     )
-    palette = px.colors.qualitative.Set2 + px.colors.qualitative.Dark24
-    sessions = ordered_unique(plot_df["session_label"])
-    session_colors = {session: palette[idx % len(palette)] for idx, session in enumerate(sessions)}
-
-    for model_idx, model in enumerate(plot_models):
-        model_df = plot_df[plot_df["model"].astype(str) == model]
-        row = model_idx // n_cols + 1
-        col = model_idx % n_cols + 1
-        for session in sessions:
-            session_df = model_df[model_df["session_label"].astype(str) == session]
-            if session_df.empty:
-                continue
+
+    condition_values = sorted(plot_df["condition"].astype(str).unique())
+    use_categorical_conditions = len(condition_values) <= 12
+    direction_palette = [
+        "#B23AEE",
+        "#3B1C54",
+        "#2DD4F6",
+        "#289285",
+        "#E3D724",
+        "#00A65A",
+        "#5B8FF9",
+        "#F97316",
+        "#E45756",
+        "#72B7B2",
+        "#54A24B",
+        "#B279A2",
+    ]
+    condition_colors = {
+        condition: direction_palette[idx % len(direction_palette)]
+        for idx, condition in enumerate(condition_values)
+    }
+    condition_prefix = "Reach direction" if dataset == "monkey" else "Condition"
+
+    for session_idx, session in enumerate(sessions):
+        session_df = plot_df[plot_df["session_label"].astype(str) == session]
+        if session_df.empty:
+            continue
+        row = session_idx // n_cols + 1
+        col = session_idx % n_cols + 1
+
+        if use_categorical_conditions:
+            for condition in condition_values:
+                cond_df = session_df[session_df["condition"].astype(str) == condition]
+                if cond_df.empty:
+                    continue
+                trace_name = f"{condition_prefix} {condition}"
+                fig.add_trace(
+                    go.Scatter3d(
+                        x=cond_df["x"],
+                        y=cond_df["y"],
+                        z=cond_df["z"],
+                        mode="markers",
+                        name=trace_name,
+                        legendgroup=condition,
+                        showlegend=session_idx == 0,
+                        marker=dict(
+                            size=2.6,
+                            opacity=0.72,
+                            color=condition_colors[condition],
+                        ),
+                        customdata=np.stack(
+                            [
+                                cond_df["session_label"].astype(str),
+                                cond_df["trial_index"].astype(str),
+                                cond_df["time_index"].astype(str),
+                            ],
+                            axis=-1,
+                        ),
+                        hovertemplate=(
+                            "session=%{customdata[0]}<br>"
+                            f"{condition_prefix.lower()}={condition}<br>"
+                            "trial=%{customdata[1]} time=%{customdata[2]}"
+                            "<extra></extra>"
+                        ),
+                    ),
+                    row=row,
+                    col=col,
+                )
+        else:
             fig.add_trace(
                 go.Scatter3d(
                     x=session_df["x"],
@@ -280,16 +347,17 @@ def latent_space_figure(dataset: str, models: list[str] | None) -> go.Figure:
                     z=session_df["z"],
                     mode="markers",
                     name=session,
-                    legendgroup=session,
-                    showlegend=model_idx == 0,
+                    showlegend=False,
                     marker=dict(
-                        size=2.4,
-                        opacity=0.58,
-                        color=session_colors[session],
+                        size=2.6,
+                        opacity=0.72,
+                        color=session_df["condition_num"],
+                        colorscale="Viridis",
+                        showscale=session_idx == 0,
+                        colorbar=dict(title=condition_prefix),
                     ),
                     customdata=np.stack(
                         [
-                            session_df["model"].astype(str),
                             session_df["condition"].astype(str),
                             session_df["trial_index"].astype(str),
                             session_df["time_index"].astype(str),
@@ -297,10 +365,9 @@ def latent_space_figure(dataset: str, models: list[str] | None) -> go.Figure:
                         axis=-1,
                     ),
                     hovertemplate=(
-                        "model=%{customdata[0]}<br>"
                         "session=%{fullData.name}<br>"
-                        "condition=%{customdata[1]}<br>"
-                        "trial=%{customdata[2]} time=%{customdata[3]}"
+                        f"{condition_prefix.lower()}=%{{customdata[0]}}<br>"
+                        "trial=%{customdata[1]} time=%{customdata[2]}"
                         "<extra></extra>"
                     ),
                 ),
@@ -308,21 +375,57 @@ def latent_space_figure(dataset: str, models: list[str] | None) -> go.Figure:
                 col=col,
             )
 
-    for idx in range(len(plot_models)):
+    extent = max(
+        float(np.nanpercentile(np.abs(plot_df["x"]), 99)),
+        float(np.nanpercentile(np.abs(plot_df["y"]), 99)),
+        float(np.nanpercentile(np.abs(plot_df["z"]), 99)),
+        1.0,
+    )
+    lim = extent * 1.08
+    for idx in range(len(sessions)):
         scene_id = "scene" if idx == 0 else f"scene{idx + 1}"
         fig.update_layout(
             **{
                 scene_id: dict(
-                    xaxis=dict(title="", showgrid=False, zeroline=False, showticklabels=False),
-                    yaxis=dict(title="", showgrid=False, zeroline=False, showticklabels=False),
-                    zaxis=dict(title="", showgrid=False, zeroline=False, showticklabels=False),
+                    xaxis=dict(
+                        title="",
+                        range=[-lim, lim],
+                        showgrid=False,
+                        zeroline=False,
+                        showticklabels=False,
+                    ),
+                    yaxis=dict(
+                        title="",
+                        range=[-lim, lim],
+                        showgrid=False,
+                        zeroline=False,
+                        showticklabels=False,
+                    ),
+                    zaxis=dict(
+                        title="",
+                        range=[-lim, lim],
+                        showgrid=False,
+                        zeroline=False,
+                        showticklabels=False,
+                    ),
+                    aspectmode="cube",
                     bgcolor="#ffffff",
-                    camera=dict(eye=dict(x=1.55, y=1.55, z=1.1)),
+                    camera=dict(eye=dict(x=1.55, y=1.45, z=1.05)),
                 )
             }
         )
-    fig.update_layout(title=f"{DATASET_LABELS.get(dataset, dataset)} latent-space samples")
-    return fig_layout(fig, height=720 if n_rows > 1 else 420)
+    score = None
+    if not score_df.empty and model in set(score_df["model"].astype(str)):
+        score = float(score_df.loc[score_df["model"].astype(str) == model, "mean_r2"].iloc[0])
+    score_text = "" if score is None or pd.isna(score) else f" (mean R2={score:.2f})"
+    fig.update_layout(
+        title=(
+            f"{DATASET_LABELS.get(dataset, dataset)} latent samples: {model}{score_text}"
+        ),
+    )
+    fig.for_each_annotation(lambda ann: ann.update(font=dict(size=11, color="#536271")))
+    fig.update_layout(legend=dict(title=condition_prefix))
+    return latent_fig_layout(fig, height=720 if n_rows > 1 else 430)
 
 
 def dataframe_table(
@@ -489,7 +592,19 @@ app.layout = html.Div(
                     children=[
                         panel(
                             "Cross-Session Latent Consistency",
-                            dcc.Graph(id="latent-space", config={"displayModeBar": False}),
+                            dcc.Graph(
+                                id="latent-space",
+                                config={
+                                    "displayModeBar": "hover",
+                                    "toImageButtonOptions": {
+                                        "format": "png",
+                                        "filename": "benchdash_latent_space",
+                                        "height": 900,
+                                        "width": 1200,
+                                        "scale": 2,
+                                    },
+                                },
+                            ),
                             dcc.Graph(id="consistency-bars", config={"displayModeBar": False}),
                             dcc.Graph(id="consistency-heatmap", config={"displayModeBar": False}),
                             dataframe_table("consistency-table"),