From 0e8accd5a1986449ca9aeee88f67ce5957004eda Mon Sep 17 00:00:00 2001 From: Gregoire CATTAN Date: Tue, 1 Sep 2026 22:01:57 +0200 Subject: [PATCH] Add reservoir dynamics benchmark and factor metrics into visualization utils hypercube_metrics (relaxation time, PCA signal-over-noise, separability ratio) now lives in p_kit.visualization.utils alongside its computation helpers, shared by hypercube_plot's statistical view. p_kit.benchmark.reservoir adds drive_reservoir/benchmark_reservoir so any p-bit reservoir model can be scored with those three numbers without plotting. --- examples/llm/llm_demo_005.py | 26 +++--- p_kit/benchmark/__init__.py | 3 + p_kit/benchmark/reservoir.py | 26 ++++++ p_kit/visualization/__init__.py | 5 +- p_kit/visualization/hypercube.py | 80 +++-------------- p_kit/visualization/utils.py | 146 +++++++++++++++++++++++++++++++ 6 files changed, 204 insertions(+), 82 deletions(-) create mode 100644 p_kit/benchmark/__init__.py create mode 100644 p_kit/benchmark/reservoir.py diff --git a/examples/llm/llm_demo_005.py b/examples/llm/llm_demo_005.py index 684d405..1c2521d 100644 --- a/examples/llm/llm_demo_005.py +++ b/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__": diff --git a/p_kit/benchmark/__init__.py b/p_kit/benchmark/__init__.py new file mode 100644 index 0000000..55739b8 --- /dev/null +++ b/p_kit/benchmark/__init__.py @@ -0,0 +1,3 @@ +from .reservoir import drive_reservoir, benchmark_reservoir + +__all__ = ["drive_reservoir", "benchmark_reservoir"] diff --git a/p_kit/benchmark/reservoir.py b/p_kit/benchmark/reservoir.py new file mode 100644 index 0000000..13d02c4 --- /dev/null +++ b/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) diff --git a/p_kit/visualization/__init__.py b/p_kit/visualization/__init__.py index 6076d2d..b6d5b4b 100644 --- a/p_kit/visualization/__init__.py +++ b/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" ] diff --git a/p_kit/visualization/hypercube.py b/p_kit/visualization/hypercube.py index 78abb70..238cdf5 100644 --- a/p_kit/visualization/hypercube.py +++ b/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) = / 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,17 +116,13 @@ 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) = /n") @@ -170,15 +130,14 @@ def _hypercube_plot_projected(output, n_pbits, labels=None): 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() diff --git a/p_kit/visualization/utils.py b/p_kit/visualization/utils.py index d78eb1d..8aa8f60 100644 --- a/p_kit/visualization/utils.py +++ b/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) = / 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)