quchip.engine.dressing

Assign dressed eigenvectors to bare or eigenstate references, and build dressed results.

label_eigensystem() supports greedy, global-greedy, and continuation assignments. BareProductReference and EigenstateReference supply vectors and labels, and compute_overlaps() computes their overlaps. The assignment kernels are inspired by SuperGrad’s compute_energy_map. DressedResult is the eager, label-keyed view of one dressed eigensystem that Chip.dress() and EngineResult.dress() return.

Functions

assign_argmax(overlaps)	One-shot argmax per row.
assign_global_greedy(overlaps)	Global descending-overlap greedy with lax.scan and masking.
assign_rowwise_greedy(overlaps)	Confidence-ordered row-greedy assignment — O(n_labels * n_dressed).
compute_overlaps(reference, evecs)	|<ref_k | psi_j>|**2 matrix, shape (n_labels, n_dressed).
dress_engine_result(result, *[, at_time, ...])	Materialize the instantaneous dressed eigensystem of an engine snapshot.
gauge_phases(amplitudes, indices)	Return the phase factor that makes each column's assigned reference amplitude real and nonnegative.
label_eigensystem(evecs, reference[, policy])	Pure-JAX labeling kernel.
labeled_eigensystem(engine_result, backend)	Diagonalize a static analysis result and assign its bare product labels.
materialize_dressed_result(eigenvalues, ...)	Build the eager label-keyed result shared by chip and resolved dressing.
reference_amplitudes(reference, evecs)	Express solver eigenvectors in the captured reference coordinates.
track_path(evals_path, evecs_path, ...[, policy])	Propagate a labeling along a parameter path of stacked eigensystems.
warn_weak_assignment(label, overlap)	Warn when the dressed state assigned to a label keeps less than 0.9 of its bare product.

Classes

BareProductReference(dims[, local_vectors])	Bare product basis states in Kronecker order.
DressedResult(eigenvalues, state_map, ...[, ...])	Store a frozen snapshot of a dressed-state diagonalization.
EigenstateReference(vectors, keys)	Reference vectors are explicit eigenvectors with attached label keys.
Labeling(keys, indices, overlaps, margins, ...)	A labeled assignment of dressed eigenstates to user-meaningful keys.
LabelingPath(keys, indices, overlaps, ...)	Stacked labelings along a parameter path.

class quchip.engine.dressing.Labeling(keys, indices, overlaps, margins, duplicates)

Bases: object

A labeled assignment of dressed eigenstates to user-meaningful keys.

The dataclass has only array fields and no embedded Reference object, so it is a clean pytree under jax.jit(), jax.vmap(), and jax.grad().

indices is the integer assignment (non-differentiable, conceptually stop_gradient’d). overlaps and margins are differentiable. Energies indexed via indices flow gradients through the gathered values, so eigvals[labeling.indices[k]] is differentiable with respect to parameters that affect eigvals.

keys: tuple[Any, ...]

indices: Array

overlaps: Array

margins: Array

duplicates: Array

class quchip.engine.dressing.LabelingPath(keys, indices, overlaps, margins, duplicates, swap_events)

Bases: object

Stacked labelings along a parameter path.

Compared with Labeling, each array field has a leading n_steps axis. A swap_events array shows where labels changed dressed-index between adjacent steps. A swap alone is only a reassignment. To distinguish avoided crossings from routine mode mixing, combine a swap with a margin collapse (margins near zero) at the same step.

keys: tuple[Any, ...]

indices: Array

overlaps: Array

margins: Array

duplicates: Array

swap_events: Array

property final: Labeling

The labeling at the end of the path.

class quchip.engine.dressing.BareProductReference(dims, local_vectors=())

Bases: object

Bare product basis states in Kronecker order.

Optional local vectors map each component’s energy states into its solver basis, and None denotes an identity factor. The full product matrix is never materialized; identity references reduce directly to |evecs|**2.

dims: tuple[int, ...]

local_vectors: tuple[Any | None, ...] = ()

property keys: tuple[tuple[int, ...], ...]

class quchip.engine.dressing.EigenstateReference(vectors, keys)

Bases: object

Reference vectors are explicit eigenvectors with attached label keys.

Used for path continuation: the previous step’s selected eigvecs become the next step’s reference. vectors is row-major: each row is one reference state in the Kronecker basis.

vectors: Array

keys: tuple[Any, ...]

quchip.engine.dressing.reference_amplitudes(reference, evecs)

Express solver eigenvectors in the captured reference coordinates.

