An implementation of the SRC algorithm introduced in Camaño, Epperly and Tropp, Quantum 10, 2022 (2026), extending the idea to other kinds of tensor networks.
The following primitives are supported:
- MPO-MPS randomized contraction-compression.
- MPO-MPO randomized contraction-compression.
- MPO randomized compression.
- MPS randomized compression.
- Randomized contraction-compression of a whole stack of trains in one sweep:
MPO^k,MPO^k . MPSandMPS . MPO^k.
src_method has no tensor-network framework dependency: it takes and returns plain lists of per-site NumPy arrays, one array per site.
from src_method import apply, compress, srcThe apply function covers cases 1 and 2 above, the compress function cases 3 and 4, and src all five: apply and compress are its two- and one-train special cases. All three functions are pure, meaning no in-place modification ever happens. The user should
manage the assignment of the returned objects, possibly overwriting the input variables.
See the reference documentation for details, and the tests or benchmarks folders for usage examples.
Whether a train is an MPS or an MPO is inferred from the rank of its first site tensor, so no wrapper type is needed.
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.
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.
The array layout follows the default quimb tensor indexing conventions, so results round-trip through Quimb without any permutation:
import quimb.tensor as qtn
result = qtn.MatrixProductOperator(apply(H1.arrays, H2.arrays, chi_out=64))-
MPO Tensors: Bulk tensors have index order
('l', 'r', 'u', 'd'). Boundary tensors (at the edges) are rank-3, dropping the outer'l'or'r'index. -
MPS Tensors: Bulk tensors have index order
('l', 'r', 'u'). Boundary tensors are rank-2, dropping the outer bond index.
Where 'l'/'r' are left/right virtual bonds and 'u'/'d' are the upper/lower physical legs.
Please keep this in mind when constructing or manipulating tensors directly.
src contracts a stack of trains and compresses the result in a single sweep:
from src_method import src
state = src(U3, U2, U1, psi, chi_out=64) # U3 U2 U1 |psi>, U1 acts firstThe stack is written in mathematical order. Each contraction joins the 'd' leg of a train with the 'u' leg (or the physical leg of an MPS) of the train to its right:
| Stack | Contraction | Result |
|---|---|---|
src(A), src(psi) |
none | compressed MPO or MPS |
src(A, B, ...) |
A.d with B.u |
MPO |
src(A, ..., psi) |
A.d with psi |
MPS on the 'u' leg of A (a ket) |
src(phi, A, ...) |
phi with A.u |
MPS on the 'd' leg of the last MPO (a bra) |
An MPS may appear only first or last, and not both. apply(A, B) and compress(A) are the two- and one-train cases; apply accepts only an MPO on the left.
A leading MPS is a row vector used without conjugation: src(phi, A, B) computes phiᵀ A B, which equals src(Bᵀ, Aᵀ, phi) with ᵀ swapping the 'u' and 'd' legs. For the physical bra <psi| A B, conjugate first:
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 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 for measurements.
src_method runs on CPU (NumPy) by default and can be accelerated on GPUs via CuPy. Both NVIDIA (CUDA) and AMD (ROCm) GPUs are supported through optional install extras.
At runtime, pass device="gpu" to use GPU acceleration. The library handles backend dispatch automatically.
Stacks that do not fit on the GPU, or in host memory, still run: the sweep reads the input cores one site at a time (from NumPy arrays, np.memmap, zarr or HDF5 datasets), batches every contraction to a memory budget, and keeps the sketched environments on the GPU, in host memory or on local disk. The budgets are detected, or set explicitly with Resources:
from src_method import Resources, src
out = src(
N,
V,
M,
U,
chi_out=2000,
device="gpu",
resources=Resources(gpu_memory="36GB", scratch_dir="/local/scratch"),
)On a node with several GPUs, Resources(devices=4) splits the sweep among four of them, in the same process.
See Large problems for the budgets, the scratch directory, several GPUs and how to read the logged plan.
# CPU only (default)
uv pip install src_method
# With NVIDIA GPU support (CUDA 13.x, driver >= 580)
uv pip install "src_method[gpu-nvidia]"
# With AMD GPU support (ROCm)
uv pip install "src_method[gpu-rocm]"When including src_method in another project's pyproject.toml:
# CPU only
dependencies = ["src_method"]
# With NVIDIA GPU support
dependencies = ["src_method[gpu-nvidia]"]
# With AMD GPU support
dependencies = ["src_method[gpu-rocm]"]The repository ships a flake.nix that provides
a development shell with uv, Git and the GitHub CLI, plus the native
libraries the binary wheels need at runtime. Python itself and all project
dependencies remain managed by uv.
With flakes enabled, run:
nix developEntering the shell runs uv sync --all-groups and activates the uv project
environment ($UV_PROJECT_ENVIRONMENT if set, otherwise .venv), so you land
in a ready-to-use environment. GPU extras are not installed by the flake: add them
explicitly with uv sync --all-groups --extra gpu-nvidia (or --all-groups --extra gpu-rocm) on a machine with the matching drivers.
If you use direnv, the provided .envrc enters the shell automatically:
direnv allowUnlike the DevContainer, the Nix shell does not install the git hooks for
you. Run prek install --prepare-hooks once after the first nix develop.
The documentation is a Fumadocs site under docs/, published at this link.
It covers concepts, features, tutorials, the API reference generated from the
docstrings, and the developer guide.