ui(viz): 3D PCA flow + diversity + loss-spectrogram tabs

Replace the flat 2D PCA tab with a Plotly Scatter3d "PCA flow" that
overlays per-replica trajectories, epoch-coloured snapshot scatter and
a final-loss-coloured terminal cloud. The 2D fallback stays available
inside an expander.

Add two adjacent tabs that visualise the same population dynamics from
complementary angles:
- "Diversity": loss-σ across replicas + median−min loss + (when ``x`` is
  recorded) genotypic variance ⟨Var_b x⟩ on a secondary axis.
- "Loss spectrogram": row-normalised histogram heatmap of replica losses
  per epoch, exposing convergence as a sharpening horizontal stripe.

Wire a GUI smoke test that exercises all three new tabs through a real
PopulationTracker run.

Made-with: Cursor

Author: Yuma-Ichikawa

app/pages/2_Visualize.py

@@ -59,6 +59,8 @@
     tab_sched,
     tab_pop,
     tab_pca,
+    tab_div,
+    tab_spec,
     tab_ridge,
     tab_fate,
 ) = st.tabs(
@@ -68,7 +70,9 @@
         "Best trajectory",
         "Schedule",
         "Parallel population",
-        "PCA trajectory",
+        "3D PCA flow",
+        "Diversity",
+        "Loss spectrogram",
         "Ridgeline",
         "Replica fate",
     ]
@@ -145,15 +149,250 @@ def _retheme(fig):
         st.info("No x-snapshots recorded (population_tracker requires record_x=True).")
     else:
         try:
-            fig = viz.plot_population_embedding(pop_tracker, backend="plotly", show=False)
-            st.plotly_chart(_retheme(fig), width="stretch")
+            import plotly.graph_objects as go
+
+            xs = pop_tracker.x  # list[T] of (B, ...) ndarrays
+            epochs = np.asarray(pop_tracker.epochs)
+            X = np.stack([np.asarray(snap).reshape(snap.shape[0], -1) for snap in xs], axis=0)
+            # X: (T, B, D). Stack into a flat (T*B, D) cloud and run a
+            # truncated SVD so PCA-3D scales to N≈10⁴ variables comfortably.
+            T_, B, D = X.shape
+            flat = X.reshape(T_ * B, D)
+            mean = flat.mean(axis=0, keepdims=True)
+            flat0 = flat - mean
+            ncomp = min(3, flat0.shape[1], flat0.shape[0])
+            try:
+                # Faster than np.linalg.svd for tall matrices.
+                _, _, vt = np.linalg.svd(flat0, full_matrices=False)
+            except np.linalg.LinAlgError:
+                vt = np.eye(D)[:ncomp]
+            comps = vt[:ncomp].T  # (D, ncomp)
+            proj = (flat0 @ comps).reshape(T_, B, ncomp)  # (T, B, 3)
+            if ncomp < 3:
+                pad = np.zeros((T_, B, 3 - ncomp))
+                proj = np.concatenate([proj, pad], axis=-1)
+
+            final_loss = np.asarray(pop_tracker.loss[-1])
+            cmin = float(final_loss.min())
+            cmax = float(final_loss.max())
+            p = palette()
+            fig = go.Figure()
+            # One faint trajectory per replica.
+            stride = max(1, B // 64)  # cap visible lines for readability
+            for i in range(0, B, stride):
+                fig.add_trace(
+                    go.Scatter3d(
+                        x=proj[:, i, 0],
+                        y=proj[:, i, 1],
+                        z=proj[:, i, 2],
+                        mode="lines",
+                        line={"color": p["muted"], "width": 1.5},
+                        opacity=0.18,
+                        showlegend=False,
+                        hoverinfo="skip",
+                    )
+                )
+            # Snapshot scatter coloured by epoch — gives the "flow" feel.
+            for k in range(T_):
+                fig.add_trace(
+                    go.Scatter3d(
+                        x=proj[k, :, 0],
+                        y=proj[k, :, 1],
+                        z=proj[k, :, 2],
+                        mode="markers",
+                        marker={
+                            "size": 2.5,
+                            "color": [int(epochs[k])] * B,
+                            "colorscale": "Viridis",
+                            "cmin": int(epochs.min()),
+                            "cmax": int(epochs.max()),
+                            "showscale": k == T_ - 1,
+                            "colorbar": {"title": "epoch"} if k == T_ - 1 else None,
+                            "opacity": 0.55,
+                        },
+                        showlegend=False,
+                        name=f"epoch {int(epochs[k])}",
+                        hovertemplate=f"epoch {int(epochs[k])}<br>PC1=%{{x:.3f}}, PC2=%{{y:.3f}}, PC3=%{{z:.3f}}<extra></extra>",
+                    )
+                )
+            # Overlay final population coloured by per-replica final loss.
+            fig.add_trace(
+                go.Scatter3d(
+                    x=proj[-1, :, 0],
+                    y=proj[-1, :, 1],
+                    z=proj[-1, :, 2],
+                    mode="markers",
+                    marker={
+                        "size": 4.2,
+                        "color": final_loss,
+                        "colorscale": "Plasma",
+                        "cmin": cmin,
+                        "cmax": cmax,
+                        "showscale": True,
+                        "colorbar": {"title": "final loss", "x": 1.12},
+                        "line": {"color": "white", "width": 0.6},
+                    },
+                    name="final",
+                    showlegend=False,
+                    hovertemplate="final<br>loss=%{marker.color:.4f}<extra></extra>",
+                )
+            )
+            fig.update_layout(
+                **plotly_layout(
+                    title={"text": "3D PCA flow of the parallel population"},
+                    height=620,
+                    showlegend=False,
+                )
+            )
+            fig.update_scenes(
+                xaxis_title="PC1",
+                yaxis_title="PC2",
+                zaxis_title="PC3",
+                bgcolor=palette()["surface"],
+            )
+            st.plotly_chart(fig, width="stretch")
             st.caption(
-                "2D PCA projection of the entire continuous-variable population over time. "
-                "Each faint grey line is one replica's trajectory; markers are coloured by epoch."
+                "Truncated-SVD PCA (3 components) of the continuous-variable "
+                "population. Each grey thread is one replica's trajectory; "
+                "snapshot dots are coloured by epoch (Viridis); the final "
+                "population is overlaid in Plasma by per-replica final loss. "
+                "Drag to rotate."
             )
+            with st.expander("Show 2D projection", expanded=False):
+                fig2 = viz.plot_population_embedding(pop_tracker, backend="plotly", show=False)
+                st.plotly_chart(_retheme(fig2), width="stretch")
         except Exception as e:
             st.info(f"PCA could not be computed: {e}")
 
+with tab_div:
+    if pop_tracker is None or not pop_tracker.loss:
+        st.info("No population snapshots recorded for this run.")
+    else:
+        import plotly.graph_objects as go
+
+        p = palette()
+        epochs = np.asarray(pop_tracker.epochs)
+        loss_mat = np.stack(pop_tracker.loss, axis=0)  # (T, B)
+        # Loss spread = std over replicas — collapses to ~0 once the
+        # population converges.
+        loss_std = loss_mat.std(axis=1)
+        loss_min = loss_mat.min(axis=1)
+        loss_med = np.median(loss_mat, axis=1)
+
+        # Genotypic diversity: mean pairwise variance of the projected
+        # bits. Cheap proxy for Hamming distance and works for
+        # categorical / spin variables alike. Falls back gracefully
+        # when ``record_x`` was off.
+        geno_div: np.ndarray | None = None
+        if pop_tracker.x:
+            try:
+                xs = [np.asarray(s).reshape(s.shape[0], -1) for s in pop_tracker.x]
+                geno_div = np.array([s.var(axis=0).mean() for s in xs])
+            except Exception:
+                geno_div = None
+
+        fig = go.Figure()
+        fig.add_trace(
+            go.Scatter(
+                x=epochs,
+                y=loss_std,
+                mode="lines",
+                line={"color": p["palette"][0], "width": 2.4},
+                name="loss σ across replicas",
+            )
+        )
+        fig.add_trace(
+            go.Scatter(
+                x=epochs,
+                y=loss_med - loss_min,
+                mode="lines",
+                line={"color": p["palette"][1], "width": 2.0, "dash": "dot"},
+                name="median − min loss",
+            )
+        )
+        if geno_div is not None and geno_div.size == epochs.size:
+            fig.add_trace(
+                go.Scatter(
+                    x=epochs,
+                    y=geno_div,
+                    mode="lines",
+                    line={
+                        "color": p["palette"][2 % len(p["palette"])],
+                        "width": 2.0,
+                    },
+                    name="genotypic variance ⟨Var_b x⟩",
+                    yaxis="y2",
+                )
+            )
+        layout_kwargs = plotly_layout(
+            title={"text": "Population diversity over time"},
+            xaxis_title="Epoch",
+            yaxis_title="Loss spread",
+            height=420,
+        )
+        if geno_div is not None and geno_div.size == epochs.size:
+            layout_kwargs["yaxis2"] = {
+                "overlaying": "y",
+                "side": "right",
+                "title": "⟨Var_b x⟩",
+                "showgrid": False,
+            }
+        fig.update_layout(**layout_kwargs)
+        st.plotly_chart(fig, width="stretch")
+        st.caption(
+            "Tracks how quickly the parallel population collapses. "
+            "Healthy CRA / repulsion runs keep the genotypic variance "
+            "above zero for longer."
+        )
+
+with tab_spec:
+    if pop_tracker is None or not pop_tracker.loss:
+        st.info("No population snapshots recorded for this run.")
+    else:
+        import plotly.graph_objects as go
+
+        epochs = np.asarray(pop_tracker.epochs)
+        loss_mat = np.stack(pop_tracker.loss, axis=0)  # (T, B)
+        # Per-row histogram → (T, n_bins) heatmap.
+        n_bins = 64
+        lo = float(loss_mat.min())
+        hi = float(loss_mat.max())
+        if not np.isfinite(lo) or not np.isfinite(hi) or hi <= lo:
+            hi = lo + 1.0
+        edges = np.linspace(lo, hi, n_bins + 1)
+        spec = np.zeros((len(epochs), n_bins), dtype=float)
+        for t in range(len(epochs)):
+            h, _ = np.histogram(loss_mat[t], bins=edges)
+            spec[t] = h
+        # Row-normalise to keep the colour scale stable when sol_size is large.
+        row_max = spec.max(axis=1, keepdims=True)
+        row_max[row_max == 0] = 1.0
+        spec = spec / row_max
+        centres = 0.5 * (edges[:-1] + edges[1:])
+        fig = go.Figure(
+            go.Heatmap(
+                z=spec.T,
+                x=epochs,
+                y=centres,
+                colorscale="Magma",
+                colorbar={"title": "density"},
+                hovertemplate="epoch %{x}<br>loss %{y:.4f}<br>p=%{z:.2f}<extra></extra>",
+            )
+        )
+        fig.update_layout(
+            **plotly_layout(
+                title={"text": "Loss-distribution spectrogram"},
+                xaxis_title="Epoch",
+                yaxis_title="Loss",
+                height=460,
+            )
+        )
+        st.plotly_chart(fig, width="stretch")
+        st.caption(
+            "Row-normalised histogram of replica losses at each recorded "
+            "epoch. A sharp horizontal stripe at the bottom = converged."
+        )
+
 with tab_ridge:
     if pop_tracker is None or not pop_tracker.loss:
         st.info("No population snapshots recorded for this run.")

tests/test_gui_apptest.py

@@ -256,6 +256,42 @@ def test_visualize_tab_order_solution_first(tmp_path):
     assert tab_labels[0] == "Solution", f"Solution must be the first tab (got order: {tab_labels})"
 
 
+def test_visualize_renders_3d_pca_and_diversity_with_population(tmp_path):
+    """The new "3D PCA flow", "Diversity" and "Loss spectrogram" tabs must
+    render without exception when a PopulationTracker is attached."""
+    import sys
+
+    sys.path.insert(0, str(APP.parent))
+    from _common import build_problem as _build  # noqa: F811
+
+    import qqa  # noqa: F811
+    from qqa.callbacks import PopulationTracker  # noqa: F811
+
+    cfg = {"kind": "mis", "size": 16, "seed": 0, "device": "cpu", "extra": {}}
+    problem = _build(cfg)
+    tracker = PopulationTracker(stride=2, record_x=True)
+    result = qqa.anneal(
+        problem,
+        sol_size=8,
+        num_epochs=12,
+        learning_rate=0.1,
+        device="cpu",
+        verbose=False,
+        callbacks=[tracker],
+    )
+
+    at = AppTest.from_file(str(PAGE_DIR / "2_Visualize.py"), default_timeout=60)
+    at.session_state["last_result"] = result
+    at.session_state["last_problem"] = problem
+    at.session_state["last_pop_tracker"] = tracker
+    at.session_state["problem_config"] = cfg
+    at.run()
+    assert not at.exception, at.exception
+    labels = [t.label for t in at.tabs]
+    for required in ("3D PCA flow", "Diversity", "Loss spectrogram"):
+        assert required in labels, f"missing tab {required!r} (got {labels})"
+
+
 def test_solve_dynamics_separates_discrete_and_relaxed_best():
     """Regression: the per-replica chart used to plot the running discrete
     best (``state.best_obj``) on the same y-axis as the relaxed mean.  For