Skip to content

Generalize the four SRC kernels to an arbitrary-depth stack #34

Description

@Panadestein

Observation

_src_mps, _src_mpo, _src_mpo_mps and _src_mpo_mpo are the same algorithm four
times: sketch the open physical legs with a Gaussian omega, sweep left-to-right
accumulating the C environments, sweep right-to-left building the output through
truncated_qr while carrying S. The only things that vary are

  • how many layers sit at each site (one or two),
  • the rank of omega, which follows the number of open physical legs of the output,
  • the contract index strings, which are written out by hand for each case.

So the kernel count is a function of the layer count, and it grows combinatorially if we
keep adding cases by hand.

Proposed generalization

Take a stack: an ordered sequence of trains contracted along their physical legs. The
composition rules are MPO . MPO -> MPO and MPO . MPS -> MPS, so the well-formed
stacks are MPO^k and MPO^k . MPS for any k >= 1. The existing kernels are the
k = 1 and k = 2 corners of that family.

A single entry point, e.g. src(*trains, chi_out=...), would cover all four of today's
cases plus arbitrary depth, with apply and compress kept as thin wrappers for
backwards compatibility.

Why it is worth doing

Beyond removing the duplication, there is a real numerical argument. Applying k MPO
layers today means k sequential apply calls, each truncating to chi_out, so
truncation error compounds layer by layer. Sketching the whole stack and compressing
once should be more accurate for the same output bond dimension. That matters directly
for Trotterized evolution, where the layers are known up front.

The cost moves the other way: the per-site contraction grows with the product of the
layer bond dimensions, and the C environments gain one bond index per layer, so they
become rank k + 1. There is a crossover in k beyond which blocking the stack and
compressing in stages wins. Finding that crossover is part of the work, and benches/
already gives us the harness for it.

Contraction paths

With more than two layers the per-site expression stops having an obvious order, which
is where cotengra and friends come in. Two separate points:

  1. Independently of this issue, we currently call contract fresh at every site of
    every sweep, so the path is re-planned each time even though the expression shape
    recurs. Building the expression once with opt_einsum.contract_expression and
    reusing it is a cheap win available today.
  2. For deep stacks, opt_einsum's default heuristics may not be enough and a
    hyper-optimizer becomes worthwhile. The sweep structure bounds the intermediate
    sizes, so this probably only pays off past some depth. It should be an optional
    dependency and a measured decision, not an assumption.

Sketch of the work

  1. A stack type or validator: equal site counts, all layers MPO except an optional
    trailing MPS, matching physical dimensions between adjacent layers.
  2. Generate the sweep index strings from the layer count instead of hard-coding them,
    and cache the compiled expressions across sites.
  3. Size omega from the output kind (one open leg for MPS, two for MPO).
  4. Reduce the existing kernels to calls into the general one, keeping the exact
    sub-MIN_SRC_SITES fallback as it is.
  5. Accuracy and cost benchmarks against sequential pairwise application, as a function
    of depth.

Open questions

  • Does the SRC error analysis carry over to k > 2? The reference results cover the
    products they analyse; how sketch quality degrades with stack depth is an empirical
    question we should answer before promising anything.
  • Deep stacks have a wider spectrum. A single Gaussian sketch may need oversampling or
    a power iteration to stay accurate.
  • Memory for the rank-k + 1 environments may cap the practical depth on its own.
  • Is the accuracy gain over sequential application large enough to justify the extra
    cost per sweep? If it is not, this is a refactor rather than a feature, which is
    still worth doing but changes the priority.

Suggest starting with a spike at k = 3 to answer the accuracy question before
committing to the general machinery.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions