Skip to content

_tensor_train

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.

Functions#

check_exact_supported #

check_exact_supported(n_sites: int) -> None

Reject sub-MIN_SRC_SITES trains the exact path cannot handle.

Called at the public boundary before the fallback is announced, so that a degenerate train raises instead of first logging a misleading warning.

PARAMETER DESCRIPTION
n_sites

The number of sites in the train.

TYPE: int

RAISES DESCRIPTION
ValueError

If the train does not have exactly two sites.
Source code in src/src_method/_tensor_train.py 
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
def check_exact_supported(n_sites: int) -> None:
    """Reject sub-``MIN_SRC_SITES`` trains the exact path cannot handle.

    Called at the public boundary before the fallback is announced, so that a
    degenerate train raises instead of first logging a misleading warning.

    Args:
        n_sites: The number of sites in the train.

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

exact_apply #

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. Site counts are validated by the caller via check_exact_supported.

PARAMETER DESCRIPTION
left_tensor

The two site tensors of the left MPO.

TYPE: Sequence[NDArray]

right_tensor

The two site tensors of the right MPS or MPO.

TYPE: Sequence[NDArray]

chi_out

The maximum bond dimension to keep.

TYPE: int

kind

Whether right_tensor is an "mps" or an "mpo".

TYPE: TrainKind

RETURNS DESCRIPTION
list[NDArray]

The compressed product, in right-canonical form, as numpy arrays.
Source code in src/src_method/_tensor_train.py 
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
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.
    Site counts are validated by the caller via `check_exact_supported`.

    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, as numpy arrays.
    """
    left_tensor = [to_numpy(arr) for arr in left_tensor]
    right_tensor = [to_numpy(arr) for arr in 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)

exact_compress #

exact_compress(
    arrays: Sequence[NDArray], chi_out: int, kind: TrainKind
) -> list[NDArray]

Compress a two-site train exactly via a single truncated SVD.

Site counts are validated by the caller via check_exact_supported.

PARAMETER DESCRIPTION
arrays

The two site tensors of the train.

TYPE: Sequence[NDArray]

chi_out

The maximum bond dimension to keep.

TYPE: int

kind

Whether the train is an "mps" or an "mpo".

TYPE: TrainKind

RETURNS DESCRIPTION
list[NDArray]

The compressed train, in right-canonical form, as numpy arrays.
Source code in src/src_method/_tensor_train.py 
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
def exact_compress(
    arrays: Sequence[NDArray], chi_out: int, kind: TrainKind
) -> list[NDArray]:
    """Compress a two-site train exactly via a single truncated SVD.

    Site counts are validated by the caller via `check_exact_supported`.

    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, as numpy arrays.
    """
    # The dense SVD is host-side, so accept device arrays like the sweep does.
    arrays = [to_numpy(arr) for arr in 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),
    ]

infer_kind #

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.

PARAMETER DESCRIPTION
arrays

The site tensors of the train.

TYPE: Sequence[NDArray]

RETURNS DESCRIPTION
TrainKind | None

"mps", "mpo", or None if the layout is unrecognised.
Source code in src/src_method/_tensor_train.py 
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
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