FazBrowse GitHub Viewer | Trending |
URL:
| Home
Tools: [Download Repo ZIP]   [Original HTTPS Page]

feat: ✨ oversample SRC sketches before exact SVD truncation by Copilot · Pull Request #63 · Algorithmiq/src-method · GitHub

Repository navigation

18 changes: 17 additions & 1 deletion README.md
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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

Expand Down
32 changes: 23 additions & 9 deletions docs/content/docs/concepts/algorithm.mdx
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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.
Expand All @@ -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$.
Expand Down
48 changes: 38 additions & 10 deletions docs/content/docs/features/cutoff.mdx
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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.
4 changes: 3 additions & 1 deletion docs/content/docs/features/stacks.mdx
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
48 changes: 44 additions & 4 deletions src/src_method/_sweep.py
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,7 @@
release_cached_blocks,
to_numpy,
truncated_qr,
use_device,
)

if TYPE_CHECKING:
Expand Down Expand Up @@ -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]:
Expand All @@ -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
Expand Down Expand Up @@ -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),
Expand All @@ -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.

Expand Down
35 changes: 34 additions & 1 deletion src/src_method/_validation.py
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
from __future__ import annotations

import numbers
from math import ceil, isfinite
from typing import TYPE_CHECKING

import numpy as np
Expand All @@ -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:
Expand Down
20 changes: 13 additions & 7 deletions src/src_method/apply.py
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand All @@ -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.
Expand All @@ -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.
Expand All @@ -88,6 +93,7 @@ def apply(
right_tensor,
chi_out=chi_out,
cutoff=cutoff,
oversample=oversample,
dtype=dtype,
seed=seed,
device=device,
Expand Down
Loading
Loading

Back | FazBrowse Home | New Git URL