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
C tensors,
- right-to-left sweep building the isometries via
truncated_qr,
- handle the first site, convert back to host, wrap in a
quimb object.
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.py or the existing utils).
- 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] = None memory release, and the final host conversion plus quimb
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
Problem
apply.pyandcompress.pycontain four implementations —_src_mpo_mps,_src_mpo_mpo,_src_mpo,_src_mps— that are structurallyidentical. Each one is roughly 80 lines with the same five phases:
Ctensors,truncated_qr,quimbobject.The only genuine differences are the
opt_einsumsubscript strings and thereshape/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_TIMEandLOG_WARN_SMALLconstants, the same imports, and the same
setup_logging()/loggerblock.The cost is real and already visible in the history: the terminal-bond
truncation bug had to be fixed in four places, and the
cutofffeature had tobe 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.
LOG_*constants, imports and logger setup into ashared private module (
src_method/_common.pyor the existingutils).arrays, the timing/logging wrapper around each sweep, the progressive
C[j - 1] = Nonememory release, and the final host conversion plusquimbwrapping.
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
LOG_*constants and logger setup exist in exactly one place.