Source code for arcade.features.path_signature.extractor

"""Path signature feature extractor (standalone)."""

from __future__ import annotations

from typing import Any

import numpy as np

from ..base import BaseFeatureExtractor, register_extractor
from .path_builder import build_weighted_path
from .signature import compute_signature_features
from .weight_function import compute_weights




@register_extractor("path_signature")
class PathSignatureExtractor(BaseFeatureExtractor):
    """Path signature feature extractor.

    Combines: compute_weights -> build_weighted_path -> compute_signature_features.
    """

    feature_kind = "signature"

    def __init__(self, signature_depth: int = 5, time_augmentation: bool = False) -> None:
        self.signature_depth = signature_depth
        self.time_augmentation = time_augmentation
        self.weights: np.ndarray | None = None



    def fit(
        self,
        train_traces: np.ndarray,
        train_labels: np.ndarray,
        **kwargs: Any,
    ) -> "PathSignatureExtractor":
        """Compute weight function from labelled training data."""
        num_levels = kwargs.get("num_levels", 2)
        self.weights = compute_weights(train_traces, train_labels, num_levels=num_levels)
        return self




    def extract(self, traces: np.ndarray, **kwargs: Any) -> np.ndarray:
        """Extract path signature features from traces."""
        traces = np.asarray(traces, dtype=np.float64)
        if self.weights is not None:
            paths = build_weighted_path(traces, self.weights)
        else:
            paths = traces
        return compute_signature_features(
            paths,
            depth=self.signature_depth,
            time_augmentation=self.time_augmentation,
        )