"""Declarative parameter sweeps for chip studies.
Sweep axes compose as Cartesian products or zipped correlated scans. Spectrum
sweeps dress the chip at each point, keeping eigenvalues and bare→dressed label
assignments.
Sweep values keep their native physical units (GHz for frequencies, ns for
times, mK for temperatures). :class:`Sweep` keeps JAX arrays verbatim, so
downstream code can differentiate through swept values. :class:`Sweep` promotes
all other values to ``np.ndarray``, so ``len`` and indexing behave predictably.
The dressed-spectrum machinery here computes the eigenvalues and overlaps of
the lab-frame chip Hamiltonian. For a general reference on the dressed-state
picture of coupled circuit qubits, see Blais, Grimsmo, Girvin & Wallraff,
*Circuit quantum electrodynamics*, Rev. Mod. Phys. 93, 025005 (2021).
"""
from __future__ import annotations
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, Callable, Mapping, Sequence
import numpy as np
from quchip.utils.batching import expand_axis_groups
from quchip.utils.labeling import bare_label_from_mapping, top_components
if TYPE_CHECKING:
from quchip.chip.chip import Chip
[docs]
class ZippedSweep:
"""Pair sweep axes element by element.
Use :meth:`Sweep.zip` to validate lengths before you construct this bundle.
Parameters
----------
sweeps : tuple of Sweep
Nonempty axes with equal lengths and distinct names.
Attributes
----------
sweeps : tuple of Sweep
Axes advanced together.
size : int
Length of the first axis.
"""
def __init__(self, sweeps: tuple[Sweep, ...]) -> None:
self.sweeps = sweeps
self.size = sweeps[0].size
def __repr__(self) -> str:
names = [s.name for s in self.sweeps]
return f"ZippedSweep({names}, size={self.size})"
[docs]
class Sweep:
"""Declare a one-dimensional parameter axis.
Parameters
----------
values : array_like, shape (n,)
Values in the parameter's units. JAX arrays are kept and other inputs
are converted to NumPy arrays. Scalar inputs are unsupported.
name : str or None, default None
Parameter path, e.g. ``"q.freq"``. Defaults to ``"unnamed"``. Names
must be unique when axes are combined.
Attributes
----------
values : array_like
Stored axis values.
name : str
Parameter path or axis name.
size : int
Number of values.
"""
def __init__(self, values: Any, *, name: str | None = None) -> None:
from quchip.utils.jax_utils import is_jax_array
self.values = values if is_jax_array(values) else np.asarray(values)
self.name = name or "unnamed"
self.size = len(self.values)
def __repr__(self) -> str:
return f"Sweep({self.name!r}, size={self.size})"
[docs]
@staticmethod
def zip(*sweeps: Sweep) -> ZippedSweep:
"""Pair axes for element-wise iteration.
Requires at least two sweeps of equal size, and raises
:class:`ValueError` otherwise.
Parameters
----------
*sweeps
Two or more :class:`Sweep` axes. Their ``name`` attributes must all
be distinct. A :class:`ZippedSweep` with repeated axis names raises
:class:`ValueError` when it is enumerated later (:meth:`expand`,
:class:`SpectrumSweep`).
Returns
-------
ZippedSweep
Bundle iterated element-wise instead of taken in Cartesian
product with other axes.
"""
if len(sweeps) < 2:
raise ValueError(f"Sweep.zip() requires at least two sweeps, got {len(sweeps)}")
sizes = {s.size for s in sweeps}
if len(sizes) != 1:
raise ValueError(f"Zipped sweeps must have equal lengths, got {sorted(sizes)}")
return ZippedSweep(sweeps)
[docs]
@staticmethod
def expand(axes: Sequence[Sweep | ZippedSweep]) -> list[dict[str, Any]]:
"""Expand independent axes as a Cartesian product and zipped axes pairwise.
Parameters
----------
axes : sequence of Sweep or ZippedSweep
Axes in output order. Names must be unique across all groups.
Returns
-------
list of dict
Parameter mappings in C order (last axis fastest). No axes gives
one empty mapping, and an empty axis gives no points.
"""
_shape, points = _iter_axis_points(axes)
return [params for _coord, params in points]
def _axis_groups(
axes: Sequence[Sweep | ZippedSweep],
) -> list[list[dict[str, Any]]]:
"""Normalize axes into independent groups of parameter dictionaries."""
groups: list[list[dict[str, Any]]] = []
for axis in axes:
if isinstance(axis, ZippedSweep):
groups.append([{s.name: s.values[i] for s in axis.sweeps} for i in range(axis.size)])
else:
groups.append([{axis.name: value} for value in axis.values])
return groups
def _check_unique_axis_names(axes: Sequence[Sweep | ZippedSweep]) -> None:
"""Reject duplicate names across independent axes and zipped members."""
names: list[str] = []
for axis in axes:
if isinstance(axis, ZippedSweep):
names.extend(s.name for s in axis.sweeps)
else:
names.append(axis.name)
duplicates = sorted({name for name in names if names.count(name) > 1})
if duplicates:
raise ValueError(f"Sweep axis names must be unique across all axes; duplicated: {duplicates}")
def _iter_axis_points(
axes: Sequence[Sweep | ZippedSweep],
) -> tuple[tuple[int, ...], list[tuple[tuple[int, ...], dict[str, Any]]]]:
"""Yield every (N-D coordinate, param dict) pair on the sweep grid."""
_check_unique_axis_names(axes)
return expand_axis_groups(_axis_groups(axes))
def _axis_metadata(
axes: Sequence[Sweep | ZippedSweep],
rename: Callable[[str], str] = str,
) -> tuple[tuple[str, Any], ...]:
"""Return named grid metadata for ordinary and zipped sweep axes."""
metadata: list[tuple[str, Any]] = []
for axis in axes:
if isinstance(axis, ZippedSweep):
names = tuple(rename(sweep.name) for sweep in axis.sweeps)
values = tuple(
{
name: sweep.values[index]
for name, sweep in zip(names, axis.sweeps, strict=True)
}
for index in range(axis.size)
)
metadata.append(("/".join(names), values))
else:
metadata.append((rename(axis.name), axis.values))
return tuple(metadata)
[docs]
@dataclass
class SpectrumSweepResult:
"""Dressed-spectrum data collected across a :class:`SpectrumSweep` grid.
Attributes
----------
device_labels
Device labels of the underlying chip, in the chip's tensor-product
order (:attr:`~quchip.chip.chip.Chip.devices`).
bare_labels
Bare-state labels tracked by the sweep (one tuple of occupation numbers
per device).
eigenvalues
``(*grid_shape, n_evals)`` array of eigenvalues in GHz, one row
per sweep point.
dressed_indices
``(*grid_shape, n_bare)`` array giving the dressed-state index
assigned to each bare label at each sweep point, or ``NaN``
when no confident assignment is found.
assignment_overlaps
``(*grid_shape, n_bare)`` overlap weights accompanying
``dressed_indices``.
overlap_threshold
Minimum overlap below which an assignment is unreliable and returned as
``NaN``.
params
Object array of parameter dicts, one per grid point.
eigenvector_matrices
Optional object array of dressed eigenvector matrices (one per
grid point), produced when ``store_eigenstates=True``.
eigenstates
Optional object array of dressed eigenstates aligned with
``eigenvector_matrices``.
"""
device_labels: tuple[str, ...]
bare_labels: tuple[tuple[int, ...], ...]
eigenvalues: np.ndarray
dressed_indices: np.ndarray
assignment_overlaps: np.ndarray
overlap_threshold: float
params: np.ndarray
eigenvector_matrices: np.ndarray | None = None
eigenstates: np.ndarray | None = None
@property
def shape(self) -> tuple[int, ...]:
"""Sweep-grid shape (excludes the trailing eigenvalue axis)."""
return self.eigenvalues.shape[:-1]
def _normalize_point(self, point: int | tuple[int, ...] | None) -> tuple[int, ...]:
if self.shape == ():
return ()
if point is None:
raise ValueError(
f"This result spans a sweep grid of shape {self.shape}; pass a "
"point index (int for 1-D sweeps, tuple of ints otherwise) to "
"select which grid point to read."
)
normalized = point if isinstance(point, tuple) else (point,)
if len(normalized) != len(self.shape):
raise ValueError(
f"Point {point!r} has {len(normalized)} coordinate(s) but this result's sweep "
f"grid has shape {self.shape} ({len(self.shape)} axes); pass a coordinate tuple "
"matching the grid's dimensionality."
)
return normalized
def _normalize_label(
self,
device_states: Mapping[str, int] | None,
device_state_kwargs: Mapping[str, int],
) -> tuple[int, ...]:
"""Build a bare-state tuple from device-or-label keys, defaulting to level zero.
Duplicate or unknown labels raise ``ValueError``.
"""
return bare_label_from_mapping(self.device_labels, device_states, device_state_kwargs)
def _bare_label_position(self, label: tuple[int, ...]) -> int:
"""Return the column index of ``label`` in ``bare_labels`` or raise."""
try:
return self.bare_labels.index(label)
except ValueError:
raise ValueError(f"Bare-state label {label} is not present in this sweep") from None
[docs]
def dressed_index(
self,
device_states: Mapping[str, int] | None = None,
/,
**device_state_kwargs: int,
) -> np.ndarray:
"""Dressed-state indices across the sweep for one bare label.
Entries where the assignment overlap falls below
``overlap_threshold`` (or was never found) are returned as
``NaN``.
Parameters
----------
device_states
Mapping from device (label or object) to occupation number, resolved
through :func:`~quchip.utils.labeling.bare_label_from_mapping`.
Unspecified devices default to Fock index 0. Do not use with
``device_state_kwargs``.
**device_state_kwargs
Keyword form of ``device_states`` (device label as keyword).
Returns
-------
numpy.ndarray
Floating-point array of shape :attr:`shape` (the sweep grid,
excluding the bare-label axis). A floating dtype is needed to
represent unreliable or missing assignments as ``NaN``. An
unreliable assignment is a grid point whose overlap with this bare
label falls below ``overlap_threshold``. A missing assignment is a
grid point where no dressed state was assigned to this bare label.
Examples
--------
>>> # Track the bare |1, 0> state across a sweep
>>> # result.dressed_index({qubit_a: 1, qubit_b: 0}) # doctest: +SKIP
"""
label = self._normalize_label(device_states, device_state_kwargs)
bare_pos = self._bare_label_position(label)
indices = np.asarray(self.dressed_indices[..., bare_pos], dtype=float)
overlaps = np.asarray(self.assignment_overlaps[..., bare_pos], dtype=float)
invalid = np.isnan(indices) | (overlaps < self.overlap_threshold)
result = indices.copy()
result[invalid] = np.nan
return result
[docs]
def energy_by_bare_label(
self,
device_states: Mapping[str, int] | None = None,
/,
**device_state_kwargs: int,
) -> np.ndarray:
"""Dressed energies (GHz) traced out for one bare label across the sweep.
Invalid points inherit the ``NaN`` mask from
:meth:`dressed_index`.
Parameters
----------
device_states
Mapping from device (label or object) to occupation number. See
:meth:`dressed_index`. Do not use with ``device_state_kwargs``.
**device_state_kwargs
Keyword form of ``device_states`` (device label as keyword).
Returns
-------
numpy.ndarray
Floating-point array of shape :attr:`shape`, in GHz. Entries
where :meth:`dressed_index` returns ``NaN`` (overlap below
``overlap_threshold``, or no assignment found) are ``NaN``.
"""
indices = self.dressed_index(device_states, **device_state_kwargs)
flat_indices = indices.reshape(-1)
flat_evals = self.eigenvalues.reshape(-1, self.eigenvalues.shape[-1])
flat_out = np.full(flat_indices.shape, np.nan, dtype=float)
valid = ~np.isnan(flat_indices)
flat_out[valid] = flat_evals[np.nonzero(valid)[0], flat_indices[valid].astype(int)]
return flat_out.reshape(indices.shape)
[docs]
def state_components_at(
self,
point: int | tuple[int, ...] | None,
state: int | Mapping[str, int] | None = None,
/,
*,
n_components: int = 5,
**device_state_kwargs: int,
) -> dict[tuple[int, ...], float]:
"""Return the largest bare-state probabilities of one dressed state.
Requires running the sweep with ``store_eigenstates=True``.
Parameters
----------
point : int, tuple of int, or None
Grid coordinate. An integer selects a one-dimensional sweep.
``None`` is valid only for a result with no sweep axes.
state : int, mapping, or None, default None
Dressed index, or a mapping from device object or label to bare
level. A mapping is resolved through the stored overlap assignment.
Omitted devices have level zero.
n_components : int, default 5
Positive maximum number of components returned.
**device_state_kwargs : int
Bare levels keyed by device label, as an alternative to a mapping.
Do not combine with a nonempty mapping.
Returns
-------
dict
Bare occupation tuples mapped to squared amplitudes, in descending order.
Raises
------
ValueError
Eigenvectors were not saved, a coordinate is invalid, the requested
count is nonpositive, or the bare label has no confident assignment.
See Also
--------
SpectrumSweep : Dressed-spectrum conventions and physics reference.
"""
if self.eigenvector_matrices is None:
raise ValueError(
"Eigenvector matrices were not stored for this sweep. "
"Re-run it with store_eigenstates=True."
)
if n_components <= 0:
raise ValueError(f"n_components must be positive, got {n_components}")
point_index = self._normalize_point(point)
eigenvector_matrix = self.eigenvector_matrices[point_index]
if isinstance(state, int) and not isinstance(state, bool):
dressed_idx: int = state
else:
if state is not None and not isinstance(state, Mapping):
raise TypeError(f"state must be an int or mapping, got {type(state).__name__}")
label = self._normalize_label(state, device_state_kwargs)
bare_pos = self._bare_label_position(label)
raw_idx = self.dressed_indices[point_index + (bare_pos,)]
overlap = self.assignment_overlaps[point_index + (bare_pos,)]
if np.isnan(raw_idx) or float(overlap) < self.overlap_threshold:
raise ValueError(
f"No confidently labeled dressed state for bare label {label} at point {point_index}"
)
dressed_idx = int(raw_idx)
return top_components(eigenvector_matrix, self.bare_labels, dressed_idx, n_components)
[docs]
class SpectrumSweep:
"""Calculate dressed spectra over a parameter grid.
Each point uses a fresh :meth:`Chip.with_params` clone and an independent
bare-state overlap assignment, so labels are not transported along the
path.
Parameters
----------
chip : Chip
Model to dress in the lab frame.
axes : sequence of Sweep or ZippedSweep
Nonempty axes named by paths in ``chip.parameters``. Independent axes
form a Cartesian product, and zipped axes advance together.
evals_count : int or None, default None
Number of lowest eigenvalues stored, from 1 to ``chip.total_dim``.
``None`` stores all. Dressing still computes the full eigensystem.
store_eigenstates : bool, default False
Keep eigenvectors and backend states for component inspection.
overlap_threshold : float, default 0.5
Minimum squared overlap for a confident bare-state assignment.
Use a value between 0 and 1. Unassigned labels yield ``NaN``.
References
----------
Blais et al., *Circuit quantum electrodynamics*, Rev. Mod. Phys. 93,
025005 (2021), https://doi.org/10.1103/RevModPhys.93.025005.
The reference describes dressed states; grid traversal is quchip policy.
"""
def __init__(
self,
chip: "Chip",
axes: Sequence[Sweep | ZippedSweep],
*,
evals_count: int | None = None,
store_eigenstates: bool = False,
overlap_threshold: float = 0.5,
) -> None:
self.chip = chip
self.axes = list(axes)
self.evals_count = evals_count
self.store_eigenstates = bool(store_eigenstates)
self.overlap_threshold = float(overlap_threshold)
[docs]
def run(self, progress: bool = True) -> SpectrumSweepResult:
"""Execute the sweep and collect the dressed-spectrum analysis.
Parameters
----------
progress
If true, display a ``tqdm`` progress bar.
Returns
-------
SpectrumSweepResult
Populated result.
"""
from tqdm import tqdm
for axis in self.axes:
if axis.size == 0:
name = axis.name if isinstance(axis, Sweep) else [s.name for s in axis.sweeps]
raise ValueError(
f"Sweep axis {name!r} has zero length; SpectrumSweep.run() cannot iterate an empty grid."
)
shape, expanded = _iter_axis_points(self.axes)
bare_labels = self.chip._canonical_bare_labels()
n_bare = len(bare_labels)
total_dim = self.chip.total_dim
if self.evals_count is None:
n_evals = total_dim
else:
n_evals = self.evals_count
if isinstance(n_evals, bool) or not isinstance(n_evals, int) or not (1 <= n_evals <= total_dim):
raise ValueError(
f"evals_count must be an integer in [1, {total_dim}] (chip.total_dim), got {n_evals!r}"
)
param_store = np.empty(shape, dtype=object)
dressed_indices = np.full(shape + (n_bare,), np.nan, dtype=float)
assignment_overlaps = np.full(shape + (n_bare,), np.nan, dtype=float)
matrices_store = np.empty(shape, dtype=object) if self.store_eigenstates else None
states_store = np.empty(shape, dtype=object) if self.store_eigenstates else None
eigenvalues_store: list[np.ndarray] = []
iterator = tqdm(expanded, desc="SpectrumSweep") if progress else expanded
for coord, params in iterator:
chip_point = self.chip.with_params(params)
dressed = chip_point.dress(overlap_threshold=self.overlap_threshold)
evals = np.asarray(dressed.eigenvalues, dtype=float)[:n_evals]
eigenvalues_store.append(evals)
param_store[coord] = dict(params)
for bare_pos, bare_label in enumerate(bare_labels):
dressed_idx = dressed.state_map.get(bare_label)
if dressed_idx is None or dressed_idx >= len(evals):
continue
dressed_indices[coord + (bare_pos,)] = dressed_idx
assignment_overlaps[coord + (bare_pos,)] = dressed.assignment_overlaps[bare_label]
if matrices_store is not None and states_store is not None:
matrices_store[coord] = np.asarray(dressed.eigenvector_matrix[:, : len(evals)])
states_store[coord] = list(dressed.eigenstates[: len(evals)])
if shape:
eigenvalues = np.asarray(eigenvalues_store, dtype=float).reshape(shape + (len(eigenvalues_store[0]),))
else:
eigenvalues = np.asarray(eigenvalues_store[0], dtype=float)
return SpectrumSweepResult(
device_labels=tuple(device.label for device in self.chip.devices),
bare_labels=bare_labels,
eigenvalues=eigenvalues,
dressed_indices=dressed_indices,
assignment_overlaps=assignment_overlaps,
overlap_threshold=self.overlap_threshold,
params=param_store,
eigenvector_matrices=matrices_store,
eigenstates=states_store,
)
