|
| 1 | +"""Array-list tensor-train conventions and exact small-system primitives. |
| 2 | +
|
| 3 | +Tensor trains are plain lists of arrays, one per site. The index ordering |
| 4 | +matches the default `quimb` layout, so a result can be handed straight to |
| 5 | +``qtn.MatrixProductState(arrays)`` / ``qtn.MatrixProductOperator(arrays)`` |
| 6 | +without any permutation: |
| 7 | +
|
| 8 | +* MPS: ``(bond_r, phys)``, ``(bond_l, bond_r, phys)``, ..., ``(bond_l, phys)`` |
| 9 | +* MPO: ``(bond_r, up, down)``, ``(bond_l, bond_r, up, down)``, ..., |
| 10 | + ``(bond_l, up, down)`` |
| 11 | +
|
| 12 | +The SRC sweep needs at least three sites, so two-site trains are handled here |
| 13 | +instead. At that size the whole network fits in a single dense matrix, and one |
| 14 | +exact SVD is both cheaper and more accurate than a randomized sketch. |
| 15 | +""" |
| 16 | + |
| 17 | +from __future__ import annotations |
| 18 | + |
| 19 | +from typing import TYPE_CHECKING, Literal |
| 20 | + |
| 21 | +import numpy as np |
| 22 | +from opt_einsum import contract |
| 23 | + |
| 24 | +from .utils import to_numpy |
| 25 | + |
| 26 | +if TYPE_CHECKING: |
| 27 | + from collections.abc import Sequence |
| 28 | + |
| 29 | + from numpy.typing import NDArray |
| 30 | + |
| 31 | +# Minimum number of sites for which the randomized SRC sweep is defined. |
| 32 | +MIN_SRC_SITES = 3 |
| 33 | + |
| 34 | +# The only sub-``MIN_SRC_SITES`` size the exact path can handle. |
| 35 | +_EXACT_SITES = 2 |
| 36 | + |
| 37 | +# Rank of a boundary (first / last) site tensor, which identifies the train type. |
| 38 | +_MPS_BOUNDARY_NDIM = 2 |
| 39 | +_MPO_BOUNDARY_NDIM = 3 |
| 40 | + |
| 41 | +TrainKind = Literal["mps", "mpo"] |
| 42 | + |
| 43 | +__all__ = [ |
| 44 | + "MIN_SRC_SITES", |
| 45 | + "TrainKind", |
| 46 | + "check_exact_supported", |
| 47 | + "exact_apply", |
| 48 | + "exact_compress", |
| 49 | + "infer_kind", |
| 50 | +] |
| 51 | + |
| 52 | + |
| 53 | +def infer_kind(arrays: Sequence[NDArray]) -> TrainKind | None: |
| 54 | + """Classify a tensor train from the rank of its first site tensor. |
| 55 | +
|
| 56 | + A boundary site carries one bond index plus either a single physical |
| 57 | + index (MPS) or an upper/lower pair (MPO), so the rank is unambiguous. |
| 58 | +
|
| 59 | + Args: |
| 60 | + arrays: The site tensors of the train. |
| 61 | +
|
| 62 | + Returns: |
| 63 | + ``"mps"``, ``"mpo"``, or ``None`` if the layout is unrecognised. |
| 64 | + """ |
| 65 | + if len(arrays) == 0: |
| 66 | + return None |
| 67 | + ndim = np.ndim(arrays[0]) |
| 68 | + if ndim == _MPS_BOUNDARY_NDIM: |
| 69 | + return "mps" |
| 70 | + if ndim == _MPO_BOUNDARY_NDIM: |
| 71 | + return "mpo" |
| 72 | + return None |
| 73 | + |
| 74 | + |
| 75 | +def check_exact_supported(n_sites: int) -> None: |
| 76 | + """Reject sub-``MIN_SRC_SITES`` trains the exact path cannot handle. |
| 77 | +
|
| 78 | + Called at the public boundary before the fallback is announced, so that a |
| 79 | + degenerate train raises instead of first logging a misleading warning. |
| 80 | +
|
| 81 | + Args: |
| 82 | + n_sites: The number of sites in the train. |
| 83 | +
|
| 84 | + Raises: |
| 85 | + ValueError: If the train does not have exactly two sites. |
| 86 | + """ |
| 87 | + if n_sites != _EXACT_SITES: |
| 88 | + msg = ( |
| 89 | + f"Expected a two-site tensor train, got {n_sites} site(s). " |
| 90 | + "Single-site trains are degenerate; use three or more sites for SRC." |
| 91 | + ) |
| 92 | + raise ValueError(msg) |
| 93 | + |
| 94 | + |
| 95 | +def exact_compress( |
| 96 | + arrays: Sequence[NDArray], chi_out: int, kind: TrainKind |
| 97 | +) -> list[NDArray]: |
| 98 | + """Compress a two-site train exactly via a single truncated SVD. |
| 99 | +
|
| 100 | + Site counts are validated by the caller via `check_exact_supported`. |
| 101 | +
|
| 102 | + Args: |
| 103 | + arrays: The two site tensors of the train. |
| 104 | + chi_out: The maximum bond dimension to keep. |
| 105 | + kind: Whether the train is an ``"mps"`` or an ``"mpo"``. |
| 106 | +
|
| 107 | + Returns: |
| 108 | + The compressed train, in right-canonical form, as numpy arrays. |
| 109 | + """ |
| 110 | + # The dense SVD is host-side, so accept device arrays like the sweep does. |
| 111 | + arrays = [to_numpy(arr) for arr in arrays] |
| 112 | + if kind == "mps": |
| 113 | + # (b, p0) x (b, p1) -> (p0, p1) |
| 114 | + theta = contract("ab,ac->bc", arrays[0], arrays[1]) |
| 115 | + left, right = _truncated_svd(theta, chi_out) |
| 116 | + return [left.T, right] |
| 117 | + |
| 118 | + # (b, u0, d0) x (b, u1, d1) -> (u0, d0, u1, d1) |
| 119 | + theta = contract("aij,akl->ijkl", arrays[0], arrays[1]) |
| 120 | + up_l, down_l, up_r, down_r = theta.shape |
| 121 | + left, right = _truncated_svd(theta.reshape(up_l * down_l, up_r * down_r), chi_out) |
| 122 | + rank = left.shape[1] |
| 123 | + return [ |
| 124 | + left.reshape(up_l, down_l, rank).transpose(2, 0, 1), |
| 125 | + right.reshape(rank, up_r, down_r), |
| 126 | + ] |
| 127 | + |
| 128 | + |
| 129 | +def exact_apply( |
| 130 | + left_tensor: Sequence[NDArray], |
| 131 | + right_tensor: Sequence[NDArray], |
| 132 | + chi_out: int, |
| 133 | + kind: TrainKind, |
| 134 | +) -> list[NDArray]: |
| 135 | + """Contract and compress two two-site trains exactly. |
| 136 | +
|
| 137 | + The MPO on the left is contracted site-wise with the right train, fusing |
| 138 | + the two bond indices, and the result is compressed with a single SVD. |
| 139 | + Site counts are validated by the caller via `check_exact_supported`. |
| 140 | +
|
| 141 | + Args: |
| 142 | + left_tensor: The two site tensors of the left MPO. |
| 143 | + right_tensor: The two site tensors of the right MPS or MPO. |
| 144 | + chi_out: The maximum bond dimension to keep. |
| 145 | + kind: Whether ``right_tensor`` is an ``"mps"`` or an ``"mpo"``. |
| 146 | +
|
| 147 | + Returns: |
| 148 | + The compressed product, in right-canonical form, as numpy arrays. |
| 149 | + """ |
| 150 | + left_tensor = [to_numpy(arr) for arr in left_tensor] |
| 151 | + right_tensor = [to_numpy(arr) for arr in right_tensor] |
| 152 | + if kind == "mps": |
| 153 | + # Contract the MPO lower leg with the MPS physical leg, fusing both bonds. |
| 154 | + product = [ |
| 155 | + contract("aij,bj->abi", left_tensor[i], right_tensor[i]).reshape( |
| 156 | + -1, left_tensor[i].shape[1] |
| 157 | + ) |
| 158 | + for i in range(2) |
| 159 | + ] |
| 160 | + else: |
| 161 | + product = [ |
| 162 | + contract("aij,bjk->abik", left_tensor[i], right_tensor[i]).reshape( |
| 163 | + -1, left_tensor[i].shape[1], right_tensor[i].shape[2] |
| 164 | + ) |
| 165 | + for i in range(2) |
| 166 | + ] |
| 167 | + return exact_compress(product, chi_out, kind) |
| 168 | + |
| 169 | + |
| 170 | +def _truncated_svd(theta: NDArray, chi_out: int) -> tuple[NDArray, NDArray]: |
| 171 | + """Split a matrix as ``(U @ diag(S), Vh)``, keeping at most ``chi_out`` values.""" |
| 172 | + U, S, Vh = np.linalg.svd(theta, full_matrices=False) |
| 173 | + rank = min(chi_out, S.size) |
| 174 | + return U[:, :rank] * S[:rank], Vh[:rank] |
0 commit comments