Fix consistency columns and latent colors

app.py

@@ -263,6 +263,9 @@ def build_dataset_labels() -> dict[str, str]:
 
 DATASET_LABELS = build_dataset_labels()
 DATASETS = ordered_unique(prediction.get("dataset", pd.Series(dtype=str)))
+CONSISTENCY_DATASETS = set(
+    active_rows(consistency).get("dataset", pd.Series(dtype=str)).dropna().astype(str)
+)
 MODEL_SET = set(prediction.get("model", pd.Series(dtype=str)).dropna().astype(str))
 MODELS = [model for model in PAPER_MODEL_ORDER if model in MODEL_SET]
 MODELS += sorted(model for model in MODEL_SET if model not in set(MODELS))
@@ -346,15 +349,54 @@ def condition_label(dataset: str, condition: object) -> str:
     return text
 
 
-def condition_axis_label(dataset: str) -> str:
+def condition_axis_label(dataset: str, color_mode: str = "condition") -> str:
+    if dataset == "ratinabox":
+        if color_mode == "x":
+            return "X position bin"
+        if color_mode == "y":
+            return "Y position bin"
+        return "Position bin"
     return {
         "monkey": "Reach direction",
         "allen_neuropixels": "Orientation",
         "speech": "Cue",
-        "ratinabox": "Position bin",
     }.get(dataset, "Condition")
 
 
+def latent_color_label(dataset: str, color_mode: str, value: object) -> str:
+    if pd.isna(value):
+        return "Unknown"
+    if dataset == "ratinabox":
+        try:
+            idx = int(float(value))
+        except ValueError:
+            return str(value)
+        if color_mode == "x":
+            return f"x{idx}"
+        if color_mode == "y":
+            return f"y{idx}"
+    return condition_label(dataset, value)
+
+
+def add_latent_color_columns(df: pd.DataFrame, dataset: str, color_mode: str) -> pd.DataFrame:
+    out = df.copy()
+    condition_num = pd.to_numeric(out["condition"], errors="coerce")
+    if condition_num.isna().any() or (condition_num < 0).any():
+        raise ValueError("Latent samples contain missing condition labels.")
+
+    if dataset == "ratinabox" and color_mode == "x":
+        values = (condition_num.astype(int) % 10).astype(str)
+    elif dataset == "ratinabox" and color_mode == "y":
+        values = (condition_num.astype(int) // 10).astype(str)
+    else:
+        values = condition_num.astype(int).astype(str)
+
+    out["color_value"] = values
+    out["color_num"] = pd.to_numeric(values, errors="coerce")
+    out["color_label"] = out["color_value"].map(lambda value: latent_color_label(dataset, color_mode, value))
+    return out
+
+
 def session_display_label(dataset: str, session: object) -> str:
     text = "" if pd.isna(session) else str(session)
     if dataset == "monkey":
@@ -409,6 +451,14 @@ def sort_table(df: pd.DataFrame, sort_by: list[dict] | None, default: list[tuple
     )
 
 
+def supported_sort(sort_by: list[dict] | None, columns: Iterable[str]) -> list[dict] | None:
+    allowed = set(columns)
+    if not sort_by:
+        return None
+    filtered = [item for item in sort_by if item.get("column_id") in allowed]
+    return filtered or None
+
+
 def records(df: pd.DataFrame) -> list[dict]:
     clean = df.astype(object).where(pd.notna(df), None)
     return clean.to_dict("records")
@@ -810,7 +860,7 @@ def consistency_heatmap(models: list[str] | None) -> go.Figure:
     return fig_layout(fig, height=max(360, 25 * len(pivot.index) + 135))
 
 
-def latent_space_figure(dataset: str, model: str | None) -> go.Figure:
+def latent_space_figure(dataset: str, model: str | None, color_mode: str = "condition") -> go.Figure:
     if not model:
         return empty_figure("No latent-space view is available for this selection.")
     if latent_samples.empty:
@@ -825,8 +875,10 @@ def latent_space_figure(dataset: str, model: 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))
+    try:
+        plot_df = add_latent_color_columns(plot_df, dataset, color_mode)
+    except ValueError as exc:
+        return empty_figure(str(exc))
     plot_df["session_display"] = plot_df["session_label"].map(lambda value: session_display_label(dataset, value))
     plot_df = plot_df.dropna(subset=["x", "y", "z"])
     if plot_df.empty:
@@ -840,9 +892,14 @@ def latent_space_figure(dataset: str, model: str | None) -> go.Figure:
         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))