Concrete inputs combine on the host, where they compile no device programs.

quchip.engine.dressing.compute_overlaps(reference, evecs)

|<ref_k | psi_j>|**2 matrix, shape (n_labels, n_dressed).

quchip.engine.dressing.assign_argmax(overlaps)

One-shot argmax per row. Allows duplicate dressed-index assignments.

Returns (indices, chosen, margins). This policy is cheap and jit-clean, and matches SuperGrad’s default greedy mode. Use it in a weak-coupling regime where each bare label has a clearly dominant dressed eigenstate.

quchip.engine.dressing.assign_global_greedy(overlaps)

Global descending-overlap greedy with lax.scan and masking.

At each step, the function assigns the (label, dressed_idx) pair with the highest remaining overlap. It then masks out that label’s row and that column, so this pure-JAX implementation assigns each label and dressed index at most once.

Margins are computed from the unmasked overlap matrix, so they describe per-label confidence independent of assignment order.

Requires n_labels <= n_dressed.

quchip.engine.dressing.assign_rowwise_greedy(overlaps)

Confidence-ordered row-greedy assignment — O(n_labels * n_dressed).

The function visits labels (rows) in descending row-max-overlap order. Each label takes its highest-overlap still-available dressed index, and the function masks that column, so the result is a permutation. Each scan step is a single masked-row argmax over n_dressed columns, so the whole assignment is O(n_labels * n_dressed). For a square D = L**k overlap matrix, this is O(L**(2k)). assign_global_greedy() is O(L**(3k)) (a full L**(2k)-entry argmax at each of L**k steps). This is the SuperGrad / scqubits confidence-ordered greedy variant.

Semantics vs. global-greedy. This function is identical to assign_global_greedy() in the near-permutation regime (dispersive / weak hybridization), where each bare label has a clearly dominant dressed eigenstate. The two can diverge on strongly-hybridized overlap matrices. Global-greedy always commits the single largest remaining entry, whereas this function visits rows in (original) row-max order. In exactly that regime the assignment is intrinsically ambiguous, and Chip.dress already warns that bare labels are approximate. For global-greedy ordering, pass policy=assign_global_greedy to label_eigensystem(). Row-max ties are broken by ascending index (argsort).

Margins (top1 - top2 per label) come from the unmasked matrix via top_k, so they match the margins of assign_global_greedy(). Requires n_labels <= n_dressed.

quchip.engine.dressing.warn_weak_assignment(label, overlap)

Warn when the dressed state assigned to a label keeps less than 0.9 of its bare product.

Every query that selects the dressed state of a label calls this. A traced overlap does not warn. The warning names the first caller outside quchip.

Parameters:

• label (tuple of int) – Product label of local energy levels.

• overlap (scalar) – Squared overlap of the assigned dressed state with the bare product.

quchip.engine.dressing.gauge_phases(amplitudes, indices)

Return the phase factor that makes each column’s assigned reference amplitude real and nonnegative.

Parameters:

• amplitudes (array_like) – (n_labels, n_dressed) amplitudes <ref_k|psi_j> of the dressed columns.

• indices (array_like) – Column assigned to each reference label.

Notes

The anchor of a column is its amplitude on the label assigned to it. A one-to-one assignment can give a column whose assigned amplitude is zero. When that amplitude is not above a floor, the square root of the floating-point epsilon, and for an unassigned column, the anchor is the first amplitude within a relative 1e-6 of the column’s largest magnitude. The factor of a column is conj(anchor) / |anchor|. The rule is piecewise: the phase can jump where the assigned amplitude crosses the floor or where the fallback selects another amplitude. A phase cannot fix a rotation, so an exactly degenerate eigenspace keeps the basis of the eigensolver.

quchip.engine.dressing.label_eigensystem(evecs, reference, policy=<function assign_rowwise_greedy>)

Pure-JAX labeling kernel.

Parameters:

• evecs (jnp.ndarray) – (dim, n_dressed) eigenvectors as columns.

• reference (object) – Data-only object (BareProductReference or EigenstateReference) defining what label keys mean.

• policy (callable) – Assignment function overlaps -> (indices, chosen, margins). Default assign_rowwise_greedy(), the confidence-ordered row-greedy policy that Chip.dress runs (its sole live caller passes it explicitly).

Notes

The caller gets energies as evals[labeling.indices[k]]. policy is a static Python callable resolved once at trace time, not a traced argument. Under jax.jit, close over policy or pass it as a static_argnames entry. Never pass it as a dynamic (traced) argument.

