eegdash 0.4.0.dev150__py3-none-any.whl → 0.4.0.dev162__py3-none-any.whl

eegdash/features/decorators.py

@@ -12,6 +12,21 @@ from .extractors import (
 
 
 class FeaturePredecessor:
+    """A decorator to specify parent extractors for a feature function.
+
+    This decorator attaches a list of parent extractor types to a feature
+    extraction function. This information can be used to build a dependency
+    graph of features.
+
+    Parameters
+    ----------
+    *parent_extractor_type : list of Type
+        A list of feature extractor classes (subclasses of
+        :class:`~eegdash.features.extractors.FeatureExtractor`) that this
+        feature depends on.
+
+    """
+
     def __init__(self, *parent_extractor_type: List[Type]):
         parent_cls = parent_extractor_type
         if not parent_cls:
@@ -20,17 +35,58 @@ class FeaturePredecessor:
             assert issubclass(p_cls, FeatureExtractor)
         self.parent_extractor_type = parent_cls
 
-    def __call__(self, func: Callable):
+    def __call__(self, func: Callable) -> Callable:
+        """Apply the decorator to a function.
+
+        Parameters
+        ----------
+        func : callable
+            The feature extraction function to decorate.
+
+        Returns
+        -------
+        callable
+            The decorated function with the `parent_extractor_type` attribute
+            set.
+
+        """
         f = _get_underlying_func(func)
         f.parent_extractor_type = self.parent_extractor_type
         return func
 
 
 class FeatureKind:
+    """A decorator to specify the kind of a feature.
+
+    This decorator attaches a "feature kind" (e.g., univariate, bivariate)
+    to a feature extraction function.
+
+    Parameters
+    ----------
+    feature_kind : MultivariateFeature
+        An instance of a feature kind class, such as
+        :class:`~eegdash.features.extractors.UnivariateFeature` or
+        :class:`~eegdash.features.extractors.BivariateFeature`.
+
+    """
+
     def __init__(self, feature_kind: MultivariateFeature):
         self.feature_kind = feature_kind
 
-    def __call__(self, func):
+    def __call__(self, func: Callable) -> Callable:
+        """Apply the decorator to a function.
+
+        Parameters
+        ----------
+        func : callable
+            The feature extraction function to decorate.
+
+        Returns
+        -------
+        callable
+            The decorated function with the `feature_kind` attribute set.
+
+        """
         f = _get_underlying_func(func)
         f.feature_kind = self.feature_kind
         return func
@@ -38,9 +94,33 @@ class FeatureKind:
 
 # Syntax sugar
 univariate_feature = FeatureKind(UnivariateFeature())
+"""Decorator to mark a feature as univariate.
+
+This is a convenience instance of :class:`FeatureKind` pre-configured for
+univariate features.
+"""
 
 
-def bivariate_feature(func, directed=False):
+def bivariate_feature(func: Callable, directed: bool = False) -> Callable:
+    """Decorator to mark a feature as bivariate.
+
+    This decorator specifies that the feature operates on pairs of channels.
+
+    Parameters
+    ----------
+    func : callable
+        The feature extraction function to decorate.
+    directed : bool, default False
+        If True, the feature is directed (e.g., connectivity from channel A
+        to B is different from B to A). If False, the feature is undirected.
+
+    Returns
+    -------
+    callable
+        The decorated function with the appropriate bivariate feature kind
+        attached.
+
+    """
     if directed:
         kind = DirectedBivariateFeature()
     else:
@@ -49,3 +129,8 @@ def bivariate_feature(func, directed=False):
 
 
 multivariate_feature = FeatureKind(MultivariateFeature())
+"""Decorator to mark a feature as multivariate.
+
+This is a convenience instance of :class:`FeatureKind` pre-configured for
+multivariate features, which operate on all channels simultaneously.
+"""

eegdash/features/extractors.py

@@ -7,7 +7,23 @@ import numpy as np
 from numba.core.dispatcher import Dispatcher
 
 
-def _get_underlying_func(func):
+def _get_underlying_func(func: Callable) -> Callable:
+    """Get the underlying function from a potential wrapper.
+
+    This helper unwraps functions that might be wrapped by `functools.partial`
+    or `numba.dispatcher.Dispatcher`.
+
+    Parameters
+    ----------
+    func : callable
+        The function to unwrap.
+
+    Returns
+    -------
+    callable
+        The underlying Python function.
+
+    """
     f = func
     if isinstance(f, partial):
         f = f.func
@@ -17,22 +33,46 @@ def _get_underlying_func(func):
 
 
 class TrainableFeature(ABC):
+    """Abstract base class for features that require training.
+
+    This ABC defines the interface for feature extractors that need to be
+    fitted on data before they can be used. It includes methods for fitting
+    the feature extractor and for resetting its state.
+    """
+
     def __init__(self):
         self._is_trained = False
         self.clear()
 
     @abstractmethod
     def clear(self):
+        """Reset the internal state of the feature extractor."""
         pass
 
     @abstractmethod
     def partial_fit(self, *x, y=None):
+        """Update the feature extractor's state with a batch of data.
+
+        Parameters
+        ----------
+        *x : tuple
+            The input data for fitting.
+        y : any, optional
+            The target data, if required for supervised training.
+
+        """
         pass
 
     def fit(self):
+        """Finalize the training of the feature extractor.
+
+        This method should be called after all data has been seen via
+        `partial_fit`. It marks the feature as fitted.
+        """
         self._is_fitted = True
 
     def __call__(self, *args, **kwargs):
+        """Check if the feature is fitted before execution."""
         if not self._is_fitted:
             raise RuntimeError(
                 f"{self.__class__} cannot be called, it has to be fitted first."
@@ -40,6 +80,22 @@ class TrainableFeature(ABC):
 
 
 class FeatureExtractor(TrainableFeature):
+    """A composite feature extractor that applies multiple feature functions.
+
+    This class orchestrates the application of a dictionary of feature
+    extraction functions to input data. It can handle nested extractors,
+    pre-processing, and trainable features.
+
+    Parameters
+    ----------
+    feature_extractors : dict[str, callable]
+        A dictionary where keys are feature names and values are the feature
+        extraction functions or other `FeatureExtractor` instances.
+    **preprocess_kwargs
+        Keyword arguments to be passed to the `preprocess` method.
+
+    """
+
     def __init__(
         self, feature_extractors: Dict[str, Callable], **preprocess_kwargs: Dict
     ):
@@ -63,30 +119,64 @@ class FeatureExtractor(TrainableFeature):
             if isinstance(fe, partial):
                 self.features_kwargs[fn] = fe.keywords
 
-    def _validate_execution_tree(self, feature_extractors):
+    def _validate_execution_tree(self, feature_extractors: dict) -> dict:
+        """Validate the feature dependency graph."""
         for fname, f in feature_extractors.items():
             f = _get_underlying_func(f)
             pe_type = getattr(f, "parent_extractor_type", [FeatureExtractor])
-            assert type(self) in pe_type
+            if type(self) not in pe_type:
+                raise TypeError(
+                    f"Feature '{fname}' cannot be a child of {type(self).__name__}"
+                )
         return feature_extractors
 
-    def _check_is_trainable(self, feature_extractors):
-        is_trainable = False
+    def _check_is_trainable(self, feature_extractors: dict) -> bool:
+        """Check if any of the contained features are trainable."""
         for fname, f in feature_extractors.items():
             if isinstance(f, FeatureExtractor):
-                is_trainable = f._is_trainable
-            else:
-                f = _get_underlying_func(f)
-                if isinstance(f, TrainableFeature):
-                    is_trainable = True
-            if is_trainable:
-                break
-        return is_trainable
+                if f._is_trainable:
+                    return True
+            elif isinstance(_get_underlying_func(f), TrainableFeature):
+                return True
+        return False
 
     def preprocess(self, *x, **kwargs):
+        """Apply pre-processing to the input data.
+
+        Parameters
+        ----------
+        *x : tuple
+            Input data.
+        **kwargs
+            Additional keyword arguments.
+
+        Returns
+        -------
+        tuple
+            The pre-processed data.
+
+        """
         return (*x,)
 
-    def __call__(self, *x, _batch_size=None, _ch_names=None):
+    def __call__(self, *x, _batch_size=None, _ch_names=None) -> dict:
+        """Apply all feature extractors to the input data.
+
+        Parameters
+        ----------
+        *x : tuple
+            Input data.
+        _batch_size : int, optional
+            The number of samples in the batch.
+        _ch_names : list of str, optional
+            The names of the channels in the input data.
+
+        Returns
+        -------
+        dict
+            A dictionary where keys are feature names and values are the
+            computed feature values.
+
+        """
         assert _batch_size is not None
         assert _ch_names is not None
         if self._is_trainable:
@@ -100,59 +190,83 @@ class FeatureExtractor(TrainableFeature):
                 r = f(*z, _batch_size=_batch_size, _ch_names=_ch_names)
             else:
                 r = f(*z)
-            f = _get_underlying_func(f)
-            if hasattr(f, "feature_kind"):
-                r = f.feature_kind(r, _ch_names=_ch_names)
+            f_und = _get_underlying_func(f)
+            if hasattr(f_und, "feature_kind"):
+                r = f_und.feature_kind(r, _ch_names=_ch_names)
             if not isinstance(fname, str) or not fname:
-                if isinstance(f, FeatureExtractor) or not hasattr(f, "__name__"):
-                    fname = ""
-                else:
-                    fname = f.__name__
+                fname = getattr(f_und, "__name__", "")
             if isinstance(r, dict):
-                if fname:
-                    fname += "_"
+                prefix = f"{fname}_" if fname else ""
                 for k, v in r.items():
-                    self._add_feature_to_dict(results_dict, fname + k, v, _batch_size)
+                    self._add_feature_to_dict(results_dict, prefix + k, v, _batch_size)
             else:
                 self._add_feature_to_dict(results_dict, fname, r, _batch_size)
         return results_dict
 
-    def _add_feature_to_dict(self, results_dict, name, value, batch_size):
-        if not isinstance(value, np.ndarray):
-            results_dict[name] = value
-        else:
+    def _add_feature_to_dict(
+        self, results_dict: dict, name: str, value: any, batch_size: int
+    ):
+        """Add a computed feature to the results dictionary."""
+        if isinstance(value, np.ndarray):
             assert value.shape[0] == batch_size
-            results_dict[name] = value
+        results_dict[name] = value
 
     def clear(self):
+        """Clear the state of all trainable sub-features."""
         if not self._is_trainable:
             return
-        for fname, f in self.feature_extractors_dict.items():
-            f = _get_underlying_func(f)
-            if isinstance(f, TrainableFeature):
-                f.clear()
+        for f in self.feature_extractors_dict.values():
+            if isinstance(_get_underlying_func(f), TrainableFeature):
+                _get_underlying_func(f).clear()
 
     def partial_fit(self, *x, y=None):
+        """Partially fit all trainable sub-features."""
         if not self._is_trainable:
             return
         z = self.preprocess(*x, **self.preprocess_kwargs)
-        for fname, f in self.feature_extractors_dict.items():
-            f = _get_underlying_func(f)
-            if isinstance(f, TrainableFeature):
-                f.partial_fit(*z, y=y)
+        if not isinstance(z, tuple):
+            z = (z,)
+        for f in self.feature_extractors_dict.values():
+            if isinstance(_get_underlying_func(f), TrainableFeature):
+                _get_underlying_func(f).partial_fit(*z, y=y)
 
     def fit(self):
+        """Fit all trainable sub-features."""
         if not self._is_trainable:
             return
-        for fname, f in self.feature_extractors_dict.items():
-            f = _get_underlying_func(f)
-            if isinstance(f, TrainableFeature):
+        for f in self.feature_extractors_dict.values():
+            if isinstance(_get_underlying_func(f), TrainableFeature):
                 f.fit()
         super().fit()
 
 
 class MultivariateFeature:
-    def __call__(self, x, _ch_names=None):
+    """A mixin for features that operate on multiple channels.
+
+    This class provides a `__call__` method that converts a feature array into
+    a dictionary with named features, where names are derived from channel
+    names.
+    """
+
+    def __call__(
+        self, x: np.ndarray, _ch_names: list[str] | None = None
+    ) -> dict | np.ndarray:
+        """Convert a feature array to a named dictionary.
+
+        Parameters
+        ----------
+        x : numpy.ndarray
+            The computed feature array.
+        _ch_names : list of str, optional
+            The list of channel names.
+
+        Returns
+        -------
+        dict or numpy.ndarray
+            A dictionary of named features, or the original array if feature
+            channel names cannot be generated.
+
+        """
         assert _ch_names is not None
         f_channels = self.feature_channel_names(_ch_names)
         if isinstance(x, dict):
@@ -163,37 +277,66 @@ class MultivariateFeature:
         return self._array_to_dict(x, f_channels)
 
     @staticmethod
-    def _array_to_dict(x, f_channels, name=""):
+    def _array_to_dict(
+        x: np.ndarray, f_channels: list[str], name: str = ""
+    ) -> dict | np.ndarray:
+        """Convert a numpy array to a dictionary with named keys."""
         assert isinstance(x, np.ndarray)
-        if len(f_channels) == 0:
-            assert x.ndim == 1
-            if name:
-                return {name: x}
-            return x
-        assert x.shape[1] == len(f_channels)
+        if not f_channels:
+            return {name: x} if name else x
+        assert x.shape[1] == len(f_channels), f"{x.shape[1]} != {len(f_channels)}"
         x = x.swapaxes(0, 1)
-        names = [f"{name}_{ch}" for ch in f_channels] if name else f_channels
+        prefix = f"{name}_" if name else ""
+        names = [f"{prefix}{ch}" for ch in f_channels]
         return dict(zip(names, x))
 
-    def feature_channel_names(self, ch_names):
+    def feature_channel_names(self, ch_names: list[str]) -> list[str]:
+        """Generate feature names based on channel names.
+
+        Parameters
+        ----------
+        ch_names : list of str
+            The names of the input channels.
+
+        Returns
+        -------
+        list of str
+            The names for the output features.
+
+        """
         return []
 
 
 class UnivariateFeature(MultivariateFeature):
-    def feature_channel_names(self, ch_names):
+    """A feature kind for operations applied to each channel independently."""
+
+    def feature_channel_names(self, ch_names: list[str]) -> list[str]:
+        """Return the channel names themselves as feature names."""
         return ch_names
 
 
 class BivariateFeature(MultivariateFeature):
-    def __init__(self, *args, channel_pair_format="{}<>{}"):
+    """A feature kind for operations on pairs of channels.
+
+    Parameters
+    ----------
+    channel_pair_format : str, default="{}<>{}"
+        A format string used to create feature names from pairs of
+        channel names.
+
+    """
+
+    def __init__(self, *args, channel_pair_format: str = "{}<>{}"):
         super().__init__(*args)
         self.channel_pair_format = channel_pair_format
 
     @staticmethod
-    def get_pair_iterators(n):
+    def get_pair_iterators(n: int) -> tuple[np.ndarray, np.ndarray]:
+        """Get indices for unique, unordered pairs of channels."""
         return np.triu_indices(n, 1)
 
-    def feature_channel_names(self, ch_names):
+    def feature_channel_names(self, ch_names: list[str]) -> list[str]:
+        """Generate feature names for each pair of channels."""
         return [
             self.channel_pair_format.format(ch_names[i], ch_names[j])
             for i, j in zip(*self.get_pair_iterators(len(ch_names)))
@@ -201,8 +344,11 @@ class BivariateFeature(MultivariateFeature):
 
 
 class DirectedBivariateFeature(BivariateFeature):
+    """A feature kind for directed operations on pairs of channels."""
+
     @staticmethod
-    def get_pair_iterators(n):
+    def get_pair_iterators(n: int) -> list[np.ndarray]:
+        """Get indices for all ordered pairs of channels (excluding self-pairs)."""
         return [
             np.append(a, b)
             for a, b in zip(np.tril_indices(n, -1), np.triu_indices(n, 1))

eegdash/features/inspect.py

@@ -5,7 +5,27 @@ from . import extractors, feature_bank
 from .extractors import FeatureExtractor, MultivariateFeature, _get_underlying_func
 
 
-def get_feature_predecessors(feature_or_extractor: Callable):
+def get_feature_predecessors(feature_or_extractor: Callable) -> list:
+    """Get the dependency hierarchy for a feature or feature extractor.
+
+    This function recursively traverses the `parent_extractor_type` attribute
+    of a feature or extractor to build a list representing its dependency
+    lineage.
+
+    Parameters
+    ----------
+    feature_or_extractor : callable
+        The feature function or :class:`FeatureExtractor` class to inspect.
+
+    Returns
+    -------
+    list
+        A nested list representing the dependency tree. For a simple linear
+        chain, this will be a flat list from the specific feature up to the
+        base `FeatureExtractor`. For multiple dependencies, it will contain
+        tuples of sub-dependencies.
+
+    """
     current = _get_underlying_func(feature_or_extractor)
     if current is FeatureExtractor:
         return [current]
@@ -20,18 +40,59 @@ def get_feature_predecessors(feature_or_extractor: Callable):
         return [current, tuple(predecessors)]
 
 
-def get_feature_kind(feature: Callable):
+def get_feature_kind(feature: Callable) -> MultivariateFeature:
+    """Get the 'kind' of a feature function.
+
+    The feature kind (e.g., univariate, bivariate) is typically attached by a
+    decorator.
+
+    Parameters
+    ----------
+    feature : callable
+        The feature function to inspect.
+
+    Returns
+    -------
+    MultivariateFeature
+        An instance of the feature kind (e.g., `UnivariateFeature()`).
+
+    """
     return _get_underlying_func(feature).feature_kind
 
 
-def get_all_features():
+def get_all_features() -> list[tuple[str, Callable]]:
+    """Get a list of all available feature functions.
+
+    Scans the `eegdash.features.feature_bank` module for functions that have
+    been decorated to have a `feature_kind` attribute.
+
+    Returns
+    -------
+    list[tuple[str, callable]]
+        A list of (name, function) tuples for all discovered features.
+
+    """
+
     def isfeature(x):
         return hasattr(_get_underlying_func(x), "feature_kind")
 
     return inspect.getmembers(feature_bank, isfeature)
 
 
-def get_all_feature_extractors():
+def get_all_feature_extractors() -> list[tuple[str, type[FeatureExtractor]]]:
+    """Get a list of all available `FeatureExtractor` classes.
+
+    Scans the `eegdash.features.feature_bank` module for all classes that
+    subclass :class:`~eegdash.features.extractors.FeatureExtractor`.
+
+    Returns
+    -------
+    list[tuple[str, type[FeatureExtractor]]]
+        A list of (name, class) tuples for all discovered feature extractors,
+        including the base `FeatureExtractor` itself.
+
+    """
+
     def isfeatureextractor(x):
         return inspect.isclass(x) and issubclass(x, FeatureExtractor)
 
@@ -41,7 +102,19 @@ def get_all_feature_extractors():
     ]
 
 
-def get_all_feature_kinds():
+def get_all_feature_kinds() -> list[tuple[str, type[MultivariateFeature]]]:
+    """Get a list of all available feature 'kind' classes.
+
+    Scans the `eegdash.features.extractors` module for all classes that
+    subclass :class:`~eegdash.features.extractors.MultivariateFeature`.
+
+    Returns
+    -------
+    list[tuple[str, type[MultivariateFeature]]]
+        A list of (name, class) tuples for all discovered feature kinds.
+
+    """
+
     def isfeaturekind(x):
         return inspect.isclass(x) and issubclass(x, MultivariateFeature)