Repository navigation
Description
Problem
apply.py and compress.py contain four implementations —
_src_mpo_mps, _src_mpo_mpo, _src_mpo, _src_mps — that are structurally
identical. Each one is roughly 80 lines with the same five phases:
- read problem dimensions and log,
- move the input arrays to the device,
- left-to-right sweep accumulating the sketched
Ctensors, - right-to-left sweep building the isometries via
truncated_qr, - handle the first site, convert back to host, wrap in a
quimbobject.
The only genuine differences are the opt_einsum subscript strings and the
reshape/transpose bookkeeping around them.
On top of that, the module preamble is duplicated verbatim between the two
files: the identical LOG_LTR, LOG_RTL, LOG_TIME and LOG_WARN_SMALL
constants, the same imports, and the same setup_logging() / logger block.
The cost is real and already visible in the history: the terminal-bond
truncation bug had to be fixed in four places, and the cutoff feature had to
be threaded through four times. Each new primitive or fix multiplies by four.
Proposed fix
Factor out the shared skeleton without obscuring the algorithm — the einsum
subscripts should stay visible and readable, because they are the algorithm.
- Move the duplicated
LOG_*constants, imports and logger setup into a
shared private module (src_method/_common.pyor the existingutils). - Extract the parts that are genuinely identical: device transfer of the input
arrays, the timing/logging wrapper around each sweep, the progressive
C[j - 1] = Nonememory release, and the final host conversion plusquimb
wrapping. - Express each primitive as its set of contraction subscripts plus the small
reshape/transpose rules, driven through the common sweep driver. A small
dataclass or a module-level table of subscripts per primitive keeps this
declarative and greppable.
Explicit non-goal: do not collapse the four primitives into one clever
generic einsum builder. That would trade a duplication problem for an
unreadable one. The aim is one sweep driver plus four short, obvious
specifications.
Acceptance criteria
- The duplicated
LOG_*constants and logger setup exist in exactly one place. - The left-to-right and right-to-left sweep scaffolding is written once.
- Each primitive's contraction subscripts remain plainly readable.
- No behaviour change: the existing test suite passes untouched.
- No measurable performance regression.
Activity
Panadestein commentedon Oct 6, 2026
#41 fixed this.