Add reservoir dynamics benchmark and factor metrics into visualizatio…

examples/llm/llm_demo_005.py

@@ -50,7 +50,8 @@
 from torch.utils.data import DataLoader
 from transformers import GPT2Config,GPT2LMHeadModel
 from p_kit.llm.llm_models import SparsePBitLMTemporalMemory
-from p_kit.visualization import hypercube_plot
+from p_kit.visualization import hypercube_plot, hypercube_metrics
+from p_kit.benchmark import drive_reservoir
 
 SEED=7
 CONTEXT=64
@@ -186,16 +187,19 @@ def main():
     print(f"{'GPT':>10} {gp:10d} {gT:5.2f} {gacc:7.3f} {gppl:8.3f} {gtrain:8.2f}s {ginfer:8.2f}s")
     print(f"\np-kit uses {100*(1-pp/gp):.1f}% fewer trainable parameters.")
 
-    # Visualize the trained reservoir's {-1,+1}^n_pbits state-space
-    # trajectory while it reads real Shakespeare text (test split). n_pbits
-    # is far above hypercube_plot's exact-enumeration limit, so this shows
-    # a relaxation curve + PCA-vs-noise spectrum, plus a labeled projection
-    # colored by the driving character (the strongest external signal we
-    # found in this state space - see the conversation this was built in).
-    model.reset()
-    sample = test[:2000]
-    char_ids = [model.char_to_id[c] for c in sample if c in model.char_to_id]
-    history = np.array([model.step(i) for i in char_ids])
+    # Benchmark the trained reservoir's {-1,+1}^n_pbits state-space
+    # trajectory while it reads real Shakespeare text (test split): a
+    # relaxation time, a PCA signal-over-noise readout, and a separability
+    # ratio for how cleanly the driving character clusters the state space
+    # (the strongest external signal we found here - see the conversation
+    # this was built in). n_pbits is far above hypercube_plot's
+    # exact-enumeration limit, so hypercube_metrics uses its statistical
+    # view under the hood; hypercube_plot shows the same numbers visually.
+    history,char_ids=drive_reservoir(model,test)
+    rel_time,pca_snr,sep_ratio=hypercube_metrics(history,labels=char_ids)
+    print(f"\nReservoir dynamics: relaxation time {rel_time:.0f} samples, "
+          f"PCA SNR {pca_snr:.0f} std devs, separability ratio {sep_ratio:.3f}")
+
     hypercube_plot(history, labels=char_ids)
 
 if __name__=="__main__":

p_kit/benchmark/__init__.py

@@ -0,0 +1,3 @@
+from .reservoir import drive_reservoir, benchmark_reservoir
+
+__all__ = ["drive_reservoir", "benchmark_reservoir"]

p_kit/benchmark/reservoir.py

@@ -0,0 +1,26 @@
+"""Benchmark a p-bit reservoir model's state-space dynamics."""
+
+import numpy as np
+
+from p_kit.visualization import hypercube_metrics
+
+
+def drive_reservoir(model, text, n_chars=2000):
+    """Reset `model` and step its reservoir through a slice of `text`,
+    returning the {-1,+1}^n_pbits state trajectory and the driving
+    character ids (one per step)."""
+    model.reset()
+    sample = text[:n_chars]
+    char_ids = [model.char_to_id[c] for c in sample if c in model.char_to_id]
+    history = np.array([model.step(i) for i in char_ids])
+    return history, char_ids
+
+
+def benchmark_reservoir(model, text, n_chars=2000):
+    """Drive `model`'s reservoir on a slice of `text` and report its
+    dynamics as three numbers: relaxation time (samples), PCA
+    signal-over-noise (std devs above the shuffled-column null), and
+    label-separability ratio (how cleanly the driving characters cluster
+    in the reservoir's top-3 PCA directions)."""
+    history, char_ids = drive_reservoir(model, text, n_chars)
+    return hypercube_metrics(history, labels=char_ids)

p_kit/visualization/__init__.py

@@ -1,4 +1,4 @@
-from .utils import m_to_string, tsp_hist
+from .utils import m_to_string, tsp_hist, hypercube_metrics
 from .histplot import histplot, energyplot
 from .vin_vout import vin_vout
 from .plot3d import plot3d
@@ -15,5 +15,6 @@
     "vin_vout",
     "plot3d",
     "visualize_tsp_route",
-    "hypercube_plot"
+    "hypercube_plot",
+    "hypercube_metrics"
 ]

p_kit/visualization/hypercube.py

@@ -5,7 +5,8 @@
 import matplotlib.pyplot as plt
 from mpl_toolkits.mplot3d.art3d import Line3DCollection
 
-from .utils import m_to_string
+from .utils import m_to_string, N_PCA_COMPONENTS, _relaxation, \
+    _pca_vs_noise, _separability
 
 # Above this many p-bits, {-1,+1}^n_pbits can no longer be enumerated (2^n
 # vertices) - e.g. a 128-p-bit reservoir has more states than atoms in the
@@ -17,9 +18,6 @@
 # avoid an unreadable plot.
 MAX_LABELED_VERTS = 64
 
-# How many leading principal components to compare against the noise null.
-N_PCA_COMPONENTS = 15
-
 
 def _nested_cube_coords(n, shrink=0.55):
     """Recursively embed the n-cube's vertices in 3D: the first 3 bits form
@@ -94,40 +92,6 @@ def _hypercube_plot_exact(output, n_pbits):
     plt.show()
 
 
-def _pca_explained_variance(X):
-    Xc = X - X.mean(axis=0, keepdims=True)
-    _, S, _ = np.linalg.svd(Xc, full_matrices=False)
-    var = S ** 2
-    return var / var.sum()
-
-
-def _column_shuffle_null(X, rng, n_components, n_trials=20):
-    """Null PCA spectrum: independently permute each column (p-bit) across
-    samples. This preserves every p-bit's own marginal (its bias/frequency)
-    exactly while destroying all cross-p-bit and temporal correlation - the
-    fair 'no structure' baseline to compare the real spectrum against."""
-    spectra = np.zeros((n_trials, n_components))
-    noise = np.empty_like(X)
-    for t in range(n_trials):
-        for j in range(X.shape[1]):
-            noise[:, j] = rng.permutation(X[:, j])
-        spectra[t] = _pca_explained_variance(noise)[:n_components]
-    return spectra
-
-
-def _correlation_function(X, max_lag, n_pbits):
-    """C(dt) = <m(t).m(t+dt)> / n_pbits, the standard two-point overlap
-    used to characterize relaxation/mixing in stochastic spin dynamics.
-    Related to the mean Hamming distance at lag dt by
-    d_H(dt) = n_pbits * (1 - C(dt)) / 2."""
-    lags = np.arange(1, max_lag + 1)
-    C = np.empty(max_lag)
-    for i, dt in enumerate(lags):
-        overlap = (X[:-dt] * X[dt:]).sum(axis=1) / n_pbits
-        C[i] = overlap.mean()
-    return lags, C
-
-
 def _hypercube_plot_projected(output, n_pbits, labels=None):
     """Too many p-bits to enumerate exactly (2^n_pbits vertices - e.g. a
     128-p-bit reservoir has ~3.4e38 of them). A spatial scatter of a lossy
@@ -152,33 +116,28 @@ def _hypercube_plot_projected(output, n_pbits, labels=None):
       by time or raw position alone.
     """
     X = np.asarray(output, dtype=float)
-    n_samples = len(X)
     rng = np.random.default_rng(0)
 
     n_panels = 3 if labels is not None else 2
     fig = plt.figure(figsize=(6.5 * n_panels, 5.5))
 
     # --- panel 1: relaxation / correlation function ---
-    max_lag = max(1, min(300, n_samples // 3))
-    lags, C = _correlation_function(X, max_lag, n_pbits)
-    shuffled = X[rng.permutation(n_samples)]
-    baseline = (X * shuffled).sum(axis=1).mean() / n_pbits
+    lags, C, baseline, tau = _relaxation(X, n_pbits, rng)
 
     ax1 = fig.add_subplot(1, n_panels, 1)
     ax1.plot(lags, C, label="C(dt) = <m(t).m(t+dt)>/n")
     ax1.axhline(baseline, color="gray", linestyle="--",
                 label=f"fully-mixed baseline ({baseline:.3f})")
     ax1.set_xlabel("lag dt (samples)")
     ax1.set_ylabel("overlap / correlation")
-    ax1.set_title("Relaxation: correlation vs. time lag")
+    ax1.set_title(f"Relaxation: correlation vs. time lag\n"
+                  f"(relaxation time: {tau:.0f} samples)")
     ax1.legend()
 
     # --- panel 2: PCA spectrum vs. noise null ---
-    n_components = min(N_PCA_COMPONENTS, n_pbits)
-    real_var = _pca_explained_variance(X)[:n_components]
-    null_spectra = _column_shuffle_null(X, rng, n_components)
-    null_mean = null_spectra.mean(axis=0)
-    null_std = null_spectra.std(axis=0)
+    real_var, null_mean, null_std, top3_z = _pca_vs_noise(
+        X, n_pbits, rng, N_PCA_COMPONENTS)
+    n_components = len(real_var)
 
     ax2 = fig.add_subplot(1, n_panels, 2)
     idx = np.arange(n_components)
@@ -187,31 +146,14 @@ def _hypercube_plot_projected(output, n_pbits, labels=None):
             label="shuffled-column noise")
     ax2.set_xlabel("principal component")
     ax2.set_ylabel("explained variance ratio")
-    top3_excess = real_var[:3].sum() - null_mean[:3].sum()
-    top3_z = top3_excess / max(null_std[:3].sum(), 1e-12)
     ax2.set_title(f"PCA spectrum vs. noise null\n"
                   f"(top-3 excess: {top3_z:.0f} std devs above noise)")
     ax2.legend()
 
     # --- panel 3 (optional): projection colored by external label ---
     if labels is not None:
-        labels = np.asarray(labels)
-        Xc = X - X.mean(axis=0, keepdims=True)
-        _, _, Vt = np.linalg.svd(Xc, full_matrices=False)
-        proj = Xc @ Vt[:3].T
-
-        grand_mean = proj.mean(axis=0)
-        between = within = 0.0
-        for g in np.unique(labels):
-            pts = proj[labels == g]
-            if len(pts) < 2:
-                continue
-            gm = pts.mean(axis=0)
-            between += len(pts) * np.sum((gm - grand_mean) ** 2)
-            within += np.sum((pts - gm) ** 2)
-        sep_ratio = between / within if within > 0 else float("nan")
-
-        _, label_codes = np.unique(labels, return_inverse=True)
+        proj, sep_ratio = _separability(X, labels)
+        _, label_codes = np.unique(np.asarray(labels), return_inverse=True)
         ax3 = fig.add_subplot(1, n_panels, 3, projection="3d")
         ax3.scatter(proj[:, 0], proj[:, 1], proj[:, 2],
                     c=label_codes, cmap="tab20", s=25, depthshade=False)
@@ -221,7 +163,7 @@ def _hypercube_plot_projected(output, n_pbits, labels=None):
 
     fig.suptitle(f"{{-1,+1}}^{n_pbits} has {2 ** n_pbits:.3g} vertices "
                  f"(too many to enumerate)\n"
-                 f"statistical view of {n_samples} sampled states")
+                 f"statistical view of {len(X)} sampled states")
     fig.tight_layout()
     plt.show()
 

p_kit/visualization/utils.py

@@ -1,6 +1,11 @@
 """Utils method for visualization"""
+from collections import namedtuple
 import numpy as np
 
+# How many leading principal components to compare against the noise null,
+# in hypercube_metrics/hypercube_plot's statistical view.
+N_PCA_COMPONENTS = 15
+
 
 def m_to_string(outputs):
     ret = ""
@@ -30,3 +35,144 @@ def tsp_hist(samples, city_graph):
             hist[key] = 1
 
     return hist
+
+
+def _pca_explained_variance(X):
+    Xc = X - X.mean(axis=0, keepdims=True)
+    _, S, _ = np.linalg.svd(Xc, full_matrices=False)
+    var = S ** 2
+    return var / var.sum()
+
+
+def _column_shuffle_null(X, rng, n_components, n_trials=20):
+    """Null PCA spectrum: independently permute each column (p-bit) across
+    samples. This preserves every p-bit's own marginal (its bias/frequency)
+    exactly while destroying all cross-p-bit and temporal correlation - the
+    fair 'no structure' baseline to compare the real spectrum against."""
+    spectra = np.zeros((n_trials, n_components))
+    noise = np.empty_like(X)
+    for t in range(n_trials):
+        for j in range(X.shape[1]):
+            noise[:, j] = rng.permutation(X[:, j])
+        spectra[t] = _pca_explained_variance(noise)[:n_components]
+    return spectra
+
+
+def _correlation_function(X, max_lag, n_pbits):
+    """C(dt) = <m(t).m(t+dt)> / n_pbits, the standard two-point overlap
+    used to characterize relaxation/mixing in stochastic spin dynamics.
+    Related to the mean Hamming distance at lag dt by
+    d_H(dt) = n_pbits * (1 - C(dt)) / 2."""
+    lags = np.arange(1, max_lag + 1)
+    C = np.empty(max_lag)
+    for i, dt in enumerate(lags):
+        overlap = (X[:-dt] * X[dt:]).sum(axis=1) / n_pbits
+        C[i] = overlap.mean()
+    return lags, C
+
+
+def _relaxation(X, n_pbits, rng, max_lag=None):
+    """Correlation curve C(dt), its fully-mixed baseline, and a scalar
+    relaxation time: the first lag at which C(dt)'s excess over baseline
+    decays to 1/e of its initial value (the usual exponential-decay-time
+    convention, read off a curve that need not itself be exponential)."""
+    n_samples = len(X)
+    if max_lag is None:
+        max_lag = max(1, min(300, n_samples // 3))
+    lags, C = _correlation_function(X, max_lag, n_pbits)
+    shuffled = X[rng.permutation(n_samples)]
+    baseline = (X * shuffled).sum(axis=1).mean() / n_pbits
+
+    excess = C - baseline
+    if excess[0] <= 0:
+        tau = float(lags[0])
+    else:
+        below = np.where(excess <= excess[0] / np.e)[0]
+        tau = float(lags[below[0]]) if len(below) else float(lags[-1])
+
+    return lags, C, baseline, tau
+
+
+def _pca_vs_noise(X, n_pbits, rng, n_components):
+    """Real PCA spectrum, its column-shuffled noise null, and a scalar
+    signal-to-noise readout: how many noise standard deviations the top-3
+    real components' combined explained variance sits above the top-3
+    noise mean - "is there structure at all" as a number instead of an
+    eyeballed bar chart."""
+    n_components = min(n_components, n_pbits)
+    real_var = _pca_explained_variance(X)[:n_components]
+    null_spectra = _column_shuffle_null(X, rng, n_components)
+    null_mean = null_spectra.mean(axis=0)
+    null_std = null_spectra.std(axis=0)
+    top3_excess = real_var[:3].sum() - null_mean[:3].sum()
+    snr = top3_excess / max(null_std[:3].sum(), 1e-12)
+    return real_var, null_mean, null_std, snr
+
+
+def _separability(X, labels):
+    """Top-3 PCA projection plus its between/within group separability
+    ratio for the given per-sample labels."""
+    labels = np.asarray(labels)
+    Xc = X - X.mean(axis=0, keepdims=True)
+    _, _, Vt = np.linalg.svd(Xc, full_matrices=False)
+    proj = Xc @ Vt[:3].T
+
+    grand_mean = proj.mean(axis=0)
+    between = within = 0.0
+    for g in np.unique(labels):
+        pts = proj[labels == g]
+        if len(pts) < 2:
+            continue
+        gm = pts.mean(axis=0)
+        between += len(pts) * np.sum((gm - grand_mean) ** 2)
+        within += np.sum((pts - gm) ** 2)
+    sep_ratio = between / within if within > 0 else float("nan")
+
+    return proj, sep_ratio
+
+
+HypercubeMetrics = namedtuple(
+    "HypercubeMetrics", ["relaxation_time", "pca_snr", "separability_ratio"]
+)
+
+
+def hypercube_metrics(output, n_pbits=None, labels=None,
+                      n_components=N_PCA_COMPONENTS, seed=0):
+    """Compute hypercube_plot's statistical-view diagnostics without
+    plotting - handy for benchmarking a reservoir/model's dynamics
+    directly (e.g. to compare configurations or models numerically).
+
+    Parameters
+    ----------
+    output : array-like, shape (n_samples, n_pbits)
+        A +/-1 state trajectory, e.g. from repeated ``model.step(...)``.
+    n_pbits : int, optional
+        Inferred from ``output``'s second dimension when omitted.
+    labels : array-like, shape (n_samples,), optional
+        Per-sample group labels (e.g. a driving token id). When omitted,
+        ``separability_ratio`` is ``None``.
+    n_components : int, optional
+        Leading PCA components to compare against the noise null.
+    seed : int, optional
+        Seed for the noise-null RNG, for reproducible metrics.
+
+    Returns
+    -------
+    HypercubeMetrics
+        ``relaxation_time``: lag (in samples) at which the correlation
+        curve's excess over its fully-mixed baseline decays to 1/e.
+        ``pca_snr``: top-3 PCA explained-variance excess over the
+        column-shuffled noise null, in noise standard deviations.
+        ``separability_ratio``: between/within-group variance ratio of
+        the top-3 PCA projection, or ``None`` if ``labels`` is omitted.
+    """
+    X = np.asarray(output, dtype=float)
+    if n_pbits is None:
+        n_pbits = X.shape[1]
+    rng = np.random.default_rng(seed)
+
+    _, _, _, tau = _relaxation(X, n_pbits, rng)
+    _, _, _, snr = _pca_vs_noise(X, n_pbits, rng, n_components)
+    sep_ratio = _separability(X, labels)[1] if labels is not None else None
+
+    return HypercubeMetrics(tau, snr, sep_ratio)