Optimize v3 native KPM energy truncation, benchmark against untruncated KPM

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.")

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,24 +3477,51 @@ class Chain
         if (nrm<1E-14) return {phi0,0.0};
         std::vector<ITensor> V;
         V.push_back(phi0/nrm);
+        // Hk_col[j][i] = Hk(i,j) = <V_i|H|V_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
+        // <V_i|H*V_j> 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<std::vector<Cplx>> Hk_col;
         for (int j=0;j<dK-1;++j)
             {
             ITensor w; PH.product(V.at(j),w);
+            std::vector<Cplx> 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<V.size();++i)
+                    {
+                    Cplx c = eltC(dag(V.at(i))*w);
+                    col[i] += c;
+                    w -= c*V.at(i);
+                    }
+            Hk_col.push_back(col); // column j, entries i=0..j
             double nw = norm(w);
             if (nw<1E-12) break; // invariant subspace found; fewer than dK vectors is fine
             V.push_back(w/nw);
             }
         int k = (int)V.size();
         auto a = Index(k,TagSet("KPMEnergyTrunc,a"));
         auto Hk = ITensor(prime(a),a);
-        for (int j=0;j<k;++j)
+        for (int j=0;j<(int)Hk_col.size();++j)
+            for (int i=0;i<(int)Hk_col[j].size();++i)
+                {
+                Hk.set(prime(a)(i+1),a(j+1),Hk_col[j][i]);
+                if (i!=j) Hk.set(prime(a)(j+1),a(i+1),std::conj(Hk_col[j][i]));
+                }
+        if ((int)Hk_col.size()<k) // ran the full dK-1 extensions: last vector's own column was never formed above
             {
-            ITensor w; PH.product(V.at(j),w);
+            ITensor w; PH.product(V.at(k-1),w);
             for (int i=0;i<k;++i)
-                Hk.set(prime(a)(i+1),a(j+1),eltC(dag(V.at(i))*w));
+                Hk.set(prime(a)(i+1),a(k),eltC(dag(V.at(i))*w));
             }
         Hk = 0.5*(Hk + dag(swapPrime(Hk,0,1))); // Hermitize away floating-point asymmetry
         ITensor U,D;

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)