feat: ✨ oversample SRC sketches before exact SVD truncation

README.md

@@ -33,6 +33,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
+
+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
 
 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:
@@ -81,7 +97,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
 

docs/content/docs/concepts/algorithm.mdx

@@ -9,7 +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 the output bond dimension $\chi_{\text{out}}$ directly.
+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
 
@@ -22,12 +24,20 @@ 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}} = \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
-   $\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 +47,20 @@ 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 (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
 
 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$.

docs/content/docs/features/cutoff.mdx

@@ -1,11 +1,10 @@
 ---
 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.
 ---
 
-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,41 @@ 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, 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`.
+- Two-site trains use the exact fallback, which ignores `cutoff` and `oversample`,
+  truncates to `chi_out` and returns right-canonical output.

docs/content/docs/features/stacks.mdx

@@ -29,11 +29,13 @@ 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.
+[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

src/src_method/_sweep.py

@@ -49,6 +49,7 @@
     release_cached_blocks,
     to_numpy,
     truncated_qr,
+    use_device,
 )
 
 if TYPE_CHECKING:
@@ -164,6 +165,7 @@ def sweep(
     *,
     device: str = "cpu",
     cutoff: float = 0.0,
+    truncate_to: int | None = None,
     dtype: DTypeLike,
     resources: Resources | None = None,
 ) -> list[np.ndarray]:
@@ -175,18 +177,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.
         device: ``"cpu"`` or ``"gpu"``, as in `src_method.stack.src`.
         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
@@ -220,7 +225,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),
@@ -230,9 +235,44 @@ 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:
+        xp = group.xp
+        with (
+            use_device(xp, group.devices[0]),
+            device_pool_limit(xp, budgets.device_cap),
+        ):
+            eta = _truncate(eta, truncate_to, cutoff, group)
     return [to_numpy(site) for site in unpad(eta, kind)]
 
 
+def _truncate(
+    sites: list[np.ndarray], chi_out: int, cutoff: float, group: DeviceGroup
+) -> list[np.ndarray]:
+    """Truncate a padded right-canonical train, staging one centre at a time."""
+    xp = group.xp
+    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)
+        # 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((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.reshape(left, up, down, rank).transpose(0, 3, 1, 2))
+        )
+        centre = xp.tensordot(
+            xp.asarray(S[:rank, None] * Vh[:rank]), xp.asarray(sites[j + 1]), axes=1
+        )
+    sites[-1] = to_numpy(centre)
+    return sites
+
+
 class _Sketches:
     """The Gaussian tensors of the left-to-right pass, shared by the devices.
 

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:

src/src_method/apply.py

@@ -29,6 +29,7 @@ def apply(
     chi_out: int | np.integer,
     *,
     cutoff: float = 0.0,
+    oversample: float | None = 1.0,
     dtype: DTypeLike | None = None,
     seed: int | None = None,
     device: str = "cpu",
@@ -49,11 +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.
-            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``; 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.
@@ -62,13 +64,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
+        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
             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 +93,7 @@ def apply(
         right_tensor,
         chi_out=chi_out,
         cutoff=cutoff,
+        oversample=oversample,
         dtype=dtype,
         seed=seed,
         device=device,