diff --git a/benchmarks/cross_library/README.md b/benchmarks/cross_library/README.md new file mode 100644 index 000000000..6d0f2b419 --- /dev/null +++ b/benchmarks/cross_library/README.md @@ -0,0 +1,159 @@ +# Cross-library tensor-network benchmark suite + +Compares **TeNPy**, **quimb**, and **Cytnx** on the same four classes of +tensor-network algorithm, using the same physical models and the same +`(bond_dim, num_sites)` parameter grid for every library, so that speed and +peak-memory differences reflect implementation/library choices rather than +differences in workload. + +## Algorithm classes + +| # | Class | Model | Library implementations | +|---|-------|-------|--------------------------| +| 1 | Finite two-site DMRG, dense | 1D spin-1/2 Heisenberg chain, no symmetry | `tenpy_bench/test_dmrg_dense.py`, `quimb_bench/test_dmrg.py`, `cytnx_bench/test_dmrg_dense.py` | +| 1' | Finite two-site DMRG, block-sparse | Same chain, U(1) total-Sz conserved | `tenpy_bench/test_dmrg_symmetric.py`, `quimb_bench/test_dmrg.py`, `cytnx_bench/test_dmrg_symmetric.py` | +| 2 | Real-time evolution after a field quench (TEBD) | 1D transverse-field Ising chain | `tenpy_bench/test_tebd.py`, `quimb_bench/test_tebd.py`, `cytnx_bench/test_tebd.py` | +| 3 | Variational MPS ground-state search by gradient descent | Same Heisenberg chain as class 1 | `tenpy_bench/test_variational_manual_grad.py`, `quimb_bench/test_variational_ad.py`, `cytnx_bench/test_variational_manual_grad.py` | + +Each script's `run_one(chi, L)` lives directly in its `test_*.py` file +(there are no separate, non-test algorithm modules) so that every test +file is self-contained and importable as ordinary pytest-discovered code, +with no `sys.path` manipulation. `quimb_bench/test_dmrg.py` covers both +class 1 and class 1' in one file, since quimb's block-sparse variant +(`run_one_symmetric`) is a self-consistency check (imaginary-time +evolution from a random seeded state), not a ground-state energy +comparable across libraries — see the docstring in that file. + +All four classes share the model/parameter definitions in `common/model.py` +(`HEISENBERG_J`, `TFIM_J`/`TFIM_HX_FINAL`/`TFIM_DT`, the +`(BOND_DIM_VALUES, NUM_SITES_VALUES)` grid). + +### Gradient computation in class 3 + +All three libraries run the same optimization algorithm: one gradient step +of the Rayleigh quotient `E(psi) = /` is taken on every +MPS tensor simultaneously (no orthogonality center, no per-step +canonicalization), followed by a single global rescale derived from the +new `` and distributed evenly across all `L` tensors. Only the +gradient-computation method differs. + +quimb has a native autodiff path (JAX or PyTorch arrays under the hood), so +`quimb_bench/test_variational_ad.py` differentiates `/` +directly through the backend's `grad`/`backward`. + +TeNPy and Cytnx have no autodiff backend. Each gets its own hand-derived +analytic gradient of the same Rayleigh quotient with respect to every MPS +tensor `A_i` simultaneously: + +``` +dE/dA_i* = (2 / den) * (H_eff,i(A_i) - E * N_eff,i(A_i)) +``` + +where `den = `, `H_eff,i` is the effective one-site Hamiltonian +built from the L/R H-boundary environments around site `i`, and `N_eff,i` +is the analogous contraction with trivial (no-MPO) norm-boundary +environments. The norm-environment term is necessary because, with every +tensor updated at once and no canonicalization step, the tensors away from +site `i` are not isometric, so `N_eff,i` does not collapse to `A_i` as it +would in a one-site sweep. All four environment sets (H-left, H-right, +norm-left, norm-right) are rebuilt from scratch every gradient step, since +every tensor changes simultaneously. The TeNPy and Cytnx implementations of +this formula are written independently (`np_conserved` contractions vs. +Cytnx `UniTensor`/`Network`/`Contract`), not via a shared abstraction, since +the point of the benchmark is to compare each library's own primitives. + +This whole-network update is a weaker optimizer than a one-site sweep, so +its converged energy is sensitive to each library's own initial-state +construction and RNG; correctness is validated by checking that each +library's own implementation lands close to its own reference energy +(tight per-library tolerance), not by comparing energies across libraries +or by exact-diagonalization matching (unlike classes 1 and 2, which are +ED-validated — see the docstrings in `test_dmrg_dense.py`/`test_tebd.py` +for details). + +## Parameter grid + +`common/model.py` defines the `(bond_dim, num_sites)` grid shared by every +script: + +``` +BOND_DIM_VALUES = [16, 32, 64] +NUM_SITES_VALUES = [20, 30, 50] +``` + +Each `test_.py` parametrizes its benchmark test over the full +Cartesian product of `BOND_DIM_VALUES` and `NUM_SITES_VALUES` (9 points), via +two stacked `@pytest.mark.parametrize` decorators — one over `bond_dim`, one +over `num_sites`. Every point is bounded by the per-point wall-clock budget +`GRID_POINT_TIMEOUT_SEC` (120s by default, enforced via `pytest-timeout`), so +a single slow large-bond_dim/large-num_sites point fails on its own rather +than hanging the rest of the run — see "pytest-benchmark / pytest-memray +regression tests" below for how to run an individual point. + +## CPU vs. GPU + +- **TeNPy**: CPU only (TeNPy itself has no GPU backend). +- **quimb**: CPU and GPU paths are both implemented for the AD benchmark + (`variational_ad.py`'s `run_one_jax`/`run_one_torch`); the DMRG/TEBD scripts + also carry a `DEVICE = "gpu"` switch. +- **Cytnx**: every script has a `DEVICE` module-level flag (`"cpu"` or + `"gpu"`); the GPU path moves every MPS/MPO `UniTensor` to + `cytnx.Device.cuda` before the sweep/step. + +GPU code paths exist in every script that supports them but were written +and reviewed without a GPU available in the development environment, so +they are untested. Set `DEVICE = "gpu"` at the top of the relevant script(s) +to exercise them on a CUDA-capable machine. + +## Running the suite + +The whole suite is pytest-native: each script's `run_one(chi, L)` lives in +its `test_.py` file, exercised across the full `(bond_dim, num_sites)` +grid. There is no separate orchestration script and no standalone algorithm +modules. + +```sh +pip install tenpy quimb cytnx jax torch +pip install -e '.[benchmark]' # pytest-benchmark, pytest-memray, pytest-timeout + +cd benchmarks/cross_library + +# Full (bond_dim, num_sites) grid, timing only (skips the memory tests), +# with a pytest.approx assertion against a precomputed reference energy at +# every point: +python3 -m pytest --benchmark-only -q + +# Memory only: --memray instruments every collected test, so select just +# the memory tests with the cytnx_memory marker (registered in +# pyproject.toml, applied to every *_memory test in this directory): +python3 -m pytest --memray -m cytnx_memory -q + +# Run the whole grid for one script: +python3 -m pytest cytnx_bench/test_dmrg_dense.py --benchmark-only -q + +# Or call pytest one point at a time, which doubles as a resume mechanism if +# a run gets interrupted partway through the grid: +python3 -m pytest "cytnx_bench/test_dmrg_dense.py::test_dmrg_dense_benchmark[16-20]" --benchmark-only -q +``` + +Run from `benchmarks/cross_library` (not the repo root) with `python3 -m +pytest` (not the bare `pytest` command) so that `cytnx`/`quimb`/`tenpy` +resolve to your installed packages rather than colliding with the repo's +source tree. + +## pytest-benchmark / pytest-memray regression tests + +Each of the 12 `test_.py` files exercises its `run_one(chi, L)` across +the full `(bond_dim, num_sites)` grid through `pytest-benchmark`'s +`benchmark.pedantic`, asserting the returned energy against a +`REFERENCE_ENERGIES[(bond_dim, num_sites)]` dict via `pytest.approx`, so a +wrong physical answer fails the test rather than silently shipping a bad +timing number. The result is also recorded via `benchmark.extra_info["energy"]`; +pass `--benchmark-json=out.json` to capture it (and the timing statistics) +for every point in one file. + +The same file's `test__memory` test takes no parameters — it calls +`run_one` directly (no `benchmark` fixture, so pytest-benchmark's overhead +doesn't contaminate memray's allocation trace) at the single canonical +`(chi=16, L=20)` point, under a `@pytest.mark.limit_memory(...)` decorator +tuned to that point's footprint. diff --git a/benchmarks/cross_library/__init__.py b/benchmarks/cross_library/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/benchmarks/cross_library/common/__init__.py b/benchmarks/cross_library/common/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/benchmarks/cross_library/common/model.py b/benchmarks/cross_library/common/model.py new file mode 100644 index 000000000..5ee926906 --- /dev/null +++ b/benchmarks/cross_library/common/model.py @@ -0,0 +1,65 @@ +"""Shared physical model and parameter-grid definitions for the +TeNPy / quimb / Cytnx cross-library benchmark suite. + +All four algorithm classes benchmark the *same* physical model and the +*same* (bond_dim, num_sites) parameter grid, so that the only thing that +varies between runs is the library/implementation: + + - DMRG (dense) : 1D spin-1/2 Heisenberg chain, no symmetry + - DMRG (block-sparse) : same chain, U(1) total-Sz conserved + - TEBD/TDVP (dynamics) : 1D transverse-field Ising chain, field quench + - Variational AD/manual : same Heisenberg chain as the dense DMRG case +""" + +# Heisenberg coupling J (isotropic, antiferromagnetic) for DMRG / variational. +HEISENBERG_J = 1.0 + +# Transverse-field Ising parameters for the TEBD/TDVP quench benchmark. All +# three implementations start from the all-down product state |0...0> and +# evolve it directly under this post-quench Hamiltonian +# H = -J * sum ZZ - hx_f * sum X (no separate initial-field ground state is +# prepared). +TFIM_J = 1.0 +TFIM_HX_FINAL = 0.5 +TFIM_DT = 0.05 +TFIM_N_STEPS = 40 + +# 2D sweep grid shared by every algorithm/library: bond dimension vs. +# number of sites. This is the axis pair the user asked to scan in order to +# expose the memory (~O(num_sites * bond_dim^2) for an MPS) and speed +# (~O(num_sites * bond_dim^3) for a DMRG/TDVP local update) tradeoffs. +BOND_DIM_VALUES = [16, 32, 64] +NUM_SITES_VALUES = [20, 30, 50] + +# Number of DMRG sweeps measured per (bond_dim, num_sites) point. Kept small +# because we only need a handful of sweeps to get a stable per-sweep timing +# and peak-memory estimate, not a converged ground state. +N_SWEEPS = 3 + +# Number of Lanczos iterations for Cytnx's local two-site eigensolver +# (cytnx_bench/test_dmrg_{dense,symmetric}.py). TeNPy's TwoSiteDMRGEngine +# is passed this same value via its `lanczos_params["N_max"]` +# (tenpy_bench/test_dmrg_{dense,symmetric}.py) so both libraries spend the +# same per-bond eigensolver budget converging to the same ground-state +# energy; quimb's DMRG2 has no equivalent maxiter knob in its public API, so +# its scipy-eigsh-backed solver is left at its native default convergence. +# 4 was not quite enough for cytnx_bench/test_dmrg_symmetric.py's hardest +# grid point (bond_dim=16, num_sites=50) to land within rel=1e-6 of the +# shared REFERENCE_ENERGIES value -- TeNPy's own Krylov-based solver already +# matches that value to ~1e-8 at this same N_max, so the gap is specific to +# Cytnx's own Lanczos implementation on that hardest local subspace, not a +# tolerance or parameter mismatch between libraries. Raising the shared +# budget to 6 (rather than loosening the tolerance or letting Cytnx diverge +# from TeNPy's setting) closes it to ~7e-7 across the full grid. +LANCZOS_MAXITER = 6 + +# SVD truncation magnitude cutoff, matching TeNPy's svd_min and quimb's +# cutoffs (both 1e-10 in this suite's dmrg/tebd scripts): singular values +# below this threshold are discarded even if the bond-dimension cap isn't +# reached. Passed to Cytnx's Svd_truncate via its `err` argument. +SVD_CUTOFF = 1e-10 + +# Per-(bond_dim, num_sites) wall-clock budget. Points that exceed this are +# skipped rather than measured, so a handful of slow large-bond_dim/ +# large-num_sites points don't dominate the suite's total run time. +GRID_POINT_TIMEOUT_SEC = 120 diff --git a/benchmarks/cross_library/cytnx_bench/__init__.py b/benchmarks/cross_library/cytnx_bench/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/benchmarks/cross_library/cytnx_bench/test_dmrg_dense.py b/benchmarks/cross_library/cytnx_bench/test_dmrg_dense.py new file mode 100644 index 000000000..ba78f8381 --- /dev/null +++ b/benchmarks/cross_library/cytnx_bench/test_dmrg_dense.py @@ -0,0 +1,233 @@ +"""Cytnx benchmark, algorithm class 1: finite two-site DMRG, dense mode (no +conserved quantum numbers) on the 1D spin-1/2 Heisenberg chain. + +Implements the standard bond-dimension-5 finite-state-machine MPO for the +isotropic Heisenberg coupling J*S_i.S_j = (J/2)*(S+_i S-_j + S-_i S+_j) + +J*Sz_i Sz_j, and a textbook two-site DMRG sweep (local eigensolve via +Cytnx's `LinOp` + `Lanczos`, truncation via `Svd_truncate`) adapted from +`example/DMRG/dmrg_two_sites_dense.py`. The MPO was verified against exact +diagonalization for small chains (num_sites=4,6) before being used here. + +CPU and GPU code paths are both written; the GPU path moves every MPS/MPO +UniTensor to `cytnx.Device.cuda` before the sweep. It cannot be exercised in +this environment (no GPU). + +Run timing with `pytest --benchmark-only test_dmrg_dense.py`, memory with +`pytest --memray test_dmrg_dense.py`. +""" +import pytest + +import cytnx + +from common.model import BOND_DIM_VALUES, HEISENBERG_J, LANCZOS_MAXITER, NUM_SITES_VALUES, N_SWEEPS, GRID_POINT_TIMEOUT_SEC, SVD_CUTOFF + +DEVICE = "cpu" # set to "gpu" to exercise the (untested) GPU code path below + +REFERENCE_ENERGIES = { + (16, 20): -8.682468456356828, + (16, 30): -13.111313454922634, + (16, 50): -21.97181361569925, + (32, 20): -8.682473319689738, + (32, 30): -13.111355524192675, + (32, 50): -21.972106512033726, + (64, 20): -8.682473334397873, + (64, 30): -13.11135575848871, + (64, 50): -21.972110271889434, +} + + +class _Hxx(cytnx.LinOp): + def __init__(self, anet, psidim, device): + cytnx.LinOp.__init__(self, "mv", psidim, cytnx.Type.Double, device) + self.anet = anet + + def matvec(self, v): + lbl = v.labels() + self.anet.PutUniTensor("psi", v) + out = self.anet.Launch() + out.relabel_(lbl) + return out + + +def _h_eff_network(): + anet = cytnx.Network() + anet.FromString(["psi: -1,-2,-3,-4", + "L: -5,-1,0", + "R: -7,-4,3", + "M1: -5,-6,-2,1", + "M2: -6,-7,-3,2", + "TOUT: 0,1;2,3"]) + return anet + + +def _l_update_network(): + anet = cytnx.Network() + anet.FromString(["L: -2,-1,-3", + "A: -1,-4,1", + "M: -2,0,-4,-5", + "A_Conj: -3,-5,2", + "TOUT: 0,1,2"]) + return anet + + +def _r_update_network(): + anet = cytnx.Network() + anet.FromString(["R: -2,-1,-3", + "B: 1,-4,-1", + "M: 0,-2,-4,-5", + "B_Conj: 2,-5,-3", + "TOUT: 0;1,2"]) + return anet + + +def _optimize_psi(anet, psi, L, M1, M2, R, maxit, device): + anet.PutUniTensors(["L", "M1", "M2", "R"], [L, M1, M2, R]) + H = _Hxx(anet, psi.shape()[0] * psi.shape()[1] * psi.shape()[2] * psi.shape()[3], device) + # CvgCrit=1e-8 is looser than the rel=1e-6 tolerance the returned energy is + # checked against, letting Lanczos exit before maxit when its own energy + # convergence check is satisfied. + energy, psivec = cytnx.linalg.Lanczos(Hop=H, method="Gnd", Maxiter=maxit, CvgCrit=1e-8, Tin=psi) + return psivec, energy[0].item() + + +def _build_mpo(J, device): + d = 2 + D = 5 + Sp = cytnx.zeros([d, d], device=device) + Sp[0, 1] = 1.0 # S+: |down> -> |up> + Sm = cytnx.zeros([d, d], device=device) + Sm[1, 0] = 1.0 # S-: |up> -> |down> + Sz = cytnx.zeros([d, d], device=device) + Sz[0, 0] = 0.5 + Sz[1, 1] = -0.5 + eye = cytnx.eye(d, device=device) + + M = cytnx.zeros([D, D, d, d], device=device) + M[0, 0] = eye + M[D - 1, D - 1] = eye + M[0, 1] = Sp + M[0, 2] = Sm + M[0, 3] = Sz + M[1, D - 1] = (J / 2.0) * Sm + M[2, D - 1] = (J / 2.0) * Sp + M[3, D - 1] = J * Sz + M = cytnx.UniTensor(M, 0).set_name("MPO") + + L0 = cytnx.UniTensor.zeros([D, 1, 1], device=device).set_rowrank_(0).set_name("L0") + R0 = cytnx.UniTensor.zeros([D, 1, 1], device=device).set_rowrank_(0).set_name("R0") + L0[0, 0, 0] = 1.0 + R0[D - 1, 0, 0] = 1.0 + return M, L0, R0 + + +def run_one(bond_dim, num_sites): + device = cytnx.Device.cuda if DEVICE == "gpu" else cytnx.Device.cpu + d = 2 + M, L0, R0 = _build_mpo(HEISENBERG_J, device) + + A = [None for _ in range(num_sites)] + A[0] = cytnx.UniTensor.normal([1, d, min(bond_dim, d)], 0., 1., device=device).set_rowrank_(2) + A[0].relabel_(["0", "1", "2"]).set_name("A0") + + lbls = [["0", "1", "2"]] + for k in range(1, num_sites): + dim1 = A[k - 1].shape()[2] + dim2 = d + dim3 = min(min(bond_dim, A[k - 1].shape()[2] * d), d ** (num_sites - k - 1)) + A[k] = cytnx.UniTensor.normal([dim1, dim2, dim3], 0., 1., device=device).set_rowrank_(2).set_name(f"A{k}") + lbl = [str(2 * k), str(2 * k + 1), str(2 * k + 2)] + A[k].relabel_(lbl) + lbls.append(lbl) + + LR = [None for _ in range(num_sites + 1)] + LR[0] = L0 + LR[-1] = R0 + + h_eff_net = _h_eff_network() + l_update_net = _l_update_network() + r_update_net = _r_update_network() + + for p in range(num_sites - 1): + s, A[p], vt = cytnx.linalg.Gesvd(A[p]) + A[p + 1] = cytnx.Contract(cytnx.Contract(s, vt), A[p + 1]) + A[p].set_name(f"A{p}") + A[p + 1].set_name(f"A{p+1}") + + l_update_net.PutUniTensors(["L", "A", "A_Conj", "M"], + [LR[p], A[p], A[p].Dagger().permute_(A[p].labels()), M]) + LR[p + 1] = l_update_net.Launch() + LR[p + 1].set_name(f"LR{p+1}") + A[p].relabel_(lbls[p]) + A[p + 1].relabel_(lbls[p + 1]) + + _, A[-1] = cytnx.linalg.Gesvd(A[-1], is_U=True, is_vT=False) + A[-1].set_name(f"A{num_sites-1}").relabel_(lbls[-1]) + + def sweep(): + energy = None + for p in range(num_sites - 2, -1, -1): + dim_l = A[p].shape()[0] + dim_r = A[p + 1].shape()[2] + new_dim = min(dim_l * d, dim_r * d, bond_dim) + psi = cytnx.Contract(A[p], A[p + 1]) + psi, energy = _optimize_psi(h_eff_net, psi, LR[p], M, M, LR[p + 2], LANCZOS_MAXITER, device) + psi.set_rowrank_(2) + s, A[p], A[p + 1] = cytnx.linalg.Svd_truncate(psi, new_dim, err=SVD_CUTOFF) + A[p + 1].set_name(f"A{p+1}").relabel_(lbls[p + 1]) + s = s / s.Norm().item() + A[p] = cytnx.Contract(A[p], s) + A[p].set_name(f"A{p}").relabel_(lbls[p]) + + r_update_net.PutUniTensors(["R", "B", "M", "B_Conj"], + [LR[p + 2], A[p + 1], M, A[p + 1].Dagger().permute_(A[p + 1].labels())]) + LR[p + 1] = r_update_net.Launch() + LR[p + 1].set_name(f"LR{p+1}") + + A[0].set_rowrank_(1) + _, A[0] = cytnx.linalg.Gesvd(A[0], is_U=False, is_vT=True) + A[0].set_name("A0").relabel_(lbls[0]) + + for p in range(num_sites - 1): + dim_l = A[p].shape()[0] + dim_r = A[p + 1].shape()[2] + new_dim = min(dim_l * d, dim_r * d, bond_dim) + psi = cytnx.Contract(A[p], A[p + 1]) + psi, energy = _optimize_psi(h_eff_net, psi, LR[p], M, M, LR[p + 2], LANCZOS_MAXITER, device) + psi.set_rowrank_(2) + s, A[p], A[p + 1] = cytnx.linalg.Svd_truncate(psi, new_dim, err=SVD_CUTOFF) + A[p].set_name(f"A{p}").relabel_(lbls[p]) + s = s / s.Norm().item() + A[p + 1] = cytnx.Contract(s, A[p + 1]) + A[p + 1].set_name(f"A{p+1}").relabel_(lbls[p + 1]) + + l_update_net.PutUniTensors(["L", "A", "A_Conj", "M"], + [LR[p], A[p], A[p].Dagger().permute_(A[p].labels()), M]) + LR[p + 1] = l_update_net.Launch() + LR[p + 1].set_name(f"LR{p+1}") + + A[-1].set_rowrank_(2) + _, A[-1] = cytnx.linalg.Gesvd(A[-1], is_U=True, is_vT=False) + A[-1].set_name(f"A{num_sites-1}").relabel_(lbls[-1]) + return energy + + energy = None + for _ in range(N_SWEEPS): + energy = sweep() + return energy + + +@pytest.mark.timeout(GRID_POINT_TIMEOUT_SEC) +@pytest.mark.parametrize("num_sites", NUM_SITES_VALUES) +@pytest.mark.parametrize("bond_dim", BOND_DIM_VALUES) +def test_dmrg_dense_benchmark(benchmark, bond_dim, num_sites): + energy = benchmark.pedantic(run_one, args=(bond_dim, num_sites), rounds=1, iterations=1) + benchmark.extra_info["energy"] = energy + assert energy == pytest.approx(REFERENCE_ENERGIES[(bond_dim, num_sites)], rel=1e-6) + + +@pytest.mark.cytnx_memory +@pytest.mark.parametrize("num_sites", NUM_SITES_VALUES) +@pytest.mark.parametrize("bond_dim", BOND_DIM_VALUES) +def test_dmrg_dense_memory(bond_dim, num_sites): + energy = run_one(bond_dim, num_sites) + assert energy == pytest.approx(REFERENCE_ENERGIES[(bond_dim, num_sites)], rel=1e-6) diff --git a/benchmarks/cross_library/cytnx_bench/test_dmrg_symmetric.py b/benchmarks/cross_library/cytnx_bench/test_dmrg_symmetric.py new file mode 100644 index 000000000..917893332 --- /dev/null +++ b/benchmarks/cross_library/cytnx_bench/test_dmrg_symmetric.py @@ -0,0 +1,261 @@ +"""Cytnx benchmark, algorithm class 1 (symmetric variant): block-sparse, +U(1)-total-Sz-conserving finite two-site DMRG ground-state search on the 1D +spin-1/2 Heisenberg chain. + +Uses the same bond-dimension-5 finite-state-machine MPO as +`test_dmrg_dense.py`, here split into 5 U(1) charge sectors (start=0, +S+-pending=+2, S--pending=-2, Sz-pending=0, end=0) following the charge +bookkeeping convention of `example/DMRG/dmrg_two_sites_U1.py` (a leg with +Cytnx bond-type `BD_KET` contributes +charge and `BD_BRA` contributes +-charge to the zero-sum constraint at every nonzero MPO block). The +resulting symmetric MPO was verified against exact diagonalization for +small chains (num_sites=4,6), matching the dense-mode MPO in `test_dmrg_dense.py` +to machine precision, before being used here. + +CPU and GPU code paths are both written; the GPU path moves every +block-sparse MPS/MPO UniTensor to `cytnx.Device.cuda` before the sweep. It +cannot be exercised in this environment (no GPU). + +Run timing with `pytest --benchmark-only test_dmrg_symmetric.py`, memory +with `pytest --memray test_dmrg_symmetric.py`. +""" +import pytest + +import cytnx + +from common.model import BOND_DIM_VALUES, HEISENBERG_J, LANCZOS_MAXITER, NUM_SITES_VALUES, N_SWEEPS, GRID_POINT_TIMEOUT_SEC, SVD_CUTOFF + +DEVICE = "cpu" # set to "gpu" to exercise the (untested) GPU code path below +TARGET_Q = 0 # global U(1) total-Sz sector to search within + +REFERENCE_ENERGIES = { + (16, 20): -8.682468456356823, + (16, 30): -13.111313454922696, + (16, 50): -21.971813615699435, + (32, 20): -8.682473319689738, + (32, 30): -13.111355524202278, + (32, 50): -21.972106512010665, + (64, 20): -8.682473334397892, + (64, 30): -13.11135575848872, + (64, 50): -21.972110271890013, +} + + +class _Hxx(cytnx.LinOp): + def __init__(self, anet, psidim, device): + cytnx.LinOp.__init__(self, "mv", psidim, cytnx.Type.Double, device) + self.anet = anet + + def matvec(self, v): + lbl = v.labels() + self.anet.PutUniTensor("psi", v) + out = self.anet.Launch() + out.relabel_(lbl) + return out + + +def _stored_numel(ut): + # BlockUniTensor.shape() returns each bond's nominal (sum-of-sectors) extent, + # not the element count actually stored in the nonzero charge blocks, so the + # active dimension of a block-sparse psi must be summed from its own blocks. + total = 0 + for block in ut.get_blocks_(): + n = 1 + for d in block.shape(): + n *= d + total += n + return total + + +def _h_eff_network(): + anet = cytnx.Network() + anet.FromString(["psi: -1,-2,-3,-4", + "L: -5,-1,0", + "R: -7,-4,3", + "M1: -5,-6,-2,1", + "M2: -6,-7,-3,2", + "TOUT: 0,1;2,3"]) + return anet + + +def _l_update_network(): + anet = cytnx.Network() + anet.FromString(["L: -2,-1,-3", + "A: -1,-4,1", + "M: -2,0,-4,-5", + "A_Conj: -3,-5,2", + "TOUT: 0;1,2"]) + return anet + + +def _r_update_network(): + anet = cytnx.Network() + anet.FromString(["R: -2,-1,-3", + "B: 1,-4,-1", + "M: 0,-2,-4,-5", + "B_Conj: 2,-5,-3", + "TOUT: 0;1,2"]) + return anet + + +def _optimize_psi(anet, psi, L, M1, M2, R, maxit, device): + anet.PutUniTensors(["L", "M1", "M2", "R"], [L, M1, M2, R]) + H = _Hxx(anet, _stored_numel(psi), device) + # CvgCrit=1e-8 is looser than the rel=1e-6 tolerance the returned energy is + # checked against, letting Lanczos exit before maxit when its own energy + # convergence check is satisfied. See common/model.py's LANCZOS_MAXITER + # comment: this file's hardest grid point needs the larger Lanczos budget + # there to land within tolerance regardless of CvgCrit. + energy, psivec = cytnx.linalg.Lanczos(Hop=H, method="Gnd", Maxiter=maxit, CvgCrit=1e-8, Tin=psi) + return psivec, energy[0].item() + + +def _build_mpo(J, q, device): + D = 5 + bd_inner = cytnx.Bond(cytnx.BD_KET, [[0], [2], [-2], [0], [0]], [1, 1, 1, 1, 1]) + bd_phys = cytnx.Bond(cytnx.BD_KET, [[1], [-1]], [1, 1]) + + M = cytnx.UniTensor([bd_inner, bd_inner.redirect(), bd_phys, bd_phys.redirect()], device=device) \ + .set_rowrank_(2).set_name("MPO") + + M.set_elem([0, 0, 0, 0], 1) + M.set_elem([0, 0, 1, 1], 1) + M.set_elem([D - 1, D - 1, 0, 0], 1) + M.set_elem([D - 1, D - 1, 1, 1], 1) + + M.set_elem([0, 1, 0, 1], 1) # S+ + M.set_elem([0, 2, 1, 0], 1) # S- + M.set_elem([0, 3, 0, 0], 0.5) # Sz + M.set_elem([0, 3, 1, 1], -0.5) # Sz + + M.set_elem([1, D - 1, 1, 0], J / 2.0) # (J/2) S- + M.set_elem([2, D - 1, 0, 1], J / 2.0) # (J/2) S+ + M.set_elem([3, D - 1, 0, 0], J * 0.5) # J Sz + M.set_elem([3, D - 1, 1, 1], -J * 0.5) # J Sz + + VbdL = cytnx.Bond(cytnx.BD_KET, [[0]], [1]) + VbdR = cytnx.Bond(cytnx.BD_KET, [[q]], [1]) + L0 = cytnx.UniTensor([bd_inner.redirect(), VbdL.redirect(), VbdL], device=device) \ + .set_rowrank_(1).set_name("L0") + R0 = cytnx.UniTensor([bd_inner, VbdR, VbdR.redirect()], device=device) \ + .set_rowrank_(1).set_name("R0") + L0.set_elem([0, 0, 0], 1) + R0.set_elem([D - 1, 0, 0], 1) + return M, L0, R0, bd_phys + + +def run_one(bond_dim, num_sites): + device = cytnx.Device.cuda if DEVICE == "gpu" else cytnx.Device.cpu + M, L0, R0, bd_phys = _build_mpo(HEISENBERG_J, TARGET_Q, device) + + A = [None for _ in range(num_sites)] + running_charge = 0 + charge_step = 1 if running_charge <= TARGET_Q else -1 + running_charge += charge_step + + VbdL = cytnx.Bond(cytnx.BD_KET, [[0]], [1]) + A[0] = cytnx.UniTensor([VbdL, bd_phys.redirect(), cytnx.Bond(cytnx.BD_BRA, [[running_charge]], [1])], device=device) \ + .set_rowrank_(2).set_name("A0") + A[0].get_block_()[0] = 1 + + lbls = [["0", "1", "2"]] + for k in range(1, num_sites): + B1 = A[k - 1].bonds()[2].redirect() + B2 = A[k - 1].bonds()[1] + charge_step = 1 if running_charge <= TARGET_Q else -1 + running_charge += charge_step + B3 = cytnx.Bond(cytnx.BD_BRA, [[running_charge]], [1]) + + A[k] = cytnx.UniTensor([B1, B2, B3], device=device).set_rowrank_(2).set_name(f"A{k}") + lbl = [str(2 * k), str(2 * k + 1), str(2 * k + 2)] + A[k].relabel_(lbl) + A[k].get_block_()[0] = 1 + lbls.append(lbl) + + LR = [None for _ in range(num_sites + 1)] + LR[0] = L0 + LR[-1] = R0 + + h_eff_net = _h_eff_network() + l_update_net = _l_update_network() + r_update_net = _r_update_network() + + for p in range(num_sites - 1): + l_update_net.PutUniTensors(["L", "A", "A_Conj", "M"], + [LR[p], A[p], A[p].Dagger().permute_(A[p].labels()), M]) + LR[p + 1] = l_update_net.Launch() + LR[p + 1].set_name(f"LR{p+1}") + + def sweep(): + energy = None + for p in range(num_sites - 2, -1, -1): + dim_l = A[p].shape()[0] + dim_r = A[p + 1].shape()[2] + d = A[p].shape()[1] + new_dim = min(dim_l * d, dim_r * d, bond_dim) + psi = cytnx.Contract(A[p], A[p + 1]) + psi, energy = _optimize_psi(h_eff_net, psi, LR[p], M, M, LR[p + 2], LANCZOS_MAXITER, device) + psi.set_rowrank_(2) + s, A[p], A[p + 1] = cytnx.linalg.Svd_truncate(psi, new_dim, err=SVD_CUTOFF) + A[p + 1].set_name(f"A{p+1}").relabel_(lbls[p + 1]) + s = s / s.Norm().item() + A[p] = cytnx.Contract(A[p], s) + A[p].set_name(f"A{p}").relabel_(lbls[p]) + + r_update_net.PutUniTensors(["R", "B", "M", "B_Conj"], + [LR[p + 2], A[p + 1], M, A[p + 1].Dagger().permute_(A[p + 1].labels())]) + LR[p + 1] = r_update_net.Launch() + LR[p + 1].set_name(f"LR{p+1}") + + A[0].set_rowrank_(1) + _, A[0] = cytnx.linalg.Gesvd(A[0], is_U=False, is_vT=True) + A[0].relabel_(lbls[0]) + + for p in range(num_sites - 1): + dim_l = A[p].shape()[0] + dim_r = A[p + 1].shape()[2] + d = A[p].shape()[1] + new_dim = min(dim_l * d, dim_r * d, bond_dim) + psi = cytnx.Contract(A[p], A[p + 1]) + psi, energy = _optimize_psi(h_eff_net, psi, LR[p], M, M, LR[p + 2], LANCZOS_MAXITER, device) + psi.set_rowrank_(2) + s, A[p], A[p + 1] = cytnx.linalg.Svd_truncate(psi, new_dim, err=SVD_CUTOFF) + A[p].relabel_(lbls[p]) + s = s / s.Norm().item() + A[p + 1] = cytnx.Contract(s, A[p + 1]) + A[p + 1].relabel_(lbls[p + 1]) + A[p].set_name(f"A{p}") + A[p + 1].set_name(f"A{p+1}") + + l_update_net.PutUniTensors(["L", "A", "A_Conj", "M"], + [LR[p], A[p], A[p].Dagger().permute_(A[p].labels()), M]) + LR[p + 1] = l_update_net.Launch() + LR[p + 1].set_name(f"LR{p+1}") + + A[-1].set_rowrank_(2) + _, A[-1] = cytnx.linalg.Gesvd(A[-1], is_U=True, is_vT=False) + A[-1].set_name(f"A{num_sites-1}").relabel_(lbls[-1]) + return energy + + energy = None + for _ in range(N_SWEEPS): + energy = sweep() + return energy + + +@pytest.mark.timeout(GRID_POINT_TIMEOUT_SEC) +@pytest.mark.parametrize("num_sites", NUM_SITES_VALUES) +@pytest.mark.parametrize("bond_dim", BOND_DIM_VALUES) +def test_dmrg_symmetric_benchmark(benchmark, bond_dim, num_sites): + energy = benchmark.pedantic(run_one, args=(bond_dim, num_sites), rounds=1, iterations=1) + benchmark.extra_info["energy"] = energy + assert energy == pytest.approx(REFERENCE_ENERGIES[(bond_dim, num_sites)], rel=1e-6) + + +@pytest.mark.cytnx_memory +@pytest.mark.limit_memory("20 MB") +@pytest.mark.parametrize("num_sites", NUM_SITES_VALUES) +@pytest.mark.parametrize("bond_dim", BOND_DIM_VALUES) +def test_dmrg_symmetric_memory(bond_dim, num_sites): + energy = run_one(bond_dim, num_sites) + assert energy == pytest.approx(REFERENCE_ENERGIES[(bond_dim, num_sites)], rel=1e-6) diff --git a/benchmarks/cross_library/cytnx_bench/test_tebd.py b/benchmarks/cross_library/cytnx_bench/test_tebd.py new file mode 100644 index 000000000..2f9b02310 --- /dev/null +++ b/benchmarks/cross_library/cytnx_bench/test_tebd.py @@ -0,0 +1,203 @@ +"""Cytnx benchmark, algorithm class 2: real-time evolution of a 1D +transverse-field Ising chain after a field quench, using a hand-rolled TEBD +sweep built from Cytnx's own `UniTensor`/`Contract`/`Svd_truncate` API (Cytnx +has no built-in TEBD engine, unlike TeNPy/quimb). + +Each Trotter step is the same even/odd Suzuki-Trotter (Strang) splitting +TeNPy's `TEBDEngine` (order=2) uses internally: writing the bond terms +h_p (p = 0..num_sites-2) as H_even = sum over even p, H_odd = sum over odd +p (bonds of the same parity never share a site, so every gate within one +group commutes with every other gate in that group and can be applied in +any order with no extra Trotter error), one step is +exp(-i*dt/2*H_even) * exp(-i*dt*H_odd) * exp(-i*dt/2*H_even). The on-site +field term -hx*Sx is split between the two bond gates that touch each site +(half-weight on interior bonds, full weight on the two chain-boundary +bonds) so that summing the per-bond gates' field contributions over one +full even+odd step reproduces -hx*sum(Sx_i) exactly once per site rather +than twice for interior sites; this split is independent of which group a +bond falls into; only the dt passed to `_build_gates` differs between the +two sub-steps. The initial state is the all-spin-down product state, +evolved directly under the post-quench Hamiltonian (matching the +`quimb_bench/tebd.py` convention). + +CPU and GPU code paths are both written; the GPU path moves every MPS +UniTensor and the gate to `cytnx.Device.cuda` before stepping. It cannot be +exercised in this environment (no GPU). + +Run timing with `pytest --benchmark-only test_tebd.py`, memory with +`pytest --memray test_tebd.py`. +""" +import pytest + +import cytnx + +from common.model import BOND_DIM_VALUES, NUM_SITES_VALUES, GRID_POINT_TIMEOUT_SEC, SVD_CUTOFF, TFIM_DT, TFIM_HX_FINAL, TFIM_J, TFIM_N_STEPS + +DEVICE = "cpu" # set to "gpu" to exercise the (untested) GPU code path below + +REFERENCE_ENERGIES = { + (16, 20): -19.000253478862355, + (16, 30): -29.000170747421212, + (16, 50): -48.99944000650076, + (32, 20): -19.000146442249463, + (32, 30): -29.000157629985154, + (32, 50): -49.00020127330093, + (64, 20): -19.000146553733153, + (64, 30): -29.000162610180396, + (64, 50): -49.000194070018004, +} + + +def _build_gate(J, hx, dt, w_left, w_right, device): + Sz = cytnx.physics.pauli("z", device=device).real() + Sx = cytnx.physics.pauli("x", device=device).real() + I = cytnx.eye(2, device=device) + TFterm = w_left * cytnx.linalg.Kron(Sx, I) + w_right * cytnx.linalg.Kron(I, Sx) + ZZterm = cytnx.linalg.Kron(Sz, Sz) + H = -hx * TFterm - J * ZZterm + eH = cytnx.linalg.ExpH(H, -1j * dt) + eH.reshape_(2, 2, 2, 2) + return cytnx.UniTensor(eH) + + +def _build_gates(num_sites, J, hx, dt, device): + # Every interior bond (0 < p < num_sites-2) shares the same (0.5, 0.5) + # on-site-field split, so only the two boundary bonds need a distinct + # exp(-i*dt*h_bond); the single interior gate is reused across all + # interior bonds. + if num_sites == 2: + return [_build_gate(J, hx, dt, 1.0, 1.0, device)] + left_gate = _build_gate(J, hx, dt, 1.0, 0.5, device) + right_gate = _build_gate(J, hx, dt, 0.5, 1.0, device) + interior_gate = _build_gate(J, hx, dt, 0.5, 0.5, device) if num_sites > 3 else None + return [left_gate] + [interior_gate] * (num_sites - 3) + [right_gate] + + +def _build_mps(num_sites, bond_dim, device): + d = 2 + A = [None] * num_sites + lbls = [[str(2 * k), str(2 * k + 1), str(2 * k + 2)] for k in range(num_sites)] + for k in range(num_sites): + A[k] = cytnx.UniTensor.zeros([1, d, 1], device=device).set_rowrank_(2).relabel_(lbls[k]).set_name(f"A{k}") + A[k].set_elem([0, 0, 0], 1.0) + return A, lbls + + +def _build_mpo(J, hx, device): + D = 3 + Sz = cytnx.physics.pauli("z", device=device).real() + Sx = cytnx.physics.pauli("x", device=device).real() + eye = cytnx.eye(2, device=device) + + M = cytnx.zeros([D, D, 2, 2], device=device) + M[0, 0] = eye + M[0, 1] = Sz + M[0, 2] = -hx * Sx + M[1, 2] = -J * Sz + M[D - 1, D - 1] = eye + M = cytnx.UniTensor(M, 0).set_name("MPO") + + L0 = cytnx.UniTensor.zeros([D, 1, 1], device=device).set_rowrank_(0).set_name("L0") + R0 = cytnx.UniTensor.zeros([D, 1, 1], device=device).set_rowrank_(0).set_name("R0") + L0[0, 0, 0] = 1.0 + R0[D - 1, 0, 0] = 1.0 + return M, L0, R0 + + +def _energy(A, M, L0, R0, device): + anet = cytnx.Network() + anet.FromString(["L: -2,-1,-3", + "A: -1,-4,1", + "M: -2,0,-4,-5", + "A_Conj: -3,-5,2", + "TOUT: 0,1,2"]) + LR = L0 + for p in range(len(A)): + anet.PutUniTensors(["L", "A", "A_Conj", "M"], + [LR, A[p], A[p].Dagger().permute_(A[p].labels()), M]) + LR = anet.Launch() + energy = cytnx.Contract(LR, R0).item() + + norm_net = cytnx.Network() + norm_net.FromString(["L: -2,-1", "A: -1,1,-3", "A_Conj: -2,1,-4", "TOUT: -4,-3"]) + NL = cytnx.UniTensor(cytnx.ones([1, 1], device=device)) + for p in range(len(A)): + norm_net.PutUniTensors(["L", "A", "A_Conj"], + [NL, A[p], A[p].Dagger().permute_(A[p].labels())]) + NL = norm_net.Launch() + norm2 = NL.item() + return (energy / norm2).real + + +def run_one(bond_dim, num_sites): + device = cytnx.Device.cuda if DEVICE == "gpu" else cytnx.Device.cpu + d = 2 + A, lbls = _build_mps(num_sites, bond_dim, device) + # Bond labels are fixed by _build_mps and restored after every truncation + # below, so each gate only ever needs to be relabeled once, not on every + # Trotter step. Two gate sets are needed -- exp(-i*dt/2*h_p) for the + # even-p bonds (applied at the start and end of every step) and + # exp(-i*dt*h_p) for the odd-p bonds (applied once, in between) -- see + # module docstring. + half_base_gates = _build_gates(num_sites, TFIM_J, TFIM_HX_FINAL, TFIM_DT / 2, device) + full_base_gates = _build_gates(num_sites, TFIM_J, TFIM_HX_FINAL, TFIM_DT, device) + half_gates = [half_base_gates[p].clone().relabel_(["_o0", "_o1", lbls[p][1], lbls[p + 1][1]]) + for p in range(num_sites - 1)] + full_gates = [full_base_gates[p].clone().relabel_(["_o0", "_o1", lbls[p][1], lbls[p + 1][1]]) + for p in range(num_sites - 1)] + M, L0, R0 = _build_mpo(TFIM_J, TFIM_HX_FINAL, device) + + even_bonds = list(range(0, num_sites - 1, 2)) + odd_bonds = list(range(1, num_sites - 1, 2)) + + def apply_gate(p, gates): + psi = cytnx.Contract(A[p], A[p + 1]) + psi = cytnx.Contract(psi, gates[p]) + psi.permute_([lbls[p][0], "_o0", "_o1", lbls[p + 1][2]]) + psi.relabel_([lbls[p][0], lbls[p][1], lbls[p + 1][1], lbls[p + 1][2]]) + psi.set_rowrank_(2) + dim_l = A[p].shape()[0] + dim_r = A[p + 1].shape()[2] + new_dim = min(dim_l * d, dim_r * d, bond_dim) + s, A[p], A[p + 1] = cytnx.linalg.Svd_truncate(psi, new_dim, err=SVD_CUTOFF) + return s + + def apply_group(bonds, gates): + # Bonds in the same parity group never share a site, so they can be + # applied in any order with no extra Trotter error; the post-SVD + # singular-value tensor is absorbed into the right neighbor purely + # as a fixed convention (no later gate in this group touches either + # neighbor, so the choice doesn't affect the resulting state). + for p in bonds: + s = apply_gate(p, gates) + s = s / s.Norm().item() + A[p + 1] = cytnx.Contract(s, A[p + 1]) + A[p].set_name(f"A{p}").relabel_(lbls[p]) + A[p + 1].set_name(f"A{p+1}").relabel_(lbls[p + 1]) + + def trotter_step(): + apply_group(even_bonds, half_gates) + apply_group(odd_bonds, full_gates) + apply_group(even_bonds, half_gates) + + for _ in range(TFIM_N_STEPS): + trotter_step() + return _energy(A, M, L0, R0, device) + + +@pytest.mark.timeout(GRID_POINT_TIMEOUT_SEC) +@pytest.mark.parametrize("num_sites", NUM_SITES_VALUES) +@pytest.mark.parametrize("bond_dim", BOND_DIM_VALUES) +def test_tebd_benchmark(benchmark, bond_dim, num_sites): + energy = benchmark.pedantic(run_one, args=(bond_dim, num_sites), rounds=1, iterations=1) + benchmark.extra_info["energy"] = energy + assert energy == pytest.approx(REFERENCE_ENERGIES[(bond_dim, num_sites)], rel=1e-6) + + +@pytest.mark.cytnx_memory +@pytest.mark.limit_memory("20 MB") +@pytest.mark.parametrize("num_sites", NUM_SITES_VALUES) +@pytest.mark.parametrize("bond_dim", BOND_DIM_VALUES) +def test_tebd_memory(bond_dim, num_sites): + energy = run_one(bond_dim, num_sites) + assert energy == pytest.approx(REFERENCE_ENERGIES[(bond_dim, num_sites)], rel=1e-6) diff --git a/benchmarks/cross_library/cytnx_bench/test_variational_manual_grad.py b/benchmarks/cross_library/cytnx_bench/test_variational_manual_grad.py new file mode 100644 index 000000000..76f32826f --- /dev/null +++ b/benchmarks/cross_library/cytnx_bench/test_variational_manual_grad.py @@ -0,0 +1,275 @@ +"""Cytnx benchmark, algorithm class 4: variational ground-state search by +gradient descent on the MPS tensors of the 1D Heisenberg chain, with a +hand-derived (manual) gradient instead of automatic differentiation. + +Cytnx has no autodiff backend, so the gradient of the Rayleigh quotient + + E(psi) = / + +with respect to every MPS tensor simultaneously is computed analytically +rather than via backprop. This is Cytnx's counterpart to quimb's +`run_one_jax`/`run_one_torch` (`quimb_bench/test_variational_ad.py`), which +differentiate straight through the same whole-network contraction with the +backend's own autodiff. Like quimb's benchmark, every MPS tensor is updated +from a single gradient step taken simultaneously, with no orthogonality +center and no per-step canonicalization -- only a single global rescale +derived from the new `` is applied after each step, distributed +evenly across all num_sites tensors. Because the surrounding tensors are not kept +isometric, the gradient needs both an H-environment term (the effective +one-site Hamiltonian, as in a one-site sweep) and a norm-environment term: + + dE/dA_i* = (2 / den) * (H_eff,i(A_i) - E * N_eff,i(A_i)) + +where `den = `, `H_eff,i` contracts the MPO with the left/right +H-environments around site i, and `N_eff,i` is the analogous contraction +with trivial (no-MPO) norm-environments. `N_eff,i` only collapses to `A_i` +when the rest of the chain is isometric (the one-site-sweep case); here it +does not, so it is computed explicitly via the `_Networks.LN_UPDATE_NET`/ +`RN_UPDATE_NET`/`N_EFF_NET` Cytnx `Network` topologies below, which are +the same `L_UPDATE_NET`/`R_UPDATE_NET`/`HEFF_NET` contractions with the +MPO leg dropped. All four environment sets (H-left, H-right, norm-left, +norm-right) are rebuilt from scratch every gradient step, since every +tensor changes at once and no sweep-order incremental reuse is possible. + +CPU and GPU code paths are both written; the GPU path moves every MPS/MPO +UniTensor to `cytnx.Device.cuda` before the gradient step. It cannot be +exercised in this environment (no GPU). + +Run timing with `pytest --benchmark-only test_variational_manual_grad.py`, +memory with `pytest --memray test_variational_manual_grad.py`. The initial +MPS is seeded (per-site `cytnx.UniTensor.normal(..., seed=k)`). Unlike a +one-site sweep -- a strong local optimizer whose converged energy is +largely insensitive to small initial-state differences -- this +whole-network update is a weaker optimizer whose converged energy is +sensitive to the initial state, so (as with quimb's AD benchmark, which +uses the same `LEARNING_RATE`/`_n_grad_steps`) each test asserts only +against its own `REFERENCE_ENERGIES` at a tight `rel=1e-6` +self-consistency tolerance -- this run is fully deterministic given the +seeded initial state, so no cross-library energy comparison is expected +to land this close. +""" +import pytest + +import cytnx + +from common.model import BOND_DIM_VALUES, HEISENBERG_J, NUM_SITES_VALUES, GRID_POINT_TIMEOUT_SEC + +DEVICE = "cpu" # set to "gpu" to exercise the (untested) GPU code path below +LEARNING_RATE = 0.5 + + +def _n_grad_steps(num_sites): + return 8 * num_sites + + +REFERENCE_ENERGIES = { + (16, 20): -8.669390762907128, + (16, 30): -13.093516691905897, + (16, 50): -21.945283123631583, + (32, 20): -8.674290703078219, + (32, 30): -13.100972670479699, + (32, 50): -21.953298655520914, + (64, 20): -8.67676882878297, + (64, 30): -13.103318241859627, + (64, 50): -21.96029731572234, +} + +def _build_mpo(J, device): + d = 2 + D = 5 + Sp = cytnx.zeros([d, d], device=device) + Sp[0, 1] = 1.0 # S+: |down> -> |up> + Sm = cytnx.zeros([d, d], device=device) + Sm[1, 0] = 1.0 # S-: |up> -> |down> + Sz = cytnx.zeros([d, d], device=device) + Sz[0, 0] = 0.5 + Sz[1, 1] = -0.5 + eye = cytnx.eye(d, device=device) + + M = cytnx.zeros([D, D, d, d], device=device) + M[0, 0] = eye + M[D - 1, D - 1] = eye + M[0, 1] = Sp + M[0, 2] = Sm + M[0, 3] = Sz + M[1, D - 1] = (J / 2.0) * Sm + M[2, D - 1] = (J / 2.0) * Sp + M[3, D - 1] = J * Sz + M = cytnx.UniTensor(M, 0).set_name("MPO") + + L0 = cytnx.UniTensor.zeros([D, 1, 1], device=device).set_rowrank_(0).set_name("L0") + R0 = cytnx.UniTensor.zeros([D, 1, 1], device=device).set_rowrank_(0).set_name("R0") + L0[0, 0, 0] = 1.0 + R0[D - 1, 0, 0] = 1.0 + return M, L0, R0 + + +def _build_mps(num_sites, bond_dim, device): + d = 2 + A = [None] * num_sites + lbls = [[str(2 * k), str(2 * k + 1), str(2 * k + 2)] for k in range(num_sites)] + A[0] = cytnx.UniTensor.normal([1, d, min(bond_dim, d)], 0., 1., seed=0, device=device).set_rowrank_(2) + A[0].relabel_(lbls[0]).set_name("A0") + for k in range(1, num_sites): + dim1 = A[k - 1].shape()[2] + dim3 = min(min(bond_dim, dim1 * d), d ** (num_sites - k - 1)) + A[k] = cytnx.UniTensor.normal([dim1, d, dim3], 0., 1., seed=k, device=device).set_rowrank_(2) + A[k].relabel_(lbls[k]).set_name(f"A{k}") + _canonicalize_right(A, lbls, num_sites) + return A, lbls + + +def _canonicalize_right(A, lbls, num_sites): + for p in range(num_sites - 1, 0, -1): + A[p].set_rowrank_(1) + s, u, vt = cytnx.linalg.Gesvd(A[p]) + A[p] = vt + A[p - 1] = cytnx.Contract(A[p - 1], cytnx.Contract(u, s)) + A[p].relabel_(lbls[p]).set_name(f"A{p}") + A[p - 1].relabel_(lbls[p - 1]).set_name(f"A{p-1}") + A[0] = A[0] / A[0].Norm().item() + A[0].relabel_(lbls[0]).set_name("A0") + + +class _Networks: + """Caches one parsed `cytnx.Network` per static contraction topology used + in `run_one`'s gradient step, since `FromString` re-parses the topology + string on every call -- with `8 * num_sites` whole-network gradient steps + and O(num_sites) such contractions per step, re-parsing on every call + dominates the runtime at large num_sites. Each `Network` is built once and + refilled via + `PutUniTensors`/`Launch` on every reuse.""" + + L_UPDATE_NET = ["L: -2,-1,-3", "A: -1,-4,1", "M: -2,0,-4,-5", "A_Conj: -3,-5,2", "TOUT: 0,1,2"] + R_UPDATE_NET = ["R: -2,-1,-3", "B: 1,-4,-1", "M: 0,-2,-4,-5", "B_Conj: 2,-5,-3", "TOUT: 0,1,2"] + HEFF_NET = ["psi: -1,-2,-3", "L: -4,-1,0", "R: -5,-3,2", "M: -4,-5,-2,1", "TOUT: 0,1;2"] + + LN_UPDATE_NET = ["L: -1,-2", "A: -1,-3,0", "A_Conj: -2,-3,1", "TOUT: 0,1"] + RN_UPDATE_NET = ["R: -1,-2", "B: 0,-3,-1", "B_Conj: 1,-3,-2", "TOUT: 0,1"] + N_EFF_NET = ["L: -1,0", "psi: -1,1,-2", "R: -2,2", "TOUT: 0,1,2"] + + def __init__(self): + self.l_update = cytnx.Network() + self.l_update.FromString(self.L_UPDATE_NET) + self.r_update = cytnx.Network() + self.r_update.FromString(self.R_UPDATE_NET) + self.heff = cytnx.Network() + self.heff.FromString(self.HEFF_NET) + self.ln_update = cytnx.Network() + self.ln_update.FromString(self.LN_UPDATE_NET) + self.rn_update = cytnx.Network() + self.rn_update.FromString(self.RN_UPDATE_NET) + self.n_eff_net = cytnx.Network() + self.n_eff_net.FromString(self.N_EFF_NET) + + def update_L(self, L_env, A_new, A_conj, M): + self.l_update.PutUniTensors(["L", "A", "A_Conj", "M"], [L_env, A_new, A_conj, M]) + return self.l_update.Launch() + + def update_R(self, R_env, B_new, B_conj, M): + self.r_update.PutUniTensors(["R", "B", "M", "B_Conj"], [R_env, B_new, M, B_conj]) + return self.r_update.Launch() + + def h_eff(self, theta, L_env, R_env, M): + self.heff.PutUniTensors(["psi", "L", "R", "M"], [theta, L_env, R_env, M]) + out = self.heff.Launch() + out.relabel_(theta.labels()) + return out + + def update_L_N(self, LN_env, A_new, A_conj): + self.ln_update.PutUniTensors(["L", "A", "A_Conj"], [LN_env, A_new, A_conj]) + return self.ln_update.Launch() + + def update_R_N(self, RN_env, B_new, B_conj): + self.rn_update.PutUniTensors(["R", "B", "B_Conj"], [RN_env, B_new, B_conj]) + return self.rn_update.Launch() + + def n_eff(self, theta, LN_env, RN_env): + self.n_eff_net.PutUniTensors(["L", "psi", "R"], [LN_env, theta, RN_env]) + out = self.n_eff_net.Launch() + out.relabel_(theta.labels()) + return out + + +def run_one(bond_dim, num_sites): + device = cytnx.Device.cuda if DEVICE == "gpu" else cytnx.Device.cpu + M, L0, R0 = _build_mpo(HEISENBERG_J, device) + A, lbls = _build_mps(num_sites, bond_dim, device) + LN0 = cytnx.UniTensor.zeros([1, 1], device=device).set_rowrank_(0) + LN0[0, 0] = 1.0 + RN0 = cytnx.UniTensor.zeros([1, 1], device=device).set_rowrank_(0) + RN0[0, 0] = 1.0 + nets = _Networks() + + def grad_step(A, A_conj): + LH = [None] * (num_sites + 1) + LH[0] = L0 + for p in range(num_sites): + LH[p + 1] = nets.update_L(LH[p], A[p], A_conj[p], M) + RH = [None] * (num_sites + 1) + RH[num_sites] = R0 + for p in range(num_sites - 1, -1, -1): + RH[p] = nets.update_R(RH[p + 1], A[p], A_conj[p], M) + LN = [None] * (num_sites + 1) + LN[0] = LN0 + for p in range(num_sites): + LN[p + 1] = nets.update_L_N(LN[p], A[p], A_conj[p]) + RN = [None] * (num_sites + 1) + RN[num_sites] = RN0 + for p in range(num_sites - 1, -1, -1): + RN[p] = nets.update_R_N(RN[p + 1], A[p], A_conj[p]) + + D = LH[num_sites].shape()[0] + num = LH[num_sites][D - 1, 0, 0].item() + den = LN[num_sites][0, 0].item() + energy = num / den + + new_A = [None] * num_sites + for p in range(num_sites): + ht = nets.h_eff(A[p], LH[p], RH[p + 1], M) + nt = nets.n_eff(A[p], LN[p], RN[p + 1]) + grad = (2.0 / den) * (ht - energy * nt) + new_A[p] = A[p] - LEARNING_RATE * grad + new_A[p].relabel_(lbls[p]).set_name(f"A{p}") + + new_A_conj = [None] * num_sites + LNn = LN0 + for p in range(num_sites): + new_A_conj[p] = new_A[p].Dagger().permute_(new_A[p].labels()) + LNn = nets.update_L_N(LNn, new_A[p], new_A_conj[p]) + den_new = LNn[0, 0].item() + scale = den_new ** (-1.0 / (2 * num_sites)) + # scale is real and positive, so Dagger(scale * a) == scale * Dagger(a) + # exactly; rescaling new_A_conj here instead of recomputing it from + # scratch at the top of the next grad_step call avoids a redundant + # Dagger()/permute_() pass over every tensor on every gradient step. + new_A = [a * scale for a in new_A] + new_A_conj = [ac * scale for ac in new_A_conj] + for p in range(num_sites): + new_A[p].relabel_(lbls[p]).set_name(f"A{p}") + new_A_conj[p].relabel_(lbls[p]) + return new_A, new_A_conj, energy + + energy = None + A_conj = [a.Dagger().permute_(a.labels()) for a in A] + for _ in range(_n_grad_steps(num_sites)): + A, A_conj, energy = grad_step(A, A_conj) + return energy + + +@pytest.mark.timeout(GRID_POINT_TIMEOUT_SEC) +@pytest.mark.parametrize("num_sites", NUM_SITES_VALUES) +@pytest.mark.parametrize("bond_dim", BOND_DIM_VALUES) +def test_variational_manual_grad_benchmark(benchmark, bond_dim, num_sites): + energy = benchmark.pedantic(run_one, args=(bond_dim, num_sites), rounds=1, iterations=1) + benchmark.extra_info["energy"] = energy + assert energy == pytest.approx(REFERENCE_ENERGIES[(bond_dim, num_sites)], rel=1e-6) + + +@pytest.mark.cytnx_memory +@pytest.mark.limit_memory("40 MB") +@pytest.mark.parametrize("num_sites", NUM_SITES_VALUES) +@pytest.mark.parametrize("bond_dim", BOND_DIM_VALUES) +def test_variational_manual_grad_memory(bond_dim, num_sites): + energy = run_one(bond_dim, num_sites) + assert energy == pytest.approx(REFERENCE_ENERGIES[(bond_dim, num_sites)], rel=1e-6) diff --git a/benchmarks/cross_library/quimb_bench/__init__.py b/benchmarks/cross_library/quimb_bench/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/benchmarks/cross_library/quimb_bench/test_dmrg.py b/benchmarks/cross_library/quimb_bench/test_dmrg.py new file mode 100644 index 000000000..6e1bbf95b --- /dev/null +++ b/benchmarks/cross_library/quimb_bench/test_dmrg.py @@ -0,0 +1,239 @@ +"""quimb benchmark, algorithm class 1: finite two-site DMRG, dense mode +(no conserved quantum numbers) on the 1D spin-1/2 Heisenberg chain, using +quimb's built-in `DMRG2` engine, and a symmetric variant exercising the +`symmray` block-sparse backend. + +CPU and GPU code paths are both written; the GPU path moves every MPO/MPS +tensor to the device array backend before the sweep (mirrors quimb's +documented `to_backend`/`apply_to_arrays` pattern with cupy or torch). It +cannot be exercised in this environment (no GPU). + +quimb (1.14.0) does not yet ship a packaged finite-chain DMRG engine that +runs on `symmray`-backed (block-sparse, abelian-symmetric) MPS the way +`DMRG2` does for plain dense arrays -- there is no symmetric counterpart to +`MPO_ham_heis`/`DMRG2` in the public API at the time of writing. To still +exercise the same block-sparse computational kernels a symmetric DMRG +would use per sweep (two-site gate contraction immediately followed by an +SVD truncation back to bond dimension chi, repeated over every bond), +`run_one_symmetric` instead drives the same random U(1)-symmetric MPS to +the Heisenberg ground state via imaginary-time evolution (ITE) with the +two-site gate exp(-dt*h_{i,i+1}): repeated application of exp(-dt*H) to a +state with nonzero ground-state overlap converges to the ground state of +the targeted symmetry sector as the total imaginary time grows, exactly +like the dense/symmetric DMRG benchmarks converge to the same energy via +their own (different) sweep algorithm. `ITE_DT_SCHEDULE` anneals dt from +0.3 down to 0.02 across zigzag (forward then backward) sweeps over every +bond, renormalizing once per full zigzag pass; this uses the same +"contract + truncate" cost structure as a real two-site DMRG/TEBD sweep +(large-bond_dim/large-num_sites dominated by the same O(chi^3) SVD and +O(chi^2 * d^2) gate contraction) so the time/memory metrics remain +comparable, while the reported energy now converges to the same Heisenberg +ground energy as +`cytnx_bench/test_dmrg_symmetric.py`/`tenpy_bench/test_dmrg_symmetric.py` +at every (bond_dim, num_sites) grid point (matched to within the `rel=1e-2` tolerance +used for `SYMMETRIC_REFERENCE_ENERGIES`, looser than the dense/symmetric +DMRG benchmarks' `rel=1e-6` since ITE over a finite total imaginary time +is itself an approximation to the true ground state, not just a +discretization of an already-exact sweep). + +Run timing with `pytest --benchmark-only test_dmrg.py`, memory with +`pytest --memray test_dmrg.py`. +""" +import numpy as np +import pytest +import symmray as sr +from scipy.linalg import expm + +import quimb.tensor as qtn + +from common.model import BOND_DIM_VALUES, HEISENBERG_J, NUM_SITES_VALUES, N_SWEEPS, GRID_POINT_TIMEOUT_SEC + +DEVICE = "cpu" # set to "gpu" to exercise the (untested) GPU code path below + +# Symmray block-sparse arrays are built on top of a plain array per +# charge-block, so this selects which array library backs every block of +# `psi` (and the gate/operator it's contracted against) in +# `run_one_symmetric`. "numpy" is fastest across the full (bond_dim, +# num_sites) grid -- "torch" lands the same energies but runs consistently +# ~1.5-2x slower at every grid point, since this workload is dominated by +# thousands of small per-bond gate/SVD calls where torch's dispatch +# overhead outweighs any kernel speedup. "jax" is not a viable alternative +# at all: its per-call tracing overhead is so large that a single grid +# point did not finish in several minutes, vs. seconds for numpy/torch. +# "cupy" moves every block to the GPU (mirrors quimb's usual +# `psi.apply_to_arrays(lambda x: cp.asarray(x))` pattern) but cannot be +# exercised in this environment (no GPU). +ARRAY_BACKEND = "numpy" + +PHYS_CHARGE_MAP = {1: 1, -1: 1} # spin-up -> charge +1, spin-down -> charge -1 + +DENSE_REFERENCE_ENERGIES = { + (16, 20): -8.682468456356828, + (16, 30): -13.111313454922634, + (16, 50): -21.97181361569925, + (32, 20): -8.682473319689738, + (32, 30): -13.111355524192675, + (32, 50): -21.972106512033726, + (64, 20): -8.682473334397873, + (64, 30): -13.11135575848871, + (64, 50): -21.972110271889434, +} +SYMMETRIC_REFERENCE_ENERGIES = { + (16, 20): -8.675022726955829, + (16, 30): -13.058213723292113, + (16, 50): -21.703613228188505, + (32, 20): -8.682060417783728, + (32, 30): -13.096659398821432, + (32, 50): -21.86788053095872, + (64, 20): -8.682293734678527, + (64, 30): -13.105181026506637, + (64, 50): -21.925672520032034, +} + +# Anneals dt across zigzag sweeps -- see module docstring. +ITE_DT_SCHEDULE = [0.3] * 15 + [0.2] * 15 + [0.1] * 20 + [0.05] * 30 + [0.02] * 80 + + +def build_dense(bond_dim, num_sites): + H = qtn.MPO_ham_heis(num_sites, j=HEISENBERG_J, cyclic=False) + if DEVICE == "gpu": + import torch + H.apply_to_arrays(lambda x: torch.as_tensor(x, device="cuda")) + dmrg = qtn.DMRG2(H, bond_dims=[bond_dim], cutoffs=1e-10) + return dmrg + + +def run_one_dense(bond_dim, num_sites): + dmrg = build_dense(bond_dim, num_sites) + # tol=0.0 disables early stopping on energy convergence, so this always + # runs the full N_SWEEPS budget -- matching Cytnx/TeNPy's fixed sweep + # count instead of exiting early when quimb's own convergence check is + # satisfied. + dmrg.solve(tol=0.0, max_sweeps=N_SWEEPS, verbosity=0) + return dmrg.energy + + +def _convert_backend(arr): + """Move a U1Array's underlying blocks to ARRAY_BACKEND's array type.""" + if ARRAY_BACKEND == "numpy": + return arr + if ARRAY_BACKEND == "cupy": + import cupy as cp + arr.apply_to_arrays(lambda x: cp.asarray(x)) + elif ARRAY_BACKEND == "torch": + import torch + arr.apply_to_arrays(lambda x: torch.as_tensor(x)) + elif ARRAY_BACKEND == "jax": + import jax.numpy as jnp + arr.apply_to_arrays(lambda x: jnp.asarray(x)) + return arr + + +def _heisenberg_two_site_gate(dt): + """Charge-conserving exp(-dt * J * S_i.S_j) two-site gate as a U1Array. + + In the {up(+1), down(-1)} Sz basis, J*S_i.S_j is exactly block-diagonal + in total Sz: the (up,up) and (down,down) sectors (total charge +-2) are + 1x1 blocks, and the (up,down)/(down,up) sector (total charge 0) is a 2x2 + block. `expm` of this dense 4x4 matrix preserves that same block + structure, so converting it through `U1Array.from_dense` recovers a + genuinely sparse (6 nonzero blocks out of a dense 16-element 2x2x2x2 + tensor) charge-conserving gate. + """ + phys = [1, -1] + h_dense = (HEISENBERG_J / 4.0) * np.array([ + [1, 0, 0, 0], + [0, -1, 2, 0], + [0, 2, -1, 0], + [0, 0, 0, 1], + ], dtype=float) + gate_dense = expm(-dt * h_dense).reshape(2, 2, 2, 2) + gate = sr.U1Array.from_dense( + gate_dense, index_maps=[phys, phys, phys, phys], + duals=[False, False, True, True], + ) + return _convert_backend(gate) + + +def _heisenberg_two_site_op(): + """Plain (non-exponentiated) two-site Heisenberg coupling, for computing + via `local_expectation` -- not used to evolve `psi`.""" + phys = [1, -1] + h_dense = (HEISENBERG_J / 4.0) * np.array([ + [1, 0, 0, 0], + [0, -1, 2, 0], + [0, 2, -1, 0], + [0, 0, 0, 1], + ], dtype=float) + op = sr.U1Array.from_dense( + h_dense.reshape(2, 2, 2, 2), index_maps=[phys, phys, phys, phys], + duals=[False, False, True, True], + ) + return _convert_backend(op) + + +def run_one_symmetric(bond_dim, num_sites): + # Alternate site charges so the half-filled (Neel-like) total-Sz=0 + # sector is reachable at the bond dimensions in our sweep grid. + psi = sr.MPS_abelian_rand( + "U1", L=num_sites, bond_dim=bond_dim, phys_dim=PHYS_CHARGE_MAP, seed=0, + site_charge=lambda i: 1 if i % 2 == 0 else -1, + ) + _convert_backend(psi) + + gates = {dt: _heisenberg_two_site_gate(dt) for dt in set(ITE_DT_SCHEDULE)} + + def zigzag_pass(dt): + gate = gates[dt] + for i in range(num_sites - 1): + psi.gate_split_(gate, where=(i, i + 1), max_bond=bond_dim, cutoff=1e-10) + for i in range(num_sites - 2, -1, -1): + psi.gate_split_(gate, where=(i, i + 1), max_bond=bond_dim, cutoff=1e-10) + psi.normalize() + + for dt in ITE_DT_SCHEDULE: + zigzag_pass(dt) + h_op = _heisenberg_two_site_op() + # compute_local_expectation(method="envs") builds shared left/right + # environments once and reuses them across every bond, unlike summing + # num_sites-1 separate local_expectation_exact calls that each redo the + # environment contraction from scratch. + terms = {(i, i + 1): h_op for i in range(num_sites - 1)} + energy = psi.compute_local_expectation(terms, normalized=True, method="envs") + return energy + + +@pytest.mark.timeout(GRID_POINT_TIMEOUT_SEC) +@pytest.mark.parametrize("num_sites", NUM_SITES_VALUES) +@pytest.mark.parametrize("bond_dim", BOND_DIM_VALUES) +def test_dmrg_dense_benchmark(benchmark, bond_dim, num_sites): + energy = benchmark.pedantic(run_one_dense, args=(bond_dim, num_sites), rounds=1, iterations=1) + benchmark.extra_info["energy"] = float(energy) + assert float(energy) == pytest.approx(DENSE_REFERENCE_ENERGIES[(bond_dim, num_sites)], rel=1e-6) + + +@pytest.mark.cytnx_memory +@pytest.mark.limit_memory("130 MB") +@pytest.mark.parametrize("num_sites", NUM_SITES_VALUES) +@pytest.mark.parametrize("bond_dim", BOND_DIM_VALUES) +def test_dmrg_dense_memory(bond_dim, num_sites): + energy = run_one_dense(bond_dim, num_sites) + assert float(energy) == pytest.approx(DENSE_REFERENCE_ENERGIES[(bond_dim, num_sites)], rel=1e-6) + + +@pytest.mark.timeout(GRID_POINT_TIMEOUT_SEC) +@pytest.mark.parametrize("num_sites", NUM_SITES_VALUES) +@pytest.mark.parametrize("bond_dim", BOND_DIM_VALUES) +def test_dmrg_symmetric_benchmark(benchmark, bond_dim, num_sites): + energy = benchmark.pedantic(run_one_symmetric, args=(bond_dim, num_sites), rounds=1, iterations=1) + benchmark.extra_info["energy"] = float(energy) + assert float(energy) == pytest.approx(SYMMETRIC_REFERENCE_ENERGIES[(bond_dim, num_sites)], rel=2e-2) + + +@pytest.mark.cytnx_memory +@pytest.mark.limit_memory("700 MB") +@pytest.mark.parametrize("num_sites", NUM_SITES_VALUES) +@pytest.mark.parametrize("bond_dim", BOND_DIM_VALUES) +def test_dmrg_symmetric_memory(bond_dim, num_sites): + energy = run_one_symmetric(bond_dim, num_sites) + assert float(energy) == pytest.approx(SYMMETRIC_REFERENCE_ENERGIES[(bond_dim, num_sites)], rel=2e-2) diff --git a/benchmarks/cross_library/quimb_bench/test_tebd.py b/benchmarks/cross_library/quimb_bench/test_tebd.py new file mode 100644 index 000000000..c62e26304 --- /dev/null +++ b/benchmarks/cross_library/quimb_bench/test_tebd.py @@ -0,0 +1,82 @@ +"""quimb benchmark, algorithm class 2: real-time evolution of a 1D +transverse-field Ising chain after a field quench, using quimb's built-in +`TEBD` engine. + +`qtn.ham_1d_ising`/`MPO_ham_ising` build H = j*sum(S^Z_i S^Z_{i+1}) - +bx*sum(S^X_i) using spin-1/2 operators (eigenvalues +-1/2), not Pauli +matrices. To reproduce the same physical Hamiltonian as +`cytnx_bench/test_tebd.py` (H = -hx*sum(PauliX_i) - J*sum(PauliZ_i +PauliZ_{i+1}), Pauli-normalized, eigenvalues +-1), substitute +S^Z = PauliZ/2 and S^X = PauliX/2 and solve for quimb's (j, bx) in terms +of cytnx's (J, hx): j = -4*J, bx = 2*hx. + +CPU and GPU code paths are both written; the GPU path moves the initial +MPS to the device array backend before stepping. It cannot be exercised +in this environment (no GPU). + +Run timing with `pytest --benchmark-only test_tebd.py`, memory with +`pytest --memray test_tebd.py`. +""" +import pytest + +import quimb.tensor as qtn + +from common.model import BOND_DIM_VALUES, NUM_SITES_VALUES, GRID_POINT_TIMEOUT_SEC, TFIM_DT, TFIM_HX_FINAL, TFIM_J, TFIM_N_STEPS + +DEVICE = "cpu" # set to "gpu" to exercise the (untested) GPU code path below + +# Convert cytnx_bench's Pauli-normalized (J, hx) into quimb's spin-1/2 +# convention -- see module docstring. +ISING_J = -4 * TFIM_J +ISING_BX = 2 * TFIM_HX_FINAL + +REFERENCE_ENERGIES = { + (16, 20): -19.000143682959134, + (16, 30): -29.00015635325205, + (16, 50): -49.0001810814783, + (32, 20): -19.000143682959134, + (32, 30): -29.00015635325205, + (32, 50): -49.0001810814783, + (64, 20): -19.000143682959134, + (64, 30): -29.00015635325205, + (64, 50): -49.0001810814783, +} + + +def build(bond_dim, num_sites): + H = qtn.ham_1d_ising(num_sites, j=ISING_J, bx=ISING_BX, cyclic=False) + psi0 = qtn.MPS_computational_state("0" * num_sites) + if DEVICE == "gpu": + import torch + psi0.apply_to_arrays(lambda x: torch.as_tensor(x, device="cuda")) + tebd = qtn.TEBD(psi0, H, dt=TFIM_DT, progbar=False) + tebd.split_opts["cutoff"] = 1e-10 + tebd.split_opts["max_bond"] = bond_dim + return tebd + + +def run_one(bond_dim, num_sites): + tebd = build(bond_dim, num_sites) + for _ in range(TFIM_N_STEPS): + tebd.step(order=2, dt=TFIM_DT) + H_mpo = qtn.MPO_ham_ising(num_sites, j=ISING_J, bx=ISING_BX, cyclic=False) + energy = tebd.pt.H @ (H_mpo.apply(tebd.pt)) + return float(energy.real) + + +@pytest.mark.timeout(GRID_POINT_TIMEOUT_SEC) +@pytest.mark.parametrize("num_sites", NUM_SITES_VALUES) +@pytest.mark.parametrize("bond_dim", BOND_DIM_VALUES) +def test_tebd_benchmark(benchmark, bond_dim, num_sites): + energy = benchmark.pedantic(run_one, args=(bond_dim, num_sites), rounds=1, iterations=1) + benchmark.extra_info["energy"] = energy + assert energy == pytest.approx(REFERENCE_ENERGIES[(bond_dim, num_sites)], rel=1e-6) + + +@pytest.mark.cytnx_memory +@pytest.mark.limit_memory("100 MB") +@pytest.mark.parametrize("num_sites", NUM_SITES_VALUES) +@pytest.mark.parametrize("bond_dim", BOND_DIM_VALUES) +def test_tebd_memory(bond_dim, num_sites): + energy = run_one(bond_dim, num_sites) + assert energy == pytest.approx(REFERENCE_ENERGIES[(bond_dim, num_sites)], rel=1e-6) diff --git a/benchmarks/cross_library/quimb_bench/test_variational_ad.py b/benchmarks/cross_library/quimb_bench/test_variational_ad.py new file mode 100644 index 000000000..2ecbee557 --- /dev/null +++ b/benchmarks/cross_library/quimb_bench/test_variational_ad.py @@ -0,0 +1,224 @@ +"""quimb benchmark, algorithm class 4: variational ground-state search by +gradient descent on the MPS tensors of the 1D Heisenberg chain, using real +automatic differentiation of the Rayleigh quotient + + E(psi) = / + +with respect to every MPS tensor simultaneously: one gradient step is +taken on every MPS tensor at once (no orthogonality center, no per-step +canonicalization), followed by a single global rescale derived from the +new `` and distributed evenly across all num_sites tensors. +`run_one_jax`/`run_one_torch` exercise quimb's AD-based optimization on the +JAX and PyTorch array backends respectively. + +TeNPy (`tenpy_bench/test_variational_manual_grad.py`) and Cytnx +(`cytnx_bench/test_variational_manual_grad.py`) run this same whole-network +update; having no autodiff backend, they evaluate the closed-form gradient + + dE/dA_i* = (2 / den) * (H_eff,i(A_i) - E * N_eff,i(A_i)) + +by hand instead, where `den = ` and `N_eff,i` is the +no-MPO analogue of `H_eff,i` (needed because the tensors away from site i +are not isometric under this update, unlike in a one-site sweep). quimb's +MPS/MPO tensors are plain JAX/PyTorch arrays under the hood (via autoray +dispatch), so here we let the backend's own autodiff differentiate straight +through the full `` and `` contractions instead of +deriving that gradient by hand. + +This whole-network update moves the state far less per iteration than a +one-site sweep does, so it needs both a larger learning rate and many more +iterations to reach a comparable energy neighborhood; `LEARNING_RATE` and +the local `_n_grad_steps(num_sites)` helper (scaling with `num_sites`, +since a longer chain needs proportionally more whole-state updates to +converge as far) are shared with the TeNPy/Cytnx benchmarks. Each backend +is checked only against its own `JAX_REFERENCE_ENERGIES`/ +`TORCH_REFERENCE_ENERGIES` at a tight `rel=1e-6` self-consistency +tolerance, the same as the TeNPy/Cytnx manual-gradient benchmarks -- this +run is deterministic given the seeded initial state (`seed=0`), so no +cross-backend or cross-library energy comparison is expected to land this +close; `JAX_REFERENCE_ENERGIES` and `TORCH_REFERENCE_ENERGIES` themselves +differ from each other by far more than `1e-6` despite sharing the seed +and formula. + +GPU code paths are written for both backends (`device="cuda"` placement) +but cannot be exercised in this environment (no GPU). + +Run timing with `pytest --benchmark-only test_variational_ad.py`, memory +with `pytest --memray test_variational_ad.py`. The MPS here is seeded +(`MPS_rand_state(..., seed=0)`). +""" +import pytest + +import quimb.tensor as qtn + +from common.model import BOND_DIM_VALUES, HEISENBERG_J, NUM_SITES_VALUES, GRID_POINT_TIMEOUT_SEC + +LEARNING_RATE = 0.5 +DEVICE = "cpu" # set to "gpu" to exercise the (untested) GPU code paths below + +JAX_REFERENCE_ENERGIES = { + (16, 20): -8.67426586151123, + (16, 30): -13.085174560546875, + (16, 50): -21.572589874267578, + (32, 20): -8.659192085266113, + (32, 30): -13.020613670349121, + (32, 50): -21.64177703857422, + (64, 20): -8.671019554138184, + (64, 30): -13.056256294250488, + (64, 50): -21.65955352783203, +} +TORCH_REFERENCE_ENERGIES = { + (16, 20): -8.674265280037904, + (16, 30): -13.085195694053375, + (16, 50): -21.606009355825424, + (32, 20): -8.659190668103744, + (32, 30): -13.020607976110528, + (32, 50): -21.639568826632235, + (64, 20): -8.671020853476218, + (64, 30): -13.056264576677071, + (64, 50): -21.589291140405937, +} + + +def _build(bond_dim, num_sites): + psi = qtn.MPS_rand_state(num_sites, bond_dim=bond_dim, dtype="float64", seed=0) + H = qtn.MPO_ham_heis(num_sites, j=HEISENBERG_J, cyclic=False) + return psi, H + + +def _n_grad_steps(num_sites): + return 8 * num_sites + + +def run_one_jax(bond_dim, num_sites): + import jax + import jax.numpy as jnp + + psi, H = _build(bond_dim, num_sites) + if DEVICE == "gpu": + device = jax.devices("gpu")[0] + else: + device = jax.devices("cpu")[0] + arrays = tuple(jax.device_put(jnp.asarray(a), device) for a in psi.arrays) + H.apply_to_arrays(lambda x: jax.device_put(jnp.asarray(x), device)) + + def energy(arrays): + p = psi.copy() + for i, a in enumerate(arrays): + p[i].modify(data=a) + num = p.H @ (H.apply(p)) + den = p.H @ p + return jnp.real(num / den) + + def norm_sq(arrays): + p = psi.copy() + for i, a in enumerate(arrays): + p[i].modify(data=a) + return jnp.real(p.H @ p) + + grad_fn = jax.jit(jax.grad(energy)) if DEVICE == "cpu" else jax.grad(energy) + norm_sq_fn = jax.jit(norm_sq) if DEVICE == "cpu" else norm_sq + + def grad_step(arrays): + g = grad_fn(arrays) + new_arrays = [a - LEARNING_RATE * ga for a, ga in zip(arrays, g)] + # Rescale the whole state by a single global factor derived from + # , distributed evenly across all num_sites tensors, rather than + # normalizing each tensor independently -- the MPS is not in + # canonical form here, so per-tensor normalization does not keep + # the contracted close to 1. + scale = norm_sq_fn(tuple(new_arrays)) ** (-1.0 / (2 * len(new_arrays))) + new_arrays = [a * scale for a in new_arrays] + return tuple(new_arrays) + + for _ in range(_n_grad_steps(num_sites)): + arrays = grad_step(arrays) + return float(energy(arrays)) + + +def run_one_torch(bond_dim, num_sites): + import torch + + psi, H = _build(bond_dim, num_sites) + torch_device = "cuda" if DEVICE == "gpu" else "cpu" + arrays = [ + torch.as_tensor(a, dtype=torch.float64, device=torch_device).clone().requires_grad_(True) + for a in psi.arrays + ] + H.apply_to_arrays(lambda x: torch.tensor(x, dtype=torch.float64, device=torch_device)) + + def energy(arrays): + p = psi.copy() + for i, a in enumerate(arrays): + p[i].modify(data=a) + num = p.H @ (H.apply(p)) + den = p.H @ p + e = num / den + return torch.real(e) if torch.is_complex(e) else e + + def norm_sq(arrays): + p = psi.copy() + for i, a in enumerate(arrays): + p[i].modify(data=a) + return p.H @ p + + def grad_step(arrays): + for a in arrays: + if a.grad is not None: + a.grad = None + e = energy(arrays) + e.backward() + new_arrays = [] + with torch.no_grad(): + for a in arrays: + a_new = a - LEARNING_RATE * a.grad + new_arrays.append(a_new) + # Rescale the whole state by a single global factor derived from + # , distributed evenly across all num_sites tensors, rather than + # normalizing each tensor independently -- the MPS is not in + # canonical form here, so per-tensor normalization does not keep + # the contracted close to 1. + scale = norm_sq(new_arrays) ** (-1.0 / (2 * len(new_arrays))) + new_arrays = [(a * scale).clone().requires_grad_(True) for a in new_arrays] + return new_arrays + + for _ in range(_n_grad_steps(num_sites)): + arrays = grad_step(arrays) + with torch.no_grad(): + return float(energy(arrays)) + + +@pytest.mark.timeout(GRID_POINT_TIMEOUT_SEC) +@pytest.mark.parametrize("num_sites", NUM_SITES_VALUES) +@pytest.mark.parametrize("bond_dim", BOND_DIM_VALUES) +def test_variational_ad_jax_benchmark(benchmark, bond_dim, num_sites): + energy = benchmark.pedantic(run_one_jax, args=(bond_dim, num_sites), rounds=1, iterations=1) + benchmark.extra_info["energy"] = energy + assert energy == pytest.approx(JAX_REFERENCE_ENERGIES[(bond_dim, num_sites)], rel=1e-6) + + +@pytest.mark.cytnx_memory +@pytest.mark.limit_memory("800 MB") +@pytest.mark.parametrize("num_sites", NUM_SITES_VALUES) +@pytest.mark.parametrize("bond_dim", BOND_DIM_VALUES) +def test_variational_ad_jax_memory(bond_dim, num_sites): + energy = run_one_jax(bond_dim, num_sites) + assert energy == pytest.approx(JAX_REFERENCE_ENERGIES[(bond_dim, num_sites)], rel=1e-6) + + +@pytest.mark.timeout(GRID_POINT_TIMEOUT_SEC) +@pytest.mark.parametrize("num_sites", NUM_SITES_VALUES) +@pytest.mark.parametrize("bond_dim", BOND_DIM_VALUES) +def test_variational_ad_torch_benchmark(benchmark, bond_dim, num_sites): + energy = benchmark.pedantic(run_one_torch, args=(bond_dim, num_sites), rounds=1, iterations=1) + benchmark.extra_info["energy"] = energy + assert energy == pytest.approx(TORCH_REFERENCE_ENERGIES[(bond_dim, num_sites)], rel=1e-6) + + +@pytest.mark.cytnx_memory +@pytest.mark.limit_memory("400 MB") +@pytest.mark.parametrize("num_sites", NUM_SITES_VALUES) +@pytest.mark.parametrize("bond_dim", BOND_DIM_VALUES) +def test_variational_ad_torch_memory(bond_dim, num_sites): + energy = run_one_torch(bond_dim, num_sites) + assert energy == pytest.approx(TORCH_REFERENCE_ENERGIES[(bond_dim, num_sites)], rel=1e-6) diff --git a/benchmarks/cross_library/tenpy_bench/__init__.py b/benchmarks/cross_library/tenpy_bench/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/benchmarks/cross_library/tenpy_bench/test_dmrg_dense.py b/benchmarks/cross_library/tenpy_bench/test_dmrg_dense.py new file mode 100644 index 000000000..39f4b265c --- /dev/null +++ b/benchmarks/cross_library/tenpy_bench/test_dmrg_dense.py @@ -0,0 +1,69 @@ +"""TeNPy benchmark, algorithm class 1: finite two-site DMRG, dense mode +(no conserved quantum numbers) on the 1D spin-1/2 Heisenberg chain. + +CPU only, per the benchmark plan (TeNPy has no GPU backend). + +Run timing with `pytest --benchmark-only test_dmrg_dense.py`, memory with +`pytest --memray test_dmrg_dense.py`. +""" +import pytest + +from tenpy.algorithms import dmrg +from tenpy.models.spins import SpinChain +from tenpy.networks.mps import MPS + +from common.model import BOND_DIM_VALUES, HEISENBERG_J, LANCZOS_MAXITER, NUM_SITES_VALUES, N_SWEEPS, GRID_POINT_TIMEOUT_SEC + +REFERENCE_ENERGIES = { + (16, 20): -8.682468456356828, + (16, 30): -13.111313454922634, + (16, 50): -21.97181361569925, + (32, 20): -8.682473319689738, + (32, 30): -13.111355524192675, + (32, 50): -21.972106512033726, + (64, 20): -8.682473334397873, + (64, 30): -13.11135575848871, + (64, 50): -21.972110271889434, +} + + +def run_one(bond_dim, num_sites, dmrg_chi_max=None): + model_params = dict( + L=num_sites, S=0.5, Jx=HEISENBERG_J, Jy=HEISENBERG_J, Jz=HEISENBERG_J, + bc_MPS="finite", conserve=None, + ) + M = SpinChain(model_params) + product_state = (["up", "down"] * (num_sites // 2 + 1))[:num_sites] + psi = MPS.from_product_state(M.lat.mps_sites(), product_state, bc=M.lat.bc_MPS) + + dmrg_params = { + # Disabled to match Cytnx/quimb's plain two-site sweep, which does no + # subspace expansion. + "mixer": False, + "trunc_params": {"chi_max": dmrg_chi_max or bond_dim, "svd_min": 1e-10}, + "max_sweeps": N_SWEEPS, + # Matches Cytnx's Lanczos(Maxiter=LANCZOS_MAXITER) local-eigensolve budget. + "lanczos_params": {"N_max": LANCZOS_MAXITER}, + "combine": True, + } + eng = dmrg.TwoSiteDMRGEngine(psi, M, dmrg_params) + E, psi = eng.run() + return E + + +@pytest.mark.timeout(GRID_POINT_TIMEOUT_SEC) +@pytest.mark.parametrize("num_sites", NUM_SITES_VALUES) +@pytest.mark.parametrize("bond_dim", BOND_DIM_VALUES) +def test_dmrg_dense_benchmark(benchmark, bond_dim, num_sites): + energy = benchmark.pedantic(run_one, args=(bond_dim, num_sites), rounds=1, iterations=1) + benchmark.extra_info["energy"] = energy + assert energy == pytest.approx(REFERENCE_ENERGIES[(bond_dim, num_sites)], rel=1e-6) + + +@pytest.mark.cytnx_memory +@pytest.mark.limit_memory("110 MB") +@pytest.mark.parametrize("num_sites", NUM_SITES_VALUES) +@pytest.mark.parametrize("bond_dim", BOND_DIM_VALUES) +def test_dmrg_dense_memory(bond_dim, num_sites): + energy = run_one(bond_dim, num_sites) + assert energy == pytest.approx(REFERENCE_ENERGIES[(bond_dim, num_sites)], rel=1e-6) diff --git a/benchmarks/cross_library/tenpy_bench/test_dmrg_symmetric.py b/benchmarks/cross_library/tenpy_bench/test_dmrg_symmetric.py new file mode 100644 index 000000000..bd28045b1 --- /dev/null +++ b/benchmarks/cross_library/tenpy_bench/test_dmrg_symmetric.py @@ -0,0 +1,71 @@ +"""TeNPy benchmark, algorithm class 1 (symmetric variant): finite two-site +DMRG with U(1) total-Sz conservation on the 1D spin-1/2 Heisenberg chain. + +This exercises TeNPy's `np_conserved` block-sparse tensor backend, which +is the library's flagship optimization for abelian symmetries. CPU only. + +Run timing with `pytest --benchmark-only test_dmrg_symmetric.py`, memory +with `pytest --memray test_dmrg_symmetric.py`. +""" +import pytest + +from tenpy.algorithms import dmrg +from tenpy.models.spins import SpinChain +from tenpy.networks.mps import MPS + +from common.model import BOND_DIM_VALUES, HEISENBERG_J, LANCZOS_MAXITER, NUM_SITES_VALUES, N_SWEEPS, GRID_POINT_TIMEOUT_SEC + +REFERENCE_ENERGIES = { + (16, 20): -8.682468456356823, + (16, 30): -13.111313454922696, + (16, 50): -21.971813615699435, + (32, 20): -8.682473319689738, + (32, 30): -13.111355524202278, + (32, 50): -21.972106512010665, + (64, 20): -8.682473334397892, + (64, 30): -13.11135575848872, + (64, 50): -21.972110271890013, +} + + +def run_one(bond_dim, num_sites): + model_params = dict( + L=num_sites, S=0.5, Jx=HEISENBERG_J, Jy=HEISENBERG_J, Jz=HEISENBERG_J, + bc_MPS="finite", conserve="Sz", + ) + M = SpinChain(model_params) + # Total Sz=0 sector (Neel state), required for a conserve='Sz' MPS. + product_state = (["up", "down"] * (num_sites // 2 + 1))[:num_sites] + psi = MPS.from_product_state(M.lat.mps_sites(), product_state, bc=M.lat.bc_MPS) + + dmrg_params = { + # Disabled to match Cytnx/quimb's plain two-site sweep, which does no + # subspace expansion. + "mixer": False, + "trunc_params": {"chi_max": bond_dim, "svd_min": 1e-10}, + "max_sweeps": N_SWEEPS, + # Matches Cytnx's Lanczos(Maxiter=LANCZOS_MAXITER) local-eigensolve budget. + "lanczos_params": {"N_max": LANCZOS_MAXITER}, + "combine": True, + } + eng = dmrg.TwoSiteDMRGEngine(psi, M, dmrg_params) + E, psi = eng.run() + return E + + +@pytest.mark.timeout(GRID_POINT_TIMEOUT_SEC) +@pytest.mark.parametrize("num_sites", NUM_SITES_VALUES) +@pytest.mark.parametrize("bond_dim", BOND_DIM_VALUES) +def test_dmrg_symmetric_benchmark(benchmark, bond_dim, num_sites): + energy = benchmark.pedantic(run_one, args=(bond_dim, num_sites), rounds=1, iterations=1) + benchmark.extra_info["energy"] = energy + assert energy == pytest.approx(REFERENCE_ENERGIES[(bond_dim, num_sites)], rel=1e-6) + + +@pytest.mark.cytnx_memory +@pytest.mark.limit_memory("50 MB") +@pytest.mark.parametrize("num_sites", NUM_SITES_VALUES) +@pytest.mark.parametrize("bond_dim", BOND_DIM_VALUES) +def test_dmrg_symmetric_memory(bond_dim, num_sites): + energy = run_one(bond_dim, num_sites) + assert energy == pytest.approx(REFERENCE_ENERGIES[(bond_dim, num_sites)], rel=1e-6) diff --git a/benchmarks/cross_library/tenpy_bench/test_tebd.py b/benchmarks/cross_library/tenpy_bench/test_tebd.py new file mode 100644 index 000000000..1ca53a61a --- /dev/null +++ b/benchmarks/cross_library/tenpy_bench/test_tebd.py @@ -0,0 +1,98 @@ +"""TeNPy benchmark, algorithm class 2: real-time evolution of a 1D +transverse-field Ising chain after a field quench, using TEBD. + +(TeNPy also ships a TDVP engine with an almost identical call signature; +TEBD is used here because it is the simpler, more universally-supported +entanglement-growth benchmark and keeps the same comparison point across +all three libraries. See the project report's algorithm-class 2 for the +TDVP variant, which would only change the `tebd.TEBDEngine` line below to +`tdvp.TDVPEngine`.) CPU only. + +`TFIChain`'s `init_terms` builds H = -J*sum(Sigmax_i Sigmax_{i+1}) - +g*sum(Sigmaz_i), i.e. the coupling and field axes are swapped relative to +`cytnx_bench/test_tebd.py`'s and `quimb_bench/test_tebd.py`'s H = +-hx*sum(PauliX_i) - J*sum(PauliZ_i PauliZ_{i+1}) (coupling on Z, field on +X). `_TFIChainZCoupling` below overrides `init_terms` to swap the two +Pauli operators back, so this script's Hamiltonian uses the same axis +convention as the other two libraries and all three can start from the +literal same Sigmaz computational-basis ("up") product state. + +Run timing with `pytest --benchmark-only test_tebd.py`, memory with +`pytest --memray test_tebd.py`. +""" +import numpy as np +import pytest + +from tenpy.algorithms import tebd +from tenpy.models.tf_ising import TFIChain +from tenpy.networks.mps import MPS + +from common.model import BOND_DIM_VALUES, NUM_SITES_VALUES, GRID_POINT_TIMEOUT_SEC, TFIM_DT, TFIM_HX_FINAL, TFIM_J, TFIM_N_STEPS + + +class _TFIChainZCoupling(TFIChain): + """TFIChain with the coupling on Sigmaz and the field on Sigmax, matching + cytnx_bench/test_tebd.py's and quimb_bench/test_tebd.py's H = + -hx*sum(X_i) - J*sum(Z_i Z_{i+1}) convention instead of TFIChain's own + H = -J*sum(X_i X_{i+1}) - g*sum(Z_i).""" + + def init_terms(self, model_params): + J = np.asarray(model_params.get('J', 1.0, 'real_or_array')) + g = np.asarray(model_params.get('g', 1.0, 'real_or_array')) + for u in range(len(self.lat.unit_cell)): + self.add_onsite(-g, u, 'Sigmax') + for u1, u2, dx in self.lat.pairs['nearest_neighbors']: + self.add_coupling(-J, u1, 'Sigmaz', u2, 'Sigmaz', dx) + + +REFERENCE_ENERGIES = { + (16, 20): -19.00014659257969, + (16, 30): -29.000162658620344, + (16, 50): -49.00019479070228, + (32, 20): -19.00014659257974, + (32, 30): -29.000162658620344, + (32, 50): -49.000194790702395, + (64, 20): -19.00014659257974, + (64, 30): -29.000162658620344, + (64, 50): -49.000194790702395, +} + + +def run_one(bond_dim, num_sites): + model_params = dict(L=num_sites, J=TFIM_J, g=TFIM_HX_FINAL, bc_MPS="finite", conserve=None) + M = _TFIChainZCoupling(model_params) + # Sigmaz "up" eigenstate -- the same computational-basis product state + # used by cytnx_bench/test_tebd.py and quimb_bench/test_tebd.py. + up = np.array([1.0, 0.0]) + psi = MPS.from_product_state(M.lat.mps_sites(), [up] * num_sites, bc=M.lat.bc_MPS) + + tebd_params = { + "N_steps": 1, + "dt": TFIM_DT, + "order": 2, + "trunc_params": {"chi_max": bond_dim, "svd_min": 1e-10}, + } + eng = tebd.TEBDEngine(psi, M, tebd_params) + + for _ in range(TFIM_N_STEPS): + eng.run() + energy = M.H_MPO.expectation_value(psi) + return energy + + +@pytest.mark.timeout(GRID_POINT_TIMEOUT_SEC) +@pytest.mark.parametrize("num_sites", NUM_SITES_VALUES) +@pytest.mark.parametrize("bond_dim", BOND_DIM_VALUES) +def test_tebd_benchmark(benchmark, bond_dim, num_sites): + energy = benchmark.pedantic(run_one, args=(bond_dim, num_sites), rounds=1, iterations=1) + benchmark.extra_info["energy"] = float(energy) + assert float(energy) == pytest.approx(REFERENCE_ENERGIES[(bond_dim, num_sites)], rel=1e-6) + + +@pytest.mark.cytnx_memory +@pytest.mark.limit_memory("90 MB") +@pytest.mark.parametrize("num_sites", NUM_SITES_VALUES) +@pytest.mark.parametrize("bond_dim", BOND_DIM_VALUES) +def test_tebd_memory(bond_dim, num_sites): + energy = run_one(bond_dim, num_sites) + assert float(energy) == pytest.approx(REFERENCE_ENERGIES[(bond_dim, num_sites)], rel=1e-6) diff --git a/benchmarks/cross_library/tenpy_bench/test_variational_manual_grad.py b/benchmarks/cross_library/tenpy_bench/test_variational_manual_grad.py new file mode 100644 index 000000000..84615d1a8 --- /dev/null +++ b/benchmarks/cross_library/tenpy_bench/test_variational_manual_grad.py @@ -0,0 +1,242 @@ +"""TeNPy benchmark, algorithm class 4: variational ground-state search by +gradient descent on the MPS tensors of the 1D Heisenberg chain, with a +hand-derived (manual) gradient instead of automatic differentiation. + +TeNPy has no autodiff backend, so the gradient of the Rayleigh quotient + + E(psi) = / + +with respect to every MPS tensor simultaneously is computed analytically +rather than via backprop. This is TeNPy's counterpart to quimb's +`run_one_jax`/`run_one_torch` (`quimb_bench/test_variational_ad.py`), which +differentiate straight through the same whole-network contraction with the +backend's own autodiff. Like quimb's benchmark, every MPS tensor is updated +from a single gradient step taken simultaneously, with no orthogonality +center and no per-step canonicalization -- only a single global rescale +derived from the new `` is applied after each step, distributed +evenly across all num_sites tensors. Because the surrounding tensors are not kept +isometric, the gradient needs both an H-environment term (the effective +one-site Hamiltonian, as in a one-site sweep) and a norm-environment term: + + dE/dA_i* = (2 / den) * (H_eff,i(A_i) - E * N_eff,i(A_i)) + +where `den = `, `H_eff,i` contracts the MPO with the left/right +H-environments around site i, and `N_eff,i` is the analogous contraction +with trivial (no-MPO) norm-environments. `N_eff,i` only collapses to `A_i` +when the rest of the chain is isometric (the one-site-sweep case); here it +does not, so it is computed explicitly via `update_L_N`/`update_R_N`/ +`n_eff`. All four environment sets (`LH`, `RH`, `LN`, `RN`) are rebuilt +from scratch every gradient step, since every tensor changes at once and +no sweep-order incremental reuse is possible. + +We build every contraction by hand with `tenpy.linalg.np_conserved.tensordot` +calls on plain lists of `Array` tensors -- mirroring the Cytnx benchmark's +`Network`-based `_update_L`/`_update_R`/`_h_eff`/`_update_L_N`/`_update_R_N`/ +`_n_eff` (`cytnx_bench/test_variational_manual_grad.py`) rather than going +through `tenpy.networks.mps.MPS`/`MPOEnvironment`, since neither of those +classes assumes the non-canonical, simultaneously-updated state this +algorithm produces. + +The initial MPS is also built the same way as the Cytnx benchmark: i.i.d. +normal-random site tensors (same per-site bond-dimension formula and +per-site `seed`), right-canonicalized via a chain of SVDs +(`canonicalize_right`) purely to fix a well-defined starting point -- not +TeNPy's `MPS.from_random_unitary_evolution` from a Neel-like product state, +which stays much closer to a product state after only a few two-site +random gates and converges far more slowly under gradient descent. +`MPOEnvironment.init_LP(0)`/`init_RP(num_sites-1)` (which depend only on the +trivial-boundary structure of `H_MPO`, not on any particular state) still +supply the H-environment boundary vectors; the norm-environment boundaries +are the corresponding bond-dimension-1 scalar identities. CPU only. + +Run timing with `pytest --benchmark-only test_variational_manual_grad.py`, +memory with `pytest --memray test_variational_manual_grad.py`. The initial +MPS is seeded (per-site `np.random.RandomState(seed)`). Unlike a one-site +sweep -- a strong local optimizer whose converged energy is largely +insensitive to small initial-state differences -- this whole-network +update is a weaker optimizer whose converged energy is sensitive to the +initial state, so (as with quimb's AD benchmark, which uses the same +`LEARNING_RATE`/`_n_grad_steps`) each test asserts only against its own +`REFERENCE_ENERGIES` at a tight `rel=1e-6` self-consistency tolerance -- +this run is fully deterministic given the seeded initial state, so no +cross-library energy comparison is expected to land this close. +""" +import numpy as np +import pytest +import tenpy.linalg.np_conserved as npc +from tenpy.linalg.charges import ChargeInfo, LegCharge +from tenpy.models.spins import SpinChain +from tenpy.networks.mps import MPS +from tenpy.networks.mpo import MPOEnvironment + +from common.model import BOND_DIM_VALUES, HEISENBERG_J, NUM_SITES_VALUES, GRID_POINT_TIMEOUT_SEC + +LEARNING_RATE = 0.5 + + +def _n_grad_steps(num_sites): + return 8 * num_sites + + +REFERENCE_ENERGIES = { + (16, 20): -8.67160693635544, + (16, 30): -13.094234734927348, + (16, 50): -21.94394809820539, + (32, 20): -8.674622771089668, + (32, 30): -13.09947631043835, + (32, 50): -21.954795152613357, + (64, 20): -8.67665261172479, + (64, 30): -13.103427625375156, + (64, 50): -21.960829029634567, +} + + +def update_L(LP, A_i, W_i): + LP = npc.tensordot(LP, A_i, axes=('vR', 'vL')) + LP = npc.tensordot(W_i, LP, axes=(['p0*', 'wL'], ['p0', 'wR'])) + LP = npc.tensordot(A_i.conj(), LP, axes=(['p0*', 'vL*'], ['p0', 'vR*'])) + return LP + + +def update_R(RP, A_i, W_i): + RP = npc.tensordot(A_i, RP, axes=('vR', 'vL')) + RP = npc.tensordot(RP, W_i, axes=(['p0', 'wL'], ['p0*', 'wR'])) + RP = npc.tensordot(RP, A_i.conj(), axes=(['p0', 'vL*'], ['p0*', 'vR*'])) + return RP + + +def h_eff(theta, LP, RP, W_i): + t = npc.tensordot(LP, theta, axes=['vR', 'vL']) + t = npc.tensordot(W_i, t, axes=[['wL', 'p0*'], ['wR', 'p0']]) + t = npc.tensordot(t, RP, axes=[['wR', 'vR'], ['wL', 'vL']]) + t.ireplace_labels(['vR*', 'vL*'], ['vL', 'vR']) + t.itranspose(['vL', 'p0', 'vR']) + return t + + +def update_L_N(LP, A_i): + LP = npc.tensordot(LP, A_i, axes=('vR', 'vL')) + LP = npc.tensordot(A_i.conj(), LP, axes=(['p0*', 'vL*'], ['p0', 'vR*'])) + return LP + + +def update_R_N(RP, A_i): + RP = npc.tensordot(A_i, RP, axes=('vR', 'vL')) + RP = npc.tensordot(RP, A_i.conj(), axes=(['p0', 'vL*'], ['p0*', 'vR*'])) + return RP + + +def n_eff(theta, LP, RP): + t = npc.tensordot(LP, theta, axes=['vR', 'vL']) + t = npc.tensordot(t, RP, axes=['vR', 'vL']) + t.ireplace_labels(['vR*', 'vL*'], ['vL', 'vR']) + t.itranspose(['vL', 'p0', 'vR']) + return t + + +def _random_trivial(shape, seed, labels, qconjs=(1, 1, -1)): + rng = np.random.RandomState(seed) + data = rng.normal(size=shape) + chinfo = ChargeInfo() + legs = [LegCharge.from_trivial(s, chinfo, qconj=q) for s, q in zip(shape, qconjs)] + return npc.Array.from_ndarray(data, legs, labels=labels) + + +def _build_mps(num_sites, bond_dim, d=2): + A = [None] * num_sites + A[0] = _random_trivial([1, d, min(bond_dim, d)], 0, ['vL', 'p0', 'vR']) + for k in range(1, num_sites): + dim1 = A[k - 1].get_leg('vR').ind_len + dim3 = min(min(bond_dim, dim1 * d), d ** (num_sites - k - 1)) + A[k] = _random_trivial([dim1, d, dim3], k, ['vL', 'p0', 'vR']) + canonicalize_right(A, num_sites) + return A + + +def canonicalize_right(A, num_sites): + for p in range(num_sites - 1, 0, -1): + mat = A[p].combine_legs(['p0', 'vR'], qconj=-1) + u, s, vh = npc.svd(mat, inner_labels=['vR', 'vL']) + vh = vh.split_legs(1) + s_arr = npc.diag(s, u.get_leg('vR'), labels=['vL', 'vR']) + A[p] = vh + A[p - 1] = npc.tensordot(A[p - 1], npc.tensordot(u, s_arr, axes=['vR', 'vL']), axes=['vR', 'vL']) + A[p - 1].itranspose(['vL', 'p0', 'vR']) + A[0] /= npc.norm(A[0]) + + +def run_one(bond_dim, num_sites): + M = SpinChain(dict( + L=num_sites, S=0.5, Jx=HEISENBERG_J, Jy=HEISENBERG_J, Jz=HEISENBERG_J, + bc_MPS="finite", conserve=None, + )) + Ws = [M.H_MPO.get_W(i).replace_labels(['p', 'p*'], ['p0', 'p0*']) for i in range(num_sites)] + sites = M.lat.mps_sites() + psi0 = MPS.from_product_state(sites, ["up"] * num_sites, bc="finite") + env0 = MPOEnvironment(psi0, M.H_MPO, psi0) + L0 = env0.init_LP(0) + R0 = env0.init_RP(num_sites - 1) + LN0 = npc.Array.from_ndarray_trivial(np.array([[1.0]]), labels=['vR*', 'vR']) + RN0 = npc.Array.from_ndarray_trivial(np.array([[1.0]]), labels=['vL*', 'vL']) + + A = _build_mps(num_sites, bond_dim) + + def grad_step(A): + LH = [None] * (num_sites + 1) + LH[0] = L0 + for p in range(num_sites): + LH[p + 1] = update_L(LH[p], A[p], Ws[p]) + RH = [None] * (num_sites + 1) + RH[num_sites] = R0 + for p in range(num_sites - 1, -1, -1): + RH[p] = update_R(RH[p + 1], A[p], Ws[p]) + LN = [None] * (num_sites + 1) + LN[0] = LN0 + for p in range(num_sites): + LN[p + 1] = update_L_N(LN[p], A[p]) + RN = [None] * (num_sites + 1) + RN[num_sites] = RN0 + for p in range(num_sites - 1, -1, -1): + RN[p] = update_R_N(RN[p + 1], A[p]) + + num = LH[num_sites].to_ndarray().reshape(-1)[-1] + den = LN[num_sites].to_ndarray().item() + energy = num / den + + new_A = [None] * num_sites + for p in range(num_sites): + ht = h_eff(A[p], LH[p], RH[p + 1], Ws[p]) + nt = n_eff(A[p], LN[p], RN[p + 1]) + grad = (2.0 / den) * (ht - energy * nt) + new_A[p] = A[p] - LEARNING_RATE * grad + + LNn = LN0 + for p in range(num_sites): + LNn = update_L_N(LNn, new_A[p]) + den_new = LNn.to_ndarray().item() + scale = den_new ** (-1.0 / (2 * num_sites)) + new_A = [a * scale for a in new_A] + return new_A, energy.real + + energy = None + for _ in range(_n_grad_steps(num_sites)): + A, energy = grad_step(A) + return energy + + +@pytest.mark.timeout(GRID_POINT_TIMEOUT_SEC) +@pytest.mark.parametrize("num_sites", NUM_SITES_VALUES) +@pytest.mark.parametrize("bond_dim", BOND_DIM_VALUES) +def test_variational_manual_grad_benchmark(benchmark, bond_dim, num_sites): + energy = benchmark.pedantic(run_one, args=(bond_dim, num_sites), rounds=1, iterations=1) + benchmark.extra_info["energy"] = float(energy) + assert float(energy) == pytest.approx(REFERENCE_ENERGIES[(bond_dim, num_sites)], rel=1e-6) + + +@pytest.mark.cytnx_memory +@pytest.mark.limit_memory("80 MB") +@pytest.mark.parametrize("num_sites", NUM_SITES_VALUES) +@pytest.mark.parametrize("bond_dim", BOND_DIM_VALUES) +def test_variational_manual_grad_memory(bond_dim, num_sites): + energy = run_one(bond_dim, num_sites) + assert float(energy) == pytest.approx(REFERENCE_ENERGIES[(bond_dim, num_sites)], rel=1e-6) diff --git a/benchmarks/cross_library/validate_correctness.py b/benchmarks/cross_library/validate_correctness.py new file mode 100644 index 000000000..094905321 --- /dev/null +++ b/benchmarks/cross_library/validate_correctness.py @@ -0,0 +1,562 @@ +"""Cross-library + exact-diagonalization correctness check. + +Not part of the timing/memory benchmark suite (`run_all.py`, `results/*.csv`): +those scripts run a fixed small number of sweeps/steps purely to get a stable +per-step timing sample and never record the physical answer (ground-state +energy, observable) they compute along the way. This script instead drives +each library's own DMRG/TEBD implementation to convergence on a small chain +(small enough for dense exact diagonalization) and checks that: + + 1. the converged dense-DMRG ground energy agrees across TeNPy, quimb, and + Cytnx, and agrees with the exact-diagonalization ground energy of the + same open-chain Heisenberg Hamiltonian; + 2. the converged U(1)-symmetric DMRG ground energy agrees between TeNPy + and Cytnx (the two libraries that actually run a symmetric ground-state + search here) and with the same exact-diagonalization value. quimb's + `dmrg_symmetric` benchmark is excluded: per its own docstring it runs + imaginary-time evolution of a *random* state to exercise the same + block-sparse contract+truncate kernel cost, not a ground-state search, + so it has no "ground energy" to compare; + 3. the post-quench TFIM energy after a short real-time evolution agrees + across TeNPy (TEBD, mirroring its tdvp.py), quimb (TEBD), and Cytnx + (hand-rolled TEBD), and agrees with the exact-diagonalization value + obtained by propagating the same initial state under the same dense + Hamiltonian. + +Uses a small num_sites and a bond_dim >= 2**(num_sites//2) so every MPS +representation in this script is numerically exact (no truncation error), +isolating algorithmic/Hamiltonian-convention bugs from ordinary truncation +error. + +Run with `--generate-references` instead to switch modes: rather than the +small-num_sites exact-diagonalization check above, this runs TeNPy's +TwoSiteDMRGEngine (the only DMRG implementation in this suite already +validated against exact diagonalization here, since Cytnx's is hand-rolled +for this benchmark and quimb's symmetric variant isn't a ground-state +search) to its own convergence across the full (bond_dim, num_sites) grid, +prints the resulting ground energies as dict literals, and reports how far +each dense/symmetric-DMRG test_*.py file's hardcoded REFERENCE_ENERGIES +(computed from only 3 sweeps, per common/model.py's N_SWEEPS) sits from +that converged value. + +This mode also reports drift for the TEBD and variational-gradient-descent +benchmarks, even though neither has a "ground truth" in the DMRG sense: a +real-time-evolved expectation value (TEBD) or a non-canonicalized +whole-network gradient descent (variational) has no single converged +answer to validate against, only the literal recipe each test_*.py file +runs. Since cytnx_bench's and tenpy_bench's TEBD scripts share the same +initial product state and Trotter order (quimb_bench's too, for TEBD), and +cytnx_bench's and tenpy_bench's variational scripts share the same +per-site-seeded initial MPS and analytic-gradient formula, those shared +recipes are expected to land on (TEBD, looser tolerance, since the bond-term +grouping still differs) or very near (variational) the same energy -- so +TeNPy's own `run_one(bond_dim, num_sites)` from `tenpy_bench/test_tebd.py` +and `tenpy_bench/test_variational_manual_grad.py` is reused here as the +comparison anchor, the same role TeNPy's DMRG plays for the dense/symmetric +case above. quimb_bench/test_variational_ad.py is excluded from the +variational comparison: per its own docstring it uses a different +random-MPS construction and differentiates via autodiff instead of the +analytic gradient cytnx/tenpy share, so it has no reason to land near their +result. +""" +import importlib.util +import os +import sys + +import numpy as np + +sys.path.insert(0, os.path.dirname(__file__)) + +from common.model import ( + BOND_DIM_VALUES, HEISENBERG_J, NUM_SITES_VALUES, TFIM_DT, TFIM_HX_FINAL, TFIM_J, +) + +_CROSS_LIBRARY_DIR = os.path.dirname(__file__) + + +def _load_module(rel_path, unique_name): + """Load a test_*.py file under a guaranteed-unique module name. + + cytnx_bench/, tenpy_bench/, and quimb_bench/ each contain a file named + e.g. test_tebd.py; importing all three via sys.path manipulation makes + every `import test_tebd` after the first resolve to the same cached + sys.modules entry instead of three distinct modules. Loading by file + path under a unique name sidesteps that collision. + """ + path = os.path.join(_CROSS_LIBRARY_DIR, rel_path) + spec = importlib.util.spec_from_file_location(unique_name, path) + module = importlib.util.module_from_spec(spec) + sys.modules[unique_name] = module + spec.loader.exec_module(module) + return module + +NUM_SITES_EXACT = 8 +BOND_DIM_EXACT = 2 ** (NUM_SITES_EXACT // 2) +N_SWEEPS_CONVERGED = 40 +N_QUENCH_STEPS = 10 +ENERGY_TOL = 1e-6 + +# Number of DMRG sweeps used by generate_reference_energies() below, far more +# than the 3 sweeps each test_*.py file's run_one() uses for its timing +# fingerprint (common/model.py's N_SWEEPS) -- enough for TwoSiteDMRGEngine's +# own convergence criteria (max_E_err/max_S_err) to halt sweeping early once +# the ground energy has actually converged at the grid's (bond_dim, num_sites) +# point, rather than reporting an under-converged energy. +GROUND_TRUTH_MAX_SWEEPS = 200 + +# --- spin-1/2 operators (eigenvalues +-1/2), basis order [up, down] --- +_SX = np.array([[0, 0.5], [0.5, 0]], dtype=complex) +_SY = np.array([[0, -0.5j], [0.5j, 0]], dtype=complex) +_SZ = np.array([[0.5, 0], [0, -0.5]], dtype=complex) +_I2 = np.eye(2, dtype=complex) +# Pauli matrices (eigenvalues +-1), same basis order. +_PX = 2 * _SX +_PZ = 2 * _SZ + + +def _kron_chain(num_sites, site_ops): + """sum over bonds/sites of site_ops, each (positions, local_op) pairs.""" + H = np.zeros((2 ** num_sites, 2 ** num_sites), dtype=complex) + for positions, op in site_ops: + term = None + for site in range(num_sites): + local = op if site in positions else _I2 + term = local if term is None else np.kron(term, local) + H += term + return H + + +def exact_heisenberg_ground_energy(num_sites, J): + bonds = [({i, i + 1}, None) for i in range(num_sites - 1)] + H = np.zeros((2 ** num_sites, 2 ** num_sites), dtype=complex) + for i in range(num_sites - 1): + for op in (_SX, _SY, _SZ): + term = None + for site in range(num_sites): + local = op if site in (i, i + 1) else _I2 + term = local if term is None else np.kron(term, local) + H += J * term + evals = np.linalg.eigvalsh(H) + return evals[0] + + +def exact_tfim_propagate(num_sites, J, hx, dt, n_steps, psi0): + H = _dense_tfim_hamiltonian(num_sites, J, hx) + evals, evecs = np.linalg.eigh(H) + coeffs = evecs.conj().T @ psi0 + phase = np.exp(-1j * evals * dt * n_steps) + psi_t = evecs @ (phase * coeffs) + energy_t = np.real(psi_t.conj() @ H @ psi_t) + return energy_t, psi_t + + +def report(name, value, reference, tol=ENERGY_TOL): + diff = abs(value - reference) + status = "OK" if diff < tol else "MISMATCH" + print(f" {name:30s} = {value:+.8f} (ref {reference:+.8f}, |diff|={diff:.2e}) [{status}]") + return status == "OK" + + +def validate_dmrg_dense(): + print(f"\n=== dense Heisenberg DMRG ground energy, num_sites={NUM_SITES_EXACT}, " + f"bond_dim={BOND_DIM_EXACT} (exact) ===") + e_ed = exact_heisenberg_ground_energy(NUM_SITES_EXACT, HEISENBERG_J) + print(f" {'exact diagonalization':30s} = {e_ed:+.8f}") + ok = True + + import tenpy.linalg.np_conserved as npc # noqa: F401 (import check only) + from tenpy.algorithms import dmrg as tenpy_dmrg + from tenpy.models.spins import SpinChain + from tenpy.networks.mps import MPS as TenpyMPS + + M = SpinChain(dict(L=NUM_SITES_EXACT, S=0.5, Jx=HEISENBERG_J, Jy=HEISENBERG_J, Jz=HEISENBERG_J, + bc_MPS="finite", conserve=None)) + product_state = (["up", "down"] * (NUM_SITES_EXACT // 2 + 1))[:NUM_SITES_EXACT] + psi = TenpyMPS.from_product_state(M.lat.mps_sites(), product_state, bc=M.lat.bc_MPS) + eng = tenpy_dmrg.TwoSiteDMRGEngine(psi, M, { + "mixer": True, "trunc_params": {"chi_max": BOND_DIM_EXACT, "svd_min": 1e-12}, + "max_sweeps": N_SWEEPS_CONVERGED, "combine": True, + }) + e_tenpy, _ = eng.run() + ok &= report("tenpy (TwoSiteDMRGEngine)", e_tenpy, e_ed) + + import quimb.tensor as qtn + H = qtn.MPO_ham_heis(NUM_SITES_EXACT, j=HEISENBERG_J, cyclic=False) + dmrg = qtn.DMRG2(H, bond_dims=[BOND_DIM_EXACT], cutoffs=1e-12) + dmrg.solve(tol=1e-10, max_sweeps=N_SWEEPS_CONVERGED, verbosity=0) + ok &= report("quimb (DMRG2)", dmrg.energy, e_ed) + + sys.path.insert(0, os.path.join(os.path.dirname(__file__), "cytnx_bench")) + import cytnx + import dmrg_dense as cytnx_dmrg_dense + e_cytnx = _cytnx_dense_dmrg_converged(cytnx, cytnx_dmrg_dense, NUM_SITES_EXACT, BOND_DIM_EXACT, N_SWEEPS_CONVERGED) + ok &= report("cytnx (hand-rolled two-site DMRG)", e_cytnx, e_ed) + return ok + + +def _cytnx_dense_dmrg_converged(cytnx, mod, num_sites, bond_dim, n_sweeps): + """Re-run cytnx_bench/dmrg_dense.py's algorithm but also return the final + local ground energy, which run_one() discards (it only returns timing).""" + d = 2 + M, L0, R0 = mod._build_mpo(HEISENBERG_J) + A = [None for _ in range(num_sites)] + A[0] = cytnx.UniTensor.normal([1, d, min(bond_dim, d)], 0., 1.).set_rowrank_(2) + A[0].relabel_(["0", "1", "2"]).set_name("A0") + lbls = [["0", "1", "2"]] + for k in range(1, num_sites): + dim1 = A[k - 1].shape()[2] + dim3 = min(min(bond_dim, A[k - 1].shape()[2] * d), d ** (num_sites - k - 1)) + A[k] = cytnx.UniTensor.normal([dim1, d, dim3], 0., 1.).set_rowrank_(2).set_name(f"A{k}") + lbl = [str(2 * k), str(2 * k + 1), str(2 * k + 2)] + A[k].relabel_(lbl) + lbls.append(lbl) + + LR = [None for _ in range(num_sites + 1)] + LR[0] = L0 + LR[-1] = R0 + for p in range(num_sites - 1): + s, A[p], vt = cytnx.linalg.Gesvd(A[p]) + A[p + 1] = cytnx.Contract(cytnx.Contract(s, vt), A[p + 1]) + A[p].set_name(f"A{p}") + A[p + 1].set_name(f"A{p+1}") + anet = cytnx.Network() + anet.FromString(["L: -2,-1,-3", "A: -1,-4,1", "M: -2,0,-4,-5", + "A_Conj: -3,-5,2", "TOUT: 0,1,2"]) + anet.PutUniTensors(["L", "A", "A_Conj", "M"], + [LR[p], A[p], A[p].Dagger().permute_(A[p].labels()), M]) + LR[p + 1] = anet.Launch() + LR[p + 1].set_name(f"LR{p+1}") + A[p].relabel_(lbls[p]) + A[p + 1].relabel_(lbls[p + 1]) + _, A[-1] = cytnx.linalg.Gesvd(A[-1], is_U=True, is_vT=False) + A[-1].set_name(f"A{num_sites-1}").relabel_(lbls[-1]) + + energy = None + device = cytnx.Device.cpu + for _ in range(n_sweeps): + for p in range(num_sites - 2, -1, -1): + dim_l, dim_r = A[p].shape()[0], A[p + 1].shape()[2] + new_dim = min(dim_l * d, dim_r * d, bond_dim) + psi = cytnx.Contract(A[p], A[p + 1]) + psi, energy = mod._optimize_psi(psi, (LR[p], M, M, LR[p + 2]), 30, device) + psi.set_rowrank_(2) + s, A[p], A[p + 1] = cytnx.linalg.Svd_truncate(psi, new_dim) + A[p + 1].set_name(f"A{p+1}").relabel_(lbls[p + 1]) + s = s / s.Norm().item() + A[p] = cytnx.Contract(A[p], s) + A[p].set_name(f"A{p}").relabel_(lbls[p]) + anet = cytnx.Network() + anet.FromString(["R: -2,-1,-3", "B: 1,-4,-1", "M: 0,-2,-4,-5", + "B_Conj: 2,-5,-3", "TOUT: 0;1,2"]) + anet.PutUniTensors(["R", "B", "M", "B_Conj"], + [LR[p + 2], A[p + 1], M, A[p + 1].Dagger().permute_(A[p + 1].labels())]) + LR[p + 1] = anet.Launch() + LR[p + 1].set_name(f"LR{p+1}") + A[0].set_rowrank_(1) + _, A[0] = cytnx.linalg.Gesvd(A[0], is_U=False, is_vT=True) + A[0].set_name("A0").relabel_(lbls[0]) + for p in range(num_sites - 1): + dim_l, dim_r = A[p].shape()[0], A[p + 1].shape()[2] + new_dim = min(dim_l * d, dim_r * d, bond_dim) + psi = cytnx.Contract(A[p], A[p + 1]) + psi, energy = mod._optimize_psi(psi, (LR[p], M, M, LR[p + 2]), 30, device) + psi.set_rowrank_(2) + s, A[p], A[p + 1] = cytnx.linalg.Svd_truncate(psi, new_dim) + A[p].set_name(f"A{p}").relabel_(lbls[p]) + s = s / s.Norm().item() + A[p + 1] = cytnx.Contract(s, A[p + 1]) + A[p + 1].set_name(f"A{p+1}").relabel_(lbls[p + 1]) + anet = cytnx.Network() + anet.FromString(["L: -2,-1,-3", "A: -1,-4,1", "M: -2,0,-4,-5", + "A_Conj: -3,-5,2", "TOUT: 0,1,2"]) + anet.PutUniTensors(["L", "A", "A_Conj", "M"], + [LR[p], A[p], A[p].Dagger().permute_(A[p].labels()), M]) + LR[p + 1] = anet.Launch() + LR[p + 1].set_name(f"LR{p+1}") + A[-1].set_rowrank_(2) + _, A[-1] = cytnx.linalg.Gesvd(A[-1], is_U=True, is_vT=False) + A[-1].set_name(f"A{num_sites-1}").relabel_(lbls[-1]) + return energy + + +def validate_dmrg_symmetric(): + print(f"\n=== U(1)-symmetric Heisenberg DMRG ground energy, num_sites={NUM_SITES_EXACT}, " + f"bond_dim={BOND_DIM_EXACT} (exact) ===") + e_ed = exact_heisenberg_ground_energy(NUM_SITES_EXACT, HEISENBERG_J) + print(f" {'exact diagonalization (global)':30s} = {e_ed:+.8f}") + print(" note: quimb's dmrg_symmetric benchmark is excluded -- it runs imaginary-time") + print(" evolution of a random state (no ground-state search), per its own docstring.") + ok = True + + from tenpy.algorithms import dmrg as tenpy_dmrg + from tenpy.models.spins import SpinChain + from tenpy.networks.mps import MPS as TenpyMPS + M = SpinChain(dict(L=NUM_SITES_EXACT, S=0.5, Jx=HEISENBERG_J, Jy=HEISENBERG_J, Jz=HEISENBERG_J, + bc_MPS="finite", conserve="Sz")) + product_state = (["up", "down"] * (NUM_SITES_EXACT // 2 + 1))[:NUM_SITES_EXACT] + psi = TenpyMPS.from_product_state(M.lat.mps_sites(), product_state, bc=M.lat.bc_MPS) + eng = tenpy_dmrg.TwoSiteDMRGEngine(psi, M, { + "mixer": True, "trunc_params": {"chi_max": BOND_DIM_EXACT, "svd_min": 1e-12}, + "max_sweeps": N_SWEEPS_CONVERGED, "combine": True, + }) + e_tenpy, _ = eng.run() + ok &= report("tenpy (U(1) TwoSiteDMRGEngine)", e_tenpy, e_ed) + print(" cytnx (U(1) DMRG): skipped -- cytnx_bench/dmrg_symmetric.py has no Python") + print(" hook to extract the converged ground energy without duplicating its full") + print(" block-sparse sweep; the dense-DMRG check above already covers the same MPO.") + return ok + + +def validate_tebd_quench(): + print(f"\n=== TFIM quench, num_sites={NUM_SITES_EXACT}, bond_dim={BOND_DIM_EXACT} (exact), " + f"{N_QUENCH_STEPS} steps of dt={TFIM_DT} ===") + psi0 = np.zeros(2 ** NUM_SITES_EXACT, dtype=complex) + psi0[-1] = 1.0 # all sites in the second basis state ("down"/"1"), matching the scripts + e_ed, _ = exact_tfim_propagate(NUM_SITES_EXACT, TFIM_J, TFIM_HX_FINAL, TFIM_DT, N_QUENCH_STEPS, psi0) + print(f" {'exact diagonalization (Pauli H)':30s} = {e_ed:+.8f}") + ok = True + + sys.path.insert(0, os.path.join(os.path.dirname(__file__), "cytnx_bench")) + import cytnx + import tebd as cytnx_tebd + e_cytnx = _cytnx_tebd_energy(cytnx, cytnx_tebd, NUM_SITES_EXACT, BOND_DIM_EXACT, N_QUENCH_STEPS) + ok &= report("cytnx (hand-rolled TEBD)", e_cytnx, e_ed, tol=1e-3) + + from tenpy.algorithms import tebd as tenpy_tebd + from tenpy.models.tf_ising import TFIChain + from tenpy.networks.mps import MPS as TenpyMPS + M = TFIChain(dict(L=NUM_SITES_EXACT, J=TFIM_J, g=TFIM_HX_FINAL, bc_MPS="finite", conserve=None)) + psi = TenpyMPS.from_product_state(M.lat.mps_sites(), ["down"] * NUM_SITES_EXACT, bc=M.lat.bc_MPS) + eng = tenpy_tebd.TEBDEngine(psi, M, { + "N_steps": 1, "dt": TFIM_DT, "order": 2, + "trunc_params": {"chi_max": BOND_DIM_EXACT, "svd_min": 1e-12}, + }) + for _ in range(N_QUENCH_STEPS): + eng.run() + e_tenpy = M.H_MPO.expectation_value(psi) + ok &= report("tenpy (TEBDEngine)", e_tenpy, e_ed, tol=1e-3) + + import quimb.tensor as qtn + H = qtn.ham_1d_ising(NUM_SITES_EXACT, j=TFIM_J, bx=TFIM_HX_FINAL, cyclic=False) + psi_q = qtn.MPS_computational_state("1" * NUM_SITES_EXACT) + tebd = qtn.TEBD(psi_q, H, dt=TFIM_DT, progbar=False) + tebd.split_opts["cutoff"] = 1e-12 + tebd.split_opts["max_bond"] = BOND_DIM_EXACT + for _ in range(N_QUENCH_STEPS): + tebd.step(order=2, dt=TFIM_DT) + vec_q = tebd.pt.to_dense().flatten() + e_quimb = float(np.real(vec_q.conj() @ _dense_tfim_hamiltonian(NUM_SITES_EXACT, TFIM_J, TFIM_HX_FINAL) @ vec_q)) + ok &= report("quimb (TEBD)", e_quimb, e_ed, tol=1e-3) + return ok + + +def _dense_tfim_hamiltonian(num_sites, J, hx): + H = np.zeros((2 ** num_sites, 2 ** num_sites), dtype=complex) + for i in range(num_sites - 1): + term = None + for site in range(num_sites): + local = _PZ if site in (i, i + 1) else _I2 + term = local if term is None else np.kron(term, local) + H += -J * term + for i in range(num_sites): + term = None + for site in range(num_sites): + local = _PX if site == i else _I2 + term = local if term is None else np.kron(term, local) + H += -hx * term + return H + + +def _cytnx_tebd_energy(cytnx, mod, num_sites, bond_dim, n_steps): + d = 2 + A, lbls = mod._build_mps(num_sites, bond_dim, "cpu") + for k in range(num_sites): + A[k].set_elem([0, 0, 0], 0.0) + A[k].set_elem([0, 1, 0], 1.0) # all spin-down ("1" component) + gates = mod._build_gates(num_sites, TFIM_J, TFIM_HX_FINAL, TFIM_DT, "cpu") + for _ in range(n_steps): + for p in range(num_sites - 1): + psi = cytnx.Contract(A[p], A[p + 1]) + g = gates[p].clone().relabel_(["_o0", "_o1", lbls[p][1], lbls[p + 1][1]]) + psi = cytnx.Contract(psi, g) + psi.permute_([lbls[p][0], "_o0", "_o1", lbls[p + 1][2]]) + psi.relabel_([lbls[p][0], lbls[p][1], lbls[p + 1][1], lbls[p + 1][2]]) + psi.set_rowrank_(2) + dim_l, dim_r = A[p].shape()[0], A[p + 1].shape()[2] + new_dim = min(dim_l * d, dim_r * d, bond_dim) + s, A[p], A[p + 1] = cytnx.linalg.Svd_truncate(psi, new_dim) + s = s / s.Norm().item() + A[p + 1] = cytnx.Contract(s, A[p + 1]) + A[p].set_name(f"A{p}").relabel_(lbls[p]) + A[p + 1].set_name(f"A{p+1}").relabel_(lbls[p + 1]) + + # Contract the full chain down to a dense state vector (num_sites is small) + # and evaluate the energy with the same exact dense Hamiltonian used for + # ED, rather than re-deriving an MPO expectation-value contraction here. + psi_full = A[0] + for k in range(1, num_sites): + psi_full = cytnx.Contract(psi_full, A[k]) + order = [lbls[0][0]] + [lbls[k][1] for k in range(num_sites)] + [lbls[-1][2]] + psi_full.permute_(order) + vec = psi_full.get_block().numpy().reshape(2 ** num_sites) + H = _dense_tfim_hamiltonian(num_sites, TFIM_J, TFIM_HX_FINAL) + return float(np.real(vec.conj() @ H @ vec)) + + +def _tenpy_dmrg_ground_truth(num_sites, bond_dim, conserve): + """Run TeNPy's TwoSiteDMRGEngine to its own convergence (not the 3-sweep + fingerprint each test_*.py file checks) at one full-grid (bond_dim, + num_sites) point. + + TeNPy is the source used here, rather than quimb or Cytnx, because it is + the only one of the three whose ground-state DMRG implementation is + validated against exact diagonalization in validate_dmrg_dense()/ + validate_dmrg_symmetric() above: Cytnx's two-site sweep is a hand-rolled + implementation written for this benchmark suite, and quimb's symmetric + variant is not a ground-state search at all (see its own docstring). + """ + from tenpy.algorithms import dmrg as tenpy_dmrg + from tenpy.models.spins import SpinChain + from tenpy.networks.mps import MPS as TenpyMPS + + M = SpinChain(dict(L=num_sites, S=0.5, Jx=HEISENBERG_J, Jy=HEISENBERG_J, Jz=HEISENBERG_J, + bc_MPS="finite", conserve=conserve)) + product_state = (["up", "down"] * (num_sites // 2 + 1))[:num_sites] + psi = TenpyMPS.from_product_state(M.lat.mps_sites(), product_state, bc=M.lat.bc_MPS) + eng = tenpy_dmrg.TwoSiteDMRGEngine(psi, M, { + "mixer": True, "trunc_params": {"chi_max": bond_dim, "svd_min": 1e-10}, + "max_sweeps": GROUND_TRUTH_MAX_SWEEPS, "combine": True, + }) + e, _ = eng.run() + return e + + +def _tenpy_tebd_canonical(num_sites, bond_dim): + """Re-run tenpy_bench/test_tebd.py's own run_one(bond_dim, num_sites): the + literal quench recipe (initial state, Trotter order, step count) shared + with cytnx_bench's and quimb_bench's TEBD benchmarks, used here as the + comparison anchor since TEBD has no converged "ground truth" to run to + instead (see module docstring).""" + mod = _load_module("tenpy_bench/test_tebd.py", "_canonical_tenpy_tebd") + return mod.run_one(bond_dim, num_sites) + + +def _tenpy_variational_canonical(num_sites, bond_dim): + """Re-run tenpy_bench/test_variational_manual_grad.py's own + run_one(bond_dim, num_sites): the same per-site-seeded initial MPS and + analytic gradient formula cytnx_bench's variational benchmark shares, + used here as the comparison anchor since this whole-network gradient + descent has no converged "ground truth" to run to instead (see module + docstring).""" + mod = _load_module("tenpy_bench/test_variational_manual_grad.py", "_canonical_tenpy_variational") + return mod.run_one(bond_dim, num_sites) + + +def generate_reference_energies(): + """Print TeNPy-converged ground energies across the full (bond_dim, + num_sites) grid, for comparison against the 3-sweep REFERENCE_ENERGIES + fingerprints hardcoded in each dense/symmetric-DMRG test_*.py file (those + use common/model.py's N_SWEEPS=3, kept small for fast, stable timing -- + not chosen for ground-state convergence). Also prints TeNPy's own TEBD + and variational-gradient-descent energies across the same grid, used as + the comparison anchor for those two algorithm classes (see module + docstring for why neither has a true "ground truth" to generate + instead). This does not overwrite any test_*.py file; it prints dict + literals plus a relative-difference report so a maintainer can judge how + far each hardcoded fingerprint sits from its anchor. + """ + print(f"=== TeNPy ground-truth reference energies (max_sweeps={GROUND_TRUTH_MAX_SWEEPS}) ===") + ground_truth = {} + for label, conserve in [("dense", None), ("symmetric", "Sz")]: + energies = {} + for num_sites in NUM_SITES_VALUES: + for bond_dim in BOND_DIM_VALUES: + energies[(bond_dim, num_sites)] = _tenpy_dmrg_ground_truth(num_sites, bond_dim, conserve) + ground_truth[label] = energies + print(f"\nGROUND_TRUTH_{label.upper()}_ENERGIES = {{") + for key, e in energies.items(): + print(f" {key}: {e!r},") + print("}") + + print("\n=== TeNPy anchor energies for TEBD/variational (not a ground truth -- see docstring) ===") + for label, canonical_fn in [("tebd", _tenpy_tebd_canonical), ("variational", _tenpy_variational_canonical)]: + energies = {} + for num_sites in NUM_SITES_VALUES: + for bond_dim in BOND_DIM_VALUES: + energies[(bond_dim, num_sites)] = canonical_fn(num_sites, bond_dim) + ground_truth[label] = energies + print(f"\nANCHOR_{label.upper()}_ENERGIES = {{") + for key, e in energies.items(): + print(f" {key}: {e!r},") + print("}") + return ground_truth + + +def _report_reference_drift(ground_truth): + """Compare each test_*.py file's hardcoded REFERENCE_ENERGIES against the + matching ground_truth dict produced by generate_reference_energies().""" + cytnx_dmrg_dense = _load_module("cytnx_bench/test_dmrg_dense.py", "_drift_cytnx_dmrg_dense") + cytnx_dmrg_symmetric = _load_module("cytnx_bench/test_dmrg_symmetric.py", "_drift_cytnx_dmrg_symmetric") + cytnx_tebd = _load_module("cytnx_bench/test_tebd.py", "_drift_cytnx_tebd") + cytnx_variational = _load_module("cytnx_bench/test_variational_manual_grad.py", "_drift_cytnx_variational") + tenpy_dmrg_dense = _load_module("tenpy_bench/test_dmrg_dense.py", "_drift_tenpy_dmrg_dense") + tenpy_dmrg_symmetric = _load_module("tenpy_bench/test_dmrg_symmetric.py", "_drift_tenpy_dmrg_symmetric") + tenpy_tebd = _load_module("tenpy_bench/test_tebd.py", "_drift_tenpy_tebd") + tenpy_variational = _load_module("tenpy_bench/test_variational_manual_grad.py", "_drift_tenpy_variational") + quimb_dmrg = _load_module("quimb_bench/test_dmrg.py", "_drift_quimb_dmrg") + quimb_tebd = _load_module("quimb_bench/test_tebd.py", "_drift_quimb_tebd") + + sources = { + "dense": [ + ("cytnx_bench/test_dmrg_dense.py", cytnx_dmrg_dense.REFERENCE_ENERGIES), + ("tenpy_bench/test_dmrg_dense.py", tenpy_dmrg_dense.REFERENCE_ENERGIES), + ("quimb_bench/test_dmrg.py", quimb_dmrg.DENSE_REFERENCE_ENERGIES), + ], + "symmetric": [ + ("cytnx_bench/test_dmrg_symmetric.py", cytnx_dmrg_symmetric.REFERENCE_ENERGIES), + ("tenpy_bench/test_dmrg_symmetric.py", tenpy_dmrg_symmetric.REFERENCE_ENERGIES), + ], + "tebd": [ + ("cytnx_bench/test_tebd.py", cytnx_tebd.REFERENCE_ENERGIES), + ("tenpy_bench/test_tebd.py", tenpy_tebd.REFERENCE_ENERGIES), + ("quimb_bench/test_tebd.py", quimb_tebd.REFERENCE_ENERGIES), + ], + # quimb_bench/test_variational_ad.py is excluded: per its own + # docstring it uses a different random-MPS construction and + # differentiates via autodiff instead of the analytic gradient + # cytnx/tenpy share, so it has no reason to land near their anchor. + "variational": [ + ("cytnx_bench/test_variational_manual_grad.py", cytnx_variational.REFERENCE_ENERGIES), + ("tenpy_bench/test_variational_manual_grad.py", tenpy_variational.REFERENCE_ENERGIES), + ], + } + print("\n=== drift vs. TeNPy ground truth/anchor ===") + for label, files in sources.items(): + for path, energies in files: + for key, e in energies.items(): + e_gt = ground_truth[label][key] + rel_diff = abs(e - e_gt) / abs(e_gt) + print(f" {path:40s} {key} = {e:+.8f} (anchor {e_gt:+.8f}, rel diff={rel_diff:.2e})") + + +def main(): + if "--generate-references" in sys.argv: + ground_truth = generate_reference_energies() + _report_reference_drift(ground_truth) + return + + results = { + "dmrg_dense": validate_dmrg_dense(), + "dmrg_symmetric": validate_dmrg_symmetric(), + "tebd_quench": validate_tebd_quench(), + } + print("\n=== summary ===") + all_ok = True + for name, ok in results.items(): + print(f" {name:20s} {'PASS' if ok else 'FAIL'}") + all_ok &= ok + sys.exit(0 if all_ok else 1) + + +if __name__ == "__main__": + main() diff --git a/pyproject.toml b/pyproject.toml index 5bfe79f03..09f0264ee 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -38,6 +38,12 @@ test = [ coverage = [ "gcovr", ] +# Timing/memory regression harness for `benchmarks/cross_library/`. +benchmark = [ + "pytest-benchmark", + "pytest-memray", + "pytest-timeout", +] # Convenience aggregate for contributors and CI: # pip install --editable .[dev] # pulls everything needed to build, run the test suite, and collect @@ -71,6 +77,11 @@ Documentation = "https://cytnx-dev.github.io/Cytnx/" Repository = "https://github.com/Cytnx-dev/Cytnx.git" Issues = "https://github.com/Cytnx-dev/Cytnx/issues" +[tool.pytest.ini_options] +markers = [ + "cytnx_memory: peak-memory measurement in benchmarks/cross_library/ (selected by --memray; see benchmarks/cross_library/conftest.py)", +] + [tool.scikit-build] minimum-version = "0.11" wheel.packages = ["cytnx"]