Source code for arcade.features.relaxation_filter

"""Relaxation matched filter (REMF) feature extractor."""

from __future__ import annotations

from typing import Any

import numpy as np

from .base import BaseFeatureExtractor, register_extractor
from .filter_set import FilterSet
from .matched_filter import apply_filters, matched_filter


def _trace_population(traces: np.ndarray, mask: np.ndarray) -> np.ndarray:
    """Return IQ traces ``(n, t, 2)`` for boolean *mask*."""
    traces = np.asarray(traces, dtype=np.float64)
    if traces.ndim == 2:
        traces = traces[:, :, np.newaxis]
    if traces.shape[-1] != 2:
        raise ValueError(f"Expected IQ traces (n, t, 2), got {traces.shape}")
    sel = np.asarray(mask, dtype=bool)
    if not np.any(sel):
        return np.empty((0, traces.shape[1], 2), dtype=np.float64)
    return traces[sel]


def compute_remf_filter_set(
    traces: np.ndarray,
    expected_labels: np.ndarray,
    post_measured_labels: np.ndarray,
    *,
    threshold_percentile: float = 99.5,
) -> FilterSet:
    """Build three REMF envelopes from post-selected training subsets."""
    expected = np.asarray(expected_labels, dtype=np.int64)
    post = np.asarray(post_measured_labels, dtype=np.float64)
    valid_post = np.isfinite(post)
    post_int = np.zeros_like(expected)
    post_int[valid_post] = post[valid_post].astype(np.int64)

    true_0 = (expected == 0) & valid_post & (post_int == 0)
    true_1 = (expected == 1) & valid_post & (post_int == 1)
    relaxation = (expected == 1) & valid_post & (post_int == 0)
    excitation = (expected == 0) & valid_post & (post_int == 1)

    fs = FilterSet(num_qubits=1, num_levels=2)
    fs.filters[0] = {}
    fs.thresholds[0] = {}

    pairs = [
        ("REMF_true01", true_0, true_1),
        ("REMF_true0_relax", true_0, relaxation),
        ("REMF_true1_exc", true_1, excitation),
    ]
    for name, mask_a, mask_b in pairs:
        ta = _trace_population(traces, mask_a)
        tb = _trace_population(traces, mask_b)
        if len(ta) == 0 or len(tb) == 0:
            continue
        env, thr = matched_filter(ta, tb, threshold_percentile=threshold_percentile)
        fs.filters[0][name] = env
        fs.thresholds[0][name] = thr

    if not fs.filters[0]:
        raise ValueError(
            "Could not build REMF filters: no post-selected training subsets"
        )

    if traces.ndim >= 2:
        fs.trace_length = traces.shape[1] if traces.ndim == 3 else None
    return fs




@register_extractor("relaxation_filter")
class RelaxationFilterExtractor(BaseFeatureExtractor):
    """Relaxation matched filter (REMF) feature extractor.

    Uses post-selection to separate true state preparations from
    relaxation/excitation events.
    """

    feature_kind = "filter_features"

    def __init__(self) -> None:
        self.filter_set: FilterSet | None = None



    def fit(
        self,
        train_traces: np.ndarray,
        train_labels: np.ndarray,
        post_measured_labels: np.ndarray | None = None,
        **kwargs: Any,
    ) -> "RelaxationFilterExtractor":
        """Compute REMF templates from post-selected training data."""
        if post_measured_labels is None:
            raise ValueError("RelaxationFilterExtractor requires post_measured_labels")
        self.filter_set = compute_remf_filter_set(
            train_traces,
            train_labels,
            post_measured_labels,
            threshold_percentile=kwargs.get("threshold_percentile", 99.5),
        )
        return self




    def extract(self, traces: np.ndarray, **kwargs: Any) -> np.ndarray:
        """Project traces onto stored REMF templates."""
        if self.filter_set is None:
            raise RuntimeError("Must call fit() before extract()")
        return apply_filters(traces, self.filter_set)