Animate feature attribution sorting

app.py

@@ -1701,6 +1701,7 @@ def feature_sorter_figure(
     n_features = len(frame)
     values = frame["signed_attribution"].to_numpy(dtype=float)
     ranks = frame["attribution_rank"].to_numpy(dtype=float)
+    feature_ids = frame["feature_index"].to_numpy(dtype=int)
     y_low = min(float(np.nanmin(values)), 0.0)
     y_high = max(float(np.nanmax(values)), 0.0)
     y_span = max(y_high - y_low, 1e-6)
@@ -1715,7 +1716,6 @@ def feature_sorter_figure(
         vmax = float(np.nanmax(validation))
         scale = max(vmax - vmin, 1e-9)
         start_x = (validation - vmin) / scale
-        feature_ids = frame["feature_index"].to_numpy(dtype=int)
         start_y = 0.5 + (
             ((feature_ids * 37) % 101) / 100.0 - 0.5
         ) * 0.42
@@ -1729,16 +1729,31 @@ def feature_sorter_figure(
             np.flatnonzero(frame["feature_group"].astype(str).eq(group))
             for group in trace_groups
         ]
-        centers = np.linspace(0.7, 0.3, len(trace_groups))
+        centers = np.linspace(0.72, 0.28, len(trace_groups))
         start_x = np.zeros(n_features, dtype=float)
         start_y = np.zeros(n_features, dtype=float)
         for center, indices in zip(centers, trace_indices):
             ordered = indices[
                 np.argsort(frame.loc[indices, "feature_index"].to_numpy(), kind="stable")
             ]
-            start_x[ordered] = np.linspace(0.0, 1.0, len(ordered))
-            lane_jitter = ((np.arange(len(ordered)) % 7) - 3) * 0.006
-            start_y[ordered] = center + lane_jitter
+            n_columns = max(1, int(np.ceil(np.sqrt(len(ordered) * 4.0))))
+            n_rows = max(1, int(np.ceil(len(ordered) / n_columns)))
+            row_y = (
+                np.asarray([center])
+                if n_rows == 1
+                else center + np.linspace(-0.075, 0.075, n_rows)
+            )
+            for row_index in range(n_rows):
+                row = ordered[row_index * n_columns : (row_index + 1) * n_columns]
+                if not len(row):
+                    continue
+                if n_columns == 1:
+                    row_x = np.asarray([0.5])
+                else:
+                    offset = (n_columns - len(row)) / 2.0
+                    row_x = (np.arange(len(row)) + offset) / (n_columns - 1)
+                start_x[row] = 0.04 + 0.92 * row_x
+                start_y[row] = row_y[row_index]
 
     marker_size = 6 if n_features > 300 else (8 if n_features > 120 else 10)
 
@@ -1788,7 +1803,7 @@ def feature_sorter_figure(
             ids=[f"{dataset}:{int(value)}" for value in subset["feature_index"]],
             mode="markers",
             name=group,
-            showlegend=dataset != "allen_neuropixels",
+            showlegend=ranked and dataset != "allen_neuropixels",
             marker=marker,
             customdata=customdata,
             hovertemplate=(
@@ -1810,7 +1825,6 @@ def feature_sorter_figure(
         ]
     )
     shapes: list[dict] = []
-    annotations: list[dict] = []
     if ranked:
         for rank_bin in ATTRIBUTION_BIN_ORDER:
             bin_ranks = frame.loc[
@@ -1832,17 +1846,6 @@ def feature_sorter_figure(
                     layer="below",
                 )
             )
-            annotations.append(
-                dict(
-                    x=float(bin_x.mean()),
-                    xref="x",
-                    y=1.035,
-                    yref="paper",
-                    text=f"{rank_bin} third",
-                    showarrow=False,
-                    font=dict(size=10, color="#334957"),
-                )
-            )
         shapes.append(
             dict(
                 type="line",
@@ -1861,35 +1864,16 @@ def feature_sorter_figure(
                     type="rect",
                     x0=display_range[0],
                     x1=display_range[1],
-                    y0=center - 0.055,
-                    y1=center + 0.055,
+                    y0=center - 0.095,
+                    y1=center + 0.095,
                     fillcolor=group_colors[group],
                     opacity=0.08,
                     line_width=0,
                     layer="below",
                 )
             )
-            annotations.append(
-                dict(
-                    x=0.012,
-                    xref="paper",
-                    y=center,
-                    yref="y",
-                    text=group,
-                    showarrow=False,
-                    xanchor="left",
-                    bgcolor="rgba(255,255,255,0.85)",
-                    font=dict(size=11, color="#334957"),
-                )
-            )
 
-    title = (
-        "Signed Kernel SHAP vs. rank"
-        if ranked
-        else "Continuous gOSI"
-        if dataset == "allen_neuropixels"
-        else "Features by reference group"
-    )
+    title = "Signed Kernel SHAP vs. rank" if ranked else None
     figure_layout(figure, height=560)
     if ranked:
         rank_ticks = np.unique(
@@ -1941,31 +1925,103 @@ def feature_sorter_figure(
             range=display_range,
             showticklabels=False,
             ticks="",
-            title=" ",
+            title="",
             showgrid=False,
             zeroline=False,
         )
         figure.update_yaxes(
             visible=True,
             range=display_range,
-            showticklabels=False,
+            tickmode="array",
+            tickvals=centers,
+            ticktext=trace_groups,
+            tickfont=dict(size=10),
             ticks="",
-            title=" ",
+            title="",
             showgrid=False,
             zeroline=False,
         )
     figure.update_xaxes(automargin=False)
     figure.update_yaxes(automargin=False)
+    animation_controls: list[dict] = []
+    if ranked:
+        animated_traces = list(range(len(trace_indices)))
+
+        def coordinate_frame(
+            name: str,
+            x_values: np.ndarray,
+            y_values: np.ndarray,
+        ) -> go.Frame:
+            return go.Frame(
+                name=name,
+                traces=animated_traces,
+                data=[
+                    go.Scatter(x=x_values[indices], y=y_values[indices])
+                    for indices in trace_indices
+                ],
+            )
+
+        figure.frames = [
+            coordinate_frame("feature-reference", start_x, start_y),
+            coordinate_frame("feature-ranked", final_x, final_y),
+        ]
+        animation_controls = [
+            dict(
+                type="buttons",
+                direction="left",
+                showactive=False,
+                x=1.0,
+                xanchor="right",
+                y=-0.38,
+                yanchor="top",
+                pad=dict(t=4, r=0),
+                bgcolor="#F3EFF8",
+                bordercolor="#C8BBDD",
+                borderwidth=1,
+                font=dict(size=11, color="#4D3880"),
+                buttons=[
+                    dict(
+                        label="↻ Replay sorting",
+                        method="animate",
+                        args=[
+                            ["feature-reference", "feature-ranked"],
+                            {
+                                "mode": "immediate",
+                                "fromcurrent": False,
+                                "frame": [
+                                    {"duration": 120, "redraw": False},
+                                    {"duration": 900, "redraw": False},
+                                ],
+                                "transition": [
+                                    {"duration": 0},
+                                    {
+                                        "duration": 900,
+                                        "easing": "cubic-in-out",
+                                    },
+                                ],
+                            },
+                        ],
+                    )
+                ],
+            )
+        ]
+    if ranked:
+        margin = dict(l=58, r=28, t=66, b=165)
+    elif dataset == "allen_neuropixels":
+        margin = dict(l=46, r=54, t=26, b=76)
+    else:
+        margin = dict(l=112, r=24, t=26, b=76)
     figure.update_layout(
         title=title,
         shapes=shapes,
-        annotations=annotations,
-        margin=dict(l=58, r=28, t=78, b=100),
+        annotations=[],
+        updatemenus=animation_controls,
+        margin=margin,
         uirevision=f"feature-sorter:{dataset}:{model}:{arrangement}",
         legend=dict(
             orientation="h",
-            yanchor="bottom",
-            y=1.04,
+            yanchor="top",
+            y=-0.17,
             xanchor="left",
             x=0,
             font=dict(size=10),
@@ -3003,7 +3059,7 @@ app.layout = html.Div(
                                     ),
                                     html.Div(
                                         [
-                                            html.Label("Arrange features", htmlFor="feature-arrangement"),
+                                            html.Label("Feature view", htmlFor="feature-arrangement"),
                                             dcc.RadioItems(
                                                 id="feature-arrangement",
                                                 options=[
@@ -3047,7 +3103,7 @@ app.layout = html.Div(
                                 ],
                                 className="download-grid panel-downloads",
                             ),
-                            subtitle="Each dot is one input feature. Switch views to follow it from the reference annotation into the signed Kernel SHAP score-versus-rank curve.",
+                            subtitle="Explore each input feature in its reference context, then replay how the features sort into the signed Kernel SHAP ranking.",
                             class_name="axis-feature",
                         ),
                         panel(
@@ -3379,8 +3435,8 @@ def update_feature_selection_detail(
     arrangement: str,
 ):
     instruction = (
-        "Switch views, then hover or click a dot to inspect its feature index, "
-        "signed Kernel SHAP value and reference annotation."
+        "Each dot is one input feature. In the ranked view, use ↻ Replay sorting "
+        "to follow every feature into its signed Kernel SHAP rank."
     )
     if not click_data or not click_data.get("points"):
         return instruction