"""Band-decompose dict-form ``e_ops`` and demodulate expectations after the solve.
The simulation runs in the resolved rotating frame, but user-facing observables
are in the control frame, where detunings ``Δ = ω_drive − ω_frame`` label the
matrix elements. This module performs the two operations that connect those
frames:
* **Pre-solve (:func:`decompose_eops`):** :mod:`quchip.engine.bands` splits
each user operator into excitation-change bands ``w = col − row``
(single-mode) or ``(Δa, Δb)`` (two-mode). Each band is embedded into the full
product Hilbert space and carries :class:`BandMeta`, so the post-solve step
knows which phase to apply.
* **Post-solve (:func:`recombine_expect`):** each band expectation ``⟨O_w⟩(t)``
is multiplied by ``exp(±i · 2π · ω_demod · w · t)`` and summed over bands.
The sign selects demodulation or remodulation. Physically, this inverse
rotating-frame transformation moves each band from the simulation frame to
the control frame.
An external-plane output request uses a parallel path. One request lowers both
the resolved ``L`` and ``L dagger L`` moments before the solve. After the
solve, quchip reconstructs ``S beta + L`` from the solve-bound input and
propagates the complete field trace to the requested reference plane.
Observable reconstruction reads per-device demodulation frequencies from
:class:`~quchip.engine.ir.ResolvedFrame` and forms the demodulation phase
``exp(i · 2π · ω · w · t)``, the inverse of the rotating-frame shift applied to
the Hamiltonian during assembly.
"""
from __future__ import annotations
from dataclasses import dataclass
from typing import Any, Literal, Mapping, Sequence
import numpy as np
from quchip.backend import Backend, SolverResult, _backend_context
from quchip.chip.chip import Chip
from quchip.declarative.expr import as_operator_expr, materialize_expr
from quchip.devices.base import BaseDevice
from quchip.engine.basis import BasisRecord
from quchip.engine.ir import CanonicalOperator, ResolvedFrame
from quchip.engine.reference import carrier_transfer, has_filter, time_shift
from quchip.observables import OutputField, is_output_field
from quchip.results.results import ObservableTrace, OutputFieldTrace
from quchip.utils.constants import TWO_PI
from quchip.utils.jax_utils import array_namespace as _array_namespace
from quchip.utils.labeling import resolve_label
from quchip.engine.bands import local_mode_bands
EOpKey = str | tuple[str, str]
BandWeight = int | tuple[int, int]
DeviceLabels = str | tuple[str, str]
[docs]
def prepare_local_op(
dev: BaseDevice,
spec: str | Any,
basis: BasisRecord,
backend: Any,
) -> Any:
"""Return one observable in the resolved local solver basis."""
with _backend_context(backend):
default_pauli = {"X": "sigma_x", "Y": "sigma_y", "Z": "sigma_z"}.get(spec) if isinstance(spec, str) else None
if (
default_pauli is not None
and type(dev).local_operator is BaseDevice.local_operator
and getattr(type(dev), default_pauli) is getattr(BaseDevice, default_pauli)
):
authored = dev._pauli_operator(default_pauli, basis=basis)
else:
authored = dev.local_operator(spec) if isinstance(spec, str) else spec
authored = as_operator_expr(
authored,
labels=(dev.label,),
dims=(basis.native_dim,),
name=rf"\hat O_{{{dev.label}}}",
owner=dev,
scope=dev.label,
)
local = materialize_expr(authored, backend)
if local.shape != (basis.native_dim, basis.native_dim):
raise ValueError(
f"Authored operator for {dev.label!r} must have shape "
f"{(basis.native_dim, basis.native_dim)}, got {local.shape}."
)
if basis.kind == "native":
return local
matrix = backend.to_array(local)
return backend.from_array(
basis.transform_operator(matrix),
dims=[[basis.resolved_dim], [basis.resolved_dim]],
)
def _resolve_eop_key(key: Any) -> EOpKey:
if isinstance(key, str):
return key
if isinstance(key, tuple) and len(key) == 2:
return (resolve_label(key[0]), resolve_label(key[1]))
return resolve_label(key)
[docs]
def decompose_eops(
e_ops_dict: Mapping[EOpKey, Any],
chip: Chip,
backend: Backend,
bases: Mapping[str, Any] | None = None,
engine_result: Any | None = None,
*,
space: Any | None = None,
) -> tuple[list[Any], list[BandMeta | OutputMeta]]:
"""Flatten dict-form ``e_ops`` into ``(ops, meta)`` ready for the solver.
Supported key shapes:
* ``"device_label"`` (or a device object resolved to its label) with a
single local operator. The operator is split into single-mode bands on
that device.
* ``("label_a", "label_b")`` with a tuple ``(op_a, op_b)``. Each operator
is split into single-mode bands on its device. All band pairs are
tensored to give two-body ``(w_a, w_b)`` bands.
* A string trace name with an ``OutputField`` value. The resolved first
moments and the normally ordered second moments are reconstructed after
the solve.
For each band, the returned :class:`BandMeta` records the original key, the
weight, and the sub-index, which identifies each entry when the user passed
a list of operators for one key. The backend solver can use the matching
``ops`` list directly. Device operators act on the retained labels of
``space``, or else of the engine result's excitation space when it has
one, and on the product space otherwise.
"""
if bases is None:
bases = chip.resolve().bases
flat_ops: list[Any] = []
meta: list[BandMeta | OutputMeta] = []
output_exposures: set[str] = set()
dims = chip.dims
if space is None and engine_result is not None:
space = engine_result.excitation_space
for raw_key, val in e_ops_dict.items():
key = _resolve_eop_key(raw_key)
if is_output_field(val):
if not isinstance(key, str):
raise ValueError("Output observable trace names must be strings.")
if val.exposure in output_exposures:
raise ValueError(
f"Output exposure {val.exposure!r} may be requested only once per solve."
)
output_exposures.add(val.exposure)
if engine_result is None:
raise ValueError(
"Output observables require a resolved transient solve; "
"pass them to QuantumSequence.simulate(..., e_ops=...)."
)
responses = {} if chip.port_network is None else chip.port_network._responses_by_output()
if val.exposure in responses:
raise NotImplementedError(
f"Output field {val.exposure!r} carries the emission of retained multiport mode responses "
f"{list(responses[val.exposure])}. Time-domain dynamics omit their poles. "
"Keep the boundary mode for time-domain output fields."
)
_append_output_eops(
flat_ops,
meta,
observable=val,
engine_result=engine_result,
backend=backend,
)
continue
if isinstance(key, str):
device_label = key
if isinstance(val, tuple):
raise ValueError(f"Single-device key {device_label!r} requires one operator, got tuple value")
dev_idx = chip.device_index(device_label)
dev = chip.device_map[device_label]
basis = bases[device_label]
dev_dim = basis.resolved_dim
if isinstance(val, list):
ops_with_idx: list[tuple[int | None, Any]] = [(i, op) for i, op in enumerate(val)]
else:
ops_with_idx = [(None, val)]
for sub_idx, op in ops_with_idx:
op = prepare_local_op(dev, op, basis, backend)
bands = local_mode_bands(backend, op, dim=dev_dim, label=device_label,
semantic_to_solver=basis.energy_to_solver()) or [(0, op)]
for weight, band in bands:
embedded = (backend.embed(band, dev_idx, dims) if space is None
else space.embed(backend, band, (dev_idx,)))
flat_ops.append(embedded)
meta.append(
BandMeta(
key=device_label,
weight=weight,
device_labels=device_label,
sub_index=sub_idx,
)
)
continue
if isinstance(key, tuple) and len(key) == 2:
label_a, label_b = key
if not (isinstance(val, tuple) and len(val) == 2):
raise ValueError(f"Two-device key {key!r} requires a tuple value (op_a, op_b)")
op_a, op_b = val
idx_a = chip.device_index(label_a)
idx_b = chip.device_index(label_b)
dev_a = chip.device_map[label_a]
dev_b = chip.device_map[label_b]
basis_a = bases[label_a]
basis_b = bases[label_b]
dim_a = basis_a.resolved_dim
dim_b = basis_b.resolved_dim
op_a = prepare_local_op(dev_a, op_a, basis_a, backend)
op_b = prepare_local_op(dev_b, op_b, basis_b, backend)
for w_a, band_a in local_mode_bands(
backend,
op_a,
dim=dim_a,
label=label_a,
semantic_to_solver=basis_a.energy_to_solver(),
) or [(0, op_a)]:
for w_b, band_b in local_mode_bands(
backend,
op_b,
dim=dim_b,
label=label_b,
semantic_to_solver=basis_b.energy_to_solver(),
) or [(0, op_b)]:
product = backend.tensor(band_a, band_b)
embedded = (backend.embed_two_body(product, idx_a, idx_b, dims) if space is None
else space.embed(backend, product, (idx_a, idx_b)))
flat_ops.append(embedded)
meta.append(
BandMeta(
key=(label_a, label_b),
weight=(w_a, w_b),
device_labels=(label_a, label_b),
)
)
continue
raise ValueError("e_ops dict keys must be device labels (str) or 2-tuples of device labels")
return flat_ops, meta
def _append_output_eops(
flat_ops: list[Any],
meta: list[BandMeta | OutputMeta],
*,
observable: OutputField,
engine_result: Any,
backend: Backend,
) -> None:
"""Lower one field recipe to the moments the backend must evaluate."""
if engine_result.slh.output_network is not None:
raise NotImplementedError("Branched reference networks currently support stationary VNA measurements. "
"Transient output fields require propagation through the output graph.")
external = {channel.key: channel for channel in engine_result.slh.external_channels}
channel = external.get(observable.exposure)
if channel is None:
raise ValueError(
f"Unknown output exposure {observable.exposure!r}. "
f"Available exposures: {list(external)}."
)
if has_filter(channel.reference.outbound) and (
engine_result.resolved_frame.mode == "lab" or channel.collapse.frame_frequency is None
):
raise ValueError(
f"Output plane {observable.exposure!r} has an outbound filter, but a lab-frame "
"channel has no carrier at which to evaluate transient output. Simulate in a "
"rotating frame at the readout carrier."
)
coupling = channel.coupling
flat_ops.append(backend.from_canonical_operator(coupling))
meta.append(
OutputMeta(
key=observable.exposure,
observable=observable,
role="amplitude",
)
)
flat_ops.append(collapse_number_operator(coupling, backend, tag=f"output_flux:{observable.exposure}"))
meta.append(
OutputMeta(
key=observable.exposure,
observable=observable,
role="flux",
)
)
[docs]
def collapse_number_operator(coupling: CanonicalOperator, backend: Backend, *, tag: str) -> Any:
"""Build the backend-native ``L†L`` operator used for collapse and output fluxes."""
values = coupling.to_dense()
xp = backend.array_module
number = xp.conj(xp.swapaxes(values, -1, -2)) @ values
return backend.from_canonical_operator(
CanonicalOperator.from_dense(
number,
dims=coupling.dims,
basis=coupling.basis,
subsystem_labels=coupling.subsystem_labels,
tag=tag,
)
)
[docs]
def recombine_expect(
flat_expect: Sequence[Any],
meta_list: Sequence[BandMeta],
tlist: Sequence[float] | Any,
frame_freqs: Mapping[str, float],
*,
direction: str = "demodulate",
) -> tuple[dict[EOpKey, Any], dict[EOpKey, Any]]:
"""Recombine flattened expectations into ``(band_sum, phase_corrected)`` dicts.
For each band with weight ``w`` on device(s) with demodulation
frequency ``ω_demod``, the correction factor is
.. math::
\\exp\\!\\big(\\pm i \\cdot 2\\pi \\cdot \\omega_{\\text{demod}}
\\cdot w \\cdot t \\big),
with the sign selected by *direction*. ``"demodulate"`` (``+1``) moves the
expectation from the simulation's rotating frame back to the control frame.
``"remodulate"`` (``−1``) does the reverse.
Returns two dicts. ``band_sum`` is the raw sum over bands for each key (no
phase correction). ``phase_corrected`` is the physically meaningful
expectation in the control frame.
"""
if direction not in ("demodulate", "remodulate"):
raise ValueError(f"direction must be 'demodulate' or 'remodulate', got {direction!r}")
if len(flat_expect) != len(meta_list):
raise ValueError(
f"flat_expect and meta_list must have the same length, got {len(flat_expect)} and {len(meta_list)}"
)
sign = +1 if direction == "demodulate" else -1
sample = flat_expect[0] if flat_expect else tlist
xp = _array_namespace(sample)
times = xp.asarray(tlist, dtype=float)
_AccKey = tuple[EOpKey, int | None]
band_sum_acc: dict[_AccKey, Any] = {}
corrected_acc: dict[_AccKey, Any] = {}
_has_sub: dict[EOpKey, bool] = {}
for series, meta in zip(flat_expect, meta_list):
values = xp.asarray(series, dtype=complex)
# Leading axes, such as the runs of a stochastic solve, share the time axis.
if values.shape[-1:] != times.shape:
raise ValueError(f"Expectation series shape {values.shape} does not end with tlist shape {times.shape}")
acc_key: _AccKey = (meta.key, meta.sub_index)
if meta.sub_index is not None:
_has_sub[meta.key] = True
if acc_key not in band_sum_acc:
band_sum_acc[acc_key] = xp.zeros_like(values, dtype=complex)
corrected_acc[acc_key] = xp.zeros_like(values, dtype=complex)
band_sum_acc[acc_key] += values
if isinstance(meta.weight, int):
if not isinstance(meta.device_labels, str):
raise ValueError("Single-device weight requires string device label")
omega_eff = frame_freqs[meta.device_labels] * meta.weight
else:
if not isinstance(meta.device_labels, tuple) or len(meta.device_labels) != 2:
raise ValueError("Two-device weight requires 2-tuple device labels")
label_a, label_b = meta.device_labels
w_a, w_b = meta.weight
omega_eff = frame_freqs[label_a] * w_a + frame_freqs[label_b] * w_b
phase = xp.exp(sign * 1j * TWO_PI * omega_eff * times)
corrected_acc[acc_key] += phase * values
band_sum: dict[EOpKey, Any] = {}
phase_corrected: dict[EOpKey, Any] = {}
for (orig_key, sub_idx), bs_arr in band_sum_acc.items():
pc_arr = corrected_acc[(orig_key, sub_idx)]
if orig_key in _has_sub:
if orig_key not in band_sum:
band_sum[orig_key] = []
phase_corrected[orig_key] = []
band_sum[orig_key].append(bs_arr)
phase_corrected[orig_key].append(pc_arr)
else:
band_sum[orig_key] = bs_arr
phase_corrected[orig_key] = pc_arr
return band_sum, phase_corrected
[docs]
def build_observable_traces(
solver_result: SolverResult,
tlist: np.ndarray,
*,
dict_meta: list[BandMeta | OutputMeta],
resolved_frame: ResolvedFrame,
engine_result: Any,
) -> tuple[
dict[EOpKey, ObservableTrace | list[ObservableTrace]],
dict[EOpKey, OutputFieldTrace],
]:
"""Wrap solver expectations into phase-corrected :class:`ObservableTrace` objects.
Gets the flat per-band expectations from *solver_result* and calls
:func:`recombine_expect` with the demodulation frequencies from
*resolved_frame*. Returns a dict keyed by the user's original e_ops keys.
Each value is an :class:`ObservableTrace`, or a list when the user supplied
multiple operators for the same key.
"""
raw_expect = solver_result.expect
if isinstance(raw_expect, dict):
flat_expect: Sequence[Any] = list(raw_expect.values())
else:
flat_expect = raw_expect or []
ordinary_expect = [
values
for values, meta in zip(flat_expect, dict_meta, strict=True)
if isinstance(meta, BandMeta)
]
ordinary_meta = [meta for meta in dict_meta if isinstance(meta, BandMeta)]
band_sum, phase_corrected = recombine_expect(
flat_expect=ordinary_expect,
meta_list=ordinary_meta,
tlist=tlist,
frame_freqs=resolved_frame.demod_freqs,
direction="demodulate",
)
traces: dict[EOpKey, ObservableTrace | list[ObservableTrace]] = {}
for key, values in phase_corrected.items():
raw = band_sum[key]
if isinstance(values, list):
traces[key] = [
ObservableTrace(values=value, raw=raw_value)
for value, raw_value in zip(values, raw)
]
else:
traces[key] = ObservableTrace(values=values, raw=raw)
return (
traces,
_build_output_traces(
flat_expect,
dict_meta,
tlist=tlist,
engine_result=engine_result,
),
)
def _build_output_traces(
flat_expect: Sequence[Any],
meta_list: Sequence[BandMeta | OutputMeta],
*,
tlist: Any,
engine_result: Any,
) -> dict[EOpKey, OutputFieldTrace]:
"""Reconstruct ``b_out = S b_in + L`` moments at reference planes."""
output_pairs = [
(values, meta)
for values, meta in zip(flat_expect, meta_list, strict=True)
if isinstance(meta, OutputMeta)
]
if not output_pairs:
return {}
sample = output_pairs[0][0]
xp = _array_namespace(sample)
times = xp.asarray(tlist, dtype=float)
external = engine_result.slh.external_channels
exposure_index = {channel.key: index for index, channel in enumerate(external)}
incident = [xp.zeros_like(times, dtype=complex) for _ in external]
from quchip.engine.ir import evaluate_signal_program
for bound in engine_result.coherent_inputs:
index = exposure_index.get(bound.exposure)
if index is not None:
incident[index] = incident[index] + evaluate_signal_program(
bound.beta,
times,
xp=xp,
)
grouped: dict[EOpKey, dict[str, Any]] = {}
specs: dict[EOpKey, OutputField] = {}
for values, meta in output_pairs:
grouped.setdefault(meta.key, {})[meta.role] = xp.asarray(values)
specs[meta.key] = meta.observable
traces: dict[EOpKey, OutputFieldTrace] = {}
for key, moments in grouped.items():
observable = specs[key]
output_index = exposure_index[observable.exposure]
channel = external[output_index]
direct = xp.zeros_like(times, dtype=complex)
for input_index, beta in enumerate(incident):
direct = direct + xp.asarray(engine_result.slh.S[output_index, input_index]) * beta
carrier = xp.exp(-1j * TWO_PI * xp.asarray(channel.carrier) * times)
lowering = carrier * moments["amplitude"]
field = direct + lowering
number = xp.real(moments["flux"])
boundary_flux = (
xp.abs(direct) ** 2
+ 2.0 * xp.real(xp.conj(direct) * lowering)
+ number
)
factor = carrier_transfer(channel.reference.outbound, channel.carrier, xp)
reported = factor * field
reported_flux = xp.abs(factor) ** 2 * boundary_flux
outbound_shift = time_shift(channel.reference.outbound)
traces[key] = OutputFieldTrace(
exposure=observable.exposure,
times=times,
amplitude=_delay_trace(reported, times, outbound_shift, xp),
photon_flux=_delay_trace(reported_flux, times, outbound_shift, xp),
raw_amplitude=field,
raw_photon_flux=boundary_flux,
)
return traces
def _delay_trace(values: Any, times: Any, delay: Any, xp: Any) -> Any:
"""Propagate one boundary trace to its reciprocal external reference plane."""
from quchip.utils.jax_utils import maybe_concrete_scalar
concrete = maybe_concrete_scalar(delay)
if concrete is not None and concrete == 0.0:
return values
query = times - xp.asarray(delay)
real = xp.interp(query, times, xp.real(values), left=0.0, right=0.0)
if not xp.iscomplexobj(values):
return real
imag = xp.interp(query, times, xp.imag(values), left=0.0, right=0.0)
return real + 1j * imag
