quchip.engine.bands

Band decomposition by excitation-change weight.

Given an operator O written in a number-basis product representation, this module splits it into the disjoint bands that connect Fock states differing by a fixed excitation change. Let |m⟩ denote a number state on a single mode; the entry O_{nm} = ⟨n|O|m⟩ contributes to the band with weight

\[w \;=\; m - n \;=\; \text{col} - \text{row}.\]

For a bilinear two-body operator on modes a and b the weight is the pair (Δa, Δb).

Physics use

• Stage 2 attaches a carrier exp(−i w ω t) to each band when it assembles a drive or coupling operator in a rotating frame, and drops counter-rotating bands under the RWA (Jaynes & Cummings 1963; for the structured cQED treatment see Gambetta et al., PRA 74, 042318 (2006); for cross-resonance specifically, Rigetti & Devoret, PRB 81, 134507 (2010), and Magesan & Gambetta, PRA 101, 052308 (2020)).

• Stage 3 uses the same weights to demodulate observable expectations back into the control frame.

Implementation

The canonical entry points preserve sparse layouts (CSR / DIA) where possible for concrete payloads. Dense or JAX-traced payloads take the dense path, where a single-band extraction is a where(weights == w, matrix, 0) mask — this is the only shape that keeps the sparsity pattern statically known under jit.

Functions

canonical_to_coo(canonical)	Flatten a canonical payload to (rows, cols, values) COO arrays.
canonical_to_dense_array(canonical)	Materialize canonical densely, preserving its array namespace (JAX-safe).
decompose_bands(op_matrix, dim)	Decompose a single-mode dense operator by weight w = col − row.
decompose_canonical_bands(canonical, dim)	Decompose a canonical single-mode operator by weight w = col − row.
decompose_two_body_canonical_bands(...)	Decompose a canonical two-body operator by (Δa, Δb) per-subsystem change.
embed_on_support(backend, op, support, dims)	Embed a component-local operator into the full space by support arity.
embed_single_mode_bands(backend, local_op, ...)	Like local_mode_bands(), but each band is embedded into dims.
local_mode_bands(backend, local_op, *, dim, ...)	Decompose local_op into ascending excitation-change bands.
prune_zero_diagonals(canonical)	Drop concretely all-zero stored diagonals from a DIA canonical operator.

quchip.engine.bands.decompose_bands(op_matrix, dim)

Decompose a single-mode dense operator by weight w = col − row.

Returns a dict keyed by integer weight w ∈ [−(dim−1), dim−1]; each value is a full dim×dim matrix containing only the entries on that diagonal band (everything else is zero). Concrete arrays drop zero-norm bands; JAX-traced arrays retain every band so the set of keys is statically known across traces.

Parameters:

• op_matrix (Any)

• dim (int)

Return type:

dict[int, Any]

quchip.engine.bands.decompose_canonical_bands(canonical, dim)

Decompose a canonical single-mode operator by weight w = col − row.

Chooses the most compact representation for each band:

• Concrete sparse payloads (CSR/DIA) go through the COO path and emit single-diagonal DIA bands.

• Dense or JAX-traced payloads take decompose_bands() and emit dense bands so the sparsity pattern stays static under jit.

Subsystem metadata (dims, basis, subsystem_labels, tag) is copied onto every band so downstream stages can continue to reason about which subsystem each band lives on.

Parameters:

• canonical (CanonicalOperator)

• dim (int)

Return type:

dict[int, CanonicalOperator]

quchip.engine.bands.decompose_two_body_canonical_bands(canonical, dims)

Decompose a canonical two-body operator by (Δa, Δb) per-subsystem change.

canonical is in the product basis |i_a⟩ ⊗ |i_b⟩ with mode b as the fast index; dims is [d_a, d_b]. Each band has a definite excitation change on each mode, so the carrier attached in stage 2 is exp(−i (Δa · ω_a + Δb · ω_b) t) — the standard rotating-frame form for a bilinear coupling (see e.g. Magesan & Gambetta, PRA 101, 052308 (2020), Eq. (2)).

Parameters:

• canonical (CanonicalOperator)

• dims (list[int])

Return type:

dict[tuple[int, int], CanonicalOperator]

quchip.engine.bands.canonical_to_coo(canonical)

Flatten a canonical payload to (rows, cols, values) COO arrays.

Dispatches on layout. Concrete dense and DIA payloads drop only exactly-zero entries (value != 0, no tolerance): whole-band drop decisions downstream use the relative _BAND_NORM_RTOL, which needs every surviving entry – an absolute per-entry cutoff here would strip an operator whose entire physical scale sits below that cutoff before the relative test ever sees a band. CSR is already explicitly sparse and its stored nnz is preserved as-is (the layout itself is the sparsity declaration, so dropping stored values would discard structurally meaningful zeros). Traced payloads keep their layout-native sparsity (CSR via indices/indptr, DIA via offsets) so the COO size stays static under jit. The fanout-everything fallback only fires when the layout’s structural metadata itself is traced (or for dense, which has no structural metadata at all).

Parameters:

canonical (CanonicalOperator)

Return type:

tuple[Any, Any, Any]

quchip.engine.bands.canonical_to_dense_array(canonical)

Materialize canonical densely, preserving its array namespace (JAX-safe).

Alias for CanonicalOperator.to_dense(), which owns the vectorized, array-namespace-preserving densification logic.

Parameters:

canonical (CanonicalOperator)

Return type:

Any

quchip.engine.bands.local_mode_bands(backend, local_op, *, dim, label)

Decompose local_op into ascending excitation-change bands.

Returns [(weight, band_op), ...] ordered by ascending weight, where band_op is a backend operator on the local dim-sized space — not embedded into the full chip space and without the 2π factor. Callers layer their own embedding / scaling / wrapping.

Parameters:

• backend (Any)

• local_op (Any)

• dim (int)

• label (str)

Return type:

list[tuple[int, Any]]

quchip.engine.bands.embed_single_mode_bands(backend, local_op, *, device_index, dim, label, dims)

Like local_mode_bands(), but each band is embedded into dims.

Returns [(weight, embedded_op), ...] where embedded_op acts on the full chip Hilbert space. Still in the lab frame and ordinary GHz (no 2π).

Parameters:

• backend (Any)

• local_op (Any)

• device_index (int)

• dim (int)

• label (str)

• dims (tuple[int, ...])

Return type:

list[tuple[int, Any]]

quchip.engine.bands.prune_zero_diagonals(canonical)

Drop concretely all-zero stored diagonals from a DIA canonical operator.

Operator algebra that cancels terms exactly (e.g. stage 2 subtracting the lab-frame coupling from H₀ before re-adding it band-by-band as dynamic terms) leaves the union of the operands’ diagonal offsets in the sum, with the cancelled diagonals stored as explicit zeros — dead payload the solver applies at every integration step. Concrete (tracer-free) payloads drop those diagonals here; traced payloads pass through untouched so the stored structure stays statically known under jit. Non-DIA layouts pass through unchanged (CSR structure is the layout’s own sparsity declaration; dense has no structural metadata to prune). If every diagonal is zero, the first one is kept so the operator stays constructible.

Uses the absolute _DIAGONAL_PRUNE_THRESHOLD, not the relative _BAND_NORM_RTOL other band-drop sites use: this removes structure that cancelled exactly to the roundoff floor, a fixed noise floor rather than a fraction of the operator’s own scale.

Parameters:

canonical (CanonicalOperator)

Return type:

CanonicalOperator