First mock-up

app.py

@@ -0,0 +1,614 @@
+from __future__ import annotations
+
+from typing import Dict, List, Tuple
+import math
+
+import numpy as np
+import pandas as pd
+import plotly.graph_objects as go
+from dash import Dash, Input, Output, callback, dash_table, dcc, html
+from plotly.subplots import make_subplots
+
+
+SIGFIGS = 3
+
+COLUMN_COLORS = {
+    "r2": "rgb(230, 245, 255)",
+    "rmse": "rgb(240, 255, 230)",
+    "mae": "rgb(255, 245, 230)",
+}
+
+BEHAVIORS = ["Behavior A", "Behavior B", "Behavior C"]
+BEHAVIOR_COLORS = {
+    "Behavior A": "#636EFA",
+    "Behavior B": "#EF553B",
+    "Behavior C": "#00CC96",
+}
+
+METRIC_DIRECTION = {
+    "r2": "high_is_good",
+    "rmse": "low_is_good",
+    "mae": "low_is_good",
+}
+
+# Universal axis limits for better comparison across panels
+XYZ_RANGE: Tuple[float, float] = (-25.0, 25.0)
+
+
+def build_toy_benchmark() -> Dict[str, pd.DataFrame]:
+    return {
+        "Dataset 1": pd.DataFrame(
+            [
+                {"model": "Model 1", "r2": 0.62, "rmse": 4.10, "mae": 3.20},
+                {"model": "Model 2", "r2": 0.71, "rmse": 3.60, "mae": 2.90},
+                {"model": "Model 3", "r2": 0.68, "rmse": 3.85, "mae": 3.05},
+            ]
+        ),
+        "Dataset 2": pd.DataFrame(
+            [
+                {"model": "Model 1", "r2": 0.55, "rmse": 4.55, "mae": 3.55},
+                {"model": "Model 2", "r2": 0.66, "rmse": 3.95, "mae": 3.10},
+                {"model": "Model 3", "r2": 0.63, "rmse": 4.05, "mae": 3.20},
+            ]
+        ),
+        "Dataset 3": pd.DataFrame(
+            [
+                {"model": "Model 1", "r2": 0.49, "rmse": 5.05, "mae": 3.95},
+                {"model": "Model 2", "r2": 0.61, "rmse": 4.25, "mae": 3.35},
+                {"model": "Model 3", "r2": 0.59, "rmse": 4.40, "mae": 3.50},
+            ]
+        ),
+    }
+
+
+def compute_mock_score(
+    df: pd.DataFrame,
+    *,
+    metric_direction: Dict[str, str],
+    metric_weights: Dict[str, float] | None = None,
+    eps: float = 1e-12,
+) -> pd.Series:
+    if metric_weights is None:
+        metric_weights = {}
+
+    metrics = [m for m in metric_direction.keys() if m in df.columns]
+    if not metrics:
+        return pd.Series([0.0] * len(df), index=df.index)
+
+    parts = []
+    weights = []
+
+    for m in metrics:
+        x = df[m].astype(float)
+        x_min = float(x.min())
+        x_max = float(x.max())
+
+        if (x_max - x_min) < eps:
+            x_norm = pd.Series([0.5] * len(df), index=df.index)
+        else:
+            x_norm = (x - x_min) / (x_max - x_min)
+
+        if metric_direction[m] == "low_is_good":
+            x_norm = 1.0 - x_norm
+
+        w = float(metric_weights.get(m, 1.0))
+        parts.append(x_norm * w)
+        weights.append(w)
+
+    return sum(parts) / (sum(weights) + eps)
+
+
+def generate_mock_embedding(
+    dataset_name: str,
+    model_name: str,
+    *,
+    noise_level: float,
+    n_points_per_behavior: int = 250,
+    # Behavior geometry knobs (at noise=0)
+    base_center_spread: float = 7.5,
+    base_noise_scale: float = 1.0,
+    transform_strength: float = 0.55,
+    shift_strength: float = 2.0,
+    warp_strength: float = 0.35,
+    # Noise injection knobs
+    max_replace_frac: float = 0.60,
+    background_scale: float = 9.0,
+) -> pd.DataFrame:
+    """
+    Mock 3D embedding with a "noise spike" control.
+
+    noise_level in [0,1]:
+      - reduces inter-behavior separation (centers collapse)
+      - increases within-cluster spread
+      - replaces a fraction of points with background noise points
+    """
+    nl = float(np.clip(noise_level, 0.0, 1.0))
+
+    seed = abs(hash((dataset_name, model_name))) % (2**32)
+    rng = np.random.default_rng(seed)
+
+    # Base behavior centers in a canonical triangle in 3D
+    base_centers = {
+        "Behavior A": np.array([0.0, 0.0, 0.0]),
+        "Behavior B": np.array([1.0, 0.2, 0.1]),
+        "Behavior C": np.array([0.5, 1.1, -0.2]),
+    }
+
+    # As noise increases, pull centers toward a common centroid
+    centroid = sum(base_centers.values()) / len(base_centers)
+    center_spread = base_center_spread * (1.0 - 0.80 * nl)
+    centers = {}
+    for beh, c0 in base_centers.items():
+        c_collapsed = (1.0 - 0.85 * nl) * c0 + (0.85 * nl) * centroid
+        centers[beh] = c_collapsed * center_spread
+
+    # Within-cluster spread increases with noise
+    noise_scale = base_noise_scale * (1.0 + 2.25 * nl)
+
+    # Stronger per-panel transforms so different (dataset, model) look different
+    angles = rng.uniform(low=-np.pi, high=np.pi, size=3)
+    cx, cy, cz = np.cos(angles)
+    sx, sy, sz = np.sin(angles)
+
+    Rx = np.array([[1, 0, 0], [0, cx, -sx], [0, sx, cx]])
+    Ry = np.array([[cy, 0, sy], [0, 1, 0], [-sy, 0, cy]])
+    Rz = np.array([[cz, -sz, 0], [sz, cz, 0], [0, 0, 1]])
+    R = Rz @ Ry @ Rx
+
+    scales = np.diag(1.0 + transform_strength * rng.normal(size=3))
+    shear = np.eye(3)
+    shear[0, 1] = 0.25 * transform_strength * rng.normal()
+    shear[0, 2] = 0.25 * transform_strength * rng.normal()
+    shear[1, 2] = 0.25 * transform_strength * rng.normal()
+
+    A = (R @ scales @ shear)
+    b = shift_strength * rng.normal(size=(3,))
+
+    # How many points get replaced by background "noise spikes"
+    replace_frac = max_replace_frac * nl
+
+    rows: List[dict] = []
+    for beh in BEHAVIORS:
+        X = centers[beh] + noise_scale * rng.normal(size=(n_points_per_behavior, 3))
+
+        # Apply affine transform
+        X = X @ A.T + b
+
+        # Nonlinear warp (twist + bend); stays even with noise
+        theta = warp_strength * (0.15 * X[:, 0] + 0.10 * X[:, 1])
+        ct = np.cos(theta)
+        st = np.sin(theta)
+        x_new = ct * X[:, 0] - st * X[:, 1]
+        y_new = st * X[:, 0] + ct * X[:, 1]
+        z_new = X[:, 2] + warp_strength * np.tanh(0.25 * X[:, 0]) * 2.0
+        X = np.stack([x_new, y_new, z_new], axis=1)
+
+        # Replace a fraction with background noise points
+        n_replace = int(round(replace_frac * n_points_per_behavior))
+        if n_replace > 0:
+            idx_replace = rng.choice(n_points_per_behavior, size=n_replace, replace=False)
+
+            # Background noise: heavy-ish tails + broad scale, centered near origin
+            # This creates global "salt-and-pepper" points that mix behaviors.
+            noise_bg = background_scale * rng.standard_t(df=3, size=(n_replace, 3))
+            X[idx_replace] = noise_bg
+
+        for i in range(n_points_per_behavior):
+            rows.append(
+                {
+                    "x": float(X[i, 0]),
+                    "y": float(X[i, 1]),
+                    "z": float(X[i, 2]),
+                    "behavior": beh,
+                }
+            )
+
+    return pd.DataFrame(rows)
+
+
+def _parse_rgb(rgb: str) -> tuple[int, int, int]:
+    vals = rgb.strip().lower().replace("rgb(", "").replace(")", "").split(",")
+    return int(vals[0]), int(vals[1]), int(vals[2])
+
+
+def _darken(rgb: str, factor: float = 0.65) -> str:
+    r, g, b = _parse_rgb(rgb)
+    r = int(max(0, min(255, round(r * factor))))
+    g = int(max(0, min(255, round(g * factor))))
+    b = int(max(0, min(255, round(b * factor))))
+    return f"rgb({r},{g},{b})"
+
+
+def _blend_rgb(rgb_a: str, rgb_b: str, t: float) -> str:
+    t = max(0.0, min(1.0, float(t)))
+    ra, ga, ba = _parse_rgb(rgb_a)
+    rb, gb, bb = _parse_rgb(rgb_b)
+    r = round(ra + (rb - ra) * t)
+    g = round(ga + (gb - ga) * t)
+    b = round(ba + (bb - ba) * t)
+    return f"rgb({r},{g},{b})"
+
+
+def make_column_gradient_styles(
+    df: pd.DataFrame,
+    *,
+    id_col: str,
+    n_bins: int,
+    column_base_colors: Dict[str, str],
+    metric_direction: Dict[str, str],
+    darken_factor: float = 0.65,
+) -> List[dict]:
+    numeric_cols = [
+        c
+        for c in df.columns
+        if c != id_col and c in column_base_colors and pd.api.types.is_numeric_dtype(df[c])
+    ]
+
+    styles: List[dict] = []
+    white = "rgb(255, 255, 255)"
+
+    for col in numeric_cols:
+        col_min = float(df[col].min())
+        col_max = float(df[col].max())
+        if not (col_max > col_min):
+            continue
+
+        base = column_base_colors[col]
+        dark_base = _darken(base, factor=darken_factor)
+        direction = metric_direction.get(col, "high_is_good")
+
+        for i in range(n_bins):
+            low = col_min + (col_max - col_min) * (i / n_bins)
+            high = col_min + (col_max - col_min) * ((i + 1) / n_bins)
+
+            t = (i + 1) / n_bins
+            if direction == "low_is_good":
+                t = 1.0 - t
+
+            styles.append(
+                {
+                    "if": {
+                        "column_id": col,
+                        "filter_query": f"{{{col}}} >= {low} && {{{col}}} < {high}",
+                    },
+                    "backgroundColor": _blend_rgb(white, dark_base, t),
+                }
+            )
+
+        best_val = col_max if direction == "high_is_good" else col_min
+        worst_val = col_min if direction == "high_is_good" else col_max
+
+        styles.append(
+            {
+                "if": {"column_id": col, "filter_query": f"{{{col}}} = {best_val}"},
+                "backgroundColor": _blend_rgb(white, dark_base, 1.0),
+            }
+        )
+        styles.append(
+            {
+                "if": {"column_id": col, "filter_query": f"{{{col}}} = {worst_val}"},
+                "backgroundColor": _blend_rgb(white, dark_base, 0.0),
+            }
+        )
+
+    return styles
+
+
+def _apply_universal_scene_ranges(fig: go.Figure, *, n_scenes: int) -> None:
+    """
+    Apply the same x/y/z ranges to every 3D scene in a subplot grid.
+    """
+    for idx in range(n_scenes):
+        scene_id = "scene" if idx == 0 else f"scene{idx + 1}"
+        fig.update_layout(
+            **{
+                scene_id: dict(
+                    xaxis=dict(title="", range=list(XYZ_RANGE)),
+                    yaxis=dict(title="", range=list(XYZ_RANGE)),
+                    zaxis=dict(title="", range=list(XYZ_RANGE)),
+                )
+            }
+        )
+
+
+def _embedding_grid_figure(
+    *,
+    titles: List[str],
+    dataset_for_each_panel: List[str],
+    model_for_each_panel: List[str],
+    noise_level: float,
+    n_cols: int = 3,
+) -> go.Figure:
+    n_panels = len(titles)
+    n_rows = int(math.ceil(n_panels / n_cols))
+
+    specs = [[{"type": "scene"} for _ in range(n_cols)] for _ in range(n_rows)]
+    fig = make_subplots(
+        rows=n_rows,
+        cols=n_cols,
+        specs=specs,
+        subplot_titles=titles,
+        horizontal_spacing=0.02,
+        vertical_spacing=0.06,
+    )
+
+    for idx in range(n_panels):
+        r = (idx // n_cols) + 1
+        c = (idx % n_cols) + 1
+
+        ds = dataset_for_each_panel[idx]
+        md = model_for_each_panel[idx]
+        df_emb = generate_mock_embedding(ds, md, noise_level=noise_level)
+
+        for beh in BEHAVIORS:
+            df_b = df_emb[df_emb["behavior"] == beh]
+            fig.add_trace(
+                go.Scatter3d(
+                    x=df_b["x"],
+                    y=df_b["y"],
+                    z=df_b["z"],
+                    mode="markers",
+                    name=beh,
+                    legendgroup=beh,
+                    showlegend=(idx == 0),
+                    marker={
+                        "size": 2.5,
+                        "opacity": 0.75,
+                        "color": BEHAVIOR_COLORS[beh],
+                    },
+                    hovertemplate=(
+                        "behavior=%{text}<br>"
+                        "x=%{x:.3f}<br>y=%{y:.3f}<br>z=%{z:.3f}<extra></extra>"
+                    ),
+                    text=[beh] * len(df_b),
+                ),
+                row=r,
+                col=c,
+            )
+
+    _apply_universal_scene_ranges(fig, n_scenes=n_panels)
+
+    fig.update_layout(
+        height=320 * n_rows,
+        margin=dict(l=10, r=10, t=40, b=10),
+    )
+    return fig
+
+
+TOY_STORE = build_toy_benchmark()
+ALL_DATASETS = list(TOY_STORE.keys())
+ALL_MODELS = sorted({m for df in TOY_STORE.values() for m in df["model"].astype(str).tolist()})
+
+app = Dash(__name__)
+
+app.layout = html.Div(
+    [
+        html.H3("Neural Decoder Benchmarking"),
+
+        html.Div(
+            [
+                html.Div("Mode"),
+                dcc.RadioItems(
+                    id="mode-radio",
+                    options=[
+                        {"label": "By dataset", "value": "by_dataset"},
+                        {"label": "By model", "value": "by_model"},
+                    ],
+                    value="by_dataset",
+                    inline=True,
+                ),
+            ],
+            style={"marginBottom": "12px"},
+        ),
+
+        html.Div(
+            [
+                html.Div(
+                    [
+                        html.Div("Dataset"),
+                        dcc.Dropdown(
+                            id="dataset-dropdown",
+                            options=[{"label": k, "value": k} for k in ALL_DATASETS],
+                            value=ALL_DATASETS[0],
+                            clearable=False,
+                            style={"width": "100%"},
+                        ),
+                    ],
+                    id="dataset-dropdown-container",
+                    style={
+                        "width": "520px",
+                        "minWidth": "520px",
+                        "flexShrink": 0,
+                        "marginBottom": "12px",
+                    },
+                ),
+                html.Div(
+                    [
+                        html.Div("Model"),
+                        dcc.Dropdown(
+                            id="model-dropdown",
+                            options=[{"label": m, "value": m} for m in ALL_MODELS],
+                            value=ALL_MODELS[0] if ALL_MODELS else None,
+                            clearable=False,
+                            style={"width": "100%"},
+                        ),
+                    ],
+                    id="model-dropdown-container",
+                    style={
+                        "width": "520px",
+                        "minWidth": "520px",
+                        "flexShrink": 0,
+                        "marginBottom": "12px",
+                        "display": "none",
+                    },
+                ),
+            ],
+            style={"display": "flex", "gap": "16px", "alignItems": "flex-end"},
+        ),
+
+        dash_table.DataTable(
+            id="metrics-table",
+            columns=[],
+            data=[],
+            sort_action="native",
+            page_action="none",
+            style_table={"overflowX": "auto"},
+            style_cell={
+                "padding": "8px",
+                "fontFamily": "Arial",
+                "fontSize": "14px",
+                "textAlign": "left",
+                "minWidth": "120px",
+                "width": "120px",
+                "maxWidth": "200px",
+            },
+            style_header={"fontWeight": "600"},
+            style_data_conditional=[],
+            sort_by=[{"column_id": "score", "direction": "desc"}],
+        ),
+
+        html.Div(
+            [
+                html.Div("Noise level (fraction of spikes replaced / corrupted)"),
+                dcc.Slider(
+                    id="noise-slider",
+                    min=0.0,
+                    max=1.0,
+                    step=0.05,
+                    value=0.0,
+                    marks={0.0: "0", 0.25: "0.25", 0.5: "0.5", 0.75: "0.75", 1.0: "1.0"},
+                    tooltip={"placement": "bottom", "always_visible": False},
+                ),
+            ],
+            style={"maxWidth": "1100px", "marginTop": "10px", "marginBottom": "10px"},
+        ),
+
+        dcc.Graph(
+            id="embeddings-grid",
+            style={"height": "800px"},
+            config={"displayModeBar": False},
+        ),
+    ],
+    style={"padding": "16px"},
+)
+
+
+@callback(
+    Output("dataset-dropdown-container", "style"),
+    Output("model-dropdown-container", "style"),
+    Input("mode-radio", "value"),
+)
+def toggle_controls(mode: str):
+    base_box = {
+        "width": "520px",
+        "minWidth": "520px",
+        "flexShrink": 0,
+        "marginBottom": "12px",
+    }
+    if mode == "by_model":
+        return {**base_box, "display": "none"}, {**base_box, "display": "block"}
+    return {**base_box, "display": "block"}, {**base_box, "display": "none"}
+
+
+@callback(
+    Output("metrics-table", "columns"),
+    Output("metrics-table", "data"),
+    Output("metrics-table", "style_data_conditional"),
+    Output("embeddings-grid", "figure"),
+    Input("mode-radio", "value"),
+    Input("dataset-dropdown", "value"),
+    Input("model-dropdown", "value"),
+    Input("noise-slider", "value"),
+)
+def update_outputs(mode: str, dataset_name: str, model_name: str, noise_level: float):
+    column_colors = dict(COLUMN_COLORS)
+    column_colors["score"] = "rgb(220, 220, 220)"
+
+    metric_direction_for_colors = dict(METRIC_DIRECTION)
+    metric_direction_for_colors["score"] = "high_is_good"
+
+    if mode == "by_model":
+        # Rows = datasets (fixed model)
+        rows = []
+        for ds_name, df_ds in TOY_STORE.items():
+            df_row = df_ds[df_ds["model"] == model_name]
+            if df_row.empty:
+                continue
+            rec = df_row.iloc[0].to_dict()
+            rows.append({"dataset": ds_name, **{k: rec[k] for k in rec if k != "model"}})
+
+        df = pd.DataFrame(rows)
+        if df.empty:
+            return [], [], [], go.Figure()
+
+        df["score"] = compute_mock_score(
+            df,
+            metric_direction=METRIC_DIRECTION,
+            metric_weights={"r2": 1.0, "rmse": 1.0, "mae": 1.0},
+        )
+        df = df.sort_values("score", ascending=False).reset_index(drop=True)
+        df = df[["dataset", "score"] + [c for c in df.columns if c not in {"dataset", "score"}]]
+
+        id_col = "dataset"
+        titles = df["dataset"].tolist()
+
+        fig_grid = _embedding_grid_figure(
+            titles=titles,
+            dataset_for_each_panel=titles,
+            model_for_each_panel=[model_name] * len(titles),
+            noise_level=float(noise_level),
+            n_cols=3,
+        )
+
+    else:
+        # Rows = models (fixed dataset)
+        df = TOY_STORE[dataset_name].copy()
+
+        df["score"] = compute_mock_score(
+            df,
+            metric_direction=METRIC_DIRECTION,
+            metric_weights={"r2": 1.0, "rmse": 1.0, "mae": 1.0},
+        )
+        df = df.sort_values("score", ascending=False).reset_index(drop=True)
+        df = df[["model", "score"] + [c for c in df.columns if c not in {"model", "score"}]]
+
+        id_col = "model"
+        titles = df["model"].tolist()
+
+        fig_grid = _embedding_grid_figure(
+            titles=titles,
+            dataset_for_each_panel=[dataset_name] * len(titles),
+            model_for_each_panel=titles,
+            noise_level=float(noise_level),
+            n_cols=3,
+        )
+
+    # DataTable columns + trailing zeros formatting
+    num_cols = [c for c in df.columns if c != id_col and pd.api.types.is_numeric_dtype(df[c])]
+
+    columns = [{"name": id_col, "id": id_col}]
+    for c in df.columns:
+        if c == id_col:
+            continue
+        col_def = {"name": c, "id": c}
+        if c in num_cols:
+            col_def["type"] = "numeric"
+            col_def["format"] = {"specifier": f".{SIGFIGS}f"}
+        columns.append(col_def)
+
+    data = df.to_dict("records")
+
+    styles = make_column_gradient_styles(
+        df,
+        id_col=id_col,
+        n_bins=12,
+        column_base_colors=column_colors,
+        metric_direction=metric_direction_for_colors,
+        darken_factor=0.65,
+    )
+
+    return columns, data, styles, fig_grid
+
+
+if __name__ == "__main__":
+    app.run_server(debug=True)