From a51b8b0286e3c940616b4b9082dbc0f33b167fd7 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 14 Oct 2026 11:16:47 +0000 Subject: [PATCH 1/5] Initial plan From 87a7a299a3b554b18b63325c9d5f2ece5088d6ba Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 14 Oct 2026 11:25:38 +0000 Subject: [PATCH 2/5] =?UTF-8?q?feat:=20=E2=9C=A8=20oversample=20SRC=20befo?= =?UTF-8?q?re=20exact=20SVD=20truncation?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Assisted-by: GitHub Copilot:unknown Co-authored-by: robertodr <3708689+robertodr@users.noreply.github.com> --- README.md | 18 +- docs/content/docs/concepts/algorithm.mdx | 30 +++- docs/content/docs/contributing/testing.mdx | 2 +- docs/content/docs/features/cutoff.mdx | 25 ++- docs/content/docs/features/stacks.mdx | 6 +- src/src_method/_sweep.py | 38 ++++- src/src_method/_validation.py | 35 +++- src/src_method/apply.py | 23 ++- src/src_method/compress.py | 23 ++- src/src_method/stack.py | 36 ++-- tests/test_gpu_backend.py | 51 +++++- tests/test_stack.py | 182 +++++++++++++++++++-- 12 files changed, 395 insertions(+), 74 deletions(-) diff --git a/README.md b/README.md index 82c2800..7528947 100644 --- a/README.md +++ b/README.md @@ -31,6 +31,22 @@ Whether a train is an MPS or an MPO is inferred from the rank of its first site **NOTE**: the current implementation targets tensor networks with 3 or more sites. For smaller networks, an exact SVD-based fallback is dispatched, with a warning. +### Oversampling and Truncation + +`apply`, `compress` and `src` sketch at +`max(ceil(1.5 * chi_out), chi_out + 10)` by default, then use a left-to-right +SVD sweep to truncate to `chi_out`. The result is **left-canonical**, with the +norm on the last site. When positive, `cutoff` discards Schmidt values below +`cutoff * sigma_max` at each cut in this final pass. + +Set `oversample` to a finite factor of at least `1.0` to use sketch size +`ceil(oversample * chi_out)`. `oversample=1.0` recovers the original behavior, +including right-canonical output and cutoff on the sketch's QR factor. Two-site +trains keep the exact, right-canonical fallback and ignore both options. + +Larger sketches improve accuracy but increase the sweep's time and memory costs; +compare factors at the same `chi_out` when choosing an accuracy/runtime trade-off. + ### Tensor Indexing Conventions The array layout follows the default `quimb` tensor indexing conventions, so results round-trip through [Quimb](https://quimb.readthedocs.io/en/latest/autoapi/quimb/tensor/index.html) without any permutation: @@ -79,7 +95,7 @@ bra = src([t.conj() for t in psi], A, B, chi_out=64) The result then pairs with a ket by plain contraction, with no further conjugation. -**Cost.** The per-site cost of the sweep grows with `chi_out**2` times the product of the bond dimensions of the layers. Compressing a whole stack at once pays off for shallow stacks of thin layers, such as two or three Trotter layers; apply anything else pairwise. See [the stack depth benchmarks](benches/stack/README.md) for measurements. +**Cost.** The per-site cost of the sweep grows with the sketch size squared times the product of the bond dimensions of the layers. Compressing a whole stack at once pays off for shallow stacks of thin layers, such as two or three Trotter layers; apply anything else pairwise. See [the stack depth benchmarks](benches/stack/README.md) for measurements. ### GPU Support diff --git a/docs/content/docs/concepts/algorithm.mdx b/docs/content/docs/concepts/algorithm.mdx index 082d4ab..b0dc876 100644 --- a/docs/content/docs/concepts/algorithm.mdx +++ b/docs/content/docs/concepts/algorithm.mdx @@ -9,7 +9,8 @@ $A$ and an MPS $|\psi\rangle$, without ever forming the exact product. The exact product of an MPO with bond dimension $\chi_A$ and an MPS with bond dimension $\chi_\psi$ has bond dimension $\chi_A \chi_\psi$; compressing it afterwards, as in the usual contract-then-truncate approach, pays for that intermediate size. -SRC targets the output bond dimension $\chi_{\text{out}}$ directly. +SRC targets a small oversampled bond dimension $\chi_{\text{os}}$ before a final +SVD truncation to the requested $\chi_{\text{out}}$. ## Randomized range finding @@ -22,12 +23,17 @@ probability, and the cost is dominated by the product $M\Omega$. SRC applies this idea one site at a time along the chain, with the matrix $M$ being an unfolding of the (never formed) product network. -## Two sweeps +## Sketching and truncation + +By default, $\chi_{\text{os}} = +\max(\lceil 1.5\chi_{\text{out}}\rceil, \chi_{\text{out}} + 10)$. +An explicit `oversample` factor sets +$\chi_{\text{os}} = \lceil \text{oversample}\,\chi_{\text{out}}\rceil$. 1. **Left to right: sketching.** At every site the open physical legs are contracted with a Gaussian tensor $\Omega_j$ of sketch size - $\chi_{\text{out}}$. Accumulating these from the left gives the sketched - environments $C_j$, each of size $\chi_{\text{out}}$ times the product of the + $\chi_{\text{os}}$. Accumulating these from the left gives the sketched + environments $C_j$, each of size $\chi_{\text{os}}$ times the product of the input bond dimensions. The same sketch index is shared by every site, which makes the sketch a Khatri-Rao product rather than a dense Gaussian matrix of the full Hilbert space. @@ -37,16 +43,22 @@ being an unfolding of the (never formed) product network. tensor as an isometry. The isometry is then folded into $S$ for the next site. The first site absorbs the remaining weight. -The output is therefore in right-canonical form, with bond dimension at most -$\chi_{\text{out}}$. Without a [cutoff](/features/cutoff) every inner bond is -$\chi_{\text{out}}$ unless the physical dimensions near the ends of the chain -force it lower. +3. **Left to right: exact truncation.** The intermediate is already + right-canonical, so no re-orthogonalisation is needed. An SVD of each centre + tensor gives the Schmidt values at that cut. Keep at most $\chi_{\text{out}}$ + values, apply the relative [cutoff](/features/cutoff), and absorb the remaining + weight into the next site. The output is **left-canonical**, with the norm on + the last site. The discarded weights add in quadrature, giving the usual + TT-SVD quasi-optimality bound for this intermediate. + +`oversample=1.0` recovers the original two sweeps without the final pass, including +right-canonical output and cutoff on the sketch's QR factor. ## Cost For an MPO with bond dimension $\chi_A$ applied to an MPS with bond dimension $\chi_\psi$, the dominant per-site cost grows like -$\chi_{\text{out}}^2 \chi_A \chi_\psi$, compared with the $(\chi_A\chi_\psi)^3$ of +$\chi_{\text{os}}^2 \chi_A \chi_\psi$, compared with the $(\chi_A\chi_\psi)^3$ of an exact contraction followed by an SVD-based truncation. The same kernel handles MPO-MPO products and plain compression; for a [stack](/features/stacks) of trains, the product of all their bond dimensions replaces $\chi_A\chi_\psi$. diff --git a/docs/content/docs/contributing/testing.mdx b/docs/content/docs/contributing/testing.mdx index ce1aa0c..d44820b 100644 --- a/docs/content/docs/contributing/testing.mdx +++ b/docs/content/docs/contributing/testing.mdx @@ -80,7 +80,7 @@ tests of that path assert on the log instead, as below. The batched contractions and the host and disk tiers of the sweep only engage when the budgets are tight, so tests force them with tiny budgets: on the CPU, -`Resources(host_memory="1MB", scratch_dir=tmp_path)` gives small batches and spills +`Resources(host_memory="1.5MB", scratch_dir=tmp_path)` gives small batches and spills the environments of a depth-4 stack with bonds of 4 to `tmp_path` (see `tests/test_stack.py`). Compare the dense operator of the result with that of a default run, not the cores: batching changes the rounding, and with it the cores of diff --git a/docs/content/docs/features/cutoff.mdx b/docs/content/docs/features/cutoff.mdx index cd8a021..22c3748 100644 --- a/docs/content/docs/features/cutoff.mdx +++ b/docs/content/docs/features/cutoff.mdx @@ -3,9 +3,8 @@ title: Adaptive truncation description: Trim bonds below chi_out with a relative singular-value cutoff. --- -By default every inner bond of the output has dimension `chi_out`, even if the -product has lower rank across that cut. Set `cutoff` to trim each bond to its -effective rank: +`chi_out` is the maximum output bond dimension. Set `cutoff` to trim each bond +further using a relative singular-value threshold: ```python import quimb.tensor as qtn @@ -22,12 +21,22 @@ print([a.shape[0] for a in fixed[1:]]) print([a.shape[0] for a in trimmed[1:]]) ``` -During the right-to-left sweep, the QR factorisation at each site is followed by -an SVD of its small `(chi_out, chi_out)` triangular factor, and singular values -below `cutoff * sigma_max` are discarded. The extra cost is negligible next to the -contractions. +By default, `apply`, `compress` and `src` use sketch size +`max(ceil(1.5 * chi_out), chi_out + 10)`, following the original paper's +oversampling recommendation. A final left-to-right SVD sweep truncates the +right-canonical sketch to `chi_out`, discarding singular values below +`cutoff * sigma_max` at each cut. These are the Schmidt values of the intermediate +train as truncated so far, not of the random sketch matrix. The result is +**left-canonical**, with the norm on the last site. Each cut is truncated optimally +for the train as truncated so far; the discarded weights add in quadrature. + +Set `oversample` to a finite multiplicative factor of at least `1.0` to use sketch +size `ceil(oversample * chi_out)`. Larger sketches cost more time and memory. +`oversample=1.0` disables oversampling and the final SVD pass, recovering the +original **right-canonical** output. In that mode, `cutoff` acts on the sketch's +QR triangular factor during the right-to-left sweep, not on Schmidt values. - `cutoff` must lie in `[0.0, 1.0)`; `0.0` (the default) disables trimming. - `chi_out` remains the upper bound: the cutoff only ever lowers a bond. - Two-site trains use the exact fallback, which ignores `cutoff` and truncates to - `chi_out`. + `chi_out`, returning right-canonical output regardless of `oversample`. diff --git a/docs/content/docs/features/stacks.mdx b/docs/content/docs/features/stacks.mdx index 1806e4e..ebb8b3a 100644 --- a/docs/content/docs/features/stacks.mdx +++ b/docs/content/docs/features/stacks.mdx @@ -29,11 +29,11 @@ form. ## When a single sweep pays off -The per-site cost of the sweep grows with $\chi_{\text{out}}^2$ times the +The per-site cost of the sweep grows with the sketch size squared times the *product* of the bond dimensions of all layers, so it grows exponentially with the depth of the stack. In exchange, a single sweep truncates once instead of once per -layer. In practice that buys at most a moderate gain in accuracy: the error of -the randomized sketch dominates the error compounded across layers. +layer. [Oversampling and exact truncation](/features/cutoff) reduce the error of +the randomized sketch for both the one-shot and pairwise approaches. Compressing a whole stack at once pays off for shallow stacks of thin layers, such as two or three Trotter layers; apply anything else pairwise. The [stack diff --git a/src/src_method/_sweep.py b/src/src_method/_sweep.py index c023cd7..d860867 100644 --- a/src/src_method/_sweep.py +++ b/src/src_method/_sweep.py @@ -48,6 +48,7 @@ release_cached_blocks, to_numpy, truncated_qr, + use_device, ) if TYPE_CHECKING: @@ -155,6 +156,7 @@ def sweep( xp: ModuleType, *, cutoff: float = 0.0, + truncate_to: int | None = None, dtype: DTypeLike, resources: Resources | None = None, ) -> list[NDArray]: @@ -166,18 +168,21 @@ def sweep( array-likes with ``shape``, ``dtype``, ``ndim`` and ``np.asarray`` support; each is read only when the sweep reaches it. kind: The kind of the contracted train. - chi_out: The sketch size, which is the maximum output bond dimension. + chi_out: The sketch size, which bounds the intermediate bond dimension. prng: The generator for the Gaussian sketches, always host-side so that a seed gives the same draws on every device. xp: Array module (``numpy`` or ``cupy``). cutoff: Relative singular-value cutoff for adaptive bond truncation. + truncate_to: If set, truncate the right-canonical intermediate to this + bond dimension with a final left-to-right SVD pass, applying + ``cutoff`` there instead of to the sketch. dtype: The data type of the sketches. resources: Memory budgets, scratch space and devices; detected when ``None``. Returns: - The site arrays of the compressed train in right-canonical form, as numpy - arrays. + The site arrays of the compressed train as numpy arrays, right-canonical + unless ``truncate_to`` is set, in which case they are left-canonical. """ shapes = padded_shapes(layers) # The dtype of the environments and the output: the sketches promoted by the @@ -212,7 +217,7 @@ def sweep( budgets, shapes, chi_out=chi_out, - cutoff=cutoff, + cutoff=cutoff if truncate_to is None else 0.0, sketches=_Sketches(prng, chi_out, dtype, group.size), work=work, depth=len(layers), @@ -222,9 +227,34 @@ def sweep( ) as source: eta = group.run(lambda rank: sweep_.run(rank, source))[0] assert eta is not None # noqa: S101 (rank 0 returns the cores) + if truncate_to is not None: + with use_device(xp, devices[0]), device_pool_limit(xp, budgets.device_cap): + eta = _truncate(eta, truncate_to, cutoff, xp) return [to_numpy(site) for site in unpad(eta, kind)] +def _truncate( + sites: list[NDArray], chi_out: int, cutoff: float, xp: ModuleType +) -> list[NDArray]: + """Truncate a padded right-canonical train, staging one centre at a time.""" + centre = xp.asarray(sites[0]) + for j in range(len(sites) - 1): + left, right, up, down = centre.shape + matrix = centre.transpose(0, 2, 3, 1).reshape(left * up * down, right) + U, S, Vh = xp.linalg.svd(matrix, full_matrices=False) + rank = min(chi_out, S.size) + if cutoff > 0: + rank = min(rank, max(1, int(xp.count_nonzero(cutoff * S[0] <= S)))) + sites[j] = to_numpy( + U[:, :rank].reshape(left, up, down, rank).transpose(0, 3, 1, 2) + ) + centre = xp.einsum( + "ab,bcud->acud", S[:rank, None] * Vh[:rank], xp.asarray(sites[j + 1]) + ) + sites[-1] = to_numpy(centre) + return sites + + class _Sketches: """The Gaussian tensors of the left-to-right pass, shared by the devices. diff --git a/src/src_method/_validation.py b/src/src_method/_validation.py index bfcb566..8e1aea1 100644 --- a/src/src_method/_validation.py +++ b/src/src_method/_validation.py @@ -7,6 +7,7 @@ from __future__ import annotations import numbers +from math import ceil, isfinite from typing import TYPE_CHECKING import numpy as np @@ -16,7 +17,39 @@ from ._tensor_train import Site -__all__ = ["validate_chi_out", "validate_cutoff", "validate_open_boundary"] +__all__ = [ + "sketch_size", + "validate_chi_out", + "validate_cutoff", + "validate_open_boundary", +] + + +def sketch_size(chi_out: int, oversample: float | None) -> int: + """Resolve the sketch size from an optional multiplicative factor. + + Args: + chi_out: The validated output bond dimension. + oversample: A finite factor of at least one, or ``None`` for the paper's + recommendation. + + Returns: + The oversampled bond dimension. + + Raises: + ValueError: If ``oversample`` is not a finite real number of at least one. + """ + if oversample is None: + return max((3 * chi_out + 1) // 2, chi_out + 10) + if ( + isinstance(oversample, bool) + or not isinstance(oversample, numbers.Real) + or not isfinite(oversample) + or oversample < 1.0 + ): + msg = f"oversample must be a finite factor >= 1.0, got {oversample!r}." + raise ValueError(msg) + return ceil(oversample * chi_out) def validate_chi_out(chi_out: object) -> int: diff --git a/src/src_method/apply.py b/src/src_method/apply.py index 1fd6969..91859d9 100644 --- a/src/src_method/apply.py +++ b/src/src_method/apply.py @@ -29,6 +29,7 @@ def apply( chi_out: int | np.integer, *, cutoff: float = 0.0, + oversample: float | None = None, dtype: DTypeLike | None = None, seed: int | None = None, device: str = "cpu", @@ -49,11 +50,15 @@ def apply( right_tensor: The site arrays of the right tensor network (MPO or MPS). chi_out: The desired maximum bond dimension of the output tensor network. cutoff: Relative singular-value cutoff for adaptive bond truncation. - When positive, bonds are trimmed to their effective rank by - discarding singular values below ``cutoff * sigma_max`` at each - site during the right-to-left sweep. The SVD operates on the - small ``(chi_out, chi_out)`` R factor from QR, so overhead is - minimal. Set to 0.0 (default) to keep all bonds at chi_out. + Discards singular values below ``cutoff * sigma_max`` in the final + SVD pass, or on the sketch's QR factor when ``oversample=1.0``. + Set to 0.0 (default) to disable cutoff trimming. Ignored for two sites. + oversample: Multiplicative sketch-size factor, finite and at least 1.0. + Defaults to ``None``, giving sketch size + ``max(ceil(1.5 * chi_out), chi_out + 10)``; an explicit factor gives + ``ceil(oversample * chi_out)``. Oversampling adds an exact SVD pass + to truncate to ``chi_out``. Set to ``1.0`` for the original behavior. + Ignored for two sites. dtype: Data type of the random sketches. Defaults to the promoted floating dtype of the inputs; an explicit dtype can promote the result. seed: An optional seed for the random number generator. @@ -62,13 +67,16 @@ def apply( resources: Memory budgets and scratch space; see `src_method.stack.src`. Returns: - The site arrays of the compressed tensor network (MPS or MPO). + The site arrays of the compressed tensor network (MPS or MPO), in + left-canonical form with oversampling, otherwise right-canonical form. + The two-site fallback is always right-canonical. Raises: TypeError: If ``chi_out`` is not an integer or the combination of input tensor types is unsupported. ValueError: If ``chi_out`` is not positive, if ``cutoff`` is not - in ``[0.0, 1.0)``, if a train is not open-boundary, if the two trains + in ``[0.0, 1.0)``, if ``oversample`` is not finite and at least 1.0, + if a train is not open-boundary, if the two trains differ in length or in the physical dimensions they join, if a sub-three-site train is not exactly two sites, or if ``device`` is not recognised. @@ -88,6 +96,7 @@ def apply( right_tensor, chi_out=chi_out, cutoff=cutoff, + oversample=oversample, dtype=dtype, seed=seed, device=device, diff --git a/src/src_method/compress.py b/src/src_method/compress.py index 46e7ffe..d531829 100644 --- a/src/src_method/compress.py +++ b/src/src_method/compress.py @@ -27,6 +27,7 @@ def compress( chi_out: int | np.integer, *, cutoff: float = 0.0, + oversample: float | None = None, dtype: DTypeLike | None = None, seed: int | None = None, device: str = "cpu", @@ -45,11 +46,15 @@ def compress( tensor: The site arrays of the tensor network to compress (MPS or MPO). chi_out: The desired maximum bond dimension of the output tensor network. cutoff: Relative singular-value cutoff for adaptive bond truncation. - When positive, bonds are trimmed to their effective rank by - discarding singular values below ``cutoff * sigma_max`` at each - site during the right-to-left sweep. The SVD operates on the - small ``(chi_out, chi_out)`` R factor from QR, so overhead is - minimal. Set to 0.0 (default) to keep all bonds at chi_out. + Discards singular values below ``cutoff * sigma_max`` in the final + SVD pass, or on the sketch's QR factor when ``oversample=1.0``. + Set to 0.0 (default) to disable cutoff trimming. Ignored for two sites. + oversample: Multiplicative sketch-size factor, finite and at least 1.0. + Defaults to ``None``, giving sketch size + ``max(ceil(1.5 * chi_out), chi_out + 10)``; an explicit factor gives + ``ceil(oversample * chi_out)``. Oversampling adds an exact SVD pass + to truncate to ``chi_out``. Set to ``1.0`` for the original behavior. + Ignored for two sites. dtype: Data type of the random sketches. Defaults to the promoted floating dtype of the inputs; an explicit dtype can promote the result. seed: An optional seed for the random number generator. @@ -58,13 +63,16 @@ def compress( resources: Memory budgets and scratch space; see `src_method.stack.src`. Returns: - The site arrays of the compressed tensor network (MPS or MPO). + The site arrays of the compressed tensor network (MPS or MPO), in + left-canonical form with oversampling, otherwise right-canonical form. + The two-site fallback is always right-canonical. Raises: TypeError: If ``chi_out`` is not an integer or the input tensor type is unsupported. ValueError: If ``chi_out`` is not positive, if ``cutoff`` is not - in ``[0.0, 1.0)``, if the train is not open-boundary, if a + in ``[0.0, 1.0)``, if ``oversample`` is not finite and at least 1.0, + if the train is not open-boundary, if a sub-three-site train is not exactly two sites, or if ``device`` is not recognised. ImportError: If ``device="gpu"`` but cupy is not installed. @@ -73,6 +81,7 @@ def compress( tensor, chi_out=chi_out, cutoff=cutoff, + oversample=oversample, dtype=dtype, seed=seed, device=device, diff --git a/src/src_method/stack.py b/src/src_method/stack.py index 631e314..40052c0 100644 --- a/src/src_method/stack.py +++ b/src/src_method/stack.py @@ -17,7 +17,7 @@ exact_stack, normalize_stack, ) -from ._validation import validate_chi_out, validate_cutoff +from ._validation import sketch_size, validate_chi_out, validate_cutoff from .utils import default_rng, get_xp, sketch_dtype if TYPE_CHECKING: @@ -41,6 +41,7 @@ def src( *trains: Sequence[Site], chi_out: int | np.integer, cutoff: float = 0.0, + oversample: float | None = None, dtype: DTypeLike | None = None, seed: int | None = None, device: str = "cpu", @@ -65,9 +66,9 @@ def src( ``[t.conj() for t in psi]``; the result then pairs with a ket by plain contraction. - The per-site cost grows with ``chi_out**2`` times the product of the layer bond - dimensions, so one sweep pays off for shallow stacks of thin layers (for example - two or three Trotter layers). Apply anything else pairwise. + The per-site cost grows with the sketch size squared times the product of the + layer bond dimensions, so one sweep pays off for shallow stacks of thin layers + (for example two or three Trotter layers). Apply anything else pairwise. Args: *trains: The site arrays of each train, in mathematical order. A site may @@ -76,11 +77,16 @@ def src( HDF5 dataset; it is then read only when the sweep reaches it. chi_out: The desired maximum bond dimension of the output train. cutoff: Relative singular-value cutoff for adaptive bond truncation. - When positive, bonds are trimmed to their effective rank by - discarding singular values below ``cutoff * sigma_max`` at each - site during the right-to-left sweep. Set to 0.0 (default) to keep - all bonds at ``chi_out``. Ignored for two-site stacks, which are - contracted and truncated to ``chi_out`` exactly. + Discards singular values below ``cutoff * sigma_max`` in the final + left-to-right SVD pass. With ``oversample=1.0``, applies instead to + the sketch's QR factor during the right-to-left sweep, as before. + Set to 0.0 (default) to disable cutoff trimming. Ignored for two-site + stacks, which are contracted and truncated to ``chi_out`` exactly. + oversample: Multiplicative sketch-size factor, finite and at least 1.0. + The sketch size is ``ceil(oversample * chi_out)``. By default (``None``), + uses ``max(ceil(1.5 * chi_out), chi_out + 10)``. An oversampled sweep + is followed by an exact SVD truncation to ``chi_out``; ``1.0`` recovers + the original sweep without the final pass. Ignored for two-site stacks. dtype: Data type of the random sketches. Defaults to the promoted floating dtype of the inputs, so single precision stays single and complex inputs get complex Ginibre sketches. An explicit dtype @@ -93,14 +99,16 @@ def src( two-site stacks. Returns: - The site arrays of the compressed train (MPS or MPO), in right-canonical - form, as numpy arrays (host-side, whatever the ``device``). + The site arrays of the compressed train (MPS or MPO), as numpy arrays + (host-side, whatever the ``device``). Oversampling returns left-canonical + form; ``oversample=1.0`` and the two-site fallback return right-canonical form. Raises: TypeError: If ``chi_out`` is not an integer, if a train has an unrecognised layout or an MPS sits anywhere other than at one end of the stack. ValueError: If ``chi_out`` is not positive, if ``cutoff`` is not in - ``[0.0, 1.0)``, if a train is not open-boundary, if the stack is empty, + ``[0.0, 1.0)``, if ``oversample`` is not finite and at least 1.0, + if a train is not open-boundary, if the stack is empty, if the trains differ in length or in the physical dimensions they join, if a sub-three-site stack is not exactly two sites, or if ``device`` is not recognised. @@ -108,6 +116,7 @@ def src( """ chi_out = validate_chi_out(chi_out) validate_cutoff(cutoff) + chi_sketch = sketch_size(chi_out, oversample) xp = get_xp(device) prng = default_rng(seed) layers, kind = normalize_stack(trains) @@ -128,10 +137,11 @@ def src( result = sweep( layers, kind, - chi_out, + chi_sketch, prng, xp, cutoff=cutoff, + truncate_to=chi_out if chi_sketch > chi_out else None, dtype=sketch_dtype(dtype, *layers), resources=resources, ) diff --git a/tests/test_gpu_backend.py b/tests/test_gpu_backend.py index 6e7c031..ca3443a 100644 --- a/tests/test_gpu_backend.py +++ b/tests/test_gpu_backend.py @@ -165,6 +165,40 @@ def test_invalid_device_raises() -> None: apply(H.arrays, psi.arrays, chi_out=4, device="tpu") +@pytest.mark.parametrize("oversample", [None, 1.0, 2.0]) +@pytest.mark.parametrize("kind", ["mps", "mpo"]) +def test_oversampled_truncation_runs_on_gpu( + oversample: float | None, kind: str, monkeypatch: pytest.MonkeyPatch +) -> None: + """The final SVD and cutoff run on the selected backend.""" + make_train = qtn.MPS_rand_state if kind == "mps" else qtn.MPO_rand + train = make_train(5, bond_dim=4, dtype=np.complex128, seed=0) + calls = [] + original = cupy.linalg.svd + + def spy(matrix: Any, **kwargs: Any) -> Any: + calls.append(matrix.device.id) + return original(matrix, **kwargs) + + monkeypatch.setattr(cupy.linalg, "svd", spy) + cpu = compress(train.arrays, 2, oversample=oversample, cutoff=0.1, seed=1) + gpu = compress( + train.arrays, + 2, + oversample=oversample, + cutoff=0.1, + seed=1, + device="gpu", + resources=Resources(devices=[0]), + ) + + assert calls == ([] if oversample == 1.0 else [0] * 4) + assert all(isinstance(site, np.ndarray) for site in gpu) + np.testing.assert_allclose( + type(train)(cpu).to_dense(), type(train)(gpu).to_dense(), atol=1e-10 + ) + + @pytest.mark.parametrize("device", ["cpu", "gpu"]) def test_src_stack_gpu_matches_reference(device: str) -> None: """A depth-3 ``src`` stack should match the exact product on both devices.""" @@ -241,15 +275,16 @@ def test_gpu_tiers_and_batches_match_cpu(tmp_path, plans: list[Plan]) -> None: rng = np.random.default_rng(7) stack = [random_mpo_arrays([4, 8, 8, 8, 4], rng) for _ in range(4)] chi = 64 + chi_sketch = 96 shapes, sizes = padded_shapes(stack), site_bytes(stack) roomy = make_plan( shapes, sizes, - chi, + chi_sketch, np.complex128, Budgets(10**10, 10**10, 10**12, tmp_path, unified=False), ) - env = chi * 8**4 * 16 # one bulk environment + env = chi_sketch * 8**4 * 16 # one bulk environment tight = Resources( gpu_memory=roomy.device_peak // 2, host_memory=roomy.host_peak + 5 * env, @@ -262,7 +297,7 @@ def test_gpu_tiers_and_batches_match_cpu(tmp_path, plans: list[Plan]) -> None: ) assert {site.tier for site in plans[-1].sites} == {"device", "host", "disk"} - assert min(site.sketch_batch for site in plans[-1].sites[1:]) < chi + assert min(site.sketch_batch for site in plans[-1].sites[1:]) < chi_sketch reference = dense_mpo(cpu) error = np.linalg.norm(dense_mpo(gpu) - reference) / np.linalg.norm(reference) assert error < 1e-10 @@ -343,6 +378,7 @@ def test_two_gpus_match_one( rng = np.random.default_rng(7) stack = [random_mpo_arrays([4, 8, 8, 8, 4], rng) for _ in range(4)] chi = 63 + chi_sketch = 95 if uploads == "full": monkeypatch.setattr( "src_method._group.enable_peer_access", lambda _xp, _devices: False @@ -351,17 +387,16 @@ def test_two_gpus_match_one( roomy = make_plan( shapes, sizes, - chi, + chi_sketch, np.complex128, Budgets(10**10, 10**10, 10**12, tmp_path, unified=False), devices=2, ) - env = (chi // 2 + 1) * 8**4 * 16 # one bulk environment on one GPU - # Checked with make_plan: one environment per GPU on the device, two in host - # memory, the oldest two on disk. + env = (chi_sketch // 2 + 1) * 8**4 * 16 # one bulk environment on one GPU + # Checked with make_plan: the oversampled environments use every tier. tight = Resources( gpu_memory=roomy.device_peak - 2 * env, - host_memory=roomy.host_peak + 12 * env, + host_memory=roomy.host_peak + 10 * env, scratch_dir=tmp_path, devices=2, ) diff --git a/tests/test_stack.py b/tests/test_stack.py index 404caa5..a0522e1 100644 --- a/tests/test_stack.py +++ b/tests/test_stack.py @@ -1,5 +1,7 @@ """Test `src` over stacks of tensor trains, and the sweep behind it.""" +from __future__ import annotations + import logging import threading from functools import reduce @@ -8,10 +10,12 @@ import pytest import src_method._sweep as sweep_module +import src_method.stack as stack_module from src_method import Resources, apply, compress, src from src_method._kernels import SiteKernels from src_method._plan import make_plan from src_method._sweep import sweep +from src_method._tensor_train import pad # ------------- # --- Utils --- @@ -207,7 +211,7 @@ def test_src_truncation_is_near_optimal(): @pytest.mark.parametrize("spec", ["A B C", "A B psi", "phi A B"]) def test_src_output_is_right_canonical(spec, rng): - out = src(*make_stack(spec, rng), chi_out=3, seed=0) + out = src(*make_stack(spec, rng), chi_out=3, oversample=1.0, seed=0) for site in out[1:]: flat = site.reshape(site.shape[0], -1) @@ -216,6 +220,141 @@ def test_src_output_is_right_canonical(spec, rng): ) +@pytest.mark.parametrize("spec", KET_STACKS + BRA_STACKS) +@pytest.mark.parametrize("oversample", [None, 1.5, 2.0]) +def test_oversampled_output_is_left_canonical(spec, oversample, rng): + out = src(*make_stack(spec, rng), chi_out=3, oversample=oversample, seed=0) + + for site in pad(out)[:-1]: + flat = site.transpose(0, 2, 3, 1).reshape(-1, site.shape[1]) + np.testing.assert_allclose( + flat.conj().T @ flat, np.eye(site.shape[1]), atol=1e-10 + ) + assert site.shape[1] <= 3 + + +@pytest.mark.parametrize("spec", KET_STACKS) +@pytest.mark.parametrize("cutoff", [0.0, 0.1]) +def test_oversample_one_matches_legacy_sweep(spec, cutoff, rng): + stack = make_stack(spec, rng) + kind = "mps" if stack[-1][0].ndim == 2 else "mpo" + expected = sweep( + stack, + kind, + 3, + np.random.default_rng(0), + np, + cutoff=cutoff, + dtype=np.complex128, + ) + out = src(*stack, chi_out=3, oversample=1.0, cutoff=cutoff, seed=0) + + assert_same_trains(out, expected) + + +@pytest.mark.parametrize( + ("chi", "oversample", "expected"), + [ + (3, None, 13), + (21, None, 32), + (3, 1.0, 3), + (3, np.float32(1.5), 5), + (3, 2.0, 6), + ], +) +def test_oversampling_sets_sweep_size(chi, oversample, expected, rng, monkeypatch): + calls = [] + original = stack_module.sweep + + def spy(*args, **kwargs): + calls.append((args[2], kwargs["truncate_to"])) + return original(*args, **kwargs) + + monkeypatch.setattr(stack_module, "sweep", spy) + src(*make_stack("A B psi", rng), chi_out=chi, oversample=oversample, seed=0) + + assert calls == [(expected, chi if expected > chi else None)] + + +@pytest.mark.parametrize( + "oversample", [0.0, -1.0, 0.99, np.nan, np.inf, -np.inf, True, "2", 2j] +) +@pytest.mark.parametrize("spec", ["A psi", "A"]) +def test_invalid_oversample_raises(spec, oversample, rng): + trains = make_stack(spec, rng) + function = apply if len(trains) == 2 else compress + + with pytest.raises(ValueError, match="oversample"): + function(*trains, chi_out=3, oversample=oversample) + + +@pytest.mark.parametrize("spec", ["A B C", "A B psi", "phi A B"]) +def test_oversampling_improves_accuracy(spec): + errors = [] + for seed in range(5): + trains = make_stack(spec, np.random.default_rng(seed)) + reference = dense_stack(*trains) + errors.append( + [ + rel_error( + src(*trains, chi_out=2, oversample=factor, seed=seed), + reference, + ) + for factor in (1.0, 2.0, None) + ] + ) + + legacy, doubled, default = np.median(errors, axis=0) + assert doubled < legacy + assert default < legacy + + +@pytest.mark.parametrize("cutoff", [0.0, 0.1]) +def test_oversampling_matches_exact_truncation_of_sketch(cutoff, rng): + qtn = pytest.importorskip("quimb.tensor") + trains = make_stack("A B C", rng) + intermediate = src(*trains, chi_out=8, oversample=1.0, seed=0) + reference = qtn.MatrixProductOperator(intermediate) + reference.compress(form="left", max_bond=3, cutoff=cutoff, cutoff_mode="rel") + + out = src(*trains, chi_out=3, oversample=8 / 3, cutoff=cutoff, seed=0) + + assert rel_error(out, reference.to_dense()) < 1e-10 + + +@pytest.mark.parametrize("kind", ["mps", "mpo"]) +@pytest.mark.parametrize("scale", [1.0, 1e-8, 1e8]) +def test_oversampled_cutoff_uses_schmidt_values(kind, scale): + sites = [ + np.diag([scale, 0.05 * scale]), + np.einsum("ab,ac->abc", np.eye(2), np.eye(2)), + np.eye(2), + ] + if kind == "mpo": + sites = [site[..., None] for site in sites] + before = [site.copy() for site in sites] + + out = compress(sites, chi_out=2, cutoff=0.1, seed=0) + + assert all(site.shape[0] == 1 for site in out[1:]) + reference = np.zeros_like(dense(sites)) + reference[0, 0] = scale + np.testing.assert_allclose(dense(out), reference, atol=1e-12 * scale) + assert_same_trains(sites, before) + + +@pytest.mark.parametrize("kind", ["mps", "mpo"]) +def test_oversampled_zero_train(kind): + sites = [np.zeros((2, 2)), np.zeros((2, 2, 2)), np.zeros((2, 2))] + if kind == "mpo": + sites = [site[..., None] for site in sites] + + out = compress(sites, chi_out=2, cutoff=0.1, seed=0) + + assert all(site.shape[0] >= 1 for site in out[1:]) + np.testing.assert_array_equal(dense(out), dense(sites)) + + def test_src_does_not_mutate_bra_stack(rng): stack = make_stack("phi A B", rng) before = [[t.copy() for t in train] for train in stack] @@ -250,10 +389,11 @@ def test_src_seed_is_deterministic(rng): @pytest.mark.parametrize("spec", ["A B psi", "A B C", "phi A B"]) -def test_src_two_site_stack_falls_back(spec, rng, caplog): +@pytest.mark.parametrize("oversample", [None, 1.0, 2.0]) +def test_src_two_site_stack_falls_back(spec, oversample, rng, caplog): stack = make_stack(spec, rng, n_sites=2) - out = src(*stack, chi_out=CHI_EXACT) + out = src(*stack, chi_out=CHI_EXACT, oversample=oversample) assert "Defaulting" in caplog.text assert rel_error(out, dense_stack(*stack)) < 1e-10 @@ -290,26 +430,44 @@ def assert_cutoff_trims(train, untrimmed): assert [t.shape for t in train] != [t.shape for t in untrimmed] +@pytest.mark.parametrize("oversample", [None, 1.0, 1.5]) @pytest.mark.parametrize("cutoff", [0.0, 0.5]) @pytest.mark.parametrize("spec", ["A psi", "A B"]) -def test_apply_matches_src(spec, cutoff, rng): +def test_apply_matches_src(spec, cutoff, oversample, rng): left, right = make_stack(spec, rng) - via_apply = apply(left, right, 4, cutoff=cutoff, dtype=np.complex128, seed=3) - via_src = src(left, right, chi_out=4, cutoff=cutoff, dtype=np.complex128, seed=3) + via_apply = apply( + left, + right, + 4, + cutoff=cutoff, + oversample=oversample, + dtype=np.complex128, + seed=3, + ) + via_src = src( + left, + right, + chi_out=4, + cutoff=cutoff, + oversample=oversample, + dtype=np.complex128, + seed=3, + ) assert_same_trains(via_apply, via_src) if cutoff: assert_cutoff_trims(via_src, src(left, right, chi_out=4, seed=3)) +@pytest.mark.parametrize("oversample", [None, 1.0, 1.5]) @pytest.mark.parametrize("cutoff", [0.0, 0.5]) @pytest.mark.parametrize("spec", ["psi", "A"]) -def test_compress_matches_src(spec, cutoff, rng): +def test_compress_matches_src(spec, cutoff, oversample, rng): (train,) = make_stack(spec, rng) - via_compress = compress(train, 4, cutoff=cutoff, seed=3) - via_src = src(train, chi_out=4, cutoff=cutoff, seed=3) + via_compress = compress(train, 4, cutoff=cutoff, oversample=oversample, seed=3) + via_src = src(train, chi_out=4, cutoff=cutoff, oversample=oversample, seed=3) assert_same_trains(via_compress, via_src) if cutoff: @@ -375,14 +533,14 @@ def spy(*args, **kwargs): def test_tiny_budget_batches_and_spills(rng, tmp_path, plans): stack = [random_mpo([3, 4, 4, 4, 3], rng) for _ in range(4)] - tight = Resources(host_memory="1MB", scratch_dir=tmp_path) + tight = Resources(host_memory="1.5MB", scratch_dir=tmp_path) default = src(*stack, chi_out=16, seed=3, dtype=np.complex128) spilled = src(*stack, chi_out=16, seed=3, dtype=np.complex128, resources=tight) assert all(site.tier == "device" for site in plans[0].sites) assert "disk" in {site.tier for site in plans[1].sites} - assert min(site.sketch_batch for site in plans[1].sites[1:]) < 16 + assert min(site.sketch_batch for site in plans[1].sites[1:]) < 26 # Rounding differs, and the output cores of an ill-conditioned sketch with it, # but not the operator they represent. assert rel_error(spilled, dense(default)) < 1e-10 @@ -501,7 +659,7 @@ def test_devices_with_cutoff_keep_the_bonds(rng): def test_devices_spill_and_clean_up(rng, tmp_path, plans): stack = [random_mpo([3, 4, 4, 4, 3], rng) for _ in range(4)] - tight = Resources(host_memory="1MB", scratch_dir=tmp_path, devices=2) + tight = Resources(host_memory="2MB", scratch_dir=tmp_path, devices=2) default = src(*stack, chi_out=16, seed=3, dtype=np.complex128) spilled = src(*stack, chi_out=16, seed=3, dtype=np.complex128, resources=tight) From bd47207a50d1f23e7015ba4cda0925ab1a5aed18 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 14 Oct 2026 11:31:02 +0000 Subject: [PATCH 3/5] =?UTF-8?q?fix:=20=F0=9F=90=9B=20release=20discarded?= =?UTF-8?q?=20SVD=20storage?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Assisted-by: GitHub Copilot:unknown Co-authored-by: robertodr <3708689+robertodr@users.noreply.github.com> --- src/src_method/_sweep.py | 4 +++- tests/test_stack.py | 17 +++++++++++++++++ 2 files changed, 20 insertions(+), 1 deletion(-) diff --git a/src/src_method/_sweep.py b/src/src_method/_sweep.py index d860867..5343502 100644 --- a/src/src_method/_sweep.py +++ b/src/src_method/_sweep.py @@ -246,7 +246,9 @@ def _truncate( if cutoff > 0: rank = min(rank, max(1, int(xp.count_nonzero(cutoff * S[0] <= S)))) sites[j] = to_numpy( - U[:, :rank].reshape(left, up, down, rank).transpose(0, 3, 1, 2) + xp.ascontiguousarray( + U[:, :rank].reshape(left, up, down, rank).transpose(0, 3, 1, 2) + ) ) centre = xp.einsum( "ab,bcud->acud", S[:rank, None] * Vh[:rank], xp.asarray(sites[j + 1]) diff --git a/tests/test_stack.py b/tests/test_stack.py index a0522e1..33e32a4 100644 --- a/tests/test_stack.py +++ b/tests/test_stack.py @@ -355,6 +355,23 @@ def test_oversampled_zero_train(kind): np.testing.assert_array_equal(dense(out), dense(sites)) +@pytest.mark.parametrize("kind", ["mps", "mpo"]) +def test_oversampled_output_does_not_retain_discarded_columns(kind, rng): + train = ( + random_mps([4, 4, 4], rng, phys=32) + if kind == "mps" + else random_mpo([4, 4, 4], rng, up=4, down=8) + ) + + out = compress(train, chi_out=1, seed=0) + + for site in out[:-1]: + storage = site + while isinstance(storage.base, np.ndarray): + storage = storage.base + assert storage.nbytes == site.nbytes + + def test_src_does_not_mutate_bra_stack(rng): stack = make_stack("phi A B", rng) before = [[t.copy() for t in train] for train in stack] From 73b0b9c1f94f4d27f7278d52a154ae8d42565d8c Mon Sep 17 00:00:00 2001 From: Panadestein Date: Thu, 15 Oct 2026 14:26:05 +0200 Subject: [PATCH 4/5] =?UTF-8?q?perf(sweep):=20=E2=9A=A1=EF=B8=8F=20reduce?= =?UTF-8?q?=20the=20final=20truncation=20to=20a=20small=20SVD=20after=20a?= =?UTF-8?q?=20QR?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Assisted-by: GitHub Copilot:claude-sonnet-5.5 --- src/src_method/_sweep.py | 15 +++++++++------ tests/test_gpu_backend.py | 15 +++------------ 2 files changed, 12 insertions(+), 18 deletions(-) diff --git a/src/src_method/_sweep.py b/src/src_method/_sweep.py index ed1db50..e9b0be4 100644 --- a/src/src_method/_sweep.py +++ b/src/src_method/_sweep.py @@ -254,17 +254,20 @@ def _truncate( for j in range(len(sites) - 1): left, right, up, down = centre.shape matrix = centre.transpose(0, 2, 3, 1).reshape(left * up * down, right) - U, S, Vh = xp.linalg.svd(matrix, full_matrices=False) + # A tall matrix is reduced by QR first, so the SVD sees a small square. It + # runs on the host, as in `truncated_qr`: device SVDs are slow at this size. + Q, R = xp.linalg.qr(matrix) if matrix.shape[0] > right else (None, matrix) + U, S, Vh = np.linalg.svd(to_numpy(R), full_matrices=False) rank = min(chi_out, S.size) if cutoff > 0: - rank = min(rank, max(1, int(xp.sum(cutoff * S[0] <= S)))) + rank = min(rank, max(1, int((cutoff * S[0] <= S).sum()))) + U = xp.asarray(U[:, :rank]) + U = U if Q is None else Q @ U sites[j] = to_numpy( - xp.ascontiguousarray( - U[:, :rank].reshape(left, up, down, rank).transpose(0, 3, 1, 2) - ) + xp.ascontiguousarray(U.reshape(left, up, down, rank).transpose(0, 3, 1, 2)) ) centre = xp.tensordot( - S[:rank, None] * Vh[:rank], xp.asarray(sites[j + 1]), axes=1 + xp.asarray(S[:rank, None] * Vh[:rank]), xp.asarray(sites[j + 1]), axes=1 ) sites[-1] = to_numpy(centre) return sites diff --git a/tests/test_gpu_backend.py b/tests/test_gpu_backend.py index ca3443a..0faece2 100644 --- a/tests/test_gpu_backend.py +++ b/tests/test_gpu_backend.py @@ -167,20 +167,12 @@ def test_invalid_device_raises() -> None: @pytest.mark.parametrize("oversample", [None, 1.0, 2.0]) @pytest.mark.parametrize("kind", ["mps", "mpo"]) -def test_oversampled_truncation_runs_on_gpu( - oversample: float | None, kind: str, monkeypatch: pytest.MonkeyPatch +def test_oversampled_truncation_matches_cpu( + oversample: float | None, kind: str ) -> None: - """The final SVD and cutoff run on the selected backend.""" + """The final truncation on the GPU agrees with the CPU one.""" make_train = qtn.MPS_rand_state if kind == "mps" else qtn.MPO_rand train = make_train(5, bond_dim=4, dtype=np.complex128, seed=0) - calls = [] - original = cupy.linalg.svd - - def spy(matrix: Any, **kwargs: Any) -> Any: - calls.append(matrix.device.id) - return original(matrix, **kwargs) - - monkeypatch.setattr(cupy.linalg, "svd", spy) cpu = compress(train.arrays, 2, oversample=oversample, cutoff=0.1, seed=1) gpu = compress( train.arrays, @@ -192,7 +184,6 @@ def spy(matrix: Any, **kwargs: Any) -> Any: resources=Resources(devices=[0]), ) - assert calls == ([] if oversample == 1.0 else [0] * 4) assert all(isinstance(site, np.ndarray) for site in gpu) np.testing.assert_allclose( type(train)(cpu).to_dense(), type(train)(gpu).to_dense(), atol=1e-10 From 8c0663e6d4f5667f1812092c3e6620bddb6f5e10 Mon Sep 17 00:00:00 2001 From: Panadestein Date: Thu, 15 Oct 2026 14:41:20 +0200 Subject: [PATCH 5/5] =?UTF-8?q?feat:=20=E2=9C=A8=20make=20oversampling=20o?= =?UTF-8?q?pt-in=20with=20a=20default=20of=201.0?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Assisted-by: GitHub Copilot:claude-opus-5.5 --- README.md | 26 +++++------ docs/content/docs/concepts/algorithm.mdx | 34 +++++++------- docs/content/docs/contributing/testing.mdx | 2 +- docs/content/docs/features/cutoff.mdx | 53 +++++++++++++++------- docs/content/docs/features/stacks.mdx | 6 ++- src/src_method/apply.py | 21 ++++----- src/src_method/compress.py | 21 ++++----- src/src_method/stack.py | 25 +++++----- tests/test_gpu_backend.py | 17 ++++--- tests/test_stack.py | 30 ++++++++---- 10 files changed, 133 insertions(+), 102 deletions(-) diff --git a/README.md b/README.md index 0781b4d..6c0d046 100644 --- a/README.md +++ b/README.md @@ -35,19 +35,19 @@ Whether a train is an MPS or an MPO is inferred from the rank of its first site ### Oversampling and Truncation -`apply`, `compress` and `src` sketch at -`max(ceil(1.5 * chi_out), chi_out + 10)` by default, then use a left-to-right -SVD sweep to truncate to `chi_out`. The result is **left-canonical**, with the -norm on the last site. When positive, `cutoff` discards Schmidt values below -`cutoff * sigma_max` at each cut in this final pass. - -Set `oversample` to a finite factor of at least `1.0` to use sketch size -`ceil(oversample * chi_out)`. `oversample=1.0` recovers the original behavior, -including right-canonical output and cutoff on the sketch's QR factor. Two-site -trains keep the exact, right-canonical fallback and ignore both options. - -Larger sketches improve accuracy but increase the sweep's time and memory costs; -compare factors at the same `chi_out` when choosing an accuracy/runtime trade-off. +By default, `apply`, `compress` and `src` sketch at `chi_out` and return +**right-canonical** output. For better accuracy at the same `chi_out`, pass +`oversample`: a finite factor of at least `1.0` sketches at +`ceil(oversample * chi_out)`, and `None` at the paper's +`max(ceil(1.5 * chi_out), chi_out + 10)`. A final left-to-right SVD sweep then +truncates to `chi_out`, and the result is **left-canonical**, with the norm on +the last site. When positive, `cutoff` discards Schmidt values below +`cutoff * sigma_max` at each cut in this final pass. Two-site trains keep the +exact, right-canonical fallback and ignore both options. + +Larger sketches improve accuracy but increase time and memory costs, often +several-fold; compare factors at the same `chi_out` when choosing an +accuracy/runtime trade-off. ### Tensor Indexing Conventions diff --git a/docs/content/docs/concepts/algorithm.mdx b/docs/content/docs/concepts/algorithm.mdx index b0dc876..ebc6040 100644 --- a/docs/content/docs/concepts/algorithm.mdx +++ b/docs/content/docs/concepts/algorithm.mdx @@ -9,8 +9,9 @@ $A$ and an MPS $|\psi\rangle$, without ever forming the exact product. The exact product of an MPO with bond dimension $\chi_A$ and an MPS with bond dimension $\chi_\psi$ has bond dimension $\chi_A \chi_\psi$; compressing it afterwards, as in the usual contract-then-truncate approach, pays for that intermediate size. -SRC targets a small oversampled bond dimension $\chi_{\text{os}}$ before a final -SVD truncation to the requested $\chi_{\text{out}}$. +SRC targets the output bond dimension $\chi_{\text{out}}$ directly, or an +oversampled sketch size $\chi_{\text{os}}$ followed by a final SVD truncation to +$\chi_{\text{out}}$. ## Randomized range finding @@ -25,10 +26,13 @@ being an unfolding of the (never formed) product network. ## Sketching and truncation -By default, $\chi_{\text{os}} = -\max(\lceil 1.5\chi_{\text{out}}\rceil, \chi_{\text{out}} + 10)$. -An explicit `oversample` factor sets -$\chi_{\text{os}} = \lceil \text{oversample}\,\chi_{\text{out}}\rceil$. +By default, $\chi_{\text{os}} = \chi_{\text{out}}$ and only the first two sweeps +run: the output is right-canonical, with bond dimension at most +$\chi_{\text{out}}$. An `oversample` factor sets +$\chi_{\text{os}} = \lceil \text{oversample}\,\chi_{\text{out}}\rceil$, and +`oversample=None` the paper's +$\max(\lceil 1.5\chi_{\text{out}}\rceil, \chi_{\text{out}} + 10)$; the third +sweep then truncates to $\chi_{\text{out}}$. 1. **Left to right: sketching.** At every site the open physical legs are contracted with a Gaussian tensor $\Omega_j$ of sketch size @@ -43,16 +47,14 @@ $\chi_{\text{os}} = \lceil \text{oversample}\,\chi_{\text{out}}\rceil$. tensor as an isometry. The isometry is then folded into $S$ for the next site. The first site absorbs the remaining weight. -3. **Left to right: exact truncation.** The intermediate is already - right-canonical, so no re-orthogonalisation is needed. An SVD of each centre - tensor gives the Schmidt values at that cut. Keep at most $\chi_{\text{out}}$ - values, apply the relative [cutoff](/features/cutoff), and absorb the remaining - weight into the next site. The output is **left-canonical**, with the norm on - the last site. The discarded weights add in quadrature, giving the usual - TT-SVD quasi-optimality bound for this intermediate. - -`oversample=1.0` recovers the original two sweeps without the final pass, including -right-canonical output and cutoff on the sketch's QR factor. +3. **Left to right: exact truncation (oversampling only).** The intermediate is + already right-canonical, so no re-orthogonalisation is needed. An SVD of each + centre tensor gives the Schmidt values at that cut. Keep at most + $\chi_{\text{out}}$ values, apply the relative [cutoff](/features/cutoff), and + absorb the remaining weight into the next site. The output is + **left-canonical**, with the norm on the last site. The discarded weights add + in quadrature, giving the usual TT-SVD quasi-optimality bound for this + intermediate. ## Cost diff --git a/docs/content/docs/contributing/testing.mdx b/docs/content/docs/contributing/testing.mdx index d44820b..ce1aa0c 100644 --- a/docs/content/docs/contributing/testing.mdx +++ b/docs/content/docs/contributing/testing.mdx @@ -80,7 +80,7 @@ tests of that path assert on the log instead, as below. The batched contractions and the host and disk tiers of the sweep only engage when the budgets are tight, so tests force them with tiny budgets: on the CPU, -`Resources(host_memory="1.5MB", scratch_dir=tmp_path)` gives small batches and spills +`Resources(host_memory="1MB", scratch_dir=tmp_path)` gives small batches and spills the environments of a depth-4 stack with bonds of 4 to `tmp_path` (see `tests/test_stack.py`). Compare the dense operator of the result with that of a default run, not the cores: batching changes the rounding, and with it the cores of diff --git a/docs/content/docs/features/cutoff.mdx b/docs/content/docs/features/cutoff.mdx index 22c3748..fee62ca 100644 --- a/docs/content/docs/features/cutoff.mdx +++ b/docs/content/docs/features/cutoff.mdx @@ -1,6 +1,6 @@ --- title: Adaptive truncation -description: Trim bonds below chi_out with a relative singular-value cutoff. +description: Trim bonds below chi_out with a relative cutoff, and oversample the sketch. --- `chi_out` is the maximum output bond dimension. Set `cutoff` to trim each bond @@ -21,22 +21,41 @@ print([a.shape[0] for a in fixed[1:]]) print([a.shape[0] for a in trimmed[1:]]) ``` -By default, `apply`, `compress` and `src` use sketch size -`max(ceil(1.5 * chi_out), chi_out + 10)`, following the original paper's -oversampling recommendation. A final left-to-right SVD sweep truncates the -right-canonical sketch to `chi_out`, discarding singular values below -`cutoff * sigma_max` at each cut. These are the Schmidt values of the intermediate -train as truncated so far, not of the random sketch matrix. The result is -**left-canonical**, with the norm on the last site. Each cut is truncated optimally -for the train as truncated so far; the discarded weights add in quadrature. - -Set `oversample` to a finite multiplicative factor of at least `1.0` to use sketch -size `ceil(oversample * chi_out)`. Larger sketches cost more time and memory. -`oversample=1.0` disables oversampling and the final SVD pass, recovering the -original **right-canonical** output. In that mode, `cutoff` acts on the sketch's -QR triangular factor during the right-to-left sweep, not on Schmidt values. +By default, the sketch size is `chi_out`. During the right-to-left sweep, the QR +factorisation at each site is followed by an SVD of its small triangular factor, +and singular values below `cutoff * sigma_max` are discarded. The output is +**right-canonical**. + +## Oversampling + +Sketching at exactly `chi_out` lets the error of the randomized range finder +dominate. Set `oversample` to a finite factor of at least `1.0` to sketch at +`ceil(oversample * chi_out)`, or to `None` for the paper's recommendation, +`max(ceil(1.5 * chi_out), chi_out + 10)`: + +```python +import quimb.tensor as qtn + +from src_method import apply + +H = qtn.MPO_rand(8, bond_dim=4, phys_dim=2, seed=0) +psi = qtn.MPS_rand_state(8, bond_dim=16, seed=0) + +out = apply(H.arrays, psi.arrays, chi_out=32, oversample=None, seed=1) +``` + +A final left-to-right SVD sweep then truncates the oversampled, right-canonical +intermediate to `chi_out`, discarding singular values below `cutoff * sigma_max` +at each cut. These are the Schmidt values of the intermediate train as truncated +so far, not of the random sketch matrix. The result is **left-canonical**, with +the norm on the last site. + +Oversampling improves accuracy at a fixed `chi_out` but costs more: the sweep +runs at the larger sketch size, `cutoff` no longer trims bonds during the sweep, +and the final pass adds one SVD per site. With `oversample=None`, expect a run +time several times that of the default. - `cutoff` must lie in `[0.0, 1.0)`; `0.0` (the default) disables trimming. - `chi_out` remains the upper bound: the cutoff only ever lowers a bond. -- Two-site trains use the exact fallback, which ignores `cutoff` and truncates to - `chi_out`, returning right-canonical output regardless of `oversample`. +- Two-site trains use the exact fallback, which ignores `cutoff` and `oversample`, + truncates to `chi_out` and returns right-canonical output. diff --git a/docs/content/docs/features/stacks.mdx b/docs/content/docs/features/stacks.mdx index ebb8b3a..5c3aa00 100644 --- a/docs/content/docs/features/stacks.mdx +++ b/docs/content/docs/features/stacks.mdx @@ -32,8 +32,10 @@ form. The per-site cost of the sweep grows with the sketch size squared times the *product* of the bond dimensions of all layers, so it grows exponentially with the depth of the stack. In exchange, a single sweep truncates once instead of once per -layer. [Oversampling and exact truncation](/features/cutoff) reduce the error of -the randomized sketch for both the one-shot and pairwise approaches. +layer. In practice that buys at most a moderate gain in accuracy: the error of +the randomized sketch dominates the error compounded across layers. +[Oversampling](/features/cutoff#oversampling) reduces that error for both the +one-shot and pairwise approaches. Compressing a whole stack at once pays off for shallow stacks of thin layers, such as two or three Trotter layers; apply anything else pairwise. The [stack diff --git a/src/src_method/apply.py b/src/src_method/apply.py index 46778e9..1c4a4e4 100644 --- a/src/src_method/apply.py +++ b/src/src_method/apply.py @@ -29,7 +29,7 @@ def apply( chi_out: int | np.integer, *, cutoff: float = 0.0, - oversample: float | None = None, + oversample: float | None = 1.0, dtype: DTypeLike | None = None, seed: int | None = None, device: str = "cpu", @@ -50,15 +50,12 @@ def apply( right_tensor: The site arrays of the right tensor network (MPO or MPS). chi_out: The desired maximum bond dimension of the output tensor network. cutoff: Relative singular-value cutoff for adaptive bond truncation. - Discards singular values below ``cutoff * sigma_max`` in the final - SVD pass, or on the sketch's QR factor when ``oversample=1.0``. - Set to 0.0 (default) to disable cutoff trimming. Ignored for two sites. - oversample: Multiplicative sketch-size factor, finite and at least 1.0. - Defaults to ``None``, giving sketch size - ``max(ceil(1.5 * chi_out), chi_out + 10)``; an explicit factor gives - ``ceil(oversample * chi_out)``. Oversampling adds an exact SVD pass - to truncate to ``chi_out``. Set to ``1.0`` for the original behavior. - Ignored for two sites. + Discards singular values below ``cutoff * sigma_max``; see + `src_method.stack.src`. Set to 0.0 (default) to disable cutoff + trimming. Ignored for two sites. + oversample: Multiplicative sketch-size factor, finite and at least 1.0, + or ``None`` for the paper's recommendation; see + `src_method.stack.src`. Defaults to 1.0, no oversampling. dtype: Data type of the random sketches. Defaults to the promoted floating dtype of the inputs; an explicit dtype can promote the result. seed: An optional seed for the random number generator. @@ -68,8 +65,8 @@ def apply( Returns: The site arrays of the compressed tensor network (MPS or MPO), in - left-canonical form with oversampling, otherwise right-canonical form. - The two-site fallback is always right-canonical. + right-canonical form, or left-canonical form when oversampling (three + sites or more). Raises: TypeError: If ``chi_out`` is not an integer or the combination of input diff --git a/src/src_method/compress.py b/src/src_method/compress.py index f645e45..605c1f8 100644 --- a/src/src_method/compress.py +++ b/src/src_method/compress.py @@ -27,7 +27,7 @@ def compress( chi_out: int | np.integer, *, cutoff: float = 0.0, - oversample: float | None = None, + oversample: float | None = 1.0, dtype: DTypeLike | None = None, seed: int | None = None, device: str = "cpu", @@ -46,15 +46,12 @@ def compress( tensor: The site arrays of the tensor network to compress (MPS or MPO). chi_out: The desired maximum bond dimension of the output tensor network. cutoff: Relative singular-value cutoff for adaptive bond truncation. - Discards singular values below ``cutoff * sigma_max`` in the final - SVD pass, or on the sketch's QR factor when ``oversample=1.0``. - Set to 0.0 (default) to disable cutoff trimming. Ignored for two sites. - oversample: Multiplicative sketch-size factor, finite and at least 1.0. - Defaults to ``None``, giving sketch size - ``max(ceil(1.5 * chi_out), chi_out + 10)``; an explicit factor gives - ``ceil(oversample * chi_out)``. Oversampling adds an exact SVD pass - to truncate to ``chi_out``. Set to ``1.0`` for the original behavior. - Ignored for two sites. + Discards singular values below ``cutoff * sigma_max``; see + `src_method.stack.src`. Set to 0.0 (default) to disable cutoff + trimming. Ignored for two sites. + oversample: Multiplicative sketch-size factor, finite and at least 1.0, + or ``None`` for the paper's recommendation; see + `src_method.stack.src`. Defaults to 1.0, no oversampling. dtype: Data type of the random sketches. Defaults to the promoted floating dtype of the inputs; an explicit dtype can promote the result. seed: An optional seed for the random number generator. @@ -64,8 +61,8 @@ def compress( Returns: The site arrays of the compressed tensor network (MPS or MPO), in - left-canonical form with oversampling, otherwise right-canonical form. - The two-site fallback is always right-canonical. + right-canonical form, or left-canonical form when oversampling (three + sites or more). Raises: TypeError: If ``chi_out`` is not an integer or the input tensor type is diff --git a/src/src_method/stack.py b/src/src_method/stack.py index 28b1ecb..5c15e95 100644 --- a/src/src_method/stack.py +++ b/src/src_method/stack.py @@ -41,7 +41,7 @@ def src( *trains: Sequence[Site], chi_out: int | np.integer, cutoff: float = 0.0, - oversample: float | None = None, + oversample: float | None = 1.0, dtype: DTypeLike | None = None, seed: int | None = None, device: str = "cpu", @@ -77,16 +77,17 @@ def src( HDF5 dataset; it is then read only when the sweep reaches it. chi_out: The desired maximum bond dimension of the output train. cutoff: Relative singular-value cutoff for adaptive bond truncation. - Discards singular values below ``cutoff * sigma_max`` in the final - left-to-right SVD pass. With ``oversample=1.0``, applies instead to - the sketch's QR factor during the right-to-left sweep, as before. - Set to 0.0 (default) to disable cutoff trimming. Ignored for two-site - stacks, which are contracted and truncated to ``chi_out`` exactly. + Discards singular values below ``cutoff * sigma_max``: on the + sketch's QR factor during the right-to-left sweep or, when + oversampling, in the final left-to-right SVD pass. Set to 0.0 + (default) to disable cutoff trimming. Ignored for two-site stacks, + which are contracted and truncated to ``chi_out`` exactly. oversample: Multiplicative sketch-size factor, finite and at least 1.0. - The sketch size is ``ceil(oversample * chi_out)``. By default (``None``), - uses ``max(ceil(1.5 * chi_out), chi_out + 10)``. An oversampled sweep - is followed by an exact SVD truncation to ``chi_out``; ``1.0`` recovers - the original sweep without the final pass. Ignored for two-site stacks. + The sketch size is ``ceil(oversample * chi_out)``; ``None`` uses the + paper's ``max(ceil(1.5 * chi_out), chi_out + 10)``. An oversampled + sweep is followed by an exact SVD truncation to ``chi_out``, which + improves accuracy at a higher cost. Defaults to 1.0, no + oversampling. Ignored for two-site stacks. dtype: Data type of the random sketches. Defaults to the promoted floating dtype of the inputs, so single precision stays single and complex inputs get complex Ginibre sketches. An explicit dtype @@ -100,8 +101,8 @@ def src( Returns: The site arrays of the compressed train (MPS or MPO), as numpy arrays - (host-side, whatever the ``device``). Oversampling returns left-canonical - form; ``oversample=1.0`` and the two-site fallback return right-canonical form. + (host-side, whatever the ``device``), in right-canonical form, or + left-canonical form when oversampling (three sites or more). Raises: TypeError: If ``chi_out`` is not an integer, if a train has an unrecognised diff --git a/tests/test_gpu_backend.py b/tests/test_gpu_backend.py index 0faece2..0e851ba 100644 --- a/tests/test_gpu_backend.py +++ b/tests/test_gpu_backend.py @@ -266,16 +266,15 @@ def test_gpu_tiers_and_batches_match_cpu(tmp_path, plans: list[Plan]) -> None: rng = np.random.default_rng(7) stack = [random_mpo_arrays([4, 8, 8, 8, 4], rng) for _ in range(4)] chi = 64 - chi_sketch = 96 shapes, sizes = padded_shapes(stack), site_bytes(stack) roomy = make_plan( shapes, sizes, - chi_sketch, + chi, np.complex128, Budgets(10**10, 10**10, 10**12, tmp_path, unified=False), ) - env = chi_sketch * 8**4 * 16 # one bulk environment + env = chi * 8**4 * 16 # one bulk environment tight = Resources( gpu_memory=roomy.device_peak // 2, host_memory=roomy.host_peak + 5 * env, @@ -288,7 +287,7 @@ def test_gpu_tiers_and_batches_match_cpu(tmp_path, plans: list[Plan]) -> None: ) assert {site.tier for site in plans[-1].sites} == {"device", "host", "disk"} - assert min(site.sketch_batch for site in plans[-1].sites[1:]) < chi_sketch + assert min(site.sketch_batch for site in plans[-1].sites[1:]) < chi reference = dense_mpo(cpu) error = np.linalg.norm(dense_mpo(gpu) - reference) / np.linalg.norm(reference) assert error < 1e-10 @@ -369,7 +368,6 @@ def test_two_gpus_match_one( rng = np.random.default_rng(7) stack = [random_mpo_arrays([4, 8, 8, 8, 4], rng) for _ in range(4)] chi = 63 - chi_sketch = 95 if uploads == "full": monkeypatch.setattr( "src_method._group.enable_peer_access", lambda _xp, _devices: False @@ -378,16 +376,17 @@ def test_two_gpus_match_one( roomy = make_plan( shapes, sizes, - chi_sketch, + chi, np.complex128, Budgets(10**10, 10**10, 10**12, tmp_path, unified=False), devices=2, ) - env = (chi_sketch // 2 + 1) * 8**4 * 16 # one bulk environment on one GPU - # Checked with make_plan: the oversampled environments use every tier. + env = (chi // 2 + 1) * 8**4 * 16 # one bulk environment on one GPU + # Checked with make_plan: one environment per GPU on the device, two in host + # memory, the oldest two on disk. tight = Resources( gpu_memory=roomy.device_peak - 2 * env, - host_memory=roomy.host_peak + 10 * env, + host_memory=roomy.host_peak + 12 * env, scratch_dir=tmp_path, devices=2, ) diff --git a/tests/test_stack.py b/tests/test_stack.py index e98aecb..d38657d 100644 --- a/tests/test_stack.py +++ b/tests/test_stack.py @@ -275,6 +275,20 @@ def spy(*args, **kwargs): assert calls == [(expected, chi if expected > chi else None)] +def test_default_does_not_oversample(rng, monkeypatch): + calls = [] + original = stack_module.sweep + + def spy(*args, **kwargs): + calls.append((args[2], kwargs["truncate_to"])) + return original(*args, **kwargs) + + monkeypatch.setattr(stack_module, "sweep", spy) + src(*make_stack("A B psi", rng), chi_out=3, seed=0) + + assert calls == [(3, None)] + + @pytest.mark.parametrize( "oversample", [0.0, -1.0, 0.99, np.nan, np.inf, -np.inf, True, "2", 2j] ) @@ -303,9 +317,9 @@ def test_oversampling_improves_accuracy(spec): ] ) - legacy, doubled, default = np.median(errors, axis=0) + legacy, doubled, recommended = np.median(errors, axis=0) assert doubled < legacy - assert default < legacy + assert recommended < legacy @pytest.mark.parametrize("cutoff", [0.0, 0.1]) @@ -333,7 +347,7 @@ def test_oversampled_cutoff_uses_schmidt_values(kind, scale): sites = [site[..., None] for site in sites] before = [site.copy() for site in sites] - out = compress(sites, chi_out=2, cutoff=0.1, seed=0) + out = compress(sites, chi_out=2, cutoff=0.1, oversample=None, seed=0) assert all(site.shape[0] == 1 for site in out[1:]) reference = np.zeros_like(dense(sites)) @@ -348,7 +362,7 @@ def test_oversampled_zero_train(kind): if kind == "mpo": sites = [site[..., None] for site in sites] - out = compress(sites, chi_out=2, cutoff=0.1, seed=0) + out = compress(sites, chi_out=2, cutoff=0.1, oversample=None, seed=0) assert all(site.shape[0] >= 1 for site in out[1:]) np.testing.assert_array_equal(dense(out), dense(sites)) @@ -362,7 +376,7 @@ def test_oversampled_output_does_not_retain_discarded_columns(kind, rng): else random_mpo([4, 4, 4], rng, up=4, down=8) ) - out = compress(train, chi_out=1, seed=0) + out = compress(train, chi_out=1, oversample=None, seed=0) for site in out[:-1]: storage = site @@ -549,14 +563,14 @@ def spy(*args, **kwargs): def test_tiny_budget_batches_and_spills(rng, tmp_path, plans): stack = [random_mpo([3, 4, 4, 4, 3], rng) for _ in range(4)] - tight = Resources(host_memory="1.5MB", scratch_dir=tmp_path) + tight = Resources(host_memory="1MB", scratch_dir=tmp_path) default = src(*stack, chi_out=16, seed=3, dtype=np.complex128) spilled = src(*stack, chi_out=16, seed=3, dtype=np.complex128, resources=tight) assert all(site.tier == "device" for site in plans[0].sites) assert "disk" in {site.tier for site in plans[1].sites} - assert min(site.sketch_batch for site in plans[1].sites[1:]) < 26 + assert min(site.sketch_batch for site in plans[1].sites[1:]) < 16 # Rounding differs, and the output cores of an ill-conditioned sketch with it, # but not the operator they represent. assert rel_error(spilled, dense(default)) < 1e-10 @@ -675,7 +689,7 @@ def test_devices_with_cutoff_keep_the_bonds(rng): def test_devices_spill_and_clean_up(rng, tmp_path, plans): stack = [random_mpo([3, 4, 4, 4, 3], rng) for _ in range(4)] - tight = Resources(host_memory="2MB", scratch_dir=tmp_path, devices=2) + tight = Resources(host_memory="1MB", scratch_dir=tmp_path, devices=2) default = src(*stack, chi_out=16, seed=3, dtype=np.complex128) spilled = src(*stack, chi_out=16, seed=3, dtype=np.complex128, resources=tight)