feat(stack): ✨ one SRC kernel for arbitrary-depth stacks

README.md

@@ -14,14 +14,16 @@ The following primitives are supported:
 2. MPO-MPO randomized contraction-compression.
 3. MPO randomized compression.
 4. MPS randomized compression.
+5. Randomized contraction-compression of a whole stack of trains in one sweep:
+   `MPO^k`, `MPO^k . MPS` and `MPS . 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.
 
 ```python
-from src_method import apply, compress
+from src_method import apply, compress, src
 ```
 
-The `apply` function covers cases 1 and 2 above, while the `compress` function covers cases 3 and 4. Both functions are pure, meaning no in-place modification ever happens. The user should
+The `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](algorithmiq.github.io/src_method/) for details, and the [tests](tests/) or [benchmarks](benches/) folders for usage examples.
 
@@ -48,6 +50,37 @@ result = qtn.MatrixProductOperator(apply(H1.arrays, H2.arrays, chi_out=64))
 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.
 
+### Contraction Conventions
+
+`src` contracts a *stack* of trains and compresses the result in a single sweep:
+
+```python
+from src_method import src
+
+state = src(U3, U2, U1, psi, chi_out=64)  # U3 U2 U1 |psi>, U1 acts first
+```
+
+The 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:
+
+```python
+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.
+
 ### GPU Support
 
 `src_method` runs on CPU (NumPy) by default and can be accelerated on GPUs via [CuPy](https://cupy.dev/). Both **NVIDIA** (CUDA) and **AMD** (ROCm) GPUs are supported through optional install extras.

benches/README.md

@@ -7,3 +7,6 @@ different SRC variants. They are grouped by the kind of workload exercised:
   SRC primitives (MPO–MPO, MPO–MPS, …). Used to assess speedup and accuracy of
   individual primitives versus `quimb` references on a fixed problem size.
   See [`primitives/README.md`](primitives/README.md).
+- [`stack/`](stack/) — One-shot SRC over stacks of trains against sequential
+  pairwise application, by stack depth: accuracy against dense references and
+  wall time. See [`stack/README.md`](stack/README.md).

benches/stack/README.md

@@ -0,0 +1,125 @@
+# Stack depth benchmarks
+
+`bench_depth.py` compares one-shot SRC over a stack (`src(A_1, ..., A_k, psi)`)
+with sequential pairwise application (`apply` from right to left, truncating to
+`chi` after every product), as a function of the number of trains in the stack.
+At depth 2 both methods run the same computation, which makes that row a sanity
+check.
+
+- `accuracy`: median relative error against the exact dense product on 10-site
+  chains over 20 random instances, next to `lower_bound`, the largest
+  best-rank-`chi` tail over all cuts, which no train of bond dimension `chi` can
+  beat. `one_shot_wins` is the fraction of instances where one-shot is more
+  accurate, ignoring ties and instances where both are exact (`nan` if none is
+  left).
+- `timing`: best-of-3 wall time on 30-site chains, MPS bond 64, `chi = 64`.
+  `ratio` is one-shot time over sequential time.
+
+Families: `random` are complex Gaussian MPOs of bond 3 (flat spectra);
+`trotter-<dt>` are brickwork layers of a mixed-field Ising model with random
+couplings (bond at most 4, decaying spectra). A `.mps` suffix ends the stack with
+an MPS. In `timing`, `x<k> . mps` is `k` layers applied to the MPS.
+
+```bash
+uv run python benches/stack/bench_depth.py accuracy --output accuracy.md
+uv run python benches/stack/bench_depth.py timing --output timing.md
+```
+
+Tuple options repeat the flag, e.g. `--depths 2 --depths 3`. Set
+`LOG_LEVEL_SRC=WARNING` to silence the per-call logs.
+
+## Results
+
+Measured on an AMD Ryzen 7 7840U, 16 threads, with light background load.
+
+### Accuracy
+
+| family | depth | chi | lower_bound | one_shot | sequential | one_shot_wins |
+|---|---|---|---|---|---|---|
+| random.mps | 2 | 4 | 0.168 | 0.499 | 0.499 | nan |
+| random.mps | 2 | 8 | 0.0298 | 0.0989 | 0.0989 | nan |
+| random.mps | 2 | 16 | 2.59e-16 | 1.58e-15 | 1.58e-15 | nan |
+| random.mps | 3 | 4 | 0.194 | 0.519 | 0.636 | 0.9 |
+| random.mps | 3 | 8 | 0.0584 | 0.188 | 0.191 | 0.75 |
+| random.mps | 3 | 16 | 0.00667 | 0.0196 | 0.0194 | 0.35 |
+| random.mps | 4 | 4 | 0.263 | 0.611 | 0.773 | 0.9 |
+| random.mps | 4 | 8 | 0.0978 | 0.275 | 0.335 | 0.85 |
+| random.mps | 4 | 16 | 0.0145 | 0.041 | 0.0533 | 0.8 |
+| random | 2 | 8 | 0.0734 | 0.318 | 0.318 | nan |
+| random | 2 | 16 | 1.3e-15 | 1.53e-15 | 1.53e-15 | nan |
+| random | 2 | 32 | 7.34e-16 | 1.61e-15 | 1.61e-15 | nan |
+| random | 3 | 8 | 0.289 | 0.75 | 0.755 | 0.7 |
+| random | 3 | 16 | 0.102 | 0.352 | 0.351 | 0.65 |
+| random | 3 | 32 | 9.17e-16 | 2.19e-15 | 2.7e-15 | nan |
+| random | 4 | 8 | 0.397 | 0.828 | 0.905 | 1 |
+| random | 4 | 16 | 0.217 | 0.579 | 0.607 | 0.95 |
+| random | 4 | 32 | 0.0819 | 0.274 | 0.27 | 0.55 |
+| trotter-0.3.mps | 2 | 2 | 0.124 | 0.265 | 0.265 | nan |
+| trotter-0.3.mps | 2 | 4 | 6.04e-16 | 1.28e-15 | 1.28e-15 | nan |
+| trotter-0.3.mps | 2 | 8 | 3.4e-16 | 1.38e-15 | 1.38e-15 | nan |
+| trotter-0.3.mps | 3 | 2 | 0.124 | 0.388 | 0.432 | 0.65 |
+| trotter-0.3.mps | 3 | 4 | 0.000264 | 0.00162 | 0.00125 | 0.5 |
+| trotter-0.3.mps | 3 | 8 | 5.53e-16 | 1.41e-15 | 1.68e-15 | nan |
+| trotter-0.3.mps | 4 | 2 | 0.15 | 0.502 | 0.517 | 0.45 |
+| trotter-0.3.mps | 4 | 4 | 0.0023 | 0.0119 | 0.0103 | 0.4 |
+| trotter-0.3.mps | 4 | 8 | 2.57e-08 | 1.16e-07 | 1.14e-07 | 0.5 |
+| trotter-0.3 | 2 | 4 | 2.87e-15 | 1.27e-15 | 1.27e-15 | nan |
+| trotter-0.3 | 2 | 8 | 2.2e-15 | 1.11e-15 | 1.11e-15 | nan |
+| trotter-0.3 | 2 | 16 | 1.38e-15 | 1.2e-15 | 1.2e-15 | nan |
+| trotter-0.3 | 3 | 4 | 0.000457 | 0.00324 | 0.00241 | 0.5 |
+| trotter-0.3 | 3 | 8 | 2.17e-07 | 2.65e-06 | 1.76e-06 | 0.35 |
+| trotter-0.3 | 3 | 16 | 1.95e-15 | 1.47e-15 | 1.65e-15 | nan |
+| trotter-0.3 | 4 | 4 | 0.000522 | 0.00471 | 0.00487 | 0.55 |
+| trotter-0.3 | 4 | 8 | 2.45e-07 | 3.78e-06 | 4.85e-06 | 0.6 |
+| trotter-0.3 | 4 | 16 | 1.98e-15 | 1.75e-15 | 2.4e-15 | nan |
+| trotter-0.8.mps | 2 | 4 | 6.44e-16 | 1.33e-15 | 1.33e-15 | nan |
+| trotter-0.8.mps | 2 | 8 | 4.41e-16 | 1.36e-15 | 1.36e-15 | nan |
+| trotter-0.8.mps | 2 | 16 | 1.93e-16 | 1.23e-15 | 1.23e-15 | nan |
+| trotter-0.8.mps | 3 | 4 | 0.00926 | 0.0393 | 0.0393 | 0.5 |
+| trotter-0.8.mps | 3 | 8 | 6.2e-16 | 1.45e-15 | 1.85e-15 | nan |
+| trotter-0.8.mps | 3 | 16 | 2.97e-16 | 1.53e-15 | 1.65e-15 | nan |
+| trotter-0.8.mps | 4 | 4 | 0.0647 | 0.223 | 0.202 | 0.4 |
+| trotter-0.8.mps | 4 | 8 | 7.87e-05 | 0.000315 | 0.00041 | 0.55 |
+| trotter-0.8.mps | 4 | 16 | 4.6e-16 | 1.76e-15 | 2.27e-15 | nan |
+| trotter-0.8 | 2 | 8 | 1.79e-15 | 1.29e-15 | 1.29e-15 | nan |
+| trotter-0.8 | 2 | 16 | 1.19e-15 | 1.4e-15 | 1.4e-15 | nan |
+| trotter-0.8 | 2 | 32 | 6.45e-16 | 1.21e-15 | 1.21e-15 | nan |
+| trotter-0.8 | 3 | 8 | 0.000394 | 0.00364 | 0.00357 | 0.5 |
+| trotter-0.8 | 3 | 16 | 1.78e-15 | 1.59e-15 | 1.97e-15 | nan |
+| trotter-0.8 | 3 | 32 | 8.17e-16 | 1.7e-15 | 1.91e-15 | nan |
+| trotter-0.8 | 4 | 8 | 0.00048 | 0.00582 | 0.00579 | 0.55 |
+| trotter-0.8 | 4 | 16 | 1.73e-15 | 2.81e-15 | 4.16e-15 | nan |
+| trotter-0.8 | 4 | 32 | 8.41e-16 | 1.88e-15 | 2.63e-15 | nan |
+
+### Timing
+
+| family | depth | chi | one_shot_s | sequential_s | ratio |
+|---|---|---|---|---|---|
+| trotter-0.3 x1 . mps | 2 | 64 | 0.0764 | 0.0642 | 1.19 |
+| trotter-0.3 x2 . mps | 3 | 64 | 0.0936 | 0.147 | 0.639 |
+| trotter-0.3 x3 . mps | 4 | 64 | 0.201 | 0.207 | 0.975 |
+| trotter-0.3 x4 . mps | 5 | 64 | 0.384 | 0.283 | 1.36 |
+| random D=8 x1 . mps | 2 | 64 | 0.283 | 0.28 | 1.01 |
+| random D=8 x2 . mps | 3 | 64 | 1.67 | 0.52 | 3.21 |
+| random D=8 x3 . mps | 4 | 64 | 24.7 | 0.837 | 29.5 |
+| random D=16 x1 . mps | 2 | 64 | 0.638 | 0.622 | 1.03 |
+| random D=16 x2 . mps | 3 | 64 | 11.8 | 1.31 | 9.05 |
+| random D=16 x3 . mps | 4 | 64 | 316 | 2.33 | 136 |
+
+## Summary
+
+- One sweep over the whole stack is at best moderately more accurate than
+  pairwise application. For random stacks ending in an MPS it has about 20 %
+  lower median error at depth 4 for every `chi` measured, and at depth 3 only
+  at the smallest `chi` (at larger `chi` the gain vanishes); it wins 75-90 % of
+  those instances. For random MPO products the gain is at most 8 % (depth 4,
+  small `chi`). For Trotter layers there is no systematic difference (wins
+  35-65 %).
+- Both methods sit 2-15x above the lower bound, so the sketch, not compounding
+  across layers, dominates the error. Oversampling
+  ([#40](https://github.com/Algorithmiq/src-method/issues/40)) is the larger
+  lever.
+- The per-site cost grows with `chi**2` times the product of the layer bonds.
+  Thin Trotter layers run at 0.6-1.4x the sequential time up to depth 5, while
+  random MPOs of bond 8 and 16 at depth 4 are 30x and 136x slower. Stack
+  shallow, thin layers; apply anything else pairwise.