"""Matched filter feature extractor."""
from __future__ import annotations
from itertools import combinations
from typing import Any, Sequence
import numpy as np
from .base import BaseFeatureExtractor, register_extractor
from .filter_set import FilterSet
def matched_filter(
traces_a: np.ndarray,
traces_b: np.ndarray,
*,
threshold_percentile: float = 99.5,
) -> tuple[np.ndarray, float]:
"""Compute a single matched filter between two trace populations.
The filter envelope is ``mean(A - B) / var(A - B)`` computed on
flattened (time x IQ) traces.
"""
a = np.asarray(traces_a, dtype=np.float64)
b = np.asarray(traces_b, dtype=np.float64)
if a.ndim == 3:
a = a.reshape(len(a), -1)
if b.ndim == 3:
b = b.reshape(len(b), -1)
n = min(len(a), len(b))
if n == 0:
dim = a.shape[1] if len(a) > 0 else (b.shape[1] if len(b) > 0 else 2)
return np.zeros(dim), 0.0
rng = np.random.default_rng(42)
if len(a) > n:
a = a[rng.choice(len(a), n, replace=False)]
if len(b) > n:
b = b[rng.choice(len(b), n, replace=False)]
diff = a - b
mean_diff = diff.mean(axis=0)
var_diff = diff.var(axis=0)
var_diff[var_diff == 0] = 1.0
envelope = mean_diff / var_diff
scores_a = a @ envelope
threshold = float(np.percentile(scores_a, threshold_percentile))
return envelope, threshold
def compute_filters(
transition_data: Any,
num_qubits: int,
num_levels: int = 2,
*,
filter_types: Sequence[str] = ("MF",),
threshold_percentile: float = 99.5,
) -> FilterSet:
"""Compute all requested filters for every qubit."""
filter_types_upper = [ft.upper() for ft in filter_types]
fs = FilterSet(num_qubits=num_qubits, num_levels=num_levels)
if hasattr(transition_data, "per_qubit"):
td = transition_data.per_qubit
else:
td = transition_data
trace_length = None
for q in range(num_qubits):
qdata = td[q]
fs.filters[q] = {}
fs.thresholds[q] = {}
if "MF" in filter_types_upper:
for si, sj in combinations(range(num_levels), 2):
ta = qdata.get(f"traces_{si}", np.empty((0,)))
tb = qdata.get(f"traces_{sj}", np.empty((0,)))
if len(ta) == 0 or len(tb) == 0:
continue
name = f"MF_{si}_{sj}"
env, thr = matched_filter(ta, tb, threshold_percentile=threshold_percentile)
fs.filters[q][name] = env
fs.thresholds[q][name] = thr
if trace_length is None and ta.ndim >= 2:
trace_length = ta.shape[1]
if "RMF" in filter_types_upper:
for si in range(num_levels):
for sj in range(si):
ta = qdata.get(f"relax_{si}_{sj}", np.empty((0,)))
tb = qdata.get(f"traces_{sj}", np.empty((0,)))
if len(ta) == 0 or len(tb) == 0:
continue
name = f"RMF_{si}_{sj}"
env, thr = matched_filter(ta, tb, threshold_percentile=threshold_percentile)
fs.filters[q][name] = env
fs.thresholds[q][name] = thr
if "EMF" in filter_types_upper:
for si in range(num_levels):
for sj in range(si + 1, num_levels):
ta = qdata.get(f"excite_{si}_{sj}", np.empty((0,)))
tb = qdata.get(f"traces_{si}", np.empty((0,)))
if len(ta) == 0 or len(tb) == 0:
continue
name = f"EMF_{si}_{sj}"
env, thr = matched_filter(ta, tb, threshold_percentile=threshold_percentile)
fs.filters[q][name] = env
fs.thresholds[q][name] = thr
fs.trace_length = trace_length
return fs
def apply_filters(
traces: np.ndarray | dict,
filter_set: FilterSet,
) -> np.ndarray:
"""Project traces through a filter set to produce scalar features."""
per_qubit = isinstance(traces, dict)
if not per_qubit:
traces_arr = np.asarray(traces, dtype=np.float64)
n_shots = len(traces_arr)
flat_shared = traces_arr.reshape(n_shots, -1)
else:
first_key = next(iter(traces))
n_shots = len(traces[first_key])
n_features = filter_set.total_features
features = np.zeros((n_shots, n_features), dtype=np.float64)
col = 0
for q in sorted(filter_set.filters.keys()):
if per_qubit:
flat_q = np.asarray(traces[q], dtype=np.float64).reshape(n_shots, -1)
else:
flat_q = flat_shared
for name in filter_set.filter_names(q):
env = filter_set.filters[q][name]
dim = min(flat_q.shape[1], len(env))
features[:, col] = flat_q[:, :dim] @ env[:dim]
col += 1
return features
def recompute_at_length(
transition_data: Any,
traces: np.ndarray | dict,
num_qubits: int,
length: int,
*,
num_levels: int = 2,
filter_types: Sequence[str] = ("MF",),
threshold_percentile: float = 99.5,
) -> tuple[FilterSet, np.ndarray]:
"""Recompute filters at a truncated trace length and apply them."""
if hasattr(transition_data, "per_qubit"):
td = transition_data.per_qubit
else:
td = transition_data
truncated_td: dict[int, dict[str, Any]] = {}
for q, qdata in td.items():
truncated_td[q] = {}
for key, val in qdata.items():
if (
isinstance(val, np.ndarray)
and val.ndim >= 2
and key.startswith(("traces_", "relax_", "excite_"))
):
truncated_td[q][key] = val[:, :length, :]
else:
truncated_td[q][key] = val
fs = compute_filters(
truncated_td,
num_qubits,
num_levels,
filter_types=filter_types,
threshold_percentile=threshold_percentile,
)
fs.trace_length = length
if isinstance(traces, dict):
truncated_traces = {q: t[:, :length, :] for q, t in traces.items()}
else:
truncated_traces = traces[:, :length, :]
features = apply_filters(truncated_traces, fs)
return fs, features