diff --git a/README.md b/README.md index 9990826..6c0d046 100644 --- a/README.md +++ b/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 diff --git a/docs/content/docs/concepts/algorithm.mdx b/docs/content/docs/concepts/algorithm.mdx index 082d4ab..ebc6040 100644 --- a/docs/content/docs/concepts/algorithm.mdx +++ b/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$. diff --git a/docs/content/docs/features/cutoff.mdx b/docs/content/docs/features/cutoff.mdx index cd8a021..fee62ca 100644 --- a/docs/content/docs/features/cutoff.mdx +++ b/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. diff --git a/docs/content/docs/features/stacks.mdx b/docs/content/docs/features/stacks.mdx index 1806e4e..5c3aa00 100644 --- a/docs/content/docs/features/stacks.mdx +++ b/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 diff --git a/src/src_method/_sweep.py b/src/src_method/_sweep.py index 114f1f9..e9b0be4 100644 --- a/src/src_method/_sweep.py +++ b/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. 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 76ff40e..1c4a4e4 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 = 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, diff --git a/src/src_method/compress.py b/src/src_method/compress.py index a2acb7e..605c1f8 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 = 1.0, dtype: DTypeLike | None = None, seed: int | None = None, device: str = "cpu", @@ -45,11 +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. - 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. @@ -58,13 +60,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 + 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 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 +78,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 3452ce7..5c15e95 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 = 1.0, 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,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. - 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``: 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)``; ``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 @@ -93,14 +100,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``), 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 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 +117,7 @@ def src( """ chi_out = validate_chi_out(chi_out) validate_cutoff(cutoff) + chi_sketch = sketch_size(chi_out, oversample) get_xp(device) # Reject an unknown device, or a missing cupy, at once. prng = default_rng(seed) layers, kind = normalize_stack(trains) @@ -128,10 +138,11 @@ def src( result = sweep( layers, kind, - chi_out, + chi_sketch, prng, device=device, 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..0e851ba 100644 --- a/tests/test_gpu_backend.py +++ b/tests/test_gpu_backend.py @@ -165,6 +165,31 @@ 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_matches_cpu( + oversample: float | None, kind: str +) -> None: + """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) + 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 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.""" diff --git a/tests/test_stack.py b/tests/test_stack.py index 4bce1c1..d38657d 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,171 @@ 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), + 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)] + + +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] +) +@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, recommended = np.median(errors, axis=0) + assert doubled < legacy + assert recommended < 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, oversample=None, 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, oversample=None, seed=0) + + assert all(site.shape[0] >= 1 for site in out[1:]) + 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, oversample=None, 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] @@ -250,10 +419,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 +460,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: