"""Boxcar averaging, truncation, IQ normalization."""
from __future__ import annotations
from typing import Any
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
import numpy as np
from ..._logging import get_logger
log = get_logger(__name__)
def preprocess(
traces: Any,
*,
boxcar_window: int | None = None,
truncate_at: int | None = None,
normalize: bool = False,
) -> np.ndarray | dict[int, np.ndarray]:
"""Apply preprocessing to ``(n_shots, time, iq)`` or per-qubit dict."""
if isinstance(traces, dict):
out: dict[int, np.ndarray] = {}
for q, ta in traces.items():
out[q] = preprocess(
ta,
boxcar_window=boxcar_window,
truncate_at=truncate_at,
normalize=normalize,
)
return out
traces_arr = np.asarray(traces, dtype=np.float64)
if boxcar_window is not None and boxcar_window > 1:
traces_arr = _boxcar(traces_arr, boxcar_window)
log.info(
"Boxcar window=%d -> new time dim=%d",
boxcar_window,
traces_arr.shape[1],
)
if truncate_at is not None:
traces_arr = traces_arr[:, :truncate_at, :]
log.info("Truncated to %d time steps", traces_arr.shape[1])
if normalize:
traces_arr = _normalize(traces_arr.copy())
log.info("Applied IQ normalization")
return traces_arr
def _boxcar(traces: np.ndarray, window: int) -> np.ndarray:
"""Average over non-overlapping windows along the time axis."""
n_shots, n_time, n_ch = traces.shape
n_bins = n_time // window
trimmed = traces[:, : n_bins * window, :]
return trimmed.reshape(n_shots, n_bins, window, n_ch).mean(axis=2)
def _normalize(traces: np.ndarray) -> np.ndarray:
"""Zero-mean, unit-variance normalization per IQ channel."""
out = traces
for ch in range(traces.shape[-1]):
channel = out[..., ch]
mean = float(channel.mean())
std = float(channel.std())
if std > 0:
out[..., ch] = (channel - mean) / std
else:
out[..., ch] = channel - mean
return out
[docs]
@register_processor("boxcar")
class BoxcarProcessor(BaseDataProcessor):
"""Non-overlapping boxcar (moving-average) downsampling."""
def __init__(self, window: int = 25) -> None:
self.window = window
[docs]
def process(self, bundle: DataBundle, **kwargs: Any) -> DataBundle:
window = kwargs.get("window", self.window)
target = bundle.demod_traces if bundle.demod_traces is not None else bundle.traces
result = preprocess(target, boxcar_window=window)
if bundle.demod_traces is not None:
bundle.demod_traces = result
else:
bundle.traces = result
return bundle
@register_processor("normalize")
class NormalizationProcessor(BaseDataProcessor):
"""Zero-mean, unit-variance IQ normalization."""
def process(self, bundle: DataBundle, **kwargs: Any) -> DataBundle:
target = bundle.demod_traces if bundle.demod_traces is not None else bundle.traces
result = preprocess(target, normalize=True)
if bundle.demod_traces is not None:
bundle.demod_traces = result
else:
bundle.traces = result
return bundle