Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion CONTRIBUTING.md
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@ comment on that thread and applies to all your future contributions.

Open an issue with:

- the versions of `src_method`, `numpy`, `quimb` and (if relevant) `cupy`,
- the versions of `src_method`, `numpy` and (if relevant) `cupy`,
- a minimal reproducer, ideally with a fixed `seed=`,
- the observed and expected behaviour.

Expand Down
14 changes: 11 additions & 3 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@ The following primitives are supported:
3. MPO randomized compression.
4. MPS randomized compression.

The package is designed to work with [Quimb](https://quimb.readthedocs.io/en/latest/autoapi/quimb/tensor/index.html) tensor network objects, as such, the API loosely follows its naming conventions:
`src_method` has no tensor-network framework dependency: it takes and returns plain lists of per-site NumPy arrays, one array per site.

```python
from src_method import apply, compress
Expand All @@ -25,11 +25,19 @@ The `apply` function covers cases 1 and 2 above, while the `compress` function c
manage the assignment of the returned objects, possibly overwriting the input variables.
See the [reference documentation](algorithmiq.github.io/src_method/) for details, and the [tests](tests/) or [benchmarks](benches/) folders for usage examples.

**NOTE**: the current implementation targets tensor networks with 3 or more sites. For smaller networks, the corresponding Quimb primitive with randomized SVD is dispatched, with a warning.
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.

### Tensor Indexing Conventions

This library follows the default `quimb` 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:

```python
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.
Expand Down
2 changes: 1 addition & 1 deletion SECURITY.md
Original file line number Diff line number Diff line change
Expand Up @@ -23,4 +23,4 @@ arrays supplied by the caller and does not parse untrusted input formats,
open network connections, or execute user-supplied code. The most likely
security-relevant issues are therefore memory-safety problems surfaced through
the optional CuPy backend, or dependency vulnerabilities. Reports about the
behaviour of `numpy`, `quimb` or `cupy` themselves should go to those projects.
behaviour of `numpy` or `cupy` themselves should go to those projects.
6 changes: 5 additions & 1 deletion benches/primitives/leonardo/bench_mpo_mpo.py
Original file line number Diff line number Diff line change
Expand Up @@ -78,7 +78,11 @@ def main(
if run == "src":
logger.info("Computing SRC's MPO-MPO contraction (with compression)...")
tms = perf_counter_ns()
H_src = apply(H1, H2, chi_out=chi_out, dtype=array_type)
# src_method takes and returns plain lists of site arrays; quimb is only
# used here to build the inputs and to measure the distance.
H_src = qtn.MatrixProductOperator(
apply(H1.arrays, H2.arrays, chi_out=chi_out, dtype=array_type)
)
tms = perf_counter_ns() - tms
logger.info(" SRC's contraction-compression took %s s", tms * 1e-9)
if compare == "yes":
Expand Down
14 changes: 11 additions & 3 deletions docs/index.md
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@ The following primitives are supported:
3. MPO randomized compression.
4. MPS randomized compression.

The package is designed to work with [Quimb](https://quimb.readthedocs.io/en/latest/autoapi/quimb/tensor/index.html) tensor network objects, as such, the API loosely follows its naming conventions:
`src_method` has no tensor-network framework dependency: it takes and returns plain lists of per-site NumPy arrays, one array per site.

```python
from src_method import apply, compress
Expand All @@ -23,11 +23,19 @@ The `apply` function covers cases 1 and 2 above, while the `compress` function c
manage the assignment of the returned objects, possibly overwriting the input variables.
See the [reference documentation](algorithmiq.github.io/src_method/) for details, and the [tests](../tests/) or [benchmarks](../benches/) folders for usage examples.

**NOTE**: the current implementation targets tensor networks with 3 or more sites. For smaller networks, the corresponding Quimb primitive with randomized SVD is dispatched, with a warning.
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.

### Tensor Indexing Conventions

This library follows the default `quimb` 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:

```python
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.
Expand Down
10 changes: 8 additions & 2 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,6 @@ dependencies = [
"numpy>=2.4.4,<3",
"opt_einsum>=3.4.0",
"structlog>=25.5.0",
"quimb>=1.14.0,<2",
]


Expand Down Expand Up @@ -67,7 +66,14 @@ Discussions = "https://github.com/Algorithmiq/src-method/discussions"


[dependency-groups]
test = ["pytest-cov>=4.0", "pytest>=9.1.0,<10.0", "pytest-benchmark>=5.2.3"]
# quimb is *not* a runtime dependency: it is only used to build reference
# tensor networks and to measure distances in the tests and benchmarks.
test = [
"pytest-cov>=4.0",
"pytest>=9.1.0,<10.0",
"pytest-benchmark>=5.2.3",
"quimb>=1.14.0,<2",
]
dev = ["scalene", "pre-commit>=3", "ruff==0.16.5", { include-group = "test" }]
interactive = [
"ipykernel>=7.2.0",
Expand Down
164 changes: 164 additions & 0 deletions src/src_method/_tensor_train.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,164 @@
"""Array-list tensor-train conventions and exact small-system primitives.

Tensor trains are plain lists of arrays, one per site. The index ordering
matches the default `quimb` layout, so a result can be handed straight to
``qtn.MatrixProductState(arrays)`` / ``qtn.MatrixProductOperator(arrays)``
without any permutation:

* MPS: ``(bond_r, phys)``, ``(bond_l, bond_r, phys)``, ..., ``(bond_l, phys)``
* MPO: ``(bond_r, up, down)``, ``(bond_l, bond_r, up, down)``, ...,
``(bond_l, up, down)``

The SRC sweep needs at least three sites, so two-site trains are handled here
instead. At that size the whole network fits in a single dense matrix, and one
exact SVD is both cheaper and more accurate than a randomized sketch.
"""

from __future__ import annotations

from typing import TYPE_CHECKING, Literal

import numpy as np
from opt_einsum import contract

if TYPE_CHECKING:
from collections.abc import Sequence

from numpy.typing import NDArray

# Minimum number of sites for which the randomized SRC sweep is defined.
MIN_SRC_SITES = 3

# Rank of a boundary (first / last) site tensor, which identifies the train type.
_MPS_BOUNDARY_NDIM = 2
_MPO_BOUNDARY_NDIM = 3

TrainKind = Literal["mps", "mpo"]

__all__ = [
"MIN_SRC_SITES",
"TrainKind",
"exact_apply",
"exact_compress",
"infer_kind",
]


def infer_kind(arrays: Sequence[NDArray]) -> TrainKind | None:
"""Classify a tensor train from the rank of its first site tensor.

A boundary site carries one bond index plus either a single physical
index (MPS) or an upper/lower pair (MPO), so the rank is unambiguous.

Args:
arrays: The site tensors of the train.

Returns:
``"mps"``, ``"mpo"``, or ``None`` if the layout is unrecognised.
"""
if len(arrays) == 0:
return None
ndim = np.ndim(arrays[0])
if ndim == _MPS_BOUNDARY_NDIM:
return "mps"
if ndim == _MPO_BOUNDARY_NDIM:
return "mpo"
return None


def exact_compress(
arrays: Sequence[NDArray], chi_out: int, kind: TrainKind
) -> list[NDArray]:
"""Compress a two-site train exactly via a single truncated SVD.

Args:
arrays: The two site tensors of the train.
chi_out: The maximum bond dimension to keep.
kind: Whether the train is an ``"mps"`` or an ``"mpo"``.

Returns:
The compressed train, in right-canonical form.

Raises:
ValueError: If the train does not have exactly two sites.
"""
_check_pair(arrays)
if kind == "mps":
# (b, p0) x (b, p1) -> (p0, p1)
theta = contract("ab,ac->bc", arrays[0], arrays[1])
left, right = _truncated_svd(theta, chi_out)
return [left.T, right]

# (b, u0, d0) x (b, u1, d1) -> (u0, d0, u1, d1)
theta = contract("aij,akl->ijkl", arrays[0], arrays[1])
up_l, down_l, up_r, down_r = theta.shape
left, right = _truncated_svd(theta.reshape(up_l * down_l, up_r * down_r), chi_out)
rank = left.shape[1]
return [
left.reshape(up_l, down_l, rank).transpose(2, 0, 1),
right.reshape(rank, up_r, down_r),
]


def exact_apply(
left_tensor: Sequence[NDArray],
right_tensor: Sequence[NDArray],
chi_out: int,
kind: TrainKind,
) -> list[NDArray]:
"""Contract and compress two two-site trains exactly.

The MPO on the left is contracted site-wise with the right train, fusing
the two bond indices, and the result is compressed with a single SVD.

Args:
left_tensor: The two site tensors of the left MPO.
right_tensor: The two site tensors of the right MPS or MPO.
chi_out: The maximum bond dimension to keep.
kind: Whether ``right_tensor`` is an ``"mps"`` or an ``"mpo"``.

Returns:
The compressed product, in right-canonical form.

Raises:
ValueError: If either train does not have exactly two sites.
"""
_check_pair(left_tensor)
_check_pair(right_tensor)
if kind == "mps":
# Contract the MPO lower leg with the MPS physical leg, fusing both bonds.
product = [
contract("aij,bj->abi", left_tensor[i], right_tensor[i]).reshape(
-1, left_tensor[i].shape[1]
)
for i in range(2)
]
else:
product = [
contract("aij,bjk->abik", left_tensor[i], right_tensor[i]).reshape(
-1, left_tensor[i].shape[1], right_tensor[i].shape[2]
)
for i in range(2)
]
return exact_compress(product, chi_out, kind)


def _truncated_svd(theta: NDArray, chi_out: int) -> tuple[NDArray, NDArray]:
"""Split a matrix as ``(U @ diag(S), Vh)``, keeping at most ``chi_out`` values."""
U, S, Vh = np.linalg.svd(theta, full_matrices=False)

Check warning on line 148 in src/src_method/_tensor_train.py

View check run for this annotation

SonarQubeCloud / SonarCloud Code Analysis

Rename this local variable "Vh" to match the regular expression ^[_a-z][a-z0-9_]*$.

See more on https://sonarcloud.io/project/issues?id=Algorithmiq_src-method&issues=AaBd6PE6qQQm2BGb4_bW&open=AaBd6PE6qQQm2BGb4_bW&pullRequest=24
rank = min(chi_out, S.size)
return U[:, :rank] * S[:rank], Vh[:rank]
Comment thread
Panadestein marked this conversation as resolved.


def _check_pair(arrays: Sequence[NDArray]) -> None:
"""Reject trains that the exact two-site path cannot handle.

Raises:
ValueError: If the train does not have exactly two sites.
"""
if len(arrays) != 2:
msg = (
f"Expected a two-site tensor train, got {len(arrays)} site(s). "
"Single-site trains are degenerate; use three or more sites for SRC."
)
raise ValueError(msg)
Loading
Loading