Source code for arcade.features.base

"""Base class and registry for feature extractors."""

from __future__ import annotations

import importlib
from abc import ABC, abstractmethod
from typing import Any

import numpy as np

_EXTRACTOR_REGISTRY: dict[str, type["BaseFeatureExtractor"]] = {}




def register_extractor(name: str) -> Any:
    """Register a :class:`BaseFeatureExtractor` subclass under *name*.

    Args:
        name: Registry key used when selecting extractors from config.

    Returns:
        A class decorator that stores the class in the registry.
    """

    def decorator(cls: type["BaseFeatureExtractor"]) -> type["BaseFeatureExtractor"]:
        _EXTRACTOR_REGISTRY[name] = cls
        cls.name = name
        return cls

    return decorator





def get_extractor(name: str) -> type["BaseFeatureExtractor"]:
    """Resolve an extractor class by registry name or dotted import path.

    Args:
        name: Registered key or ``package.module.ClassName``.

    Returns:
        The extractor class.

    Raises:
        KeyError: Unknown registry name.
        TypeError: Dotted path is not a ``BaseFeatureExtractor`` subclass.
    """
    if name in _EXTRACTOR_REGISTRY:
        return _EXTRACTOR_REGISTRY[name]

    if "." in name:
        module_path, cls_name = name.rsplit(".", 1)
        mod = importlib.import_module(module_path)
        cls = getattr(mod, cls_name)
        if isinstance(cls, type) and issubclass(cls, BaseFeatureExtractor):
            return cls
        raise TypeError(f"{name} resolved to {cls!r}, not a BaseFeatureExtractor subclass.")

    raise KeyError(
        f"Unknown extractor {name!r}. "
        f"Registered: {sorted(_EXTRACTOR_REGISTRY)}. "
        f"For external extractors, use a dotted import path."
    )





def list_extractors() -> list[str]:
    """Return sorted names of all registered feature extractors."""
    return sorted(_EXTRACTOR_REGISTRY)





class BaseFeatureExtractor(ABC):
    """Contract for turning IQ / ADC traces into a feature matrix.

    Matched filters, NG-RC polynomials, and path signatures all implement this
    interface. Subclass and decorate with :func:`register_extractor`.

    Attributes:
        name: Registry name (set by the decorator).
        feature_kind: Short tag describing the output feature family.

    Example::

        @register_extractor("my_features")
        class MyExtractor(BaseFeatureExtractor):
            feature_kind = "custom"

            def extract(self, traces, **kwargs):
                return np.zeros((len(traces), 10))
    """

    name: str = "base"
    feature_kind: str = "generic"



    def fit(self, train_traces: np.ndarray, train_labels: np.ndarray, **kwargs: Any) -> "BaseFeatureExtractor":
        """Optionally fit templates / statistics on training traces.

        Args:
            train_traces: Training IQ or ADC array (layout is extractor-specific).
            train_labels: Integer labels aligned with *train_traces*.
            **kwargs: Extractor-specific options.

        Returns:
            ``self``. Default implementation is a no-op.
        """
        return self




    @abstractmethod
    def extract(self, traces: np.ndarray, **kwargs: Any) -> np.ndarray:
        """Extract features from traces.

        Args:
            traces: Input traces; shape depends on extractor type.
            **kwargs: Extractor-specific options.

        Returns:
            Feature matrix shaped ``(n_shots, feature_dim)``.
        """
        ...