quchip.engine.dressing.track_path(evals_path, evecs_path, initial_reference, policy=<function assign_global_greedy>)

Propagate a labeling along a parameter path of stacked eigensystems.

Step 0 is labeled against initial_reference. At each later step, the reference becomes the previous step’s selected eigvecs (an EigenstateReference). This continuation method accepts a stacked eigvec tensor (usually the output of vmap(eigh) over a parameter grid), so it composes with any JAX-side parameter sweep.

Parameters:

• evals_path (jnp.ndarray) – (n_steps, n_dressed).

• evecs_path (jnp.ndarray) – (n_steps, dim, n_dressed).

• initial_reference (object) – Bootstrap reference at step 0.

• policy (callable) – Per-step assignment function. Defaults to assign_global_greedy().

Returns:

Stacked diagnostics including swap_events per step boundary.

Return type:

LabelingPath

class quchip.engine.dressing.DressedResult(eigenvalues, state_map, dressed_eigenvalues, assignment_overlaps, hybridized_labels, bare_labels, bare_labels_by_dressed_index, eigenvector_matrix=None, overlap_threshold=0.5, labeling='DE', _eigensystem=None)

Bases: object

Store a frozen snapshot of a dressed-state diagonalization.

Variables:

• eigenvalues (array-like) – Sorted dressed eigenvalues in GHz.

• eigenstates (array-like) – Backend eigenstate objects in the same order as eigenvalues, with the phases of eigenvector_matrix. They are materialized lazily on first access from the underlying EigensystemData. The dressing and sweep hot path never touches them, so the backend builds per-column kets only on request.

• state_map (dict[tuple[int, ...], int]) – Map from each bare product-basis label (one int per device) to the dressed eigenstate index with the largest overlap.

• dressed_eigenvalues (dict[tuple[int, ...], Any]) – Dressed eigenvalue for each assigned bare label, the direct lookup path for Chip.energy().

• assignment_overlaps (dict[tuple[int, ...], float]) – |⟨bare|dressed⟩|² of each assignment. Values below overlap_threshold flag hybridization.

• hybridized_labels (tuple[tuple[int, ...], ...]) – Bare labels whose assignment quality is below overlap_threshold. A non-empty tuple triggers a user warning at dress time.

• bare_labels (tuple[tuple[int, ...], ...]) – Labels of the product energy levels exposed by the chip’s resolved local dimensions, in chip order.

• bare_labels_by_dressed_index (dict[int, tuple[int, ...]]) – Inverse of state_map: dressed index → assigned bare label.

• eigenvector_matrix (array-like or None) – Dressed eigenvectors as columns, in the resolved solver basis, used by ChipAnalysis.operator_in_dressed_basis() and ChipAnalysis.state_components(). Each column has a real, nonnegative amplitude on its assigned bare state (gauge_phases()).

• overlap_threshold (float) – Minimum overlap for a confident assignment.

• labeling (str) – Algorithm id (currently only "DE"). It resolves to the confidence-ordered row-greedy overlap-matching policy run by ChipAnalysis.dress() (assign_rowwise_greedy()).

eigenvalues: Any

state_map: dict[tuple[int, ...], int]

dressed_eigenvalues: dict[tuple[int, ...], Any]

assignment_overlaps: dict[tuple[int, ...], float]

hybridized_labels: tuple[tuple[int, ...], ...]

bare_labels: tuple[tuple[int, ...], ...]

bare_labels_by_dressed_index: dict[int, tuple[int, ...]]

eigenvector_matrix: Any = None

overlap_threshold: float = 0.5

labeling: str = 'DE'

property eigenstates: Any

Backend eigenstate kets, materialized lazily from the eigensystem.

quchip.engine.dressing.materialize_dressed_result(eigenvalues, eigenvector_matrix, eigensystem, kernel_labeling, *, overlap_threshold, labeling, warning_stacklevel)

Build the eager label-keyed result shared by chip and resolved dressing.

quchip.engine.dressing.dress_engine_result(result, *, at_time=None, overlap_threshold=0.5, labeling='DE')

Materialize the instantaneous dressed eigensystem of an engine snapshot.

quchip.engine.dressing.labeled_eigensystem(engine_result, backend)

Diagonalize a static analysis result and assign its bare product labels.

Returns (eigenvalues, eigenvector_matrix, eigensystem, labeling); traced Hamiltonians stay traced. Each eigenvector follows the phase convention of gauge_phases().