From 38c698a2d2ab9a63533835c310d1c7765b9db662 Mon Sep 17 00:00:00 2001 From: joselado Date: Thu, 30 Jul 2026 14:26:27 +0300 Subject: [PATCH 1/2] Optimize v3 native KPM energy truncation: halve redundant matvecs, add benchmark Chain::kpm_local_krylov_projection() (mpscpp3/chain_session.h) built its dense per-site Krylov Hamiltonian Hk in two passes: a Gram-Schmidt extension loop that computes H*V_j via LocalMPO::product() while orthogonalizing against existing vectors, then a second, separate loop that called LocalMPO::product() again for every vector just to rebuild Hk from scratch. That second pass is redundant -- the orthogonalization coefficients the first loop already computes and discards (summed across its two reorthogonalization passes) are exactly Hk's off/diagonal entries to working precision, since a Hermitian operator's projected matrix satisfies Hk(j,i) = conj(Hk(i,j)). Only the last accepted vector's own row/column is genuinely never produced by the extension loop (which stops one short by design), so exactly one extra matvec covers it, not k. This cuts LocalMPO::product() calls -- the dominant cost of this method, a full local-effective-Hamiltonian contraction -- from ~2k-1 down to ~k per site, with no change to the numerics (same dense Hermitian Hk, same Hermitization/diagHermitian/thresholding). Verified against tests/test_kpm_energy_truncation_v3.py, tests/test_kpm_energy_truncation_v3_accuracy.py, and tests/test_kpm_divergence_guard_catchable.py (all still pass), plus a full pytest tests/ run (293 passed, 0 failed). Added examples/dynamical_correlator/kpm_energy_truncation_v3_benchmark, turning what was previously an ad hoc, unsaved timing observation into a saved, rerunnable regression. Measured on this system (4- and 6-site Heisenberg chains, dK=30/nsweeps=10, back-to-back before/after the optimization): n=4: truncated KPM 0.93s -> 0.56s (~1.7x faster) n=6: truncated KPM 13.9s -> 7.2s (~1.9x faster) Energy truncation remains a net slowdown vs. plain KPM on these small systems (untruncated: 0.06s / 0.14s) -- this optimization narrows that gap, it doesn't eliminate it; the resolution win truncation buys still doesn't pay for itself until a system is large enough for a correlator's spectral weight to genuinely separate from the full bandwidth (see the example's own module docstring). Co-Authored-By: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_01DoZKWzhMo7qKoiP12gbcvv --- .../main.py | 90 +++++++++++++++++++ src/dmrgpy/mpscpp3/chain_session.h | 46 ++++++++-- 2 files changed, 128 insertions(+), 8 deletions(-) create mode 100644 examples/dynamical_correlator/kpm_energy_truncation_v3_benchmark/main.py diff --git a/examples/dynamical_correlator/kpm_energy_truncation_v3_benchmark/main.py b/examples/dynamical_correlator/kpm_energy_truncation_v3_benchmark/main.py new file mode 100644 index 00000000..bc282baa --- /dev/null +++ b/examples/dynamical_correlator/kpm_energy_truncation_v3_benchmark/main.py @@ -0,0 +1,90 @@ +# Timing comparison: ITensor v3's native KPM (mpscpp3/chain_session.h) +# with and without energy truncation (Holzner, Weichselbaum, McCulloch & +# von Delft, "Chebyshev matrix product state approach for spectral +# functions", PRB 83, 195115 (2011), Sec. III-B). +# +# Energy truncation (kpm_energy_truncate=True) lets the KPM window be +# narrower than the full many-body bandwidth (see +# examples/dynamical_correlator/dynamical_correlator_kpm_energy_truncation +# for the accuracy-focused demo, and +# src/dmrgpy/mpscpp3/chain_session.h's kpm_energy_truncate()/ +# scaled_hamiltonian_gs_anchored() for the implementation) -- but the +# extra truncation sweeps it runs (kpm_truncate_dK Krylov vectors x +# kpm_truncate_nsweeps sweeps, once per Chebyshev moment) are themselves +# real cost. This script turns what was previously an ad hoc, unsaved +# measurement into a saved, rerunnable timing regression: it is a net +# *slowdown* on small test systems like the ones below (the narrower +# window's resolution gain doesn't pay for itself until a system is +# large enough for a correlator's spectral weight to genuinely separate +# from the full bandwidth) -- this script exists to quantify that +# slowdown concretely, not to claim truncation is faster here. + +# Add the root path of the dmrgpy library +import os ; import sys ; sys.path.append(os.getcwd()+'/../../../src') + +import time +import numpy as np +from dmrgpy import spinchain + +DELTA = 0.08 +ES = np.linspace(0.3, 1.2, 41) +NARROW_KPM_SCALE = 0.65 # matches tests/test_kpm_energy_truncation_v3_accuracy.py +TRUNC_DK = 30 # paper's Table I recommended values +TRUNC_NSWEEPS = 10 + + +def make_chain(n): + sc = spinchain.Spin_Chain(["S=1/2" for _ in range(n)]) + sc.setup_cpp(3) + h = 0 + for i in range(n - 1): + h = h + sc.Sx[i]*sc.Sx[i+1] + sc.Sy[i]*sc.Sy[i+1] + sc.Sz[i]*sc.Sz[i+1] + sc.set_hamiltonian(h) + return sc + + +def run_untruncated(n): + sc = make_chain(n) + name = (sc.Sz[0], sc.Sz[0]) + t0 = time.time() + x, y = sc.get_dynamical_correlator(mode="DMRG", submode="KPM", name=name, + es=ES, delta=DELTA) + dt = time.time() - t0 + y = np.array(y).real + peak = np.array(x)[np.argmax(y)] + return dt, peak, y + + +def run_truncated(n): + sc = make_chain(n) + sc.kpm_scale = NARROW_KPM_SCALE + sc.kpm_energy_truncate = True + sc.kpm_truncate_dK = TRUNC_DK + sc.kpm_truncate_nsweeps = TRUNC_NSWEEPS + name = (sc.Sz[0], sc.Sz[0]) + t0 = time.time() + x, y = sc.get_dynamical_correlator(mode="DMRG", submode="KPM", name=name, + es=ES, delta=DELTA) + dt = time.time() - t0 + y = np.array(y).real + peak = np.array(x)[np.argmax(y)] + return dt, peak, y + + +print(f"{'n':>3} {'untrunc [s]':>12} {'trunc [s]':>10} {'slowdown':>9} " + f"{'peak untrunc':>13} {'peak trunc':>11}") +for n in (4, 6): + dt_u, peak_u, y_u = run_untruncated(n) + dt_t, peak_t, y_t = run_truncated(n) + slowdown = dt_t / dt_u + print(f"{n:3d} {dt_u:12.3f} {dt_t:10.3f} {slowdown:8.1f}x " + f"{peak_u:13.3f} {peak_t:11.3f}") + + assert np.all(np.isfinite(y_u)) and np.all(np.isfinite(y_t)) + # both must locate the same resonance -- energy truncation (and any + # optimization of its own internals) must not move the physics, + # only the wall-clock time + assert abs(peak_u - peak_t) <= 2 * (ES[1] - ES[0]), \ + f"n={n}: truncated/untruncated peaks disagree ({peak_t} vs {peak_u})" + +print("OK: truncated and untruncated v3 KPM agree on the resonance location.") diff --git a/src/dmrgpy/mpscpp3/chain_session.h b/src/dmrgpy/mpscpp3/chain_session.h index 899fe082..0a471b86 100644 --- a/src/dmrgpy/mpscpp3/chain_session.h +++ b/src/dmrgpy/mpscpp3/chain_session.h @@ -3458,9 +3458,12 @@ class Chain // - Unlike arnoldi_smallest_real()'s Hessenberg-only bookkeeping // (sufficient there since only one Ritz pair is ever extracted), // every Krylov component below threshold must survive here, so - // the *full* dense Hermitian projected matrix is built (one extra - // matvec pass over the already-orthonormalized basis) rather than - // reusing only the entries encountered during Gram-Schmidt. + // the *full* dense Hermitian projected matrix is built -- but + // reusing the entries already computed while orthogonalizing + // below (see the comment on Hk_col), rather than a second, + // separate PH.product() pass over the whole basis: PH.product() + // (a full local-effective-Hamiltonian contraction) dominates this + // method's cost, so avoiding a redundant pass roughly halves it. // - The projector keeps |eps_alpha| < threshold (both signs), not // just eps_alpha < threshold as in Eq. (38): scaled_hamiltonian_ // gs_anchored() pins the ground state near -1 by construction, so @@ -3474,12 +3477,33 @@ class Chain if (nrm<1E-14) return {phi0,0.0}; std::vector V; V.push_back(phi0/nrm); + // Hk_col[j][i] = Hk(i,j) = for i<=j: while extending + // from V_j, w starts out as H*V_j, and each orthogonalization + // pass below subtracts eltC(dag(Vi)*w)*Vi from it -- the *sum* of + // those coefficients across both passes already equals + // to working precision (w's own residual component + // along V_i is at machine-epsilon level after two passes), so + // capturing them here needs no extra matvec. This only ever + // yields entries for i<=j (V_j's own extension only ever + // orthogonalizes against vectors already present, i.e. i<=j); + // the missing i>j half is filled by Hermitian symmetry below, + // and the one column this loop can never produce -- j=k-1, since + // the loop stops one short of extending past the last accepted + // vector -- gets exactly one extra PH.product() call instead of + // the k calls the original full second pass used. + std::vector> Hk_col; for (int j=0;j col(V.size(),0.0); for (int pass=0;pass<2;++pass) - for (auto const& Vi : V) - w -= eltC(dag(Vi)*w)*Vi; + for (size_t i=0;i Date: Thu, 30 Jul 2026 14:26:40 +0300 Subject: [PATCH 2/2] Fix order-dependent flake in test_energy_truncate_noop_when_window_is_wide Found while running the full suite to verify an unrelated v3 KPM optimization: this test builds an excited state via _excited_state(sc), which consumes draws from Python's global np.random state through pyitensor's randomized excited-state search -- but unlike its two sibling tests in the same file, it never called np.random.seed(0) first. Whichever tests ran earlier in the same pytest process therefore determined its random draws, making it pass or fail depending on test collection order rather than on its own logic (reproduced deterministically twice in a row against the same commit: identical failure value 0.10780681457731987 both times, matching pytest's fixed, non-randomized default collection order). Added the same np.random.seed(0) call its sibling tests already use, in the same position (before any chain/DMRG construction). Verified with a full pytest tests/ run (293 passed, 0 failed) and a separate run of this file after a broad slice of other random-draw-consuming test files (190 passed, 0 failed). Co-Authored-By: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_01DoZKWzhMo7qKoiP12gbcvv --- tests/test_kpm_energy_truncation.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_kpm_energy_truncation.py b/tests/test_kpm_energy_truncation.py index 7cf6d173..91fdf5b9 100644 --- a/tests/test_kpm_energy_truncation.py +++ b/tests/test_kpm_energy_truncation.py @@ -62,6 +62,7 @@ def test_energy_truncate_noop_when_window_is_wide(): """With Ws~W (today's safe default kpm_scale), every site's local Krylov spectrum should already sit inside the [-1,1] threshold, so truncation must leave the state numerically unchanged.""" + np.random.seed(0) sc = _heisenberg_chain(6) session = sc._session scaled_H, emin, emax, scale = session._scaled_hamiltonian(kpm_scale=0.7)