Source code for arcade.features.ngrc.extractor

"""NG-RC polynomial feature extractor."""

from __future__ import annotations

import itertools
from typing import Any

import numpy as np

from ..base import BaseFeatureExtractor, register_extractor


def apply_boxcar_masks(
    traces: np.ndarray,
    boxcar_ends: list[int] | None,
) -> np.ndarray:
    """Zero IQ samples after per-qubit end indices."""
    if boxcar_ends is None:
        return np.asarray(traces, dtype=np.float64)

    out = np.asarray(traces, dtype=np.float64).copy()
    if out.ndim == 3:
        end = int(boxcar_ends[0]) if boxcar_ends else out.shape[1]
        if end < out.shape[1]:
            out[:, end:, :] = 0.0
        return out

    if out.ndim != 4:
        raise ValueError(f"Expected 3D or 4D traces, got {out.shape}")

    n_q = out.shape[1]
    ends = list(boxcar_ends)
    if len(ends) < n_q:
        ends = ends + [out.shape[2]] * (n_q - len(ends))
    for q in range(n_q):
        end = min(int(ends[q]), out.shape[2])
        if end < out.shape[2]:
            out[:, q, end:, :] = 0.0
    return out


def window_average_iq(traces: np.ndarray, window_size: int) -> np.ndarray:
    """Average I,Q within fixed non-overlapping windows."""
    traces = np.asarray(traces, dtype=np.float64)
    if window_size < 1:
        raise ValueError("window_size must be >= 1")

    if traces.ndim == 3:
        n_time = traces.shape[1]
        n_win = n_time // window_size
        if n_win < 1:
            raise ValueError(f"Trace too short for window_size={window_size}")
        trimmed = traces[:, : n_win * window_size, :]
        reshaped = trimmed.reshape(traces.shape[0], n_win, window_size, 2)
        return reshaped.mean(axis=2)

    if traces.ndim == 4:
        return np.stack(
            [window_average_iq(traces[:, q, :, :], window_size) for q in range(traces.shape[1])],
            axis=1,
        )

    raise ValueError(f"Unsupported trace shape {traces.shape}")


def _shot_features(averaged: np.ndarray, degree: int) -> np.ndarray:
    """Build feature vector for one shot from window-averaged IQ."""
    flat = averaged.reshape(-1)
    parts: list[np.ndarray] = [np.ones(1, dtype=np.float64), flat]
    if degree >= 2:
        for i, j in itertools.combinations_with_replacement(range(len(flat)), 2):
            parts.append(np.asarray([flat[i] * flat[j]], dtype=np.float64))
    if degree >= 3:
        for i, j, k in itertools.combinations_with_replacement(range(len(flat)), 3):
            parts.append(np.asarray([flat[i] * flat[j] * flat[k]], dtype=np.float64))
    return np.concatenate(parts)


def expand_features(averaged_traces: np.ndarray, *, degree: int = 2) -> np.ndarray:
    """Expand window-averaged IQ to NG-RC feature matrix."""
    averaged_traces = np.asarray(averaged_traces, dtype=np.float64)
    if averaged_traces.ndim in (3, 4):
        rows = [_shot_features(averaged_traces[i], degree) for i in range(averaged_traces.shape[0])]
        return np.stack(rows, axis=0)
    raise ValueError(f"Unsupported shape {averaged_traces.shape}")




@register_extractor("ngrc")
class NGRCFeatureExtractor(BaseFeatureExtractor):
    """NG-RC polynomial feature extractor.

    Combines: apply_boxcar_masks -> window_average_iq -> expand_features.
    """

    feature_kind = "ngrc_expansion"

    def __init__(
        self,
        window_size: int = 10,
        polynomial_degree: int = 2,
        boxcar_ends: list[int] | None = None,
    ) -> None:
        self.window_size = window_size
        self.polynomial_degree = polynomial_degree
        self.boxcar_ends = boxcar_ends



    def extract(self, traces: np.ndarray, **kwargs: Any) -> np.ndarray:
        """Extract NG-RC polynomial features from traces."""
        masked = apply_boxcar_masks(traces, self.boxcar_ends)
        averaged = window_average_iq(masked, self.window_size)
        return expand_features(averaged, degree=self.polynomial_degree)