"""IQ demodulation — functional API and processor."""
from __future__ import annotations
from typing import Any, Sequence
import numpy as np
from ..._logging import get_logger
from ..base import BaseDataProcessor, register_processor
from ..bundle import DataBundle
log = get_logger(__name__)
from typing import Any, Sequence
import numpy as np
from ..._logging import get_logger
log = get_logger(__name__)
def demodulate(
traces_raw: np.ndarray,
frequencies: Sequence[float],
sampling_rate: float,
*,
time_bin_width: float | None = None,
skip_samples: int = 0,
correct_iq_offset: bool = True,
correct_iq_amplitude: bool = True,
) -> dict[int, np.ndarray]:
"""Digitally demodulate multiplexed IQ data at listed readout tones.
Returns:
``{qubit_index: iq}`` with each array shaped ``(n_shots, n_time, 2)``.
"""
traces_raw = np.asarray(traces_raw, dtype=np.float64)
if traces_raw.ndim != 3 or traces_raw.shape[2] != 2:
raise ValueError(
f"Expected shape (n_shots, n_samples, 2), got {traces_raw.shape}",
)
n_shots = traces_raw.shape[0]
data_i = traces_raw[:, skip_samples:, 0]
data_q = traces_raw[:, skip_samples:, 1]
dlen = data_i.shape[1]
if correct_iq_offset:
data_i = data_i - data_i.mean()
data_q = data_q - data_q.mean()
if correct_iq_amplitude:
std_i = data_i.std()
std_q = data_q.std()
if std_q > 0:
data_q = data_q * (std_i / std_q)
dt = 1.0 / sampling_rate
time_vec = dt * (skip_samples + np.arange(dlen, dtype=np.float64))
log.info(
"Demodulating %d shots x %d samples at %d frequencies",
n_shots,
dlen,
len(frequencies),
)
result: dict[int, np.ndarray] = {}
for qi, freq in enumerate(frequencies):
cos_vec = np.cos(2.0 * np.pi * time_vec * freq)
sin_vec = np.sin(2.0 * np.pi * time_vec * freq)
demod_i = 2.0 * (cos_vec * data_i + sin_vec * data_q)
demod_q = 2.0 * (cos_vec * data_q - sin_vec * data_i)
iq = np.stack([demod_i, demod_q], axis=-1)
if time_bin_width is not None:
bin_size = max(1, int(sampling_rate * time_bin_width))
n_bins = dlen // bin_size
iq = iq[:, : n_bins * bin_size, :]
iq = iq.reshape(n_shots, n_bins, bin_size, 2).mean(axis=2)
result[qi] = iq.astype(np.float64)
shapes = ", ".join(f"q{k}={tuple(v.shape)}" for k, v in sorted(result.items()))
log.info("Demodulated: %s", shapes or "(empty)")
return result
def load_or_demodulate(
traces_raw: np.ndarray,
frequencies: Sequence[float] | None,
sampling_rate: float | None,
*,
pre_demodulated: bool = False,
num_qubits: int | None = None,
**kwargs: Any,
) -> dict[int, np.ndarray]:
"""Demodulate or wrap already-demodulated per-qubit arrays."""
traces_raw = np.asarray(traces_raw, dtype=np.float64)
if pre_demodulated:
if traces_raw.ndim == 3:
return {0: traces_raw}
if traces_raw.ndim == 4:
return {qi: traces_raw[qi] for qi in range(traces_raw.shape[0])}
raise ValueError(
"pre_demodulated=True expects 3D (single qubit) or 4D (stacked)",
)
if frequencies is None or sampling_rate is None:
raise ValueError(
"frequencies and sampling_rate required unless pre_demodulated=True",
)
dm = demodulate(traces_raw, frequencies, sampling_rate, **kwargs)
if num_qubits is not None and len(dm) != num_qubits:
log.warning(
"Got %d demod tones but expected num_qubits=%d",
len(dm),
num_qubits,
)
return dm
[docs]
@register_processor("demodulation")
class DemodulationProcessor(BaseDataProcessor):
"""Demodulate raw ADC traces into per-qubit IQ data."""
def __init__(
self,
frequencies: dict | None = None,
sampling_rate: float | None = None,
) -> None:
self.frequencies = frequencies
self.sampling_rate = sampling_rate
[docs]
def process(self, bundle: DataBundle, **kwargs: Any) -> DataBundle:
freqs = kwargs.get("frequencies") or self.frequencies
sr = kwargs.get("sampling_rate") or self.sampling_rate
if freqs is None or sr is None:
raise ValueError(
"frequencies and sampling_rate are required for demodulation"
)
freq_list: list[float]
if isinstance(freqs, dict):
freq_list = [freqs[k] for k in sorted(freqs)]
else:
freq_list = list(freqs)
demod_kwargs = {
k: kwargs[k]
for k in ("time_bin_width", "skip_samples", "correct_iq_offset", "correct_iq_amplitude")
if k in kwargs
}
bundle.demod_traces = demodulate(
bundle.traces, freq_list, sr, **demod_kwargs
)
return bundle