+        try:
+            trajectory_df = add_latent_color_columns(trajectory_df, dataset, color_mode)
+        except ValueError:
+            trajectory_df = trajectory_df.iloc[0:0].copy()
         trajectory_df["session_display"] = trajectory_df["session_label"].map(lambda value: session_display_label(dataset, value))
         trajectory_df = trajectory_df.dropna(subset=["x", "y", "z"])
+    if "color_value" not in trajectory_df.columns:
+        trajectory_df["color_value"] = pd.Series(dtype=str)
 
     sessions = ordered_unique(plot_df["session_label"])
     session_titles = [session_display_label(dataset, session) for session in sessions]
@@ -858,8 +915,8 @@ def latent_space_figure(dataset: str, model: str | None) -> go.Figure:
         vertical_spacing=0.12,
     )
 
-    condition_values = sorted(plot_df["condition"].astype(str).unique(), key=condition_sort_key)
-    use_categorical = len(condition_values) <= 12
+    condition_values = sorted(plot_df["color_value"].astype(str).unique(), key=condition_sort_key)
+    use_categorical = len(condition_values) <= 12 and dataset != "ratinabox"
     if dataset == "monkey":
         condition_colors = {
             condition: DIRECTION_PALETTE[int(float(condition)) % len(DIRECTION_PALETTE)]
@@ -875,7 +932,7 @@ def latent_space_figure(dataset: str, model: str | None) -> go.Figure:
             condition: CATEGORICAL_PALETTE[idx % len(CATEGORICAL_PALETTE)]
             for idx, condition in enumerate(condition_values)
         }
-    condition_name = condition_axis_label(dataset)
+    condition_name = condition_axis_label(dataset, color_mode)
 
     for session_idx, session in enumerate(sessions):
         session_df = plot_df[plot_df["session_label"].astype(str) == str(session)]
@@ -887,13 +944,13 @@ def latent_space_figure(dataset: str, model: str | None) -> go.Figure:
 
         if use_categorical:
             for condition in condition_values:
-                cond_df = session_df[session_df["condition"].astype(str) == condition]
+                cond_df = session_df[session_df["color_value"].astype(str) == condition]
                 if cond_df.empty:
                     continue
-                trace_name = condition_label(dataset, condition)
+                trace_name = latent_color_label(dataset, color_mode, condition)
                 session_traj = trajectory_df[
                     (trajectory_df["session_label"].astype(str) == str(session))
-                    & (trajectory_df["condition"].astype(str) == condition)
+                    & (trajectory_df["color_value"].astype(str) == condition)
                 ].sort_values("time_index")
                 fig.add_trace(
                     go.Scatter3d(
@@ -912,7 +969,7 @@ def latent_space_figure(dataset: str, model: str | None) -> go.Figure:
                         customdata=np.stack(
                             [
                                 np.repeat(display_session, len(cond_df)),
-                                cond_df["condition_label"].astype(str),
+                                cond_df["color_label"].astype(str),
                                 cond_df["trial_index"].astype(str),
                                 cond_df["time_index"].astype(str),
                             ],
@@ -960,15 +1017,17 @@ def latent_space_figure(dataset: str, model: str | None) -> go.Figure:
                     marker=dict(
                         size=2.8,
                         opacity=0.72,
-                        color=session_df["condition_num"],
+                        color=session_df["color_num"],
                         colorscale=RATINABOX_SCALE,
+                        cmin=0,
+                        cmax=9 if color_mode in {"x", "y"} else 99,
                         showscale=session_idx == 0,
                         colorbar=dict(title=condition_name, thickness=12),
                     ),
                     customdata=np.stack(
                         [
                             np.repeat(display_session, len(session_df)),
-                            session_df["condition"].map(lambda value: condition_label(dataset, value)).astype(str),
+                            session_df["color_label"].astype(str),
                             session_df["trial_index"].astype(str),
                             session_df["time_index"].astype(str),
                         ],
@@ -1332,18 +1391,34 @@ app.layout = html.Div(
                                         ],
                                         className="consistency-table-wrap",
                                     ),
-                                    dcc.Graph(
-                                        id="latent-space",
-                                        config={
-                                            "displayModeBar": "hover",
-                                            "toImageButtonOptions": {
-                                                "format": "png",
-                                                "filename": "benchdash_latent_space",
-                                                "height": 900,
-                                                "width": 1200,
-                                                "scale": 2,
-                                            },
-                                        },
+                                    html.Div(
+                                        [
+                                            html.Div(
+                                                [
+                                                    html.Label("Color by"),
+                                                    dcc.Dropdown(
+                                                        id="latent-color-mode",
+                                                        clearable=False,
+                                                    ),
+                                                ],
+                                                id="latent-color-control",
+                                                className="control latent-color-control",
+                                            ),
+                                            dcc.Graph(
+                                                id="latent-space",
+                                                config={
+                                                    "displayModeBar": "hover",
+                                                    "toImageButtonOptions": {
+                                                        "format": "png",
+                                                        "filename": "benchdash_latent_space",
+                                                        "height": 900,
+                                                        "width": 1200,
+                                                        "scale": 2,
+                                                    },
+                                                },
+                                            ),
+                                        ],
+                                        className="latent-panel",
                                     ),
                                 ],
                                 className="latent-grid",
@@ -1449,11 +1524,13 @@ def update_leaderboard(dataset: str, sort_by: list[dict] | None):
         "method",
         "task_score",
         "robustness_auc",
-        "alignment_score",
         "training_time_sec",
         "peak_ram_gb",
         "peak_vram_gb",
     ]
+    if str(dataset) in CONSISTENCY_DATASETS:
+        visible_cols.insert(3, "alignment_score")
+    sort_by = supported_sort(sort_by, visible_cols)
     sorted_df = sort_table(df, sort_by, [("task_score", False), ("model_order", True)])
     table_df = sorted_df[[c for c in visible_cols + ["id", "model"] if c in sorted_df.columns]]
     metric = df["metric"].dropna().iloc[0] if df["metric"].notna().any() else "score"
@@ -1467,6 +1544,26 @@ def update_leaderboard(dataset: str, sort_by: list[dict] | None):
     )
 
 
+@app.callback(
+    Output("latent-color-mode", "options"),
+    Output("latent-color-mode", "value"),
+    Output("latent-color-control", "style"),
+    Input("dataset-filter", "value"),
+)
+def update_latent_color_control(dataset: str):
+    if str(dataset) == "ratinabox":
+        return (
+            [
+                {"label": "Position bin", "value": "condition"},
+                {"label": "X position", "value": "x"},
+                {"label": "Y position", "value": "y"},
+            ],
+            "condition",
+            {},
+        )
+    return ([{"label": condition_axis_label(str(dataset)), "value": "condition"}], "condition", {"display": "none"})
+
+
 @app.callback(
     Output("consistency-table", "columns"),
     Output("consistency-table", "data"),
@@ -1477,18 +1574,36 @@ def update_leaderboard(dataset: str, sort_by: list[dict] | None):
     Input("dataset-filter", "value"),
     Input("consistency-table", "active_cell"),
     Input("consistency-table", "sort_by"),
+    Input("latent-color-mode", "value"),
 )
-def update_consistency(dataset: str, active_cell: dict | None, sort_by: list[dict] | None):
+def update_consistency(
+    dataset: str,
+    active_cell: dict | None,
+    sort_by: list[dict] | None,
+    color_mode: str | None,
+):
+    if str(dataset) not in CONSISTENCY_DATASETS:
+        label = DATASET_LABELS.get(dataset, dataset)
+        message = f"Cross-session alignment is not defined for {label}."
+        return (
+            [],
+            [],
+            message,
+            empty_figure(message),
+            empty_figure(message),
+            consistency_heatmap(None),
+        )
     df = consistency_frame(dataset, None)
     model = selected_consistency_model(df, active_cell)
     visible_cols = ["method", "alignment_score", "n_sessions", "latent_dim", "n_pairwise"]
+    sort_by = supported_sort(sort_by, visible_cols)
     sorted_df = sort_table(df, sort_by, [("alignment_score", False), ("model_order", True)])
     table_df = sorted_df[[c for c in visible_cols + ["id", "model"] if c in sorted_df.columns]]
     return (
         column_defs([c for c in visible_cols if c in table_df.columns]),
         records(table_df),
         sort_status(sort_by, ("alignment_score", False)),
-        latent_space_figure(dataset, model),
+        latent_space_figure(dataset, model, color_mode or "condition"),
         consistency_bar_figure(dataset, df),
         consistency_heatmap(None),
     )

assets/styles.css

@@ -228,6 +228,15 @@ h2 {
   align-items: start;
 }
 
+.latent-panel {
+  min-width: 0;
+}
+
+.latent-color-control {
+  width: min(260px, 100%);
+  margin-bottom: 8px;
+}
+
 .chart-grid {
   display: grid;
   gap: 16px;