feat: add interactive plot to view rope rotation angle as a function of dimension pair index (d) and token position index (k)

app.py

@@ -30,6 +30,8 @@ from src.plots import (
     norm_compare_add_vs_rope,
     norms_and_cosine,
     position_sweep,
+    rope_angle_heatmap,
+    rope_angle_slices,
     rotation_2d,
     theta_heatmap,
 )
@@ -282,6 +284,57 @@ def update_bulk(data, which, head, mod_2pi):
     )
 
 
+def update_angle_explorer(data, head, token, pair, display_mode):
+    if not data:
+        return (
+            PLACEHOLDER,
+            PLACEHOLDER,
+            "Compute on the Setup tab first.",
+            gr.update(maximum=1, value=0),
+            gr.update(maximum=1, value=0),
+        )
+    before = select_head(data["q_before"], int(head))
+    seq, dim = before.shape
+    n_pairs = dim // 2
+    token = int(np.clip(token, 0, seq - 1))
+    pair = int(np.clip(pair, 0, n_pairs - 1))
+    base = float(data["base"])
+    theta = float(theta_grid(seq, dim, base)[token, pair])
+    if display_mode == "turns":
+        shown_theta = theta / (2 * np.pi)
+        display_label = "θ / 2π (turns)"
+    elif display_mode == "wrapped":
+        shown_theta = float(np.mod(theta, 2 * np.pi))
+        display_label = "θ mod 2π (radians)"
+    else:
+        shown_theta = theta
+        display_label = "θ (radians)"
+    omega = float(pair_frequencies(dim, base)[pair])
+    d0, d1 = pair_dim_labels(pair, dim, style=data["style"])
+    detail = fr"""
+### Selected rotation angle
+
+- **Head:** `{int(head)}`
+- **Token position:** `k = {token}`
+- **Pair:** `i = {pair}` → ({d0}, {d1})
+- **Frequency:** `ωᵢ = {omega:.8f}` radians per token position
+- **Raw angle:** `θ({token}, {pair}) = {theta:.8f}` radians
+- **Displayed angle ({display_label}):** `{shown_theta:.8f}`
+
+$$\theta(k,i) = k\,base^{{-2i/d}} = {token}\,({base:g})^{{-2\times{pair}/{dim}}}$$
+
+Every increase of one token position adds `ωᵢ` radians for this pair. Lower-frequency
+pairs change more slowly as `k` increases.
+"""
+    return (
+        rope_angle_heatmap(seq, dim, base, token, pair, display_mode),
+        rope_angle_slices(seq, dim, base, token, pair, display_mode),
+        detail,
+        gr.update(maximum=_safe_slider_max(seq - 1), value=token),
+        gr.update(maximum=_safe_slider_max(n_pairs - 1), value=pair),
+    )
+
+
 def update_individual(data, which, head, token, pair, sweep):
     if not data:
         return "Compute on the Setup tab first.", PLACEHOLDER, PLACEHOLDER
@@ -449,6 +502,28 @@ with gr.Blocks(title="RoPE Explorer") as demo:
             bulk_theta = gr.Plot(label="θ(k, i)")
             bulk_freq = gr.Plot(label="ω_i")
 
+        with gr.Tab("RoPE angles"):
+            gr.Markdown(
+                "Explore how the rotation angle changes with token position `k` and "
+                "dimension pair `i`. The selected head is shared with the Setup tab. "
+                "A marker identifies the selected `(k, i)` cell in the heatmap."
+            )
+            with gr.Row():
+                angle_token = gr.Slider(0, 15, step=1, value=0, label="Token position k")
+                angle_pair = gr.Slider(0, 15, step=1, value=0, label="Pair index i")
+                angle_display = gr.Radio(
+                    [
+                        ("Absolute angle (radians)", "absolute"),
+                        ("Angle / 2π (turns)", "turns"),
+                        ("Wrapped angle mod 2π", "wrapped"),
+                    ],
+                    value="absolute",
+                    label="Angle display",
+                )
+            angle_heatmap = gr.Plot(label="RoPE angle heatmap")
+            angle_slices = gr.Plot(label="Selected pair/token slices")
+            angle_detail = gr.Markdown("Compute on the Setup tab first.")
+
         with gr.Tab("Individual changes"):
             with gr.Row():
                 token_k = gr.Slider(0, 15, step=1, value=0, label="Token index k")
@@ -552,6 +627,11 @@ score does not represent a universal threshold of importance.
     for ctrl in bulk_inputs:
         ctrl.change(update_bulk, inputs=bulk_inputs, outputs=bulk_outputs)
 
+    angle_inputs = [state, head, angle_token, angle_pair, angle_display]
+    angle_outputs = [angle_heatmap, angle_slices, angle_detail, angle_token, angle_pair]
+    for ctrl in angle_inputs:
+        ctrl.change(update_angle_explorer, inputs=angle_inputs, outputs=angle_outputs)
+
     ind_inputs = [state, which, head, token_k, pair_i, sweep]
     ind_outputs = [pair_table, pair_plot, sweep_plot]
     for ctrl in ind_inputs:

src/plots.py

@@ -115,6 +115,128 @@ def theta_heatmap(seq_len: int, dim: int, base: float, mod_2pi: bool = False) ->
     return fig
 
 
+def rope_angle_heatmap(
+    seq_len: int,
+    dim: int,
+    base: float,
+    selected_token: int = 0,
+    selected_pair: int = 0,
+    display_mode: str = "absolute",
+) -> go.Figure:
+    """Interactive θ(k, i) map with token/pair selection marker."""
+    grid = theta_grid(seq_len, dim, base=base)
+    if display_mode == "turns":
+        displayed = grid / (2 * np.pi)
+        value_label = "θ / 2π (turns)"
+        shown_label = "θ / 2π"
+    elif display_mode == "wrapped":
+        displayed = np.mod(grid, 2 * np.pi)
+        value_label = "θ mod 2π (radians)"
+        shown_label = "θ mod 2π"
+    else:
+        displayed = grid
+        value_label = "θ (radians)"
+        shown_label = "θ"
+    pairs = grid.shape[1]
+    tokens = np.arange(seq_len)
+    pair_indices = np.arange(pairs)
+    selected_token = int(np.clip(selected_token, 0, seq_len - 1))
+    selected_pair = int(np.clip(selected_pair, 0, pairs - 1))
+    fig = go.Figure(
+        go.Heatmap(
+            z=displayed.T,
+            x=tokens,
+            y=pair_indices,
+            customdata=grid.T,
+            colorbar=dict(title=value_label),
+            hovertemplate=(
+                "token position k=%{x}<br>pair index i=%{y}<br>"
+                "raw θ=%{customdata:.6f} rad<br>"
+                + f"shown {shown_label}=%{{z:.6f}}" 
+                + "<extra></extra>"
+            ),
+        )
+    )
+    fig.add_trace(
+        go.Scatter(
+            x=[selected_token],
+            y=[selected_pair],
+            mode="markers",
+            name="selected (k, i)",
+            marker=dict(size=11, color="white", line=dict(color="black", width=2)),
+            hovertemplate="selected k=%{x}, i=%{y}<extra></extra>",
+        )
+    )
+    fig.update_layout(
+        title="RoPE angle θ(k, i) across token positions and dimension pairs",
+        xaxis_title="token position k",
+        yaxis_title="pair index i",
+        template="plotly_white",
+        height=480,
+        margin=dict(t=65, b=55),
+    )
+    return fig
+
+
+def rope_angle_slices(
+    seq_len: int,
+    dim: int,
+    base: float,
+    selected_token: int = 0,
+    selected_pair: int = 0,
+    display_mode: str = "absolute",
+) -> go.Figure:
+    """Show θ for one pair over positions and one position over pairs."""
+    grid = theta_grid(seq_len, dim, base=base)
+    if display_mode == "turns":
+        displayed = grid / (2 * np.pi)
+        value_label = "θ / 2π (turns)"
+    elif display_mode == "wrapped":
+        displayed = np.mod(grid, 2 * np.pi)
+        value_label = "θ mod 2π (radians)"
+    else:
+        displayed = grid
+        value_label = "θ (radians)"
+    selected_token = int(np.clip(selected_token, 0, seq_len - 1))
+    selected_pair = int(np.clip(selected_pair, 0, grid.shape[1] - 1))
+    fig = make_subplots(
+        rows=1,
+        cols=2,
+        subplot_titles=[
+            f"Pair i={selected_pair}: angle over token position k",
+            f"Token k={selected_token}: angle over pair index i",
+        ],
+    )
+    fig.add_trace(
+        go.Scatter(
+            x=np.arange(seq_len),
+            y=displayed[:, selected_pair],
+            mode="lines",
+            name=f"pair {selected_pair}",
+            hovertemplate="k=%{x}<br>θ=%{y:.6f} rad<extra></extra>",
+        ),
+        row=1,
+        col=1,
+    )
+    fig.add_trace(
+        go.Scatter(
+            x=np.arange(grid.shape[1]),
+            y=displayed[selected_token, :],
+            mode="lines+markers",
+            name=f"token {selected_token}",
+            hovertemplate="i=%{x}<br>θ=%{y:.6f} rad<extra></extra>",
+        ),
+        row=1,
+        col=2,
+    )
+    fig.update_xaxes(title_text="token position k", row=1, col=1)
+    fig.update_yaxes(title_text=value_label, row=1, col=1)
+    fig.update_xaxes(title_text="pair index i", row=1, col=2)
+    fig.update_yaxes(title_text=value_label, row=1, col=2)
+    fig.update_layout(template="plotly_white", height=360, showlegend=False)
+    return fig
+
+
 def frequency_strip(dim: int, base: float) -> go.Figure:
     omega = pair_frequencies(dim, base=base)
     fig = go.Figure(data=go.Bar(x=list(range(len(omega))), y=omega, name="ω_i"))