Discrete labels -> continuous labels

app.py

@@ -18,13 +18,6 @@ COLUMN_COLORS = {
     "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",
@@ -34,6 +27,12 @@ METRIC_DIRECTION = {
 # Universal axis limits for better comparison across panels
 XYZ_RANGE: Tuple[float, float] = (-25.0, 25.0)
 
+# Universal behavior-value range for color normalization (keeps colors comparable)
+BEHAVIOR_RANGE: Tuple[float, float] = (0.0, 1.0)
+
+# Continuous colorscale for regression target (Plotly built-in)
+CONTINUOUS_COLORSCALE = "Viridis"
+
 
 def build_toy_benchmark() -> Dict[str, pd.DataFrame]:
     return {
@@ -103,49 +102,51 @@ def generate_mock_embedding(
     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
+    n_points: int = 750,
+    # Manifold geometry (at noise=0)
+    base_radius: float = 9.0,
+    base_thickness: float = 0.8,
+    # Per-panel transforms
+    transform_strength: float = 0.70,
+    shift_strength: float = 2.5,
+    warp_strength: float = 0.45,
+    # Noise injection (spike replacement)
     max_replace_frac: float = 0.60,
-    background_scale: float = 9.0,
+    background_scale: float = 10.0,
 ) -> pd.DataFrame:
     """
-    Mock 3D embedding with a "noise spike" control.
+    Continuous target embedding.
+
+    Returns columns:
+      - x, y, z
+      - behavior_value in [0,1] (regression-style target)
 
     noise_level in [0,1]:
-      - reduces inter-behavior separation (centers collapse)
-      - increases within-cluster spread
-      - replaces a fraction of points with background noise points
+      - increases thickness/noise around the manifold
+      - replaces a fraction of points with background noise spikes
     """
     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]),
-    }
+    # Continuous behavior value (regression target)
+    t = rng.uniform(low=BEHAVIOR_RANGE[0], high=BEHAVIOR_RANGE[1], size=n_points)
+
+    # Base 1D manifold (a warped helix / loop) parameterized by t
+    # You can interpret t as the "behavioral state" / continuous label.
+    phase = 2.0 * np.pi * t
+    x0 = base_radius * (2.0 * t - 1.0)
+    y0 = base_radius * np.sin(phase)
+    z0 = 0.6 * base_radius * np.cos(phase)
 
-    # 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
+    X = np.stack([x0, y0, z0], axis=1)
 
-    # Within-cluster spread increases with noise
-    noise_scale = base_noise_scale * (1.0 + 2.25 * nl)
+    # Thickness around the manifold increases with noise_level
+    thickness = base_thickness * (1.0 + 2.5 * nl)
+    X = X + thickness * rng.normal(size=(n_points, 3))
 
-    # Stronger per-panel transforms so different (dataset, model) look different
+    # Strong per-panel affine transform
     angles = rng.uniform(low=-np.pi, high=np.pi, size=3)
     cx, cy, cz = np.cos(angles)
     sx, sy, sz = np.sin(angles)
@@ -157,53 +158,44 @@ def generate_mock_embedding(
 
     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()
+    shear[0, 1] = 0.30 * transform_strength * rng.normal()
+    shear[0, 2] = 0.30 * transform_strength * rng.normal()
+    shear[1, 2] = 0.30 * transform_strength * rng.normal()
 
-    A = (R @ scales @ shear)
+    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
+    X = X @ A.T + b
 
-    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,
-                }
-            )
+    # Nonlinear warp (keeps panels visually distinct)
+    theta = warp_strength * (0.10 * X[:, 0] + 0.08 * 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.20 * X[:, 0]) * 2.2
+    X = np.stack([x_new, y_new, z_new], axis=1)
 
-    return pd.DataFrame(rows)
+    # Replace a fraction with background noise spikes (increases mixing)
+    replace_frac = max_replace_frac * nl
+    n_replace = int(round(replace_frac * n_points))
+    if n_replace > 0:
+        idx_replace = rng.choice(n_points, size=n_replace, replace=False)
+        noise_bg = background_scale * rng.standard_t(df=3, size=(n_replace, 3))
+        X[idx_replace] = noise_bg
+
+        # Optional: also partially destroy label structure for replaced points
+        # (push their behavior_value toward random)
+        t[idx_replace] = rng.uniform(low=BEHAVIOR_RANGE[0], high=BEHAVIOR_RANGE[1], size=n_replace)
+
+    return pd.DataFrame(
+        {
+            "x": X[:, 0].astype(float),
+            "y": X[:, 1].astype(float),
+            "z": X[:, 2].astype(float),
+            "behavior_value": t.astype(float),
+        }
+    )
 
 
 def _parse_rgb(rgb: str) -> tuple[int, int, int]:
@@ -261,9 +253,9 @@ def make_column_gradient_styles(
             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
+            intensity = (i + 1) / n_bins
             if direction == "low_is_good":
-                t = 1.0 - t
+                intensity = 1.0 - intensity
 
             styles.append(
                 {
@@ -271,7 +263,7 @@ def make_column_gradient_styles(
                         "column_id": col,
                         "filter_query": f"{{{col}}} >= {low} && {{{col}}} < {high}",
                     },
-                    "backgroundColor": _blend_rgb(white, dark_base, t),
+                    "backgroundColor": _blend_rgb(white, dark_base, intensity),
                 }
             )
 
@@ -295,9 +287,6 @@ def make_column_gradient_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(
@@ -340,31 +329,39 @@ def _embedding_grid_figure(
         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>"
+        # Single trace: continuous color encodes regression target
+        # Use cmin/cmax so all panels share the same color normalization. :contentReference[oaicite:2]{index=2}
+        show_scale = (idx == 0)
+
+        fig.add_trace(
+            go.Scatter3d(
+                x=df_emb["x"],
+                y=df_emb["y"],
+                z=df_emb["z"],
+                mode="markers",
+                showlegend=False,
+                marker=dict(
+                    size=2.5,
+                    opacity=0.75,
+                    color=df_emb["behavior_value"],
+                    colorscale=CONTINUOUS_COLORSCALE,
+                    cmin=float(BEHAVIOR_RANGE[0]),
+                    cmax=float(BEHAVIOR_RANGE[1]),
+                    showscale=show_scale,
+                    colorbar=dict(  # colorbar config supported via marker.colorbar :contentReference[oaicite:3]{index=3}
+                        title=dict(text="behavior"),
+                        tickformat=".2f",
+                        len=0.70,
                     ),
-                    text=[beh] * len(df_b),
                 ),
-                row=r,
-                col=c,
-            )
+                hovertemplate=(
+                    "behavior=%{marker.color:.3f}<br>"
+                    "x=%{x:.3f}<br>y=%{y:.3f}<br>z=%{z:.3f}<extra></extra>"
+                ),
+            ),
+            row=r,
+            col=c,
+        )
 
     _apply_universal_scene_ranges(fig, n_scenes=n_panels)
 
@@ -384,7 +381,6 @@ app = Dash(__name__)
 app.layout = html.Div(
     [
         html.H3("Neural Decoder Benchmarking"),
-
         html.Div(
             [
                 html.Div("Mode"),
@@ -400,7 +396,6 @@ app.layout = html.Div(
             ],
             style={"marginBottom": "12px"},
         ),
-
         html.Div(
             [
                 html.Div(
@@ -445,7 +440,6 @@ app.layout = html.Div(
             ],
             style={"display": "flex", "gap": "16px", "alignItems": "flex-end"},
         ),
-
         dash_table.DataTable(
             id="metrics-table",
             columns=[],
@@ -466,7 +460,6 @@ app.layout = html.Div(
             style_data_conditional=[],
             sort_by=[{"column_id": "score", "direction": "desc"}],
         ),
-
         html.Div(
             [
                 html.Div("Noise level (fraction of spikes replaced / corrupted)"),
@@ -482,7 +475,6 @@ app.layout = html.Div(
             ],
             style={"maxWidth": "1100px", "marginTop": "10px", "marginBottom": "10px"},
         ),
-
         dcc.Graph(
             id="embeddings-grid",
             style={"height": "800px"},
@@ -528,7 +520,6 @@ def update_outputs(mode: str, dataset_name: str, model_name: str, noise_level: f
     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]
@@ -559,9 +550,7 @@ def update_outputs(mode: str, dataset_name: str, model_name: str, noise_level: f
             noise_level=float(noise_level),
             n_cols=3,
         )
-
     else:
-        # Rows = models (fixed dataset)
         df = TOY_STORE[dataset_name].copy()
 
         df["score"] = compute_mock_score(
@@ -583,7 +572,6 @@ def update_outputs(mode: str, dataset_name: str, model_name: str, noise_level: f
             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}]