From 7341a70aeb1d4c5e0ea52da0e4e19365ceb8a13d Mon Sep 17 00:00:00 2001 From: Hamza Abdelhedi Date: Sun, 22 Mar 2026 11:31:04 -0600 Subject: [PATCH 1/7] feat(descriptors): initial implementation of the descriptors module - Implement config-driven descriptor extraction with strict Pydantic validation. - Support spectral (Welch/Multitaper), parametric (specparam), and complexity (antropy/neurokit2) families. - Implement shared PSD and Parametric Fit planning for multi-extractor reuse. - Support parallel execution (obs-batch, work-unit, consumer-parallel) using joblib. - Achieve 96% test coverage across the module. --- .readthedocs.yml | 10 +- coco_pipe/__init__.py | 6 + coco_pipe/descriptors/__init__.py | 7 + coco_pipe/descriptors/configs.py | 446 +++++++++ coco_pipe/descriptors/core.py | 796 ++++++++++++++++ coco_pipe/descriptors/extractors/__init__.py | 3 + .../descriptors/extractors/_parametric_fit.py | 303 ++++++ coco_pipe/descriptors/extractors/_psd.py | 139 +++ coco_pipe/descriptors/extractors/base.py | 384 ++++++++ .../descriptors/extractors/complexity.py | 404 ++++++++ .../descriptors/extractors/parametric.py | 296 ++++++ coco_pipe/descriptors/extractors/spectral.py | 602 ++++++++++++ coco_pipe/descriptors/extractors/utils.py | 189 ++++ coco_pipe/descriptors/validation.py | 148 +++ coco_pipe/io/structures.py | 342 +++++-- configs/run_descriptors_eeg.yml | 70 ++ examples/.DS_Store | Bin 6148 -> 0 bytes examples/demo_structures.py | 2 +- examples/descriptors_example.py | 67 ++ pyproject.toml | 8 + requirements-dev.txt | 22 - requirements.txt | 1 - scripts/run_descriptors.py | 133 +++ tests/conftest.py | 15 + tests/test_descriptors_configs.py | 313 +++++++ tests/test_descriptors_core.py | 881 ++++++++++++++++++ tests/test_descriptors_extractors.py | 492 ++++++++++ tests/test_io_dataset.py | 2 +- tests/test_io_structures.py | 401 +++++++- 29 files changed, 6367 insertions(+), 115 deletions(-) create mode 100644 coco_pipe/descriptors/__init__.py create mode 100644 coco_pipe/descriptors/configs.py create mode 100644 coco_pipe/descriptors/core.py create mode 100644 coco_pipe/descriptors/extractors/__init__.py create mode 100644 coco_pipe/descriptors/extractors/_parametric_fit.py create mode 100644 coco_pipe/descriptors/extractors/_psd.py create mode 100644 coco_pipe/descriptors/extractors/base.py create mode 100644 coco_pipe/descriptors/extractors/complexity.py create mode 100644 coco_pipe/descriptors/extractors/parametric.py create mode 100644 coco_pipe/descriptors/extractors/spectral.py create mode 100644 coco_pipe/descriptors/extractors/utils.py create mode 100644 coco_pipe/descriptors/validation.py create mode 100644 configs/run_descriptors_eeg.yml delete mode 100644 examples/.DS_Store create mode 100644 examples/descriptors_example.py delete mode 100644 requirements-dev.txt delete mode 100644 requirements.txt create mode 100644 scripts/run_descriptors.py create mode 100644 tests/test_descriptors_configs.py create mode 100644 tests/test_descriptors_core.py create mode 100644 tests/test_descriptors_extractors.py diff --git a/.readthedocs.yml b/.readthedocs.yml index 6d76a88..737e1d4 100644 --- a/.readthedocs.yml +++ b/.readthedocs.yml @@ -8,7 +8,7 @@ version: 2 build: os: ubuntu-24.04 tools: - python: "3.9" + python: "3.10" # Build documentation in the "docs/" directory with Sphinx sphinx: @@ -18,5 +18,9 @@ sphinx: # declare the Python requirements required to build your documentation # See https://docs.readthedocs.io/en/stable/guides/reproducible-builds.html python: - install: - - requirements: requirements-dev.txt + install: + - method: pip + path: . + extra_requirements: + - docs + - full diff --git a/coco_pipe/__init__.py b/coco_pipe/__init__.py index a5bbe17..d9eb2a5 100644 --- a/coco_pipe/__init__.py +++ b/coco_pipe/__init__.py @@ -2,6 +2,10 @@ Package initializer for the coco_pipe package. """ +from .descriptors import ( + DescriptorConfig, + DescriptorPipeline, +) from .dim_reduction import ( METHODS, BaseReducer, @@ -22,6 +26,8 @@ # Core exports __all__ = [ + "DescriptorConfig", + "DescriptorPipeline", "DimReduction", "METHODS", "interpret_features", diff --git a/coco_pipe/descriptors/__init__.py b/coco_pipe/descriptors/__init__.py new file mode 100644 index 0000000..d8a0eef --- /dev/null +++ b/coco_pipe/descriptors/__init__.py @@ -0,0 +1,7 @@ +from .configs import DescriptorConfig +from .core import DescriptorPipeline + +__all__ = [ + "DescriptorConfig", + "DescriptorPipeline", +] diff --git a/coco_pipe/descriptors/configs.py b/coco_pipe/descriptors/configs.py new file mode 100644 index 0000000..6d0a52e --- /dev/null +++ b/coco_pipe/descriptors/configs.py @@ -0,0 +1,446 @@ +""" +Descriptor Configuration +======================== + +Strict Pydantic configuration models for the descriptors module. + +This module defines the static, typed configuration surface for descriptor +extraction: + +- explicit runtime input requirements +- family-specific configs for bands, parametric fitting, and complexity +- output formatting controls +- runtime execution controls + +These models validate local field structure and family-local constraints. The +remaining cross-family compatibility rule for corrected spectral outputs is +enforced by :class:`coco_pipe.descriptors.core.DescriptorPipeline` after config +parsing, because it depends on how multiple family configs interact. + +Author: Hamza Abdelhedi (hamza.abdelhedi@umontreal.ca) +""" + +from __future__ import annotations + +from typing import Any, Literal + +from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator + +__all__ = [ + "DescriptorInputConfig", + "BandDescriptorConfig", + "ParametricDescriptorConfig", + "ComplexityDescriptorConfig", + "DescriptorFamiliesConfig", + "DescriptorOutputConfig", + "DescriptorRuntimeConfig", + "DescriptorConfig", +] + + +CANONICAL_BANDS = { + "delta": (1.0, 4.0), + "theta": (4.0, 8.0), + "alpha": (8.0, 13.0), + "beta": (13.0, 30.0), + "gamma": (30.0, 45.0), +} + +_BAND_OUTPUTS = ( + "absolute_power", + "relative_power", + "ratios", + "corrected_absolute_power", + "corrected_relative_power", + "corrected_ratios", +) +_PARAM_OUTPUTS = ("aperiodic", "fit_quality", "peak_summary") +_COMPLEXITY_MEASURES = ( + "sample_entropy", + "perm_entropy", + "spectral_entropy", + "hjorth_mobility", + "hjorth_complexity", + "lziv_complexity", +) + + +class _StrictConfigModel(BaseModel): + """Shared strict Pydantic behavior.""" + + model_config = ConfigDict(extra="forbid") + + +class DescriptorInputConfig(_StrictConfigModel): + """ + Explicit runtime input requirements for descriptor extraction. + + Parameters + ---------- + require_sfreq : bool, default=True + Whether extraction requires an explicit sampling frequency input. + require_channel_names : bool, default=False + Whether extraction requires explicit channel names at runtime. + + Notes + ----- + The descriptors module accepts only explicit NumPy-like arrays with shape + ``(n_obs, n_channels, n_times)`` in observation-channel-time order. That + structural contract is fixed by the module and enforced at runtime; this + config only controls which additional runtime inputs must also be passed. + """ + + require_sfreq: bool = True + require_channel_names: bool = False + + +class BandDescriptorConfig(_StrictConfigModel): + """ + Configuration for PSD-based band summary descriptors. + + Parameters + ---------- + enabled : bool, default=False + Whether the band family is enabled. + psd_method : {"welch", "multitaper"}, default="welch" + PSD estimator used before computing band summaries. + fmin, fmax : float + Global frequency window within which PSDs and bands are evaluated. + bands : dict of str to tuple of float, default=canonical EEG bands + Mapping from band name to ``(low, high)`` boundaries. + outputs : list of {"absolute_power", "relative_power", "ratios", \ +"corrected_absolute_power", "corrected_relative_power", "corrected_ratios"} + Band descriptors to emit. + ratio_pairs : list of tuple of str, default=[] + Explicit numerator and denominator band names for ratio outputs. + log_power : bool, default=False + Whether to emit log-transformed absolute band power in addition to + absolute power when that output is enabled. + + Notes + ----- + Corrected band outputs are configured here, but their cross-family + compatibility with the parametric fit range is checked later by the + descriptor pipeline because that rule depends on both the band and + parametric family configs together. + """ + + enabled: bool = False + psd_method: Literal["welch", "multitaper"] = "welch" + fmin: float = Field(1.0, ge=0.0) + fmax: float = Field(45.0, gt=0.0) + bands: dict[str, tuple[float, float]] = Field( + default_factory=lambda: dict(CANONICAL_BANDS) + ) + outputs: list[ + Literal[ + "absolute_power", + "relative_power", + "ratios", + "corrected_absolute_power", + "corrected_relative_power", + "corrected_ratios", + ] + ] = Field(default_factory=lambda: ["absolute_power"]) + ratio_pairs: list[tuple[str, str]] = Field(default_factory=list) + log_power: bool = False + + @field_validator("bands", mode="before") + @classmethod + def _coerce_bands(cls, value: Any) -> dict[str, tuple[float, float]]: + if value is None: + return dict(CANONICAL_BANDS) + return {str(key): tuple(bounds) for key, bounds in dict(value).items()} + + @field_validator("outputs", mode="before") + @classmethod + def _validate_outputs(cls, value: list[str]) -> list[str]: + if len(set(value)) != len(value): + raise ValueError("Band outputs must not contain duplicates.") + invalid = sorted(set(value) - set(_BAND_OUTPUTS)) + if invalid: + raise ValueError(f"Unknown band outputs: {invalid}") + return value + + @field_validator("ratio_pairs", mode="before") + @classmethod + def _coerce_ratio_pairs(cls, value: Any) -> list[tuple[str, str]]: + if value is None: + return [] + return [tuple(pair) for pair in value] + + @model_validator(mode="after") + def _validate_model(self) -> "BandDescriptorConfig": + if self.fmin >= self.fmax: + raise ValueError("Band descriptor config requires fmin < fmax.") + for name, (low, high) in self.bands.items(): + if low >= high: + raise ValueError(f"Band '{name}' requires low < high.") + if low < self.fmin or high > self.fmax: + raise ValueError( + "Band " + f"'{name}' must stay within the configured " + f"[{self.fmin}, {self.fmax}] range." + ) + if ( + "ratios" in self.outputs or "corrected_ratios" in self.outputs + ) and not self.ratio_pairs: + raise ValueError("Band ratios require explicit ratio_pairs.") + return self + + +class ParametricDescriptorConfig(_StrictConfigModel): + """ + Configuration for specparam-based spectral summary descriptors. + + Parameters + ---------- + enabled : bool, default=False + Whether the parametric family is enabled. + backend : {"specparam"}, default="specparam" + Parametric modeling backend. + psd_method : {"welch", "multitaper"}, default="welch" + PSD estimator used before fitting the parametric model. + freq_range : tuple of float, default=(1.0, 45.0) + Frequency range passed to the parametric model. + peak_width_limits : tuple of float, default=(1.0, 12.0) + Peak width bounds forwarded to the model backend. + max_n_peaks : int, default=6 + Maximum number of periodic peaks to fit. + aperiodic_mode : {"fixed", "knee"}, default="fixed" + Aperiodic model form used by specparam. + outputs : list of {"aperiodic", "fit_quality", "peak_summary"} + Parametric descriptor groups to emit. + + Notes + ----- + This config describes how the shared parametric fit is produced. The same + fit can be reused by the parametric family itself and by corrected spectral + outputs when the planner detects compatible requests. + """ + + enabled: bool = False + backend: Literal["specparam"] = "specparam" + psd_method: Literal["welch", "multitaper"] = "welch" + freq_range: tuple[float, float] = (1.0, 45.0) + peak_width_limits: tuple[float, float] = (1.0, 12.0) + max_n_peaks: int = Field(6, ge=0) + aperiodic_mode: Literal["fixed", "knee"] = "fixed" + outputs: list[Literal["aperiodic", "fit_quality", "peak_summary"]] = Field( + default_factory=lambda: ["aperiodic", "fit_quality", "peak_summary"] + ) + + @field_validator("outputs", mode="before") + @classmethod + def _validate_outputs(cls, value: list[str]) -> list[str]: + if len(set(value)) != len(value): + raise ValueError("Parametric outputs must not contain duplicates.") + invalid = sorted(set(value) - set(_PARAM_OUTPUTS)) + if invalid: + raise ValueError(f"Unknown parametric outputs: {invalid}") + return value + + @model_validator(mode="after") + def _validate_model(self) -> "ParametricDescriptorConfig": + if self.freq_range[0] >= self.freq_range[1]: + raise ValueError("Parametric freq_range requires low < high.") + if self.peak_width_limits[0] >= self.peak_width_limits[1]: + raise ValueError("peak_width_limits requires low < high.") + return self + + +class ComplexityDescriptorConfig(_StrictConfigModel): + """ + Configuration for signal-complexity descriptors. + + Parameters + ---------- + enabled : bool, default=False + Whether the complexity family is enabled. + backend : {"antropy", "neurokit2", "auto"}, default="antropy" + Complexity backend used for supported measures. + measures : list of str + Complexity measures to compute. + measure_kwargs : dict of str to dict, default={} + Per-measure keyword arguments forwarded to the backend implementation. + + Notes + ----- + Complexity measures are signal-domain descriptors. Unlike the PSD-based + families, they do not participate in shared PSD planning. + """ + + enabled: bool = False + backend: Literal["antropy", "neurokit2", "auto"] = "antropy" + measures: list[str] = Field(default_factory=lambda: list(_COMPLEXITY_MEASURES)) + measure_kwargs: dict[str, dict[str, Any]] = Field(default_factory=dict) + + @field_validator("measures", mode="before") + @classmethod + def _validate_measures(cls, value: list[str]) -> list[str]: + if len(set(value)) != len(value): + raise ValueError("Complexity measures must not contain duplicates.") + invalid = sorted(set(value) - set(_COMPLEXITY_MEASURES)) + if invalid: + raise ValueError(f"Unknown complexity measures: {invalid}") + return value + + +class DescriptorFamiliesConfig(_StrictConfigModel): + """ + Group descriptor-family configuration under one top-level field. + + Attributes + ---------- + bands : BandDescriptorConfig + Configuration for PSD-based band summaries. + parametric : ParametricDescriptorConfig + Configuration for specparam-based summaries. + complexity : ComplexityDescriptorConfig + Configuration for complexity measures. + """ + + bands: BandDescriptorConfig = Field(default_factory=BandDescriptorConfig) + parametric: ParametricDescriptorConfig = Field( + default_factory=ParametricDescriptorConfig + ) + complexity: ComplexityDescriptorConfig = Field( + default_factory=ComplexityDescriptorConfig + ) + + +class DescriptorOutputConfig(_StrictConfigModel): + """ + Controls output precision and descriptor-level channel pooling. + + Parameters + ---------- + precision : {"float32", "float64"}, default="float32" + Output dtype used for the final descriptor matrix. + channel_pooling : {"none", "all"} or dict of str to list of str, default="none" + Descriptor-level channel pooling policy applied after per-channel + descriptors are computed. ``"none"`` keeps one descriptor per sensor, + ``"all"`` averages descriptor values across all sensors, and a mapping + averages descriptor values within each named group while leaving + ungrouped sensors unchanged. + + Notes + ----- + Output config is intentionally small. The descriptors module now returns a + minimal result object, so output controls are limited to matrix precision + and descriptor-level channel pooling. + """ + + precision: Literal["float32", "float64"] = "float32" + channel_pooling: Literal["none", "all"] | dict[str, list[str]] = "none" + + @field_validator("channel_pooling", mode="before") + @classmethod + def _coerce_channel_pooling( + cls, value: Any + ) -> Literal["none", "all"] | dict[str, list[str]]: + if value in (None, {}): + return "none" + if isinstance(value, str): + return value + return { + str(group_name): [str(member) for member in members] + for group_name, members in dict(value).items() + } + + @field_validator("channel_pooling") + @classmethod + def _validate_channel_pooling( + cls, value: Literal["none", "all"] | dict[str, list[str]] + ) -> Literal["none", "all"] | dict[str, list[str]]: + if isinstance(value, str): + if value not in {"none", "all"}: + raise ValueError("channel_pooling must be 'none', 'all', or a mapping.") + return value + for group_name, members in value.items(): + if not group_name: + raise ValueError( + "channel_pooling mapping keys must be non-empty strings." + ) + if not members: + raise ValueError( + f"channel_pooling['{group_name}'] must define at least one channel." + ) + if len(set(members)) != len(members): + raise ValueError( + f"channel_pooling['{group_name}'] must not contain duplicates." + ) + return value + + +class DescriptorRuntimeConfig(_StrictConfigModel): + """ + Runtime execution controls for descriptor extraction. + + Parameters + ---------- + execution_backend : {"joblib", "sequential"}, default="joblib" + Execution backend used by the pipeline. + n_jobs : int, default=1 + Number of worker slots requested for supported parallel paths. + ``-1`` means "use as much useful parallelism as the current stage can + use", while positive integers request an explicit worker count. + obs_chunk : int, default=128 + Number of observations processed per batch. + on_error : {"raise", "warn", "collect"}, default="collect" + Failure policy applied during extraction. + + Notes + ----- + Runtime config controls execution only. It does not add provenance, + reporting, or persistence metadata to the returned descriptor result. + """ + + execution_backend: Literal["joblib", "sequential"] = "joblib" + n_jobs: int = 1 + obs_chunk: int = Field(128, gt=0) + on_error: Literal["raise", "warn", "collect"] = Field( + "collect", + description=( + "Policies: " + "'raise' re-raises the first exception immediately; " + "'warn' collects all failures and emits one aggregate warning; " + "'collect' stores failures silently for inspection in result['failures']." + ), + ) + + @field_validator("n_jobs") + @classmethod + def _validate_n_jobs(cls, value: int) -> int: + if value == 0 or value < -1: + raise ValueError("n_jobs must be -1 or a positive integer.") + return value + + +class DescriptorConfig(_StrictConfigModel): + """ + Top-level descriptors configuration object. + + Attributes + ---------- + input : DescriptorInputConfig + Runtime input requirements for explicit array extraction. + families : DescriptorFamiliesConfig + Enabled descriptor families and their typed configs. + output : DescriptorOutputConfig + Output precision and formatting settings. + runtime : DescriptorRuntimeConfig + Runtime execution and error-handling settings. + + Notes + ----- + This object is the stable config boundary for + :class:`coco_pipe.descriptors.core.DescriptorPipeline`. Parsing this config + validates local structure here, then the pipeline applies the remaining + cross-family compatibility checks when it builds the execution plan. + """ + + input: DescriptorInputConfig = Field(default_factory=DescriptorInputConfig) + families: DescriptorFamiliesConfig = Field(default_factory=DescriptorFamiliesConfig) + output: DescriptorOutputConfig = Field(default_factory=DescriptorOutputConfig) + runtime: DescriptorRuntimeConfig = Field(default_factory=DescriptorRuntimeConfig) diff --git a/coco_pipe/descriptors/core.py b/coco_pipe/descriptors/core.py new file mode 100644 index 0000000..8ed408b --- /dev/null +++ b/coco_pipe/descriptors/core.py @@ -0,0 +1,796 @@ +""" +Descriptor extraction planner and execution pipeline. + +This module owns the config-bound runtime orchestration for descriptor +extraction. It does not implement family-specific descriptor math; instead it: + +- validates the explicit runtime inputs accepted by the module +- instantiates enabled descriptor families from typed config +- plans shared PSD computation for compatible PSD consumers +- executes one observation batch at a time with controlled parallelism +- merges aligned family outputs into one flat descriptor matrix + +Author: Hamza Abdelhedi (hamza.abdelhedi@umontreal.ca) +""" + +from __future__ import annotations + +import warnings +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from typing import Any + +import numpy as np + +from .configs import DescriptorConfig, ParametricDescriptorConfig +from .extractors._parametric_fit import fit_parametric_batch +from .extractors._psd import compute_psd +from .extractors.base import ( + BaseDescriptorExtractor, + BasePSDDescriptorExtractor, + _DescriptorBlock, +) +from .extractors.complexity import ComplexityDescriptorExtractor +from .extractors.parametric import ParametricDescriptorExtractor +from .extractors.spectral import BandDescriptorExtractor +from .validation import ( + validate_runtime_inputs, +) + +__all__ = ["DescriptorPipeline"] + + +@dataclass(slots=True) +class _PSDGroup: + """Plan one shared PSD compute for a compatible set of consumers. + + Attributes + ---------- + method : str + Shared PSD estimator name. + fmin, fmax : float + Union frequency window required by all consumers in the group. + consumers : list of BasePSDDescriptorExtractor + PSD-consuming extractors that will reuse the shared PSD output. + needs_parametric_fit : bool, default=False + Whether the group also needs one shared parametric fit. + need_parametric_metrics : bool, default=False + Whether the shared fit must expose scalar parametric metrics. + need_periodic_psds : bool, default=False + Whether the shared fit must reconstruct periodic-only PSDs. + fit_config : ParametricDescriptorConfig or None, default=None + Shared parametric fit configuration when a fit is required. + """ + + method: str + fmin: float + fmax: float + consumers: list[BasePSDDescriptorExtractor] + needs_parametric_fit: bool = False + need_parametric_metrics: bool = False + need_periodic_psds: bool = False + fit_config: ParametricDescriptorConfig | None = None + + +def _parallel_jobs(n_jobs: int, limit: int) -> int: + """Clamp worker count to the amount of useful parallel work. + + Parameters + ---------- + n_jobs : int + Requested worker count from runtime config. ``-1`` means "use as much + parallelism as this stage can use". + limit : int + Number of parallel tasks available in the current stage. + + Returns + ------- + int + Worker count capped at ``limit``. + """ + return limit if n_jobs == -1 else min(n_jobs, limit) + + +def _cast_precision(values: np.ndarray, precision: str) -> np.ndarray: + """Cast the final descriptor matrix to the configured floating precision. + + Parameters + ---------- + values : np.ndarray + Floating descriptor matrix to cast. + precision : {"float32", "float64"} + Requested output precision. + + Returns + ------- + np.ndarray + ``values`` cast in-place when possible. + """ + return values.astype( + np.float32 if precision == "float32" else np.float64, copy=False + ) + + +def _merge_descriptor_blocks( + blocks: list[_DescriptorBlock], + n_obs: int, + precision: str, +) -> tuple[np.ndarray, list[str], list[dict[str, Any]]]: + """Merge family descriptor blocks column-wise on the descriptor axis. + + Parameters + ---------- + blocks : list of _DescriptorBlock + Family-specific descriptor block objects to concatenate column-wise. + n_obs : int + Expected number of observations in each block. + precision : {"float32", "float64"} + Precision applied to the merged descriptor matrix. + + Returns + ------- + tuple + ``(X, descriptor_names, failures)`` where ``X`` is the merged + descriptor matrix, ``descriptor_names`` is the deterministic merged + column order, and ``failures`` concatenates all family failure records. + + Raises + ------ + ValueError + If any block is misaligned on the observation axis. + """ + if not blocks: + empty = np.empty( + (n_obs, 0), + dtype=np.float32 if precision == "float32" else np.float64, + ) + return empty, [], [] + + matrices = [] + names: list[str] = [] + failures: list[dict[str, Any]] = [] + + for block in blocks: + if block.X.shape[0] != n_obs: + raise ValueError( + "Descriptor block " + f"'{block.family}' is misaligned: expected {n_obs} rows, " + f"got {block.X.shape[0]}." + ) + matrices.append(block.X) + names.extend(block.descriptor_names) + failures.extend(block.failures) + + if len(matrices) == 1: + X = _cast_precision(matrices[0], precision) + else: + X = _cast_precision(np.concatenate(matrices, axis=1), precision) + + return ( + X, + names, + failures, + ) + + +def _sequential_runtime(runtime): + """Return a sequential runtime copy for nested work. + + Parameters + ---------- + runtime : DescriptorRuntimeConfig + Runtime configuration for the current extraction stage. + + Returns + ------- + DescriptorRuntimeConfig + Copy with nested parallelism disabled. + """ + return runtime.model_copy(update={"execution_backend": "sequential", "n_jobs": 1}) + + +def _build_psd_groups( + extractors: list[BaseDescriptorExtractor], +) -> list[_PSDGroup]: + """Plan shared PSD groups for the enabled PSD-consuming extractors. + + Parameters + ---------- + extractors : list of BaseDescriptorExtractor + Config-bound extractors in deterministic family order. + + Returns + ------- + list of _PSDGroup + Shared PSD execution groups keyed by compatible PSD method and merged + fit requirements. + + Raises + ------ + ValueError + If consumers that would share one parametric fit disagree on the fit + configuration. + """ + groups_by_method: dict[str, _PSDGroup] = {} + for extractor in extractors: + if not isinstance(extractor, BasePSDDescriptorExtractor): + continue + request = extractor.psd_request() + method = str(request["method"]) + if method not in groups_by_method: + groups_by_method[method] = _PSDGroup( + method=method, + fmin=float(request["fmin"]), + fmax=float(request["fmax"]), + consumers=[extractor], + ) + else: + current = groups_by_method[method] + current.fmin = min(current.fmin, float(request["fmin"])) + current.fmax = max(current.fmax, float(request["fmax"])) + current.consumers.append(extractor) + current = groups_by_method[method] + fit_req = extractor.parametric_fit_requirements() + if fit_req["needed"]: + current.needs_parametric_fit = True + current.need_parametric_metrics = current.need_parametric_metrics or bool( + fit_req["metrics"] + ) + current.need_periodic_psds = current.need_periodic_psds or bool( + fit_req["periodic_psds"] + ) + if current.fit_config is None: + current.fit_config = fit_req["config"] + elif current.fit_config.model_dump() != fit_req["config"].model_dump(): + raise ValueError( + "PSD consumers sharing one parametric fit must use the same " + "parametric fit configuration." + ) + return list(groups_by_method.values()) + + +def _merge_family_blocks( + batch_results: list[dict[str, _DescriptorBlock]], + family_order: list[str], +) -> list[_DescriptorBlock]: + """Merge per-batch results row-wise within each descriptor family. + + Parameters + ---------- + batch_results : list of dict + One family-block mapping per processed observation batch. + family_order : list of str + Deterministic family order used for the final merged output. + + Returns + ------- + list of _DescriptorBlock + One merged block per family, still separated by family but aligned + across all processed batches. + + Raises + ------ + ValueError + If descriptor names drift across batches for the same family. + """ + merged_blocks: list[_DescriptorBlock] = [] + for family_name in family_order: + family_blocks = [ + batch_result[family_name] + for batch_result in batch_results + if family_name in batch_result + ] + if not family_blocks: + continue + + reference_names = family_blocks[0].descriptor_names + for block in family_blocks[1:]: + if block.descriptor_names != reference_names: + raise ValueError( + "Descriptor names changed across batches for family " + f"'{family_name}'." + ) + + merged_blocks.append( + _DescriptorBlock( + family=family_name, + X=np.concatenate([block.X for block in family_blocks], axis=0), + descriptor_names=list(reference_names), + meta={}, + failures=[ + failure for block in family_blocks for failure in block.failures + ], + ) + ) + return merged_blocks + + +def _process_psd_group( + group: _PSDGroup, + X_batch: np.ndarray, + sfreq: float, + channel_names: list[str] | None, + channel_pooling: str | dict[str, list[str]], + ids_batch: np.ndarray | None, + runtime, + obs_offset: int, + joblib=None, + consumer_parallel: bool = False, + psd_n_jobs: int | None = None, +) -> dict[str, _DescriptorBlock]: + """Execute one shared PSD group for a single observation batch. + + Parameters + ---------- + group : _PSDGroup + Planned PSD reuse group for the current batch. + X_batch : np.ndarray + Observation batch with shape ``(n_obs_batch, n_channels, n_times)``. + sfreq : float + Sampling frequency in Hertz. + channel_names : list of str or None + Runtime channel labels. + channel_pooling : {"none", "all"} or dict + Descriptor-level channel pooling policy. + ids_batch : np.ndarray or None + Observation identifiers aligned with ``X_batch``. + runtime : DescriptorRuntimeConfig + Runtime policy used for this stage. + obs_offset : int + Absolute observation offset of the batch in the full input array. + joblib : module, optional + Imported ``joblib`` module when the selected strategy uses it. + consumer_parallel : bool, default=False + Whether compatible PSD consumers should run in parallel after the PSD + has been computed. + psd_n_jobs : int or None, default=None + Worker count forwarded to the PSD backend when the selected strategy is + PSD-level parallelism. + + Returns + ------- + dict[str, _DescriptorBlock] + Family-name mapping for all consumers in the PSD group. + """ + if psd_n_jobs is not None and joblib is not None: + with joblib.parallel_backend("threading", n_jobs=psd_n_jobs): + psds, freqs = compute_psd( + X_batch, + sfreq=sfreq, + method=group.method, + fmin=group.fmin, + fmax=group.fmax, + n_jobs=psd_n_jobs, + ) + else: + psds, freqs = compute_psd( + X_batch, + sfreq=sfreq, + method=group.method, + fmin=group.fmin, + fmax=group.fmax, + n_jobs=psd_n_jobs, + ) + + fit_batch = None + if group.needs_parametric_fit: + fit_batch = fit_parametric_batch( + psds, + freqs, + group.fit_config, + runtime, + need_periodic_psd=group.need_periodic_psds, + include_metrics=group.need_parametric_metrics, + ) + + consumer_runtime = _sequential_runtime(runtime) if consumer_parallel else runtime + if consumer_parallel and joblib is not None and len(group.consumers) > 1: + blocks = joblib.Parallel( + n_jobs=_parallel_jobs(runtime.n_jobs, len(group.consumers)), + prefer="threads", + )( + joblib.delayed(consumer.extract_psd)( + psds, + freqs, + channel_names=channel_names, + channel_pooling=channel_pooling, + ids=ids_batch, + runtime=consumer_runtime, + obs_offset=obs_offset, + fit_batch=fit_batch, + ) + for consumer in group.consumers + ) + else: + blocks = [ + consumer.extract_psd( + psds, + freqs, + channel_names=channel_names, + channel_pooling=channel_pooling, + ids=ids_batch, + runtime=consumer_runtime, + obs_offset=obs_offset, + fit_batch=fit_batch, + ) + for consumer in group.consumers + ] + return {block.family: block for block in blocks} + + +def _process_batch( + obs_slice: slice, + X: np.ndarray, + sfreq: float | None, + channel_names: list[str] | None, + channel_pooling: str | dict[str, list[str]], + ids: np.ndarray | None, + signal_extractors: list[BaseDescriptorExtractor], + psd_groups: list[_PSDGroup], + runtime, + strategy: str, + joblib=None, +) -> dict[str, _DescriptorBlock]: + """Execute one observation batch under the selected planner strategy. + + Parameters + ---------- + obs_slice : slice + Observation slice for the current batch. + X : np.ndarray + Full validated input array with shape ``(n_obs, n_channels, n_times)``. + sfreq : float or None + Sampling frequency in Hertz. + channel_names : list of str or None + Runtime channel labels. + channel_pooling : {"none", "all"} or dict + Descriptor-level channel pooling policy. + ids : np.ndarray or None + Observation identifiers aligned with ``X``. + signal_extractors : list of BaseDescriptorExtractor + Non-PSD families that consume raw signal batches directly. + psd_groups : list of _PSDGroup + Planned PSD reuse groups for this pipeline instance. + runtime : DescriptorRuntimeConfig + Runtime policy for the current extraction call. + strategy : str + Selected parallelization strategy for this execution path. + joblib : module, optional + Imported ``joblib`` module when the selected strategy uses it. + + Returns + ------- + dict[str, _DescriptorBlock] + Family-name mapping for all blocks produced from the batch. + """ + X_batch = X[obs_slice] + ids_batch = None if ids is None else ids[obs_slice] + obs_offset = obs_slice.start or 0 + family_blocks: dict[str, _DescriptorBlock] = {} + + if strategy == "work-unit" and joblib is not None: + + def _signal_unit(extractor): + block = extractor.extract( + X_batch, + sfreq=sfreq, + channel_names=channel_names, + channel_pooling=channel_pooling, + ids=ids_batch, + runtime=_sequential_runtime(runtime), + obs_offset=obs_offset, + ) + return {block.family: block} + + def _psd_unit(group): + return _process_psd_group( + group, + X_batch, + sfreq=sfreq, + channel_names=channel_names, + channel_pooling=channel_pooling, + ids_batch=ids_batch, + runtime=_sequential_runtime(runtime), + obs_offset=obs_offset, + joblib=None, + consumer_parallel=False, + ) + + work_units = [ + joblib.delayed(_signal_unit)(extractor) for extractor in signal_extractors + ] + [joblib.delayed(_psd_unit)(group) for group in psd_groups] + for unit_result in joblib.Parallel( + n_jobs=_parallel_jobs( + runtime.n_jobs, + len(signal_extractors) + len(psd_groups), + ), + prefer="threads", + )(work_units): + family_blocks.update(unit_result) + else: + signal_runtime = _sequential_runtime(runtime) + for extractor in signal_extractors: + block = extractor.extract( + X_batch, + sfreq=sfreq, + channel_names=channel_names, + channel_pooling=channel_pooling, + ids=ids_batch, + runtime=signal_runtime, + obs_offset=obs_offset, + ) + family_blocks[block.family] = block + + for group in psd_groups: + consumer_parallel = ( + strategy == "psd-consumer" + and joblib is not None + and len(group.consumers) > 1 + ) + psd_n_jobs = None + if strategy == "psd-n_jobs": + psd_n_jobs = runtime.n_jobs + family_blocks.update( + _process_psd_group( + group, + X_batch, + sfreq=sfreq, + channel_names=channel_names, + channel_pooling=channel_pooling, + ids_batch=ids_batch, + runtime=runtime + if strategy == "parametric-inner" and group.needs_parametric_fit + else signal_runtime, + obs_offset=obs_offset, + joblib=joblib + if consumer_parallel or strategy == "psd-n_jobs" + else None, + consumer_parallel=consumer_parallel, + psd_n_jobs=psd_n_jobs, + ) + ) + + return family_blocks + + +class DescriptorPipeline: + """Run config-driven descriptor extraction on explicit arrays. + + Parameters + ---------- + config : DescriptorConfig or Mapping[str, Any] + Typed descriptors configuration or a mapping accepted by + :class:`DescriptorConfig`. + + Attributes + ---------- + config : DescriptorConfig + Parsed descriptors configuration. + extractors : list of BaseDescriptorExtractor + Enabled family extractors in deterministic family order. + signal_extractors : list of BaseDescriptorExtractor + Enabled non-PSD extractors that consume raw signal batches directly. + psd_groups : list of _PSDGroup + Planned PSD reuse groups derived once from the enabled extractors. + family_order : list of str + Deterministic family order used when merging batch-local outputs. + + Notes + ----- + The pipeline is config-bound but runtime-stateless. Construction performs + config parsing, corrected-band compatibility checks, and planner setup once. + Each call to :meth:`extract` then validates the explicit runtime inputs, + executes the planned families, and returns one flat descriptor matrix plus + any collected failures. + """ + + def __init__(self, config: DescriptorConfig | Mapping[str, Any]): + """Create a config-bound descriptor extraction pipeline. + + Parameters + ---------- + config : DescriptorConfig or Mapping[str, Any] + Typed descriptors configuration or a mapping accepted by + :class:`DescriptorConfig`. + + Raises + ------ + ValueError + If corrected band outputs are enabled but the parametric fit range + does not cover the requested band PSD window. + """ + self.config = ( + config + if isinstance(config, DescriptorConfig) + else DescriptorConfig.model_validate(config) + ) + corrected_outputs = { + "corrected_absolute_power", + "corrected_relative_power", + "corrected_ratios", + } + if any( + output in corrected_outputs for output in self.config.families.bands.outputs + ): + fit_low, fit_high = self.config.families.parametric.freq_range + band_low = self.config.families.bands.fmin + band_high = self.config.families.bands.fmax + if fit_low > band_low or fit_high < band_high: + raise ValueError( + "Corrected band outputs require families.parametric.freq_range " + f"to cover the band PSD window [{band_low}, {band_high}]." + ) + self.extractors: list[BaseDescriptorExtractor] = [] + if self.config.families.bands.enabled: + self.extractors.append( + BandDescriptorExtractor( + self.config.families.bands, + fit_config=self.config.families.parametric, + ) + ) + if self.config.families.parametric.enabled: + self.extractors.append( + ParametricDescriptorExtractor(self.config.families.parametric) + ) + if self.config.families.complexity.enabled: + self.extractors.append( + ComplexityDescriptorExtractor(self.config.families.complexity) + ) + self.signal_extractors = [ + extractor + for extractor in self.extractors + if not isinstance(extractor, BasePSDDescriptorExtractor) + ] + self.psd_groups = _build_psd_groups(self.extractors) + self.family_order = [extractor.family_name for extractor in self.extractors] + + def extract( + self, + X: np.ndarray, + ids: Sequence[Any] | np.ndarray | None = None, + sfreq: float | None = None, + channel_names: Sequence[str] | np.ndarray | None = None, + ) -> dict[str, Any]: + """Extract descriptors from explicit NumPy inputs. + + Parameters + ---------- + X : np.ndarray + Signal array with shape ``(n_obs, n_channels, n_times)``. + ids : sequence or np.ndarray, optional + Observation identifiers aligned with ``X``. + sfreq : float, optional + Sampling frequency in Hertz. Required when enabled families depend + on spectral estimates or spectral entropy. + channel_names : sequence of str or np.ndarray, optional + Channel labels. Required for channel-resolved outputs. + + Returns + ------- + dict[str, Any] + Dictionary with keys ``X``, ``descriptor_names``, and ``failures``. + + Raises + ------ + ValueError + If the explicit input contract is not satisfied. + ImportError + If an optional backend required by the enabled families is missing. + + Notes + ----- + When ``runtime.on_error="warn"``, extraction still completes and stores + failures in ``result["failures"]`` before emitting one aggregate + warning at the pipeline level. + + The returned row order always matches the input observation order. + """ + inputs = validate_runtime_inputs( + self.config, + X=X, + ids=ids, + channel_names=channel_names, + sfreq=sfreq, + ) + + planner_runtime = self.config.runtime + n_obs = inputs["X"].shape[0] + obs_chunk = planner_runtime.obs_chunk + if not obs_chunk or obs_chunk >= n_obs: + batch_slices = [slice(0, n_obs)] + else: + batch_slices = [ + slice(start, min(start + obs_chunk, n_obs)) + for start in range(0, n_obs, obs_chunk) + ] + + if ( + planner_runtime.execution_backend == "sequential" + or planner_runtime.n_jobs == 1 + ): + parallel_strategy = "sequential" + elif len(batch_slices) > 1: + parallel_strategy = "obs-batch" + else: + work_units = len(self.signal_extractors) + len(self.psd_groups) + if work_units > 1: + parallel_strategy = "work-unit" + elif len(self.psd_groups) == 1 and len(self.psd_groups[0].consumers) > 1: + parallel_strategy = "psd-consumer" + elif ( + len(self.psd_groups) == 1 + and len(self.psd_groups[0].consumers) == 1 + and self.psd_groups[0].needs_parametric_fit + ): + parallel_strategy = "parametric-inner" + elif len(self.psd_groups) == 1 and len(self.psd_groups[0].consumers) == 1: + parallel_strategy = "psd-n_jobs" + else: + parallel_strategy = "sequential" + + if parallel_strategy == "obs-batch": + import joblib + + batch_results = joblib.Parallel( + n_jobs=_parallel_jobs(planner_runtime.n_jobs, len(batch_slices)), + prefer="threads", + )( + joblib.delayed(_process_batch)( + obs_slice, + X=inputs["X"], + sfreq=inputs["sfreq"], + channel_names=inputs["channel_names"], + channel_pooling=inputs["channel_pooling"], + ids=inputs["ids"], + signal_extractors=self.signal_extractors, + psd_groups=self.psd_groups, + runtime=_sequential_runtime(planner_runtime), + strategy="sequential", + joblib=None, + ) + for obs_slice in batch_slices + ) + else: + if parallel_strategy != "sequential": + import joblib + else: + joblib = None + batch_results = [ + _process_batch( + obs_slice, + X=inputs["X"], + sfreq=inputs["sfreq"], + channel_names=inputs["channel_names"], + channel_pooling=inputs["channel_pooling"], + ids=inputs["ids"], + signal_extractors=self.signal_extractors, + psd_groups=self.psd_groups, + runtime=planner_runtime, + strategy=parallel_strategy, + joblib=joblib, + ) + for obs_slice in batch_slices + ] + + blocks = _merge_family_blocks( + batch_results, + family_order=self.family_order, + ) + + X_desc, descriptor_names, failures = _merge_descriptor_blocks( + blocks, + n_obs=inputs["X"].shape[0], + precision=self.config.output.precision, + ) + + if self.config.runtime.on_error == "warn" and failures: + warnings.warn( + f"Collected {len(failures)} descriptor failures during extract().", + stacklevel=2, + ) + + return { + "X": X_desc, + "descriptor_names": descriptor_names, + "failures": failures, + } diff --git a/coco_pipe/descriptors/extractors/__init__.py b/coco_pipe/descriptors/extractors/__init__.py new file mode 100644 index 0000000..ee25218 --- /dev/null +++ b/coco_pipe/descriptors/extractors/__init__.py @@ -0,0 +1,3 @@ +from .base import BaseDescriptorExtractor, BasePSDDescriptorExtractor + +__all__ = ["BaseDescriptorExtractor", "BasePSDDescriptorExtractor"] diff --git a/coco_pipe/descriptors/extractors/_parametric_fit.py b/coco_pipe/descriptors/extractors/_parametric_fit.py new file mode 100644 index 0000000..0d46755 --- /dev/null +++ b/coco_pipe/descriptors/extractors/_parametric_fit.py @@ -0,0 +1,303 @@ +""" +Shared specparam fitting for PSD-consuming descriptor paths. + +This module holds the reusable fitting step used by the descriptors planner and +by extractors that consume explicit parametric-fit intermediates. It does not +define descriptor names or output pooling. It only: + +- fit specparam models on PSD batches +- collect scalar fit metrics in aligned arrays +- optionally reconstruct periodic-only PSDs for corrected band outputs +- return one batch-scoped payload that downstream extractors can consume + +Author: Hamza Abdelhedi (hamza.abdelhedi@umontreal.ca) +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +import numpy as np + +from ...utils import import_optional_dependency +from ..configs import ParametricDescriptorConfig + + +@dataclass(slots=True) +class _ParametricFitBatch: + """ + Batch-scoped parametric fit payload shared across PSD consumers. + + Attributes + ---------- + freqs : np.ndarray + Frequency grid used for the fitted spectra. + metrics : dict of str to np.ndarray + Scalar metric arrays aligned to ``(n_obs, n_channels)`` for each + requested parametric metric. + errors : list of tuple + Collected fit failures as ``(obs_index, channel_index, + exception_type, message)``. + periodic_psds : np.ndarray | None + Periodic-only PSDs aligned to ``(n_obs, n_channels, n_freqs)`` when + corrected spectral outputs request them. + meta : dict + Lightweight fit metadata propagated into downstream descriptor blocks. + """ + + freqs: np.ndarray + metrics: dict[str, np.ndarray] + errors: list[tuple[int, int, str, str]] = field(default_factory=list) + periodic_psds: np.ndarray | None = None + meta: dict[str, Any] = field(default_factory=dict) + + +def fit_single_spectrum( + freqs: np.ndarray, + spectrum: np.ndarray, + config: ParametricDescriptorConfig, + need_periodic_psd: bool = False, +) -> tuple[dict[str, float], np.ndarray | None]: + """ + Fit one specparam model to one PSD spectrum. + + Parameters + ---------- + freqs : np.ndarray + Frequency grid for the input spectrum. + spectrum : np.ndarray + One PSD spectrum aligned with ``freqs``. + config : ParametricDescriptorConfig + Parsed parametric fit configuration. + need_periodic_psd : bool, default=False + Whether to reconstruct the periodic-only PSD from the fitted model. + + Returns + ------- + tuple[dict[str, float], np.ndarray | None] + Scalar fit metrics and, when requested, the periodic-only PSD on the + same frequency grid. + + Raises + ------ + ValueError + If the spectrum is constant or entirely non-finite. + RuntimeError + If specparam fails to produce a usable model or if reconstructed model + components become non-finite. + """ + finite = spectrum[np.isfinite(spectrum)] + if finite.size == 0 or np.ptp(finite) < np.finfo(float).eps: + raise ValueError("Parametric fitting requires a non-constant spectrum.") + + SpectralModel = import_optional_dependency( + lambda: ( + __import__( + "specparam.models", + fromlist=["SpectralModel"], + ).SpectralModel + ), + feature="parametric descriptor extraction", + dependency="specparam", + install_hint="pip install coco-pipe[descriptors]", + ) + model = SpectralModel( + aperiodic_mode=config.aperiodic_mode, + peak_width_limits=config.peak_width_limits, + max_n_peaks=config.max_n_peaks, + verbose=False, + ) + model.fit(freqs, spectrum, list(config.freq_range)) + + if not model.results.has_model: + raise RuntimeError("Specparam fitting was unsuccessful.") + + aperiodic = np.asarray(model.results.get_params("aperiodic")) + periodic = np.asarray(model.results.get_params("periodic")) + error = float(np.asarray(model.results.get_metrics("error")).squeeze()) + r_squared = float( + np.asarray(model.results.get_metrics("gof", "rsquared")).squeeze() + ) + + if periodic.size == 0 or np.all(np.isnan(periodic)): + peak_count = 0.0 + dominant_freq = np.nan + dominant_power = np.nan + else: + periodic = np.atleast_2d(periodic) + peak_count = float(periodic.shape[0]) + dominant_idx = int(np.nanargmax(periodic[:, 1])) + dominant_freq = float(periodic[dominant_idx, 0]) + dominant_power = float(periodic[dominant_idx, 1]) + + offset = float(aperiodic[0]) if aperiodic.size >= 1 else np.nan + knee = float(aperiodic[1]) if aperiodic.size == 3 else np.nan + exponent = float(aperiodic[-1]) if aperiodic.size >= 2 else np.nan + + periodic_psd = None + if need_periodic_psd: + full_log = np.asarray(model.results.model.get_component("full"), dtype=float) + aperiodic_log = np.asarray( + model.results.model.get_component("aperiodic"), + dtype=float, + ) + if not np.all(np.isfinite(full_log)) or not np.all(np.isfinite(aperiodic_log)): + raise RuntimeError( + "Specparam model components became non-finite for corrected bands." + ) + periodic_psd = np.clip( + np.power(10.0, full_log) - np.power(10.0, aperiodic_log), + 0.0, + None, + ) + + return { + "offset": offset, + "knee": knee, + "exponent": exponent, + "fit_error": error, + "r_squared": r_squared, + "peak_count": peak_count, + "peak_freq_dom": dominant_freq, + "peak_power_dom": dominant_power, + }, periodic_psd + + +def fit_parametric_batch( + psds: np.ndarray, + freqs: np.ndarray, + config: ParametricDescriptorConfig, + runtime, + need_periodic_psd: bool = False, + include_metrics: bool = True, +) -> _ParametricFitBatch: + """ + Fit parametric models over one PSD batch. + + Parameters + ---------- + psds : np.ndarray + PSD batch with shape ``(n_obs, n_channels, n_freqs)``. + freqs : np.ndarray + Frequency grid aligned with the last axis of ``psds``. + config : ParametricDescriptorConfig + Parsed parametric fit configuration. + runtime : DescriptorRuntimeConfig + Runtime execution controls. Only the inner fitting parallelism path + uses this object here. + need_periodic_psd : bool, default=False + Whether to reconstruct periodic-only PSDs for each fitted spectrum. + include_metrics : bool, default=True + Whether to materialize scalar metric arrays in the returned payload. + + Returns + ------- + _ParametricFitBatch + Batch-scoped fit payload aligned to the input observation and channel + axes after restricting the PSD to ``config.freq_range``. + """ + freq_mask = (freqs >= config.freq_range[0]) & (freqs <= config.freq_range[1]) + local_freqs = freqs[freq_mask] + local_psds = psds[..., freq_mask] + + metric_names: list[str] = [] + if include_metrics: + if "aperiodic" in config.outputs: + if config.aperiodic_mode == "knee": + metric_names.extend(["offset", "knee", "exponent"]) + else: + metric_names.extend(["offset", "exponent"]) + if "fit_quality" in config.outputs: + metric_names.extend(["fit_error", "r_squared"]) + if "peak_summary" in config.outputs: + metric_names.extend(["peak_count", "peak_freq_dom", "peak_power_dom"]) + metric_arrays = { + metric_name: np.full( + (local_psds.shape[0], local_psds.shape[1]), + np.nan, + dtype=float, + ) + for metric_name in metric_names + } + periodic_psds = ( + np.full(local_psds.shape, np.nan, dtype=float) if need_periodic_psd else None + ) + + def fit_one( + obs_rel: int, + unit_idx: int, + ) -> tuple[ + int, + int, + dict[str, float] | None, + np.ndarray | None, + dict[str, str] | None, + ]: + try: + metrics, periodic = fit_single_spectrum( + local_freqs, + local_psds[obs_rel, unit_idx], + config, + need_periodic_psd=need_periodic_psd, + ) + return obs_rel, unit_idx, metrics, periodic, None + except Exception as exc: # pragma: no cover - exercised via callers + return ( + obs_rel, + unit_idx, + None, + None, + { + "exception_type": type(exc).__name__, + "message": str(exc), + }, + ) + + if runtime.execution_backend != "sequential" and runtime.n_jobs != 1: + import joblib + + tasks = [ + (obs_rel, unit_idx) + for obs_rel in range(local_psds.shape[0]) + for unit_idx in range(local_psds.shape[1]) + ] + fit_results = joblib.Parallel( + n_jobs=runtime.n_jobs, + prefer="threads", + )(joblib.delayed(fit_one)(obs_rel, unit_idx) for obs_rel, unit_idx in tasks) + else: + fit_results = [ + fit_one(obs_rel, unit_idx) + for obs_rel in range(local_psds.shape[0]) + for unit_idx in range(local_psds.shape[1]) + ] + + errors: list[tuple[int, int, str, str]] = [] + for obs_rel, unit_idx, metrics, periodic, error in fit_results: + if metrics is not None: + for metric_name in metric_names: + metric_arrays[metric_name][obs_rel, unit_idx] = metrics[metric_name] + if periodic_psds is not None and periodic is not None: + periodic_psds[obs_rel, unit_idx] = periodic + continue + errors.append( + ( + obs_rel, + unit_idx, + error["exception_type"], + error["message"], + ) + ) + + return _ParametricFitBatch( + freqs=np.asarray(local_freqs, dtype=float), + metrics=metric_arrays, + errors=errors, + periodic_psds=periodic_psds, + meta={ + "backend": config.backend, + "freq_range": list(config.freq_range), + "aperiodic_mode": config.aperiodic_mode, + }, + ) diff --git a/coco_pipe/descriptors/extractors/_psd.py b/coco_pipe/descriptors/extractors/_psd.py new file mode 100644 index 0000000..387ef3d --- /dev/null +++ b/coco_pipe/descriptors/extractors/_psd.py @@ -0,0 +1,139 @@ +""" +Shared PSD computation for PSD-consuming descriptor paths. + +This module holds the reusable PSD step used by the descriptors planner and by +PSD-consuming extractors when they need a standalone spectral input. It does +not define descriptor semantics. It only: + +- prepare a writable runtime environment for MNE-backed PSD helpers +- lazily import the MNE PSD functions used by descriptors +- compute Welch or multitaper PSD batches on explicit NumPy inputs + +Author: Hamza Abdelhedi (hamza.abdelhedi@umontreal.ca) +""" + +from __future__ import annotations + +import os +import tempfile + +import numpy as np + +from ...utils import import_optional_dependency + + +def load_mne_psd_functions(): + """Lazily import MNE PSD helpers with writable runtime cache locations. + + Returns + ------- + tuple + ``(psd_array_welch, psd_array_multitaper)`` imported from + `mne.time_frequency`. + + Notes + ----- + MNE may write cache or config files during import/use. The descriptors + module keeps those paths inside the system temp directory so PSD + computation remains sandbox-friendly. + """ + tmp_root = os.path.join(tempfile.gettempdir(), "coco_pipe_descriptors") + mpl_dir = os.path.join(tmp_root, "mplconfig") + mne_dir = os.path.join(tmp_root, "mne") + os.makedirs(mpl_dir, exist_ok=True) + os.makedirs(mne_dir, exist_ok=True) + os.environ.setdefault("MPLCONFIGDIR", mpl_dir) + os.environ.setdefault("MNE_HOME", mne_dir) + os.environ.setdefault("MNE_DONTWRITE_HOME", "true") + + return import_optional_dependency( + lambda: ( + __import__( + "mne.time_frequency", + fromlist=["psd_array_welch", "psd_array_multitaper"], + ).psd_array_welch, + __import__( + "mne.time_frequency", + fromlist=["psd_array_welch", "psd_array_multitaper"], + ).psd_array_multitaper, + ), + feature="descriptor spectral extraction", + dependency="mne", + install_hint="pip install coco-pipe[descriptors,eeg]", + ) + + +def compute_psd( + X: np.ndarray, + sfreq: float, + method: str, + fmin: float, + fmax: float, + n_jobs: int | None = None, +) -> tuple[np.ndarray, np.ndarray]: + """ + Compute PSD values for one batch of segmented signals. + + Parameters + ---------- + X : np.ndarray + Input array with shape ``(n_obs, n_channels, n_times)``. + sfreq : float + Sampling frequency in Hertz. + method : {"welch", "multitaper"} + PSD estimator to use. + fmin : float + Lower frequency bound passed to the PSD backend. + fmax : float + Upper frequency bound passed to the PSD backend. + n_jobs : int, optional + Parallel worker count forwarded to the MNE PSD backend when the caller + enables PSD-level parallelism. `None` leaves the backend default in + place. + + Returns + ------- + tuple[np.ndarray, np.ndarray] + PSD values with shape ``(n_obs, n_channels, n_freqs)`` and the aligned + frequency grid with shape ``(n_freqs,)``. + + Notes + ----- + For Welch PSDs, the descriptors module uses: + + - ``n_fft = min(n_times, 256)`` + - ``n_per_seg = n_fft`` + + while enforcing a minimum of `8` for both values. This keeps Welch + behavior bounded and deterministic across the current descriptor tests and + examples. + """ + psd_array_welch, psd_array_multitaper = load_mne_psd_functions() + + if method == "welch": + n_fft = min(int(X.shape[-1]), 256) + psd, freqs = psd_array_welch( + X, + sfreq=sfreq, + fmin=fmin, + fmax=fmax, + n_fft=max(n_fft, 8), + n_per_seg=max(n_fft, 8), + average="mean", + n_jobs=n_jobs, + verbose=False, + ) + return np.asarray(psd, dtype=float), np.asarray(freqs, dtype=float) + + if method == "multitaper": + psd, freqs = psd_array_multitaper( + X, + sfreq=sfreq, + fmin=fmin, + fmax=fmax, + n_jobs=n_jobs, + verbose=False, + ) + return np.asarray(psd, dtype=float), np.asarray(freqs, dtype=float) + + raise ValueError(f"Unknown PSD method: {method}") diff --git a/coco_pipe/descriptors/extractors/base.py b/coco_pipe/descriptors/extractors/base.py new file mode 100644 index 0000000..6ccefa6 --- /dev/null +++ b/coco_pipe/descriptors/extractors/base.py @@ -0,0 +1,384 @@ +""" +Base interfaces for descriptor extraction backends. + +This module defines the internal contracts shared by built-in descriptor +extractors. The module exposes: + +- `BaseDescriptorExtractor` for families that consume validated raw signal + batches +- `BasePSDDescriptorExtractor` for families that consume shared PSD batches +- `_DescriptorBlock` as the private family output payload + +The surrounding descriptors stack uses these interfaces to provide: + +- explicit runtime dispatch from `DescriptorPipeline` +- deterministic descriptor naming and channel reduction helpers +- family-wise metadata and failure collection +- safe merging of family outputs into one stable result dictionary + +Notes +----- +`BaseDescriptorExtractor` is an internal extension point for descriptor +families. Unlike dim-reduction reducers, descriptor extractors are stateless at +runtime and do not expose `fit`, persistence, or model objects. + +Examples +-------- +The shared finalization helper converts per-channel descriptor values into the +public column naming convention based on ``output.channel_pooling``: + +- ``channel_pooling="none"``: + ``band_abs_alpha_ch-Fz``, ``band_abs_alpha_ch-Cz`` +- ``channel_pooling="all"``: + ``band_abs_alpha_ch-all`` +- ``channel_pooling={"Frontal": ["Fz", "Cz"]}``: + ``band_abs_alpha_chgrp-Frontal`` plus any ungrouped channels such as + ``band_abs_alpha_ch-Pz`` + +Author: Hamza Abdelhedi (hamza.abdelhedi@umontreal.ca) +""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from dataclasses import dataclass, field +from typing import Any + +import numpy as np + +from ..configs import DescriptorRuntimeConfig +from .utils import pool_channel_descriptor_matrix + +__all__ = ["BaseDescriptorExtractor", "BasePSDDescriptorExtractor"] + + +@dataclass(slots=True) +class _DescriptorBlock: + """Private in-memory descriptor payload for one family. + + Attributes + ---------- + family : str + Canonical family name that produced the block. + X : np.ndarray + Family-specific descriptor matrix aligned on the observation axis. + descriptor_names : list of str + Deterministic column names aligned with the columns of ``X``. + meta : dict + Family-specific metadata to preserve under the merged result. + failures : list of dict + Normalized failure records collected during extraction. + + Notes + ----- + ``X.shape[0]`` must always match the input observation count seen by the + extractor, and ``len(descriptor_names)`` must always match ``X.shape[1]``. + The pipeline depends on these alignment guarantees when merging family + outputs. + """ + + family: str + X: np.ndarray + descriptor_names: list[str] + meta: dict[str, Any] = field(default_factory=dict) + failures: list[dict[str, Any]] = field(default_factory=list) + + +class BaseDescriptorExtractor(ABC): + """ + Abstract base class for descriptor extraction families. + + Subclasses receive already validated NumPy inputs and must return one + `_DescriptorBlock` aligned on the observation axis. The base class keeps + the extractor API narrow and provides a shared helper for channel + finalization and deterministic descriptor naming. + + Parameters + ---------- + config : Any + Typed family configuration parsed by `DescriptorConfig`. + + Attributes + ---------- + config : Any + Stored family-specific configuration object. + family_name : str + Stable family identifier used in failure records and merged metadata. + + Notes + ----- + Extractors are stateless at runtime. They do not learn parameters across + calls; all runtime state is provided explicitly through `extract()`. + + Concrete extractors are expected to: + + 1. compute family-specific values with shape ``(n_obs, n_channels)`` for + each metric + 2. pass those values through :meth:`_finalize_descriptor` + 3. return one `_DescriptorBlock` with aligned names, metadata, and failures + + Examples + -------- + A minimal concrete extractor typically looks like: + + >>> class MeanOverTimeExtractor(BaseDescriptorExtractor): + ... family_name = "toy" + ... + ... def extract( + ... self, + ... X, + ... sfreq, + ... channel_names, + ... channel_pooling, + ... ids, + ... runtime, + ... ): + ... values = X.mean(axis=-1) + ... X_out, names = self._finalize_descriptor( + ... values, + ... family_prefix="toy", + ... metric_name="mean", + ... channel_names=channel_names, + ... channel_pooling=channel_pooling, + ... ) + ... return _DescriptorBlock( + ... family=self.family_name, + ... X=X_out, + ... descriptor_names=names, + ... ) + """ + + family_name = "base" + + def __init__(self, config: Any): + """Store the typed family configuration.""" + self.config = config + + @property + def capabilities(self) -> dict[str, Any]: + """Return static extractor capability metadata. + + Returns + ------- + dict[str, Any] + Static metadata describing optional dependencies and general + execution properties for the extractor. + + Notes + ----- + The descriptors pipeline currently uses this mapping only as lightweight + backend metadata. It is intentionally much smaller than the reducer + capability surface in `dim_reduction`. + """ + return { + "requires_sfreq": False, + "supports_batching": True, + "supports_channelwise": True, + "deterministic": True, + "optional_dependencies": [], + } + + @abstractmethod + def extract( + self, + X: np.ndarray, + sfreq: float | None, + channel_names: list[str] | None, + channel_pooling: str | dict[str, list[str]], + ids: np.ndarray | None, + runtime: DescriptorRuntimeConfig, + obs_offset: int = 0, + ) -> _DescriptorBlock: + """Extract descriptors from a validated input array. + + Parameters + ---------- + X : np.ndarray + Input array with shape ``(n_obs, n_channels, n_times)``. + sfreq : float, optional + Sampling frequency in Hertz. + channel_names : list of str, optional + Explicit channel labels aligned with axis 1 of ``X``. + channel_pooling : {"none", "all"} or dict + Descriptor-level channel pooling policy applied after per-channel + descriptors are computed. + ids : np.ndarray, optional + Observation identifiers aligned with axis 0 of ``X``. + runtime : DescriptorRuntimeConfig + Runtime execution controls shared across extractors. + obs_offset : int, default=0 + Global observation offset applied to any collected failure records. + + Returns + ------- + _DescriptorBlock + Family-specific descriptor matrix plus metadata and failures. + + Raises + ------ + ImportError + If an optional backend required by the extractor is unavailable. + ValueError + If the extractor encounters an invalid runtime condition and the + configured error policy requires raising. + + Notes + ----- + The recommended pattern is to keep family-specific computation local to + the extractor and delegate all public channel naming and channel pooling + behavior to :meth:`_finalize_descriptor`. + """ + + def _finalize_descriptor( + self, + values: np.ndarray, + family_prefix: str, + metric_name: str, + channel_names: list[str] | None, + channel_pooling: str | dict[str, list[str]] = "none", + ) -> tuple[np.ndarray, list[str]]: + """Pool channels and build deterministic descriptor names. + + Parameters + ---------- + values : np.ndarray + Family metric values with shape ``(n_obs, n_channels)`` or + ``(n_obs,)``. + family_prefix : str + Stable family prefix, for example ``"band"`` or ``"param"``. + metric_name : str + Family-local metric identifier used in the descriptor name. + channel_names : list of str, optional + Channel labels used when building channel-resolved descriptor names. + channel_pooling : {"none", "all"} or dict, default="none" + Descriptor-level channel pooling policy. + + Returns + ------- + tuple + ``(X_metric, names)`` where ``X_metric`` is the finalized metric + matrix and ``names`` is the aligned list of descriptor names. + + Notes + ----- + This helper assumes ``values`` already represents descriptor values, not + raw signals. Pooling therefore always happens at the descriptor level: + + - ``"none"`` keeps one column per sensor + - ``"all"`` averages descriptor values across all sensors + - a mapping averages descriptor values within each named group and keeps + ungrouped sensors as individual columns + + Examples + -------- + Given ``channel_names=["Fz", "Cz", "Pz"]`` and + ``metric_name="abs_alpha"``: + + - ``channel_pooling="none"`` yields + ``["band_abs_alpha_ch-Fz", "band_abs_alpha_ch-Cz", "band_abs_alpha_ch-Pz"]`` + - ``channel_pooling="all"`` yields + ``["band_abs_alpha_ch-all"]`` + - ``channel_pooling={"Frontal": ["Fz", "Cz"]}`` yields + ``["band_abs_alpha_chgrp-Frontal", "band_abs_alpha_ch-Pz"]`` + """ + if values.ndim == 1: + values = values[:, None] + pooled_values, scopes = pool_channel_descriptor_matrix( + values, + channel_names=channel_names or [], + channel_pooling=channel_pooling, + ) + names = ["_".join((family_prefix, metric_name, scope)) for scope in scopes] + return pooled_values, names + + +class BasePSDDescriptorExtractor(BaseDescriptorExtractor): + """ + Abstract base class for descriptor families that consume PSD batches. + + PSD-consuming families still participate in the shared descriptor contract, + but they expose one additional explicit entry point: + + - `extract_psd(...)` consumes precomputed `psds, freqs` + - `psd_request()` tells the planner which PSD range and method is needed + + This keeps the generic raw-signal interface narrow while still giving the + planner one formal PSD-consumer contract shared by spectral and parametric + families. + + Notes + ----- + PSD consumers may still expose `extract()` to satisfy the generic family + interface, but the shared planner uses `psd_request()` and `extract_psd()` + exclusively once PSD intermediates have been materialized. + """ + + @abstractmethod + def psd_request(self) -> dict[str, Any]: + """Describe the PSD requirements for the shared planner. + + Returns + ------- + dict[str, Any] + Minimal request payload containing the PSD method and the required + frequency range for this family. + """ + + def parametric_fit_requirements(self) -> dict[str, Any]: + """Describe whether this PSD consumer needs a shared parametric fit. + + Returns + ------- + dict[str, Any] + Shared-fit requirements with the keys: + + - `needed` + - `metrics` + - `periodic_psds` + - `config` + """ + return { + "needed": False, + "metrics": False, + "periodic_psds": False, + "config": None, + } + + @abstractmethod + def extract_psd( + self, + psds: np.ndarray, + freqs: np.ndarray, + channel_names: list[str] | None, + channel_pooling: str | dict[str, list[str]], + ids: np.ndarray | None, + runtime: DescriptorRuntimeConfig, + obs_offset: int = 0, + fit_batch: Any | None = None, + ) -> _DescriptorBlock: + """Extract descriptors from explicit PSD intermediates. + + Parameters + ---------- + psds : np.ndarray + PSD batch with shape ``(n_obs, n_channels, n_freqs)``. + freqs : np.ndarray + Frequency grid aligned with the last axis of ``psds``. + channel_names : list of str, optional + Explicit channel labels aligned with the channel axis. + channel_pooling : {"none", "all"} or dict + Descriptor-level channel pooling policy. + ids : np.ndarray, optional + Observation identifiers aligned with the observation axis. + runtime : DescriptorRuntimeConfig + Runtime execution controls shared across extractors. + obs_offset : int, default=0 + Global observation offset applied to collected failure records. + fit_batch : Any, optional + Additional shared fit payload required by some PSD consumers. + + Returns + ------- + _DescriptorBlock + Family-specific descriptor block aligned with the input PSD batch. + """ diff --git a/coco_pipe/descriptors/extractors/complexity.py b/coco_pipe/descriptors/extractors/complexity.py new file mode 100644 index 0000000..c0a7fd7 --- /dev/null +++ b/coco_pipe/descriptors/extractors/complexity.py @@ -0,0 +1,404 @@ +""" +Complexity descriptor extraction backend. + +This module implements the built-in complexity family for +`coco_pipe.descriptors`. The extractor operates on already segmented NumPy +inputs with shape ``(n_obs, n_channels, n_times)`` and computes one or more +complexity measures per sensor, per observation. + +Notes +----- +The complexity family prefers batched backend calls when the selected library +supports them. In the current implementation: + +- `spectral_entropy`, `hjorth_mobility`, and `hjorth_complexity` use batched + `antropy` calls over flattened observation-channel units +- `sample_entropy`, `perm_entropy`, and `lziv_complexity` are still evaluated + one 1D signal at a time + +Author: Hamza Abdelhedi (hamza.abdelhedi@umontreal.ca) +""" + +from __future__ import annotations + +from typing import Any + +import numpy as np + +from ...utils import import_optional_dependency +from ..configs import ComplexityDescriptorConfig +from .base import BaseDescriptorExtractor, _DescriptorBlock +from .utils import make_failure_record + + +class ComplexityDescriptorExtractor(BaseDescriptorExtractor): + """ + Complexity descriptor extractor. + + This extractor computes scalar complexity measures for each observation and + sensor in a validated descriptor input array. It is intended for signals + that are already segmented upstream, such as epochs, windows, or trial + blocks. + + Parameters + ---------- + config : ComplexityDescriptorConfig + Parsed family configuration controlling the selected measures, backend, + and any per-measure keyword arguments. + + Attributes + ---------- + config : ComplexityDescriptorConfig + Stored typed configuration for the complexity family. + family_name : str + Stable family identifier used in metadata and failure records. + + Notes + ----- + The extractor always computes descriptor values per sensor first. Public + output pooling, such as `channel_pooling="all"` or grouped channel pooling, + is applied afterward through :meth:`BaseDescriptorExtractor._finalize_descriptor`. + + When `antropy` is selected, the extractor uses batched calls where the + backend supports them and falls back to scalar loops for measures that are + inherently one-signal-at-a-time in the current backend API. + """ + + family_name = "complexity" + + def __init__(self, config: ComplexityDescriptorConfig): + super().__init__(config) + self.config = config + + @property + def capabilities(self) -> dict[str, Any]: + """Return static complexity extractor capability metadata. + + Returns + ------- + dict[str, Any] + Capability metadata describing sampling-rate requirements and the + optional backends used by the complexity family. + + Notes + ----- + `spectral_entropy` requires an explicit sampling rate, while the other + currently supported measures do not. + """ + return { + **super().capabilities, + "requires_sfreq": "spectral_entropy" in self.config.measures, + "optional_dependencies": ["antropy", "neurokit2"], + } + + def _load_antropy(self): + """Import `antropy` lazily when the configured backend needs it. + + Returns + ------- + module + Imported `antropy` module. + + Raises + ------ + ImportError + If `antropy` is not installed. + """ + return import_optional_dependency( + lambda: __import__("antropy"), + feature="complexity descriptor extraction", + dependency="antropy", + install_hint="pip install coco-pipe[descriptors]", + ) + + def _load_neurokit(self): + """Import `neurokit2` lazily when the configured backend needs it. + + Returns + ------- + module + Imported `neurokit2` module. + + Raises + ------ + ImportError + If `neurokit2` is not installed. + """ + return import_optional_dependency( + lambda: __import__("neurokit2"), + feature="neurokit complexity descriptor extraction", + dependency="neurokit2", + install_hint="pip install coco-pipe[descriptors]", + ) + + def extract( + self, + X: np.ndarray, + sfreq: float | None, + channel_names: list[str] | None, + channel_pooling: str | dict[str, list[str]], + ids: np.ndarray | None, + runtime, + obs_offset: int = 0, + ) -> _DescriptorBlock: + """Extract complexity descriptors from segmented multi-channel data. + + Parameters + ---------- + X : np.ndarray + Input array with shape ``(n_obs, n_channels, n_times)``. Each row + already represents one observation segment produced upstream. + sfreq : float, optional + Sampling frequency in Hertz. Required when + `spectral_entropy` is requested. + channel_names : list of str, optional + Explicit channel labels aligned with axis 1 of ``X``. If omitted, + fallback names ``"ch-0"``, ``"ch-1"``, ... are used internally. + channel_pooling : {"none", "all"} or dict + Descriptor-level channel pooling policy applied after per-sensor + complexity values are computed. + ids : np.ndarray, optional + Observation identifiers aligned with axis 0 of ``X``. + runtime : DescriptorRuntimeConfig + Runtime execution controls shared across descriptor families. + obs_offset : int, default=0 + Global observation offset added to any collected failure records + when this extractor is called on one observation batch. + + Returns + ------- + _DescriptorBlock + Complexity-family descriptor block aligned with the input + observation axis. + + Raises + ------ + ImportError + If the configured optional backend is unavailable. + ValueError + If a requested measure is unsupported by the selected backend, or + if runtime error handling is configured to raise on a numerical or + backend failure. + + Notes + ----- + The extractor uses a mixed execution strategy: + + - batched `antropy` calls for `spectral_entropy`, + `hjorth_mobility`, and `hjorth_complexity` + - scalar calls for `sample_entropy`, `perm_entropy`, and + `lziv_complexity` + + Non-finite outputs are converted to `NaN` and recorded under + ``failures`` unless `runtime.on_error == "raise"`, in which case the + extractor fails immediately. + + Examples + -------- + With ``channel_pooling="none"`` and + ``channel_names=["Fz", "Cz"]``, a requested measure such as + ``perm_entropy`` yields channel-resolved names like + ``complexity_perm_entropy_ch-Fz`` and + ``complexity_perm_entropy_ch-Cz``. + + With ``channel_pooling="all"``, the same metric yields one pooled + column named ``complexity_perm_entropy_ch-all``. + """ + channel_names = channel_names or [f"ch-{idx}" for idx in range(X.shape[1])] + + descriptor_names: list[str] | None = None + failures: list[dict[str, Any]] = [] + metric_arrays = { + measure: np.full((X.shape[0], X.shape[1]), np.nan, dtype=float) + for measure in self.config.measures + } + measure_kwargs = { + measure: dict(self.config.measure_kwargs.get(measure, {})) + for measure in self.config.measures + } + flat_signals = X.reshape(-1, X.shape[-1]) + batched_outputs: dict[str, np.ndarray] = {} + scalar_dispatch: dict[str, Any] = {} + + if self.config.backend in {"antropy", "auto"}: + ant = self._load_antropy() + + if "spectral_entropy" in self.config.measures: + batched_outputs["spectral_entropy"] = np.asarray( + ant.spectral_entropy( + flat_signals, + sf=sfreq, + axis=-1, + **measure_kwargs["spectral_entropy"], + ), + dtype=float, + ) + + if ( + "hjorth_mobility" in self.config.measures + or "hjorth_complexity" in self.config.measures + ): + mobility, complexity = ant.hjorth_params( + flat_signals, + axis=-1, + ) + if "hjorth_mobility" in self.config.measures: + batched_outputs["hjorth_mobility"] = np.asarray( + mobility, + dtype=float, + ) + if "hjorth_complexity" in self.config.measures: + batched_outputs["hjorth_complexity"] = np.asarray( + complexity, + dtype=float, + ) + + scalar_dispatch = { + "sample_entropy": lambda signal, kwargs, sfreq: float( + ant.sample_entropy(signal, **kwargs) + ), + "perm_entropy": lambda signal, kwargs, sfreq: float( + ant.perm_entropy(signal, **kwargs) + ), + "lziv_complexity": lambda signal, kwargs, sfreq: float( + ant.lziv_complexity( + (signal > np.median(signal)).astype(int), + **kwargs, + ) + ), + } + else: + nk = self._load_neurokit() + scalar_dispatch = { + "sample_entropy": lambda signal, kwargs, sfreq: float( + nk.entropy_sample(signal, **kwargs)[0] + ), + "perm_entropy": lambda signal, kwargs, sfreq: float( + nk.entropy_permutation(signal, **kwargs)[0] + ), + "spectral_entropy": lambda signal, kwargs, sfreq: float( + nk.entropy_spectral( + signal, + sampling_rate=sfreq, + **kwargs, + )[0] + ), + } + + for measure, flat_values in batched_outputs.items(): + values = np.asarray(flat_values, dtype=float).reshape( + X.shape[0], + X.shape[1], + ) + metric_arrays[measure][:] = np.where(np.isfinite(values), values, np.nan) + bad_positions = np.argwhere(~np.isfinite(values)) + if bad_positions.size == 0: + continue + + message = f"Complexity measure '{measure}' produced a non-finite result." + if runtime.on_error == "raise": + raise ValueError(message) + + for obs_rel, unit_idx in bad_positions: + failures.append( + make_failure_record( + family=self.family_name, + obs_index=obs_offset + int(obs_rel), + obs_id=None if ids is None else ids[int(obs_rel)], + channel_index=int(unit_idx), + channel_name=channel_names[int(unit_idx)], + exception_type="NumericalIssue", + message=message, + ) + ) + + scalar_measures = [ + measure + for measure in self.config.measures + if measure not in batched_outputs + ] + unsupported = sorted(set(scalar_measures) - set(scalar_dispatch)) + if unsupported: + raise ValueError( + f"Measures {unsupported} are not supported by backend " + f"'{self.config.backend}'." + ) + + for obs_rel in range(X.shape[0]): + unit_signals = X[obs_rel] + obs_id = None if ids is None else ids[obs_rel] + + for unit_idx, signal in enumerate(unit_signals): + for measure in scalar_measures: + try: + value = scalar_dispatch[measure]( + signal, + measure_kwargs[measure], + sfreq, + ) + if np.isfinite(value): + metric_arrays[measure][obs_rel, unit_idx] = float(value) + else: + if runtime.on_error == "raise": + raise ValueError( + "Complexity measure produced a non-finite result." + ) + failures.append( + make_failure_record( + family=self.family_name, + obs_index=obs_offset + obs_rel, + obs_id=obs_id, + channel_index=unit_idx, + channel_name=channel_names[unit_idx], + exception_type="NumericalIssue", + message=( + "Complexity measure " + f"'{measure}' produced a non-finite result." + ), + ) + ) + except Exception as exc: # pragma: no cover - hit via failure tests + if isinstance(exc, ImportError): + raise + if runtime.on_error == "raise": + raise + failures.append( + make_failure_record( + family=self.family_name, + obs_index=obs_offset + obs_rel, + obs_id=obs_id, + channel_index=unit_idx, + channel_name=channel_names[unit_idx], + exception_type=type(exc).__name__, + message=str(exc), + ) + ) + + chunk_features: list[np.ndarray] = [] + chunk_names: list[str] = [] + for measure in self.config.measures: + feature, names = self._finalize_descriptor( + metric_arrays[measure], + family_prefix="complexity", + metric_name=measure, + channel_names=channel_names, + channel_pooling=channel_pooling, + ) + chunk_features.append(feature) + chunk_names.extend(names) + + descriptor_names = chunk_names + + return _DescriptorBlock( + family=self.family_name, + X=np.concatenate(chunk_features, axis=1) + if chunk_features + else np.empty((X.shape[0], 0)), + descriptor_names=descriptor_names or [], + meta={ + "backend": self.config.backend, + "measures": list(self.config.measures), + "batched_measures": sorted(batched_outputs), + }, + failures=failures, + ) diff --git a/coco_pipe/descriptors/extractors/parametric.py b/coco_pipe/descriptors/extractors/parametric.py new file mode 100644 index 0000000..8bfaa60 --- /dev/null +++ b/coco_pipe/descriptors/extractors/parametric.py @@ -0,0 +1,296 @@ +""" +Parametric spectral descriptor extraction backend. + +This module implements the built-in parametric spectral family for +`coco_pipe.descriptors`. The extractor operates on already segmented NumPy +inputs with shape ``(n_obs, n_channels, n_times)`` and computes one or more +specparam-derived summary descriptors per sensor, per observation. + +Notes +----- +The parametric family is a PSD consumer. When used through +`DescriptorPipeline.extract()`, it can share one batch-scoped PSD computation +with other compatible PSD consumers such as the spectral band family. The +actual descriptor outputs are then derived from that shared `psds, freqs` pair. + +Model fitting itself still happens one spectrum at a time. When runtime +parallelism is enabled and the planner allows it, those per-spectrum fits can +run in parallel across observation-channel units. + +Author: Hamza Abdelhedi (hamza.abdelhedi@umontreal.ca) +""" + +from __future__ import annotations + +from typing import Any + +import numpy as np + +from ..configs import ParametricDescriptorConfig +from ._parametric_fit import _ParametricFitBatch, fit_parametric_batch +from ._psd import compute_psd +from .base import BasePSDDescriptorExtractor, _DescriptorBlock +from .utils import make_failure_record + + +class ParametricDescriptorExtractor(BasePSDDescriptorExtractor): + """ + Parametric spectral descriptor extractor. + + This extractor fits one specparam model per observation and sensor in a + validated descriptor input array, then exposes scalar summaries such as + aperiodic parameters, fit quality, and dominant peak statistics. + + Parameters + ---------- + config : ParametricDescriptorConfig + Parsed family configuration controlling the PSD method, fit range, + specparam settings, and requested output groups. + + Attributes + ---------- + config : ParametricDescriptorConfig + Stored typed configuration for the parametric family. + family_name : str + Stable family identifier used in metadata and failure records. + + Notes + ----- + The extractor always computes descriptor values per sensor first. Public + output pooling, such as `channel_pooling="all"` or grouped channel pooling, + is applied afterward through + :meth:`BaseDescriptorExtractor._finalize_descriptor`. + + When the pipeline provides a precomputed PSD batch through + :meth:`extract_psd`, the extractor reuses that shared spectral input + and expects an explicit shared `fit_batch`. Standalone :meth:`extract` + remains available for family-local PSD and fit computation. + """ + + family_name = "parametric" + + def __init__(self, config: ParametricDescriptorConfig): + super().__init__(config) + self.config = config + + @property + def capabilities(self) -> dict[str, Any]: + """Return static parametric extractor capability metadata. + + Returns + ------- + dict[str, Any] + Capability metadata describing sampling-rate requirements and the + optional backends used by the parametric family. + """ + return { + **super().capabilities, + "requires_sfreq": True, + "optional_dependencies": ["specparam", "mne"], + } + + def psd_request(self) -> dict[str, Any]: + """Describe the PSD requirements for the shared planner.""" + return { + "method": self.config.psd_method, + "fmin": self.config.freq_range[0], + "fmax": self.config.freq_range[1], + } + + def parametric_fit_requirements(self) -> dict[str, Any]: + """Describe whether this family needs shared parametric-fit outputs.""" + return { + "needed": True, + "metrics": True, + "periodic_psds": False, + "config": self.config, + } + + def extract_psd( + self, + psds: np.ndarray, + freqs: np.ndarray, + channel_names: list[str] | None, + channel_pooling: str | dict[str, list[str]], + ids: np.ndarray | None, + runtime, + obs_offset: int = 0, + fit_batch: _ParametricFitBatch | None = None, + ) -> _DescriptorBlock: + """Extract parametric descriptors from a precomputed PSD batch. + + Parameters + ---------- + psds : np.ndarray + Power spectral density array with shape + ``(n_obs, n_channels, n_freqs)``. + freqs : np.ndarray + Frequency grid aligned with the last axis of ``psds``. + channel_names : list of str, optional + Explicit channel labels aligned with axis 1 of ``psds``. If + omitted, fallback names ``"ch-0"``, ``"ch-1"``, ... are used + internally. + channel_pooling : {"none", "all"} or dict + Descriptor-level channel pooling policy applied after per-sensor + parametric values are computed. + ids : np.ndarray, optional + Observation identifiers aligned with axis 0 of ``psds``. + runtime : DescriptorRuntimeConfig + Runtime execution controls shared across descriptor families. + obs_offset : int, default=0 + Global observation offset added to any collected failure records + when this extractor is called on one observation batch. + + Returns + ------- + _DescriptorBlock + Parametric-family descriptor block aligned with the input + observation axis. + + Raises + ------ + ValueError + If ``fit_batch`` is not supplied. + + Notes + ----- + This method consumes explicit shared intermediates. It does not compute + PSDs or fit models on its own. + """ + channel_names = channel_names or [f"ch-{idx}" for idx in range(psds.shape[1])] + if fit_batch is None: + raise ValueError("Parametric extract_psd() requires a supplied fit_batch.") + + chunk_metric_arrays = fit_batch.metrics + metrics: list[str] = [] + if "aperiodic" in self.config.outputs: + metrics.extend(["offset", "exponent"]) + if "knee" in chunk_metric_arrays: + metrics.append("knee") + if "fit_quality" in self.config.outputs: + metrics.extend(["fit_error", "r_squared"]) + if "peak_summary" in self.config.outputs: + metrics.extend(["peak_count", "peak_freq_dom", "peak_power_dom"]) + failures: list[dict[str, Any]] = [] + for obs_rel, unit_idx, exception_type, message in fit_batch.errors: + if runtime.on_error == "raise": + raise RuntimeError(message) + failures.append( + make_failure_record( + family=self.family_name, + obs_index=obs_offset + obs_rel, + obs_id=None if ids is None else ids[obs_rel], + channel_index=unit_idx, + channel_name=channel_names[unit_idx], + exception_type=exception_type, + message=message, + ) + ) + + chunk_features: list[np.ndarray] = [] + descriptor_names: list[str] = [] + for metric_name in metrics: + feature, names = self._finalize_descriptor( + chunk_metric_arrays[metric_name], + family_prefix="param", + metric_name=metric_name, + channel_names=channel_names, + channel_pooling=channel_pooling, + ) + chunk_features.append(feature) + descriptor_names.extend(names) + + return _DescriptorBlock( + family=self.family_name, + X=np.concatenate(chunk_features, axis=1) + if chunk_features + else np.empty((psds.shape[0], 0)), + descriptor_names=descriptor_names, + meta={ + **fit_batch.meta, + "psd_method": self.config.psd_method, + }, + failures=failures, + ) + + def extract( + self, + X: np.ndarray, + sfreq: float | None, + channel_names: list[str] | None, + channel_pooling: str | dict[str, list[str]], + ids: np.ndarray | None, + runtime, + obs_offset: int = 0, + ) -> _DescriptorBlock: + """Extract parametric descriptors from segmented multi-channel data. + + Parameters + ---------- + X : np.ndarray + Input array with shape ``(n_obs, n_channels, n_times)``. Each row + already represents one observation segment produced upstream. + sfreq : float, optional + Sampling frequency in Hertz. + channel_names : list of str, optional + Explicit channel labels aligned with axis 1 of ``X``. + channel_pooling : {"none", "all"} or dict + Descriptor-level channel pooling policy applied after per-sensor + parametric values are computed. + ids : np.ndarray, optional + Observation identifiers aligned with axis 0 of ``X``. + runtime : DescriptorRuntimeConfig + Runtime execution controls shared across descriptor families. + obs_offset : int, default=0 + Global observation offset added to any collected failure records. + + Returns + ------- + _DescriptorBlock + Parametric-family descriptor block aligned with the input + observation axis. + + Raises + ------ + ImportError + If the optional `mne` or `specparam` backend is unavailable. + ValueError + If PSD computation encounters an invalid runtime condition. + RuntimeError + If shared fit materialization encounters a runtime failure and + ``runtime.on_error == "raise"``. + + Notes + ----- + This standalone path computes a PSD for the current batch, fits one + explicit parametric batch payload for this family, and then delegates + to :meth:`extract_psd`. When the family is executed through + `DescriptorPipeline`, the shared planner supplies the PSD and fit + payload instead. + """ + psds, freqs = compute_psd( + X, + sfreq=sfreq, + method=self.config.psd_method, + fmin=self.config.freq_range[0], + fmax=self.config.freq_range[1], + n_jobs=None, + ) + fit_batch = fit_parametric_batch( + psds, + freqs, + self.config, + runtime, + need_periodic_psd=False, + include_metrics=True, + ) + return self.extract_psd( + psds, + freqs, + channel_names=channel_names, + channel_pooling=channel_pooling, + ids=ids, + runtime=runtime, + obs_offset=obs_offset, + fit_batch=fit_batch, + ) diff --git a/coco_pipe/descriptors/extractors/spectral.py b/coco_pipe/descriptors/extractors/spectral.py new file mode 100644 index 0000000..7a15130 --- /dev/null +++ b/coco_pipe/descriptors/extractors/spectral.py @@ -0,0 +1,602 @@ +""" +Band summary descriptor extraction backend. + +This module implements the built-in spectral band family for +`coco_pipe.descriptors`. The extractor operates on already segmented NumPy +inputs with shape ``(n_obs, n_channels, n_times)`` and computes PSD-derived +band summaries per sensor, per observation. + +Notes +----- +The spectral family is a PSD consumer. When used through +`DescriptorPipeline.extract()`, it can share one batch-scoped PSD computation +with other compatible PSD consumers such as the parametric family. The actual +descriptor outputs are then derived from that shared `psds, freqs` pair. + +Within one extracted PSD batch, the family computes band integrals once and +reuses them for all requested outputs: + +- absolute power +- optional log absolute power +- relative power +- band ratios +- corrected absolute power +- corrected relative power +- corrected band ratios + +Corrected outputs are derived from periodic-only PSDs produced by a shared +parametric fit batch. They are therefore only available through the shared +planner path or an explicit ``fit_batch`` passed to :meth:`extract_psd`. + +Author: Hamza Abdelhedi (hamza.abdelhedi@umontreal.ca) +""" + +from __future__ import annotations + +from typing import Any + +import numpy as np + +from ..configs import BandDescriptorConfig +from ._parametric_fit import _ParametricFitBatch +from ._psd import compute_psd +from .base import BasePSDDescriptorExtractor, _DescriptorBlock +from .utils import make_failure_record + + +class BandDescriptorExtractor(BasePSDDescriptorExtractor): + """ + Spectral band descriptor extractor. + + This extractor computes PSD-derived band summaries for each observation and + sensor in a validated descriptor input array. It is intended for signals + that are already segmented upstream, such as epochs, windows, or trial + blocks. + + Parameters + ---------- + config : BandDescriptorConfig + Parsed family configuration controlling the PSD method, frequency + window, band definitions, and requested spectral outputs. + + Attributes + ---------- + config : BandDescriptorConfig + Stored typed configuration for the spectral band family. + family_name : str + Stable family identifier used in metadata and failure records. + + Notes + ----- + The extractor always computes descriptor values per sensor first. Public + output pooling, such as `channel_pooling="all"` or grouped channel pooling, + is applied afterward through + :meth:`BaseDescriptorExtractor._finalize_descriptor`. + + When the pipeline provides a precomputed PSD batch through + :meth:`extract_psd`, the extractor reuses that shared spectral input + instead of computing its own PSD. Corrected spectral outputs additionally + require a shared parametric fit batch and are therefore only available + through the shared planner path or an explicit `fit_batch`. + """ + + family_name = "bands" + + def __init__(self, config: BandDescriptorConfig, fit_config=None): + super().__init__(config) + self.config = config + self.fit_config = fit_config + + @property + def capabilities(self) -> dict[str, Any]: + """Return static spectral extractor capability metadata. + + Returns + ------- + dict[str, Any] + Capability metadata describing sampling-rate requirements and the + optional backend used by the spectral family. + + Notes + ----- + Spectral band extraction always requires an explicit sampling rate + because the PSD frequency axis depends on it. + """ + return { + **super().capabilities, + "requires_sfreq": True, + "optional_dependencies": ["mne"], + } + + def psd_request(self) -> dict[str, Any]: + """Describe the PSD requirements for the shared planner. + + Returns + ------- + dict[str, Any] + Minimal PSD request containing the PSD method and the required + frequency range for this family. + + Notes + ----- + `DescriptorPipeline` uses this request to group compatible PSD + consumers and decide when one batch-scoped PSD can be reused across + families. + """ + if self.needs_parametric_fit(): + if self.fit_config is None: + raise ValueError( + "Corrected band outputs require parametric fit settings." + ) + fit_low, fit_high = self.fit_config.freq_range + return { + "method": self.config.psd_method, + "fmin": min(self.config.fmin, fit_low), + "fmax": max(self.config.fmax, fit_high), + } + return { + "method": self.config.psd_method, + "fmin": self.config.fmin, + "fmax": self.config.fmax, + } + + def needs_parametric_fit(self) -> bool: + """Whether corrected spectral outputs require a shared parametric fit.""" + return any( + output + in { + "corrected_absolute_power", + "corrected_relative_power", + "corrected_ratios", + } + for output in self.config.outputs + ) + + def parametric_fit_requirements(self) -> dict[str, Any]: + """Describe whether this family needs shared parametric-fit outputs.""" + return { + "needed": self.needs_parametric_fit(), + "metrics": False, + "periodic_psds": self.needs_parametric_fit(), + "config": self.fit_config, + } + + def extract_psd( + self, + psds: np.ndarray, + freqs: np.ndarray, + channel_names: list[str] | None, + channel_pooling: str | dict[str, list[str]], + ids: np.ndarray | None, + runtime, + obs_offset: int = 0, + fit_batch: _ParametricFitBatch | None = None, + ) -> _DescriptorBlock: + """Extract band descriptors from a precomputed PSD batch. + + Parameters + ---------- + psds : np.ndarray + Power spectral density array with shape + ``(n_obs, n_channels, n_freqs)``. + freqs : np.ndarray + Frequency grid aligned with the last axis of ``psds``. + channel_names : list of str, optional + Explicit channel labels aligned with axis 1 of ``psds``. If + omitted, fallback names ``"ch-0"``, ``"ch-1"``, ... are used + internally. + channel_pooling : {"none", "all"} or dict + Descriptor-level channel pooling policy applied after per-sensor + band values are computed. + ids : np.ndarray, optional + Observation identifiers aligned with axis 0 of ``psds``. + runtime : DescriptorRuntimeConfig + Runtime execution controls shared across descriptor families. + obs_offset : int, default=0 + Global observation offset added to any collected failure records + when this extractor is called on one observation batch. + + Returns + ------- + _DescriptorBlock + Spectral-family descriptor block aligned with the input + observation axis. + + Raises + ------ + ValueError + If a configured band has no overlap with the computed PSD range and + runtime error handling is configured to raise. Also raised when + corrected outputs are requested without a supplied parametric + ``fit_batch``. + + Notes + ----- + The extractor first restricts the incoming PSD to the configured + frequency window, then integrates one power value per configured band + and sensor. Those band integrals are reused for all enabled outputs, + such as absolute power, relative power, log power, and ratios. + + Examples + -------- + With ``channel_pooling="none"`` and + ``channel_names=["Fz", "Cz"]``, an absolute alpha-band request yields + names such as ``band_abs_alpha_ch-Fz`` and ``band_abs_alpha_ch-Cz``. + + With ``channel_pooling="all"``, the same metric yields one pooled + column named ``band_abs_alpha_ch-all``. + """ + channel_names = channel_names or [f"ch-{idx}" for idx in range(psds.shape[1])] + eps = np.finfo(float).eps + corrected_failed_pairs: set[tuple[int, int]] = set() + + freq_mask = (freqs >= self.config.fmin) & (freqs <= self.config.fmax) + local_freqs = freqs[freq_mask] + local_psds = psds[..., freq_mask] + + total_power = None + if "relative_power" in self.config.outputs: + if local_freqs.size == 0: + total_power = np.full(psds.shape[:-1], np.nan, dtype=float) + else: + total_power = np.trapezoid(local_psds, local_freqs, axis=-1) + + band_power: dict[str, np.ndarray] = {} + missing_bands: set[str] = set() + failures: list[dict[str, Any]] = [] + descriptor_names: list[str] = [] + chunk_features: list[np.ndarray] = [] + + def integrate_band_power( + spectra: np.ndarray, + band_freqs: np.ndarray, + band_label_prefix: str, + ) -> tuple[dict[str, np.ndarray], set[str]]: + computed_band_power: dict[str, np.ndarray] = {} + computed_missing_bands: set[str] = set() + range_label = ( + "computed PSD range" + if band_label_prefix == "Raw" + else "parametric fit range" + ) + for band_name, (low, high) in self.config.bands.items(): + mask = (band_freqs >= low) & (band_freqs <= high) + if not np.any(mask): + message = ( + f"{band_label_prefix} band '{band_name}' does not overlap " + f"the {range_label}." + ) + if runtime.on_error == "raise": + raise ValueError(message) + computed_missing_bands.add(band_name) + computed_band_power[band_name] = np.full( + spectra.shape[:-1], + np.nan, + dtype=float, + ) + for obs_rel, ch_idx in np.argwhere( + ~np.isfinite(computed_band_power[band_name]) + ): + failures.append( + make_failure_record( + family=self.family_name, + obs_index=obs_offset + int(obs_rel), + obs_id=None if ids is None else ids[obs_rel], + channel_index=int(ch_idx), + channel_name=channel_names[ch_idx], + exception_type="BandResolutionError", + message=message, + ) + ) + continue + computed_band_power[band_name] = np.trapezoid( + spectra[..., mask], + band_freqs[mask], + axis=-1, + ) + return computed_band_power, computed_missing_bands + + def append_band_outputs( + band_power_dict: dict[str, np.ndarray], + total_power_array: np.ndarray | None, + missing_band_names: set[str], + output_prefix: str | None, + enabled_absolute_output: str, + enabled_relative_output: str, + enabled_ratio_output: str, + relative_message_prefix: str, + ratio_message_prefix: str, + failed_pairs_to_skip: set[tuple[int, int]] | None = None, + ) -> None: + metric_prefix = [] if output_prefix is None else [output_prefix] + + if enabled_absolute_output in self.config.outputs: + for band_name, values in band_power_dict.items(): + feature, names = self._finalize_descriptor( + values, + family_prefix="band", + metric_name="_".join(metric_prefix + ["abs", band_name]), + channel_names=channel_names, + channel_pooling=channel_pooling, + ) + chunk_features.append(feature) + descriptor_names.extend(names) + + if self.config.log_power: + log_values = np.log10(np.clip(values, eps, None)) + feature, names = self._finalize_descriptor( + log_values, + family_prefix="band", + metric_name="_".join( + metric_prefix + ["log", "abs", band_name] + ), + channel_names=channel_names, + channel_pooling=channel_pooling, + ) + chunk_features.append(feature) + descriptor_names.extend(names) + + if enabled_relative_output in self.config.outputs: + for band_name, values in band_power_dict.items(): + relative = np.divide( + values, + total_power_array, + out=np.full_like(values, np.nan, dtype=float), + where=total_power_array > 0, + ) + if band_name not in missing_band_names: + for obs_rel, ch_idx in np.argwhere(~np.isfinite(relative)): + if ( + failed_pairs_to_skip + and ( + int(obs_rel), + int(ch_idx), + ) + in failed_pairs_to_skip + ): + continue + failures.append( + make_failure_record( + family=self.family_name, + obs_index=obs_offset + int(obs_rel), + obs_id=None if ids is None else ids[obs_rel], + channel_index=int(ch_idx), + channel_name=channel_names[ch_idx], + exception_type="NumericalIssue", + message=( + f"{relative_message_prefix} for band " + f"'{band_name}' became non-finite." + ), + ) + ) + feature, names = self._finalize_descriptor( + relative, + family_prefix="band", + metric_name="_".join(metric_prefix + ["rel", band_name]), + channel_names=channel_names, + channel_pooling=channel_pooling, + ) + chunk_features.append(feature) + descriptor_names.extend(names) + + if enabled_ratio_output in self.config.outputs: + for numerator, denominator in self.config.ratio_pairs: + ratio = np.divide( + band_power_dict[numerator], + band_power_dict[denominator], + out=np.full_like( + band_power_dict[numerator], + np.nan, + dtype=float, + ), + where=band_power_dict[denominator] > 0, + ) + if ( + numerator not in missing_band_names + and denominator not in missing_band_names + ): + for obs_rel, ch_idx in np.argwhere(~np.isfinite(ratio)): + if ( + failed_pairs_to_skip + and ( + int(obs_rel), + int(ch_idx), + ) + in failed_pairs_to_skip + ): + continue + failures.append( + make_failure_record( + family=self.family_name, + obs_index=obs_offset + int(obs_rel), + obs_id=None if ids is None else ids[obs_rel], + channel_index=int(ch_idx), + channel_name=channel_names[ch_idx], + exception_type="NumericalIssue", + message=( + f"{ratio_message_prefix} " + f"'{numerator}/{denominator}' " + "became non-finite." + ), + ) + ) + feature, names = self._finalize_descriptor( + ratio, + family_prefix="band", + metric_name="_".join( + metric_prefix + ["ratio", numerator, denominator] + ), + channel_names=channel_names, + channel_pooling=channel_pooling, + ) + chunk_features.append(feature) + descriptor_names.extend(names) + + band_power, missing_bands = integrate_band_power(local_psds, local_freqs, "Raw") + + corrected_band_power: dict[str, np.ndarray] = {} + corrected_missing_bands: set[str] = set() + corrected_total_power = None + corrected_outputs_requested = self.needs_parametric_fit() + if corrected_outputs_requested: + if fit_batch is None: + raise ValueError( + "Corrected band outputs require a supplied parametric " + "fit_batch in extract_psd()." + ) + + for obs_rel, ch_idx, exception_type, message in fit_batch.errors: + corrected_failed_pairs.add((obs_rel, ch_idx)) + failures.append( + make_failure_record( + family=self.family_name, + obs_index=obs_offset + obs_rel, + obs_id=None if ids is None else ids[obs_rel], + channel_index=ch_idx, + channel_name=channel_names[ch_idx], + exception_type=exception_type, + message=f"Corrected band estimation unavailable: {message}", + ) + ) + + corrected_freq_mask = (fit_batch.freqs >= self.config.fmin) & ( + fit_batch.freqs <= self.config.fmax + ) + corrected_freqs = fit_batch.freqs[corrected_freq_mask] + corrected_psds = fit_batch.periodic_psds[..., corrected_freq_mask] + + if "corrected_relative_power" in self.config.outputs: + if corrected_freqs.size == 0: + corrected_total_power = np.full( + psds.shape[:-1], + np.nan, + dtype=float, + ) + else: + corrected_total_power = np.trapezoid( + corrected_psds, + corrected_freqs, + axis=-1, + ) + + corrected_band_power, corrected_missing_bands = integrate_band_power( + corrected_psds, + corrected_freqs, + "Corrected", + ) + + append_band_outputs( + band_power, + total_power, + missing_bands, + None, + "absolute_power", + "relative_power", + "ratios", + "Relative power", + "Band ratio", + ) + append_band_outputs( + corrected_band_power, + corrected_total_power, + corrected_missing_bands, + "corr", + "corrected_absolute_power", + "corrected_relative_power", + "corrected_ratios", + "Corrected relative power", + "Corrected band ratio", + corrected_failed_pairs, + ) + + if chunk_features: + X_out = np.concatenate(chunk_features, axis=1) + else: + X_out = np.empty((psds.shape[0], 0), dtype=float) + + return _DescriptorBlock( + family=self.family_name, + X=X_out, + descriptor_names=descriptor_names, + meta={ + "psd_method": self.config.psd_method, + "bands": self.config.bands, + "freq_range": [self.config.fmin, self.config.fmax], + "n_freqs": int(local_freqs.size), + "corrected_outputs": [ + output + for output in self.config.outputs + if output.startswith("corrected_") + ], + }, + failures=failures, + ) + + def extract( + self, + X: np.ndarray, + sfreq: float | None, + channel_names: list[str] | None, + channel_pooling: str | dict[str, list[str]], + ids: np.ndarray | None, + runtime, + obs_offset: int = 0, + ) -> _DescriptorBlock: + """Extract band descriptors from segmented multi-channel data. + + Parameters + ---------- + X : np.ndarray + Input array with shape ``(n_obs, n_channels, n_times)``. Each row + already represents one observation segment produced upstream. + sfreq : float, optional + Sampling frequency in Hertz. + channel_names : list of str, optional + Explicit channel labels aligned with axis 1 of ``X``. + channel_pooling : {"none", "all"} or dict + Descriptor-level channel pooling policy applied after per-sensor + band values are computed. + ids : np.ndarray, optional + Observation identifiers aligned with axis 0 of ``X``. + runtime : DescriptorRuntimeConfig + Runtime execution controls shared across descriptor families. + obs_offset : int, default=0 + Global observation offset added to any collected failure records. + + Returns + ------- + _DescriptorBlock + Spectral-family descriptor block aligned with the input + observation axis. + + Notes + ----- + This is the standalone extraction path for raw spectral outputs. It + computes a PSD for the provided batch and then delegates to + :meth:`extract_psd`. Corrected spectral outputs are not supported + here because they depend on an explicit shared parametric fit batch. + When the family is executed through `DescriptorPipeline`, the shared + planner provides that batch automatically. + """ + if self.needs_parametric_fit(): + raise ValueError( + "Corrected band outputs are only available through " + "DescriptorPipeline or extract_psd(..., fit_batch=...)." + ) + psds, freqs = compute_psd( + X, + sfreq=sfreq, + method=self.config.psd_method, + fmin=self.config.fmin, + fmax=self.config.fmax, + n_jobs=None, + ) + return self.extract_psd( + psds, + freqs, + channel_names=channel_names, + channel_pooling=channel_pooling, + ids=ids, + runtime=runtime, + obs_offset=obs_offset, + ) diff --git a/coco_pipe/descriptors/extractors/utils.py b/coco_pipe/descriptors/extractors/utils.py new file mode 100644 index 0000000..b36b76c --- /dev/null +++ b/coco_pipe/descriptors/extractors/utils.py @@ -0,0 +1,189 @@ +""" +Utilities for descriptor extractors. + +This module contains small pure helpers shared across descriptor extractors. +They cover: + +- normalized failure-record creation +- descriptor-level channel pooling after per-channel values are computed + +Notes +----- +Pooling helpers in this module are descriptor-level only: + +- ``"none"`` keeps one descriptor column per input channel +- ``"all"`` averages descriptor values across all channels +- a mapping pools descriptor values within named groups and leaves ungrouped + channels unchanged + +Author: Hamza Abdelhedi (hamza.abdelhedi@umontreal.ca) +""" + +from __future__ import annotations + +from typing import Any + +import numpy as np + + +def make_failure_record( + family: str, + obs_index: int, + obs_id: Any = None, + channel_index: int | None = None, + channel_name: str | None = None, + exception_type: str | None = None, + message: str | None = None, +) -> dict[str, Any]: + """Create one normalized extractor failure record. + + Parameters + ---------- + family : str + Canonical family name that raised or collected the failure. + obs_index : int + Global observation index in the original input array. + obs_id : Any, optional + Optional user-provided observation identifier aligned with + ``obs_index``. + channel_index : int, optional + Channel index associated with the failure. + channel_name : str, optional + Explicit channel label associated with the failure. + exception_type : str + Exception class name or normalized failure type. + message : str + Stable human-readable failure description. + + Returns + ------- + dict[str, Any] + Failure record compatible with ``result["failures"]``. + """ + return { + "family": family, + "obs_index": obs_index, + "obs_id": obs_id, + "channel_index": channel_index, + "channel_name": channel_name, + "exception_type": exception_type, + "message": message, + } + + +def average_channel_matrix(values: np.ndarray) -> np.ndarray: + """Average a ``(n_obs, n_channels)`` descriptor matrix across channels. + + Parameters + ---------- + values : np.ndarray + Descriptor matrix with shape ``(n_obs, n_channels)``. A 1D vector is + returned unchanged. + + Returns + ------- + np.ndarray + Vector with shape ``(n_obs,)`` containing the NaN-aware mean across the + channel axis for each observation. + + Notes + ----- + Rows containing no finite values yield ``NaN``. + """ + if values.ndim == 1: + return values + out = np.empty(values.shape[0], dtype=float) + for idx, row in enumerate(values): + finite = row[np.isfinite(row)] + out[idx] = np.nan if finite.size == 0 else float(finite.mean()) + return out + + +def pool_channel_descriptor_matrix( + values: np.ndarray, + channel_names: list[str], + channel_pooling: str | dict[str, list[str]], +) -> tuple[np.ndarray, list[str]]: + """Pool per-channel descriptor values into the public output layout. + + Parameters + ---------- + values : np.ndarray + Descriptor matrix with shape ``(n_obs, n_channels)``. + channel_names : list of str + Channel labels aligned with the columns of ``values``. + channel_pooling : {"none", "all"} or dict of str to list of str + Pooling specification coming from ``output.channel_pooling``. + + Returns + ------- + tuple + ``(X_pooled, scopes)`` where ``X_pooled`` is the pooled descriptor + matrix and ``scopes`` is the aligned list of channel-scope tokens used + by the extractor base class to build final descriptor names. + + Raises + ------ + ValueError + If ``values`` is not 2D. + + Notes + ----- + Grouped outputs preserve first-sensor order: the first channel belonging to + a group determines where that group appears in the output column order. + Channels not assigned to a group remain as standalone outputs. + Channel labels that already carry the public ``"ch-"`` scope prefix are + preserved as-is instead of receiving a second prefix. + + Examples + -------- + Given ``channel_names=["Fz", "Cz", "Pz"]``: + + - ``channel_pooling="none"`` returns scopes + ``["ch-Fz", "ch-Cz", "ch-Pz"]`` + - ``channel_pooling="all"`` returns scopes ``["ch-all"]`` + - ``channel_pooling={"Frontal": ["Fz", "Cz"]}`` returns scopes + ``["chgrp-Frontal", "ch-Pz"]`` + """ + if values.ndim != 2: + raise ValueError("pool_channel_descriptor_matrix expects a 2D matrix.") + + def channel_scope(channel_name: str) -> str: + return channel_name if channel_name.startswith("ch-") else f"ch-{channel_name}" + + if channel_pooling == "none": + return values, [channel_scope(channel_name) for channel_name in channel_names] + if channel_pooling == "all": + return average_channel_matrix(values)[:, None], ["ch-all"] + + channel_to_index = {name: idx for idx, name in enumerate(channel_names)} + member_to_group = { + member: group_name + for group_name, members in channel_pooling.items() + for member in members + } + + grouped_columns: list[np.ndarray] = [] + scopes: list[str] = [] + emitted_groups: set[str] = set() + + for channel_name in channel_names: + group_name = member_to_group.get(channel_name) + if group_name is None: + grouped_columns.append(values[:, channel_to_index[channel_name]][:, None]) + scopes.append(channel_scope(channel_name)) + continue + if group_name in emitted_groups: + continue + member_indices = [ + channel_to_index[member] for member in channel_pooling[group_name] + ] + grouped_columns.append( + average_channel_matrix(values[:, member_indices])[:, None] + ) + scopes.append(f"chgrp-{group_name}") + emitted_groups.add(group_name) + + if not grouped_columns: + return np.empty((values.shape[0], 0), dtype=float), [] + return np.concatenate(grouped_columns, axis=1), scopes diff --git a/coco_pipe/descriptors/validation.py b/coco_pipe/descriptors/validation.py new file mode 100644 index 0000000..468850c --- /dev/null +++ b/coco_pipe/descriptors/validation.py @@ -0,0 +1,148 @@ +"""Runtime input validation helpers for descriptor extraction.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from typing import Any + +import numpy as np + +from .configs import DescriptorConfig + + +def _normalize_channel_pooling( + channel_pooling: str | Mapping[str, Sequence[str]], + channel_names: list[str] | None, +) -> str | dict[str, list[str]]: + if channel_pooling == "none": + return "none" + if channel_pooling == "all": + return "all" + if channel_names is None: + raise ValueError( + "`channel_names` must be passed explicitly when " + "output.channel_pooling uses named groups." + ) + if len(set(channel_names)) != len(channel_names): + raise ValueError( + "`channel_names` must be unique when output.channel_pooling uses " + "named groups." + ) + + known_channels = set(channel_names) + assigned: dict[str, str] = {} + normalized: dict[str, list[str]] = {} + for group_name, members in channel_pooling.items(): + normalized_members = [str(member) for member in members] + for member in normalized_members: + if member not in known_channels: + raise ValueError( + f"output.channel_pooling['{group_name}'] references unknown " + f"channel '{member}'." + ) + if member in assigned: + raise ValueError( + f"Channel '{member}' is assigned to multiple channel_pooling " + "groups: " + f"'{assigned[member]}' and '{group_name}'." + ) + assigned[member] = group_name + normalized[str(group_name)] = normalized_members + return normalized + + +def validate_runtime_inputs( + config: DescriptorConfig, + *, + X: Any, + ids: Sequence[Any] | np.ndarray | None = None, + channel_names: Sequence[str] | np.ndarray | None = None, + sfreq: float | None = None, +) -> dict[str, Any]: + """Validate explicit runtime inputs against the descriptor contract. + + Parameters + ---------- + config : DescriptorConfig + Parsed descriptor config defining the expected runtime contract. + X : Any + Candidate signal array expected to coerce to shape + ``(n_obs, n_channels, n_times)``. + ids, channel_names, sfreq + Optional runtime inputs aligned with the observation or channel axes. + + Returns + ------- + dict[str, Any] + Normalized runtime inputs ready for pipeline and extractor dispatch. + + Raises + ------ + ValueError + If array dimensionality, identifier alignment, sampling frequency, or + explicit channel-name requirements are violated. + """ + X_arr = np.asarray(X, dtype=float) + if X_arr.ndim != 3: + raise ValueError( + "Descriptors expect 3D input in 'obs_channel_time' layout; " + f"got shape {X_arr.shape}." + ) + + n_obs, n_channels, _ = X_arr.shape + + sfreq_required = ( + config.input.require_sfreq + or config.families.bands.enabled + or config.families.parametric.enabled + or ( + config.families.complexity.enabled + and "spectral_entropy" in config.families.complexity.measures + ) + ) + if sfreq_required: + if sfreq is None: + raise ValueError( + "`sfreq` must be passed explicitly for the enabled descriptor families." + ) + if sfreq <= 0: + raise ValueError("`sfreq` must be positive.") + + channel_names_out = None + channel_names_required = config.input.require_channel_names or ( + any( + getattr(config.families, family_name).enabled + for family_name in ("bands", "parametric", "complexity") + ) + and config.output.channel_pooling != "all" + ) + if channel_names is not None: + channel_names_out = [str(name) for name in np.asarray(channel_names).tolist()] + if len(channel_names_out) != n_channels: + raise ValueError( + f"`channel_names` must align with n_channels={n_channels}; " + f"got {len(channel_names_out)}." + ) + elif channel_names_required: + raise ValueError( + "`channel_names` must be passed explicitly for channel-resolved output." + ) + + ids_out = None + if ids is not None: + ids_out = np.asarray(ids) + if ids_out.shape[0] != n_obs: + raise ValueError( + f"`ids` must align with n_obs={n_obs}; got shape {ids_out.shape}." + ) + + return { + "X": X_arr, + "ids": ids_out, + "channel_names": channel_names_out, + "channel_pooling": _normalize_channel_pooling( + config.output.channel_pooling, + channel_names_out, + ), + "sfreq": sfreq, + } diff --git a/coco_pipe/io/structures.py b/coco_pipe/io/structures.py index d706c60..6445ced 100644 --- a/coco_pipe/io/structures.py +++ b/coco_pipe/io/structures.py @@ -41,6 +41,7 @@ import itertools import logging import re +import warnings from copy import deepcopy from dataclasses import dataclass, field, replace from typing import Any, Dict, List, Optional, Sequence, Tuple, Union @@ -1221,26 +1222,48 @@ def baseline_correction( return self.center(dim=dim, inplace=inplace) def aggregate( - self, by: Union[str, np.ndarray, List[Any]], method: str = "mean" + self, + by: Union[str, np.ndarray, List[Any]], + stats: Union[str, Sequence[str]] = "mean", + min_count: int = 1, + on_insufficient: str = "raise", ) -> "DataContainer": """ - Aggregate observations by groups (vectorized implementation). + Aggregate observations into grouped summaries along the ``obs`` axis. Parameters ---------- by : str or array-like - The key defining the groups. - - If str: It looks up the key in `self.coords` (e.g., 'subject_id') - or checks `self.y` if by='y'. - - If array: A sequence of labels matching the length of the 'obs' - dimension. - method : str, default='mean' - Aggregation method. Options: 'mean', 'median', 'std'. + Group definition for the observation axis. + - If str: resolve the key from ``self.coords`` or from ``self.y`` + when ``by == "y"``. + - If array-like: explicit group labels aligned with ``obs``. + stats : str or sequence of str, default="mean" + Aggregation statistic or ordered list of statistics. Supported + tokens are ``"mean"``, ``"median"``, ``"std"``, ``"var"``, + ``"sem"``, ``"min"``, ``"max"``, ``"count"``, and ``"first"``. + Legacy ``"obs-*"`` aliases are accepted and normalized. + min_count : int, default=1 + Minimum number of valid observations required per group. A valid + observation is one with at least one finite value across the + non-observation axes. + on_insufficient : {"raise", "warn", "collect"}, default="raise" + Policy applied when a group has fewer than ``min_count`` valid + observations. Returns ------- DataContainer - DataContainer with aggregated 'obs' dimension. + Aggregated container with grouped observations on the ``obs`` axis. + When multiple stats are requested, a ``stat`` dimension is inserted + immediately after ``obs``. + + Raises + ------ + ValueError + If the container has no ``obs`` dimension, grouping is invalid, + requested stats are unsupported, or ``min_count`` / + ``on_insufficient`` are invalid. """ if "obs" not in self.dims: raise ValueError("Aggregation requires 'obs' dimension.") @@ -1248,107 +1271,266 @@ def aggregate( obs_idx = self.dims.index("obs") n_obs = self.X.shape[obs_idx] - # Resolve 'by' to labels + if min_count < 1: + raise ValueError("`min_count` must be at least 1.") + if on_insufficient not in {"raise", "warn", "collect"}: + raise ValueError("`on_insufficient` must be one of: raise, warn, collect.") + + stat_aliases = { + "obs-mean": "mean", + "obs-median": "median", + "obs-std": "std", + "obs-var": "var", + "obs-sem": "sem", + "obs-min": "min", + "obs-max": "max", + "obs-count": "count", + } + supported_stats = { + "mean", + "median", + "std", + "var", + "sem", + "min", + "max", + "count", + "first", + } + if isinstance(stats, str): + stats_out = [stat_aliases.get(stats, stats)] + else: + stats_out = [stat_aliases.get(str(stat), str(stat)) for stat in stats] + if not stats_out: + raise ValueError("`stats` must not be empty.") + invalid_stats = sorted(set(stats_out) - supported_stats) + if invalid_stats: + raise ValueError( + f"Unknown stats: {invalid_stats}. Supported stats are: " + f"{sorted(supported_stats)}" + ) + if isinstance(by, str): if by == "y" and self.y is not None: - groups = self.y + groups_raw = self.y elif by in self.coords: - groups = self.coords[by] + groups_raw = self.coords[by] else: raise ValueError(f"Grouping key '{by}' not found in coords or y.") else: - groups = np.array(by) + groups_raw = by + + labels_list = list(groups_raw) + groups = np.empty(len(labels_list), dtype=object) + groups[:] = labels_list if len(groups) != n_obs: raise ValueError( f"Grouping array length {len(groups)} must match obs length {n_obs}." ) - # Transform inputs for DataFrame-based GroupBy (Vectorized) - # 1. Flatten X to (n_obs, n_features_flat) - # We need to reshape specifically so obs is axis 0, and rest is flattened if obs_idx != 0: - # Move obs to front if not already X_moved = np.moveaxis(self.X, obs_idx, 0) else: X_moved = self.X - original_shape = X_moved.shape - X_flat = X_moved.reshape(original_shape[0], -1) - - # 2. Create DataFrame - # Using numeric index for columns to avoid overhead - df = pd.DataFrame(X_flat) - df["__group__"] = groups - - # 3. Groupby & Aggregate - grouped = df.groupby("__group__") - - if method == "mean": - agg_df = grouped.mean() - elif method == "median": - agg_df = grouped.median() - elif method == "std": - agg_df = grouped.std() - elif method == "first": - agg_df = grouped.first() - else: - raise ValueError(f"Unknown method '{method}'") + other_dims = tuple(dim for dim in self.dims if dim != "obs") + group_positions: Dict[Any, List[int]] = {} + ordered_groups: List[Any] = [] + for obs_position, group_id in enumerate(groups.tolist()): + if group_id not in group_positions: + ordered_groups.append(group_id) + group_positions[group_id] = [] + group_positions[group_id].append(obs_position) + + def _reshape_reduced(values_flat: np.ndarray) -> np.ndarray | np.float64: + if rest_shape: + return np.asarray(values_flat, dtype=np.float64).reshape(rest_shape) + return np.asarray(values_flat, dtype=np.float64)[0] + + def _reduce_group( + group_X: np.ndarray, + group_X_flat: np.ndarray, + counts_flat: np.ndarray, + stat: str, + ) -> np.ndarray | np.float64: + if stat == "count": + return _reshape_reduced(counts_flat) + if stat == "first": + return np.asarray(group_X[0], dtype=np.float64) + + with warnings.catch_warnings(): + warnings.simplefilter("ignore", category=RuntimeWarning) + if stat == "mean": + values_flat = np.nanmean(group_X_flat, axis=0) + elif stat == "median": + values_flat = np.nanmedian(group_X_flat, axis=0) + elif stat == "std": + values_flat = np.nanstd(group_X_flat, axis=0) + elif stat == "var": + values_flat = np.nanvar(group_X_flat, axis=0) + elif stat == "sem": + values_flat = np.nanstd(group_X_flat, axis=0) / np.sqrt( + counts_flat.astype(np.float64) + ) + elif stat == "min": + values_flat = np.nanmin(group_X_flat, axis=0) + elif stat == "max": + values_flat = np.nanmax(group_X_flat, axis=0) + else: # pragma: no cover - guarded above + raise ValueError(f"Unknown stat '{stat}'") + + values_flat = np.asarray(values_flat, dtype=np.float64) + if counts_flat.size: + values_flat = np.where(counts_flat == 0, np.nan, values_flat) + return _reshape_reduced(values_flat) + + def _failure_record( + group_id: Any, + group_index: int, + row_count: int, + valid_row_count: int, + message: str, + ) -> Dict[str, Any]: + return { + "group_id": group_id, + "group_index": group_index, + "row_count": row_count, + "valid_row_count": valid_row_count, + "exception_type": "InsufficientObservations", + "message": message, + } - # 4. Extract Result - # agg_df index is the unique groups (sorted) - unique_groups = agg_df.index.to_numpy() - X_agg_flat = agg_df.values # (n_groups, n_features_flat) + n_groups = len(ordered_groups) + rest_shape = X_moved.shape[1:] + reduced_shape = (n_groups, len(stats_out)) + rest_shape + agg_moved = np.empty(reduced_shape, dtype=np.float64) + epoch_counts = np.empty(n_groups, dtype=np.int64) + failures: List[Dict[str, Any]] = [] + + for group_index, group_id in enumerate(ordered_groups): + obs_positions = np.asarray(group_positions[group_id], dtype=int) + group_X = X_moved[obs_positions] + row_count = int(obs_positions.size) + epoch_counts[group_index] = row_count + + if rest_shape: + group_X_flat = group_X.reshape(row_count, -1) + else: + group_X_flat = group_X.reshape(row_count, 1) + if group_X_flat.shape[1] == 0: + valid_row_count = row_count + else: + valid_row_count = int(np.isfinite(group_X_flat).any(axis=1).sum()) + if valid_row_count < min_count: + message = ( + f"Group {group_id!r} has {valid_row_count} valid rows, " + f"requires at least {min_count}." + ) + failure = _failure_record( + group_id=group_id, + group_index=group_index, + row_count=row_count, + valid_row_count=valid_row_count, + message=message, + ) + if on_insufficient == "raise": + raise ValueError(message) + if on_insufficient == "warn": + warnings.warn(message, stacklevel=2) + failures.append(failure) + agg_moved[group_index] = np.full((len(stats_out),) + rest_shape, np.nan) + continue - # 5. Reshape back - new_shape = (len(unique_groups),) + original_shape[1:] - X_agg_moved = X_agg_flat.reshape(new_shape) + counts_flat = np.isfinite(group_X_flat).sum(axis=0, dtype=np.int64) + for stat_index, stat in enumerate(stats_out): + agg_moved[group_index, stat_index] = _reduce_group( + group_X=group_X, + group_X_flat=group_X_flat, + counts_flat=counts_flat, + stat=stat, + ) - if obs_idx != 0: - X_agg = np.moveaxis(X_agg_moved, 0, obs_idx) + if len(stats_out) == 1: + moved_dims = ("obs",) + other_dims + final_dims = self.dims + agg_values = agg_moved[:, 0, ...] else: - X_agg = X_agg_moved + moved_dims = ("obs", "stat") + other_dims + final_dims_list: List[str] = [] + for dim in self.dims: + final_dims_list.append(dim) + if dim == "obs": + final_dims_list.append("stat") + final_dims = tuple(final_dims_list) + agg_values = agg_moved + + permutation = [moved_dims.index(dim) for dim in final_dims] + X_agg = np.transpose(agg_values, axes=permutation) + + unique_groups = np.empty(n_groups, dtype=object) + unique_groups[:] = ordered_groups - # 6. Metadata Handling (y consistency) new_y = None if self.y is not None: - # Check if y is consistent per group - y_df = pd.DataFrame({"g": groups, "y": self.y}) - y_nunique = y_df.groupby("g")["y"].nunique() - if y_nunique.max() == 1: - # Take first - new_y = y_df.groupby("g")["y"].first().reindex(unique_groups).values - - # 7. Update Coords - new_coords = self.coords.copy() + grouped_y: List[Any] = [] + y_consistent = True + for group_id in ordered_groups: + values = np.asarray(self.y)[group_positions[group_id]] + if len(set(values.tolist())) != 1: + y_consistent = False + break + grouped_y.append(values[0]) + if y_consistent: + new_y = np.asarray(grouped_y) + + new_coords = { + dim: deepcopy(values) + for dim, values in self.coords.items() + if dim in self.dims and dim != "obs" + } new_coords["obs"] = unique_groups + if len(stats_out) > 1: + new_coords["stat"] = np.asarray(stats_out, dtype=object) + new_coords["epoch_count"] = epoch_counts - for k, v in self.coords.items(): - if k == "obs": + for key, values in self.coords.items(): + if key == "obs" or key in self.dims: continue - if len(v) == n_obs and k not in self.dims: - coord_df = pd.DataFrame({"g": groups, "value": np.array(v)}) - coord_nunique = coord_df.groupby("g")["value"].nunique(dropna=False) - if coord_nunique.max() == 1: - new_coords[k] = ( - coord_df.groupby("g")["value"] - .first() - .reindex(unique_groups) - .values - ) - else: - del new_coords[k] + if len(values) != n_obs: + continue + grouped_values: List[Any] = [] + consistent = True + values_array = np.asarray(values, dtype=object) + for group_id in ordered_groups: + group_values = values_array[group_positions[group_id]] + if len(set(group_values.tolist())) != 1: + consistent = False + break + grouped_values.append(group_values[0]) + if consistent: + coord_out = np.empty(n_groups, dtype=object) + coord_out[:] = grouped_values + new_coords[key] = coord_out + + meta = deepcopy(self.meta) + meta.update( + { + "aggregated": True, + "agg_by": by if isinstance(by, str) else None, + "agg_stats": list(stats_out), + "min_count": int(min_count), + } + ) + if failures: + meta["aggregate_failures"] = failures return replace( self, X=X_agg, y=new_y, - ids=unique_groups if method != "std" else None, + dims=final_dims, + ids=unique_groups, coords=new_coords, - meta={ - **self.meta, - "aggregated": True, - "agg_by": str(by), - "agg_method": method, - }, + meta=meta, ) diff --git a/configs/run_descriptors_eeg.yml b/configs/run_descriptors_eeg.yml new file mode 100644 index 0000000..723d664 --- /dev/null +++ b/configs/run_descriptors_eeg.yml @@ -0,0 +1,70 @@ +data: + path: PhysioNet_EEGBCI/BIDS + type: bids + # Run 03 is shared across the first 10 local subjects. + task: motorimagery + session: "01" + runs: + - "03" + loading_mode: epochs + # Fixed windows keep shapes aligned across subjects for this dataset copy. + window_length: 2.0 + stride: 2.0 + subjects: + - "001" + - "002" + - "003" + - "004" + - "005" + - "006" + - "007" + - "008" + - "009" + - "010" + +descriptors: + input: + require_sfreq: true + require_channel_names: false + output: + precision: float32 + channel_pooling: all + include_failure_summary: true + include_runtime_meta: true + families: + bands: + enabled: true + psd_method: welch + fmin: 1.0 + fmax: 45.0 + outputs: + - absolute_power + - relative_power + parametric: + enabled: true + backend: specparam + psd_method: welch + freq_range: [1.0, 45.0] + outputs: + - aperiodic + - fit_quality + - peak_summary + complexity: + enabled: true + backend: antropy + measures: + - sample_entropy + - perm_entropy + - spectral_entropy + - hjorth_mobility + - hjorth_complexity + runtime: + execution_backend: sequential + n_jobs: 1 + chunking: + obs_chunk: 64 + channel_chunk: full + time_chunk: 2048 + on_error: collect + +save_path: outputs/descriptors_eeg.npz diff --git a/examples/.DS_Store b/examples/.DS_Store deleted file mode 100644 index bf1c00faf8e73b2de58b7829087d2e38e39adc0d..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 6148 zcmeHK%}T>S5Z-O8O(;SR3VI88E!fm5f|n5M3mDOZN=-=7V9b^#F^5vfU0=u-@p+ut z-5jbpcoMNQu=~yI%0Bnn(G)^%uj4clr6-Cs%4Dvu9ENF{6Kij=a0a8kJsO@Y z`}V>)a)w|W9xs=sv9o`0dNFxQo>TFvxa2^)lr4=lyn*titQUWtCL(y(*4ohSp`l=0g$fAhE0+K;a39%~PUW{z ahd5Vbr4VPqahVQC7Xe8Kb;Q6gFz^X0T1`#> diff --git a/examples/demo_structures.py b/examples/demo_structures.py index 45d99e0..4791e48 100644 --- a/examples/demo_structures.py +++ b/examples/demo_structures.py @@ -85,7 +85,7 @@ def section(title): # just globally by condition. # Let's say we group by 'y' (Conditions A, B) print("\n--- Aggregation Test ---") -agg_cond = container_eeg.aggregate(by=container_eeg.y, method="mean") +agg_cond = container_eeg.aggregate(by=container_eeg.y, stats="mean") print(f"Aggregated by Condition (A, B): {agg_cond.shape} ids={agg_cond.ids}") # Test Selection of Specific Epochs (Wildcard) diff --git a/examples/descriptors_example.py b/examples/descriptors_example.py new file mode 100644 index 0000000..a0257ab --- /dev/null +++ b/examples/descriptors_example.py @@ -0,0 +1,67 @@ +""" +Minimal descriptors example with explicit NumPy inputs. +""" + +from __future__ import annotations + +import numpy as np + +from coco_pipe.descriptors import DescriptorPipeline +from coco_pipe.io import DataContainer + + +def main() -> None: + rng = np.random.default_rng(42) + X = rng.normal(size=(12, 3, 256)) + t = np.linspace(0, 1, 256, endpoint=False) + X[:, 0, :] += np.sin(2 * np.pi * 10 * t) + X[:, 1, :] += np.sin(2 * np.pi * 6 * t) + ids = np.asarray([f"obs-{idx:02d}" for idx in range(12)]) + + config = { + "output": {"channel_pooling": "all"}, + "families": { + "bands": { + "enabled": True, + "outputs": ["absolute_power", "corrected_absolute_power"], + }, + "parametric": { + "enabled": True, + "outputs": ["aperiodic"], + }, + "complexity": { + "enabled": True, + "measures": ["sample_entropy", "hjorth_mobility"], + }, + }, + } + + pipe = DescriptorPipeline(config) + result = pipe.extract( + X=X, + ids=ids, + sfreq=256.0, + channel_names=["Fz", "Cz", "Pz"], + ) + + print("Descriptor matrix shape:", result["X"].shape) + print("First five names:", result["descriptor_names"][:5]) + print("Failure count:", len(result["failures"])) + + container = DataContainer( + X=result["X"], + dims=("obs", "feature"), + coords={"feature": result["descriptor_names"]}, + ) + grouped = container.aggregate( + by=["sub-01"] * 6 + ["sub-02"] * 6, + stats=["mean", "std"], + ) + + print("Grouped descriptor shape:", grouped.X.shape) + print("Grouped dims:", grouped.dims) + print("Grouped stats:", grouped.coords["stat"].tolist()) + + +if __name__ == "__main__": + main() diff --git a/pyproject.toml b/pyproject.toml index cb590c7..24f32e9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -97,6 +97,11 @@ eeg = [ "mne-bids>=0.10.0", "meegkit", ] +descriptors = [ + "specparam", + "antropy", + "neurokit2", +] dim-red = [ "dask[array]", "dask-ml", @@ -134,6 +139,9 @@ full = [ "meegkit", "mne>=1.9.0", "mne-bids>=0.10.0", + "specparam", + "antropy", + "neurokit2", ] [project.urls] diff --git a/requirements-dev.txt b/requirements-dev.txt deleted file mode 100644 index 61e38eb..0000000 --- a/requirements-dev.txt +++ /dev/null @@ -1,22 +0,0 @@ -pytest-cov -pyyaml -numpy -pandas -scikit-learn -matplotlib -sphinx>=4.0.0 -insegel -furo -sphinx_gallery -sphinx-copybutton -sphinxcontrib-mermaid -sphinx-autoapi -myst-parser -umap-learn -trimap -phate -mne>=1.0.0 -mne-bids>=0.10.0 -requests -openpyxl --e . diff --git a/requirements.txt b/requirements.txt deleted file mode 100644 index 091b1c1..0000000 --- a/requirements.txt +++ /dev/null @@ -1 +0,0 @@ --e .[dev,dim-red,eeg] diff --git a/scripts/run_descriptors.py b/scripts/run_descriptors.py new file mode 100644 index 0000000..ce36897 --- /dev/null +++ b/scripts/run_descriptors.py @@ -0,0 +1,133 @@ +#!/usr/bin/env python3 +""" +Run the descriptors pipeline from a YAML configuration. +""" + +from __future__ import annotations + +import argparse +from collections.abc import Mapping +from pathlib import Path + +import numpy as np +import yaml + +from coco_pipe.descriptors import DescriptorConfig, DescriptorPipeline +from coco_pipe.io import load_data + +_ALLOWED_DATA_TYPES = {"auto", "tabular", "bids", "embedding"} + + +def _normalize_data_config(data_cfg): + normalized = dict(data_cfg) + if "tasks" in normalized: + if "task" in normalized: + raise ValueError("Use only one of `task` or `tasks` in the data config.") + tasks = normalized.pop("tasks") + if isinstance(tasks, str): + normalized["task"] = tasks + else: + tasks = list(tasks) + if len(tasks) != 1: + raise ValueError( + "run_descriptors.py supports exactly one task in `data.tasks`." + ) + normalized["task"] = tasks[0] + return normalized + + +def _validate_data_config(data_cfg): + if not isinstance(data_cfg, Mapping): + raise ValueError("The YAML `data` section must be a mapping.") + + normalized = _normalize_data_config(data_cfg) + path = normalized.get("path") + if not path: + raise ValueError("The YAML `data` section must define a non-empty `path`.") + + data_type = normalized.get("type", "auto") + if data_type not in _ALLOWED_DATA_TYPES: + raise ValueError( + "`data.type` must be one of " + f"{sorted(_ALLOWED_DATA_TYPES)}; got {data_type!r}." + ) + + sfreq_override = normalized.get("sfreq_override") + if sfreq_override is not None and float(sfreq_override) <= 0: + raise ValueError("`data.sfreq_override` must be positive when provided.") + + return normalized + + +def _extract_explicit_inputs(container): + sfreq = getattr(container, "meta", {}).get("sfreq") + channel_names = None + coords = getattr(container, "coords", {}) or {} + if "channel" in coords: + channel_names = list(np.asarray(coords["channel"]).tolist()) + + return { + "X": np.asarray(container.X), + "ids": getattr(container, "ids", None), + "sfreq": sfreq, + "channel_names": channel_names, + } + + +def _save_result(save_path: Path, result): + save_path.parent.mkdir(parents=True, exist_ok=True) + np.savez_compressed( + save_path, + X=result["X"], + descriptor_names=np.asarray(result["descriptor_names"], dtype=object), + failures=np.asarray(result["failures"], dtype=object), + ) + + +def main() -> None: + parser = argparse.ArgumentParser(description="Run descriptor extraction.") + parser.add_argument("--config", required=True, help="Path to YAML configuration.") + args = parser.parse_args() + + config_path = Path(args.config) + payload = yaml.safe_load(config_path.read_text()) + if not isinstance(payload, Mapping): + raise ValueError("The YAML config must define a top-level mapping.") + if "data" not in payload or "descriptors" not in payload: + raise ValueError("The YAML config must contain `data` and `descriptors`.") + + data_cfg = _validate_data_config(payload["data"]) + data_path = data_cfg.pop("path") + mode = data_cfg.pop("type", "auto") + sfreq_override = data_cfg.pop("sfreq_override", None) + save_path = payload.get("save_path") + + container = load_data(data_path, mode=mode, **data_cfg) + explicit_inputs = _extract_explicit_inputs(container) + if explicit_inputs["sfreq"] is None and sfreq_override is not None: + explicit_inputs["sfreq"] = float(sfreq_override) + if explicit_inputs["sfreq"] is None: + raise ValueError( + "Could not find `sfreq` in the loaded data container's metadata. " + "Ensure the data is loaded with sampling frequency metadata, or add " + "`sfreq_override` to the YAML configuration." + ) + + descriptor_cfg = DescriptorConfig.model_validate(payload["descriptors"]) + pipe = DescriptorPipeline(descriptor_cfg) + result = pipe.extract(**explicit_inputs) + + if save_path: + _save_result(Path(save_path), result) + + print( + { + "shape": result["X"].shape, + "n_descriptors": len(result["descriptor_names"]), + "n_failures": len(result["failures"]), + } + ) + + +if __name__ == "__main__": + main() diff --git a/tests/conftest.py b/tests/conftest.py index bcb14e8..8dc4a42 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,3 +1,5 @@ +import os +import tempfile from unittest.mock import MagicMock import pytest @@ -14,3 +16,16 @@ def mock_visualizations(): import matplotlib.pyplot as plt plt.show = MagicMock() + + +@pytest.fixture(scope="session", autouse=True) +def _sandbox_runtime_env(): + """Keep third-party cache/config writes inside writable temp dirs.""" + tmp_root = os.path.join(tempfile.gettempdir(), "coco_pipe_test_runtime") + mpl_dir = os.path.join(tmp_root, "mplconfig") + mne_dir = os.path.join(tmp_root, "mne") + os.makedirs(mpl_dir, exist_ok=True) + os.makedirs(mne_dir, exist_ok=True) + os.environ.setdefault("MPLCONFIGDIR", mpl_dir) + os.environ.setdefault("MNE_HOME", mne_dir) + os.environ.setdefault("MNE_DONTWRITE_HOME", "true") diff --git a/tests/test_descriptors_configs.py b/tests/test_descriptors_configs.py new file mode 100644 index 0000000..93d413d --- /dev/null +++ b/tests/test_descriptors_configs.py @@ -0,0 +1,313 @@ +import pytest +from pydantic import ValidationError + +from coco_pipe.descriptors import DescriptorPipeline +from coco_pipe.descriptors.configs import DescriptorConfig + + +def test_descriptor_config_is_strict(): + with pytest.raises(ValidationError): + DescriptorConfig(unknown_field=1) + + +def test_band_ratios_require_pairs(): + with pytest.raises(ValidationError, match="ratio_pairs"): + DescriptorConfig( + families={ + "bands": { + "enabled": True, + "outputs": ["absolute_power", "ratios"], + } + } + ) + + +def test_corrected_band_ratios_require_pairs(): + with pytest.raises(ValidationError, match="ratio_pairs"): + DescriptorConfig( + families={ + "bands": { + "enabled": True, + "outputs": ["corrected_absolute_power", "corrected_ratios"], + } + } + ) + + +def test_corrected_bands_require_parametric_fit_range_to_cover_band_window(): + with pytest.raises(ValueError, match="Corrected band outputs require"): + DescriptorPipeline( + { + "families": { + "bands": { + "enabled": True, + "fmin": 1.0, + "fmax": 45.0, + "outputs": ["corrected_absolute_power"], + }, + "parametric": { + "freq_range": [4.0, 30.0], + }, + } + } + ) + + +def test_channel_pooling_accepts_none_all_or_mapping(): + assert ( + DescriptorConfig(output={"channel_pooling": "none"}).output.channel_pooling + == "none" + ) + assert ( + DescriptorConfig(output={"channel_pooling": "all"}).output.channel_pooling + == "all" + ) + assert DescriptorConfig( + output={"channel_pooling": {"Frontal": ["Fz", "Cz"]}} + ).output.channel_pooling == {"Frontal": ["Fz", "Cz"]} + + +def test_runtime_and_output_flags_parse_strictly(): + config = DescriptorConfig.model_validate( + { + "input": { + "require_sfreq": False, + "require_channel_names": True, + }, + "output": { + "precision": "float64", + "channel_pooling": "all", + }, + "runtime": { + "execution_backend": "joblib", + "n_jobs": -1, + "obs_chunk": 16, + "on_error": "collect", + }, + } + ) + + assert config.input.require_sfreq is False + assert config.input.require_channel_names is True + assert config.output.precision == "float64" + assert config.output.channel_pooling == "all" + assert config.runtime.execution_backend == "joblib" + assert config.runtime.n_jobs == -1 + assert config.runtime.obs_chunk == 16 + + +def test_runtime_n_jobs_must_be_minus_one_or_positive(): + with pytest.raises(ValidationError, match="n_jobs"): + DescriptorConfig.model_validate({"runtime": {"n_jobs": 0}}) + + with pytest.raises(ValidationError, match="n_jobs"): + DescriptorConfig.model_validate({"runtime": {"n_jobs": -2}}) + + +def test_removed_ceremonial_fields_are_rejected(): + with pytest.raises(ValidationError): + DescriptorConfig.model_validate( + { + "input": {"expected_ndim": 3}, + } + ) + + with pytest.raises(ValidationError): + DescriptorConfig.model_validate( + { + "runtime": {"batch_size": 32}, + } + ) + + with pytest.raises(ValidationError): + DescriptorConfig.model_validate( + { + "output": {"return_format": "dict"}, + } + ) + + with pytest.raises(ValidationError): + DescriptorConfig.model_validate( + { + "output": {"include_failure_summary": False}, + } + ) + + with pytest.raises(ValidationError): + DescriptorConfig.model_validate( + { + "output": {"include_runtime_meta": False}, + } + ) + + with pytest.raises(ValidationError): + DescriptorConfig.model_validate( + { + "output": {"channel_resolved": False}, + } + ) + + with pytest.raises(ValidationError): + DescriptorConfig.model_validate( + { + "output": {"channel_groups": {"Frontal": ["Fz", "Cz"]}}, + } + ) + + with pytest.raises(ValidationError): + DescriptorConfig.model_validate( + { + "families": {"bands": {"reduce_channels": "mean"}}, + } + ) + + with pytest.raises(ValidationError): + DescriptorConfig.model_validate( + { + "families": {"bands": {"per_channel": True}}, + } + ) + + with pytest.raises(ValidationError): + DescriptorConfig.model_validate( + { + "families": {"parametric": {"store_failures": True}}, + } + ) + + with pytest.raises(ValidationError): + DescriptorConfig.model_validate( + { + "runtime": {"chunking": {"obs_chunk": 16}}, + } + ) + + with pytest.raises(ValidationError): + DescriptorConfig.model_validate( + { + "runtime": {"channel_chunk": 4}, + } + ) + + with pytest.raises(ValidationError): + DescriptorConfig.model_validate( + { + "runtime": {"time_chunk": 128}, + } + ) + + with pytest.raises(ValidationError): + DescriptorConfig.model_validate( + { + "runtime": {"random_state": 7}, + } + ) + + +def test_band_validation_edge_cases(): + # Duplicate outputs + with pytest.raises(ValidationError, match="duplicates"): + DescriptorConfig( + families={"bands": {"outputs": ["absolute_power", "absolute_power"]}} + ) + + # Invalid fmin/fmax + with pytest.raises(ValidationError, match="fmin < fmax"): + DescriptorConfig(families={"bands": {"fmin": 10, "fmax": 5}}) + + # Band low >= high + with pytest.raises(ValidationError, match="low < high"): + DescriptorConfig(families={"bands": {"bands": {"bad": (10, 5)}}}) + + # Band out of range + with pytest.raises(ValidationError, match="stay within"): + DescriptorConfig( + families={"bands": {"fmin": 1, "fmax": 10, "bands": {"out": (5, 15)}}} + ) + + # Unknown outputs + with pytest.raises(ValidationError, match="Unknown band outputs"): + DescriptorConfig(families={"bands": {"outputs": ["non_existent"]}}) + + +def test_parametric_validation_edge_cases(): + # Duplicate outputs + with pytest.raises(ValidationError, match="duplicates"): + DescriptorConfig( + families={"parametric": {"outputs": ["aperiodic", "aperiodic"]}} + ) + + # Invalid freq_range + with pytest.raises(ValidationError, match="low < high"): + DescriptorConfig(families={"parametric": {"freq_range": (10, 5)}}) + + # Invalid peak_width_limits + with pytest.raises(ValidationError, match="low < high"): + DescriptorConfig(families={"parametric": {"peak_width_limits": (5, 2)}}) + + # Unknown outputs + with pytest.raises(ValidationError, match="Unknown parametric outputs"): + DescriptorConfig(families={"parametric": {"outputs": ["non_existent"]}}) + + +def test_complexity_validation_edge_cases(): + # Duplicate measures + with pytest.raises(ValidationError, match="duplicates"): + DescriptorConfig( + families={"complexity": {"measures": ["sample_entropy", "sample_entropy"]}} + ) + + # Unknown measures + with pytest.raises(ValidationError, match="Unknown complexity measures"): + DescriptorConfig(families={"complexity": {"measures": ["non_existent"]}}) + + +def test_channel_pooling_validation_edge_cases(): + # Invalid string + with pytest.raises( + ValidationError, + match="Input should be 'none' or 'all'|Input should be a valid dictionary", + ): + DescriptorConfig(output={"channel_pooling": "some_string"}) + + # Empty group name + with pytest.raises(ValidationError, match="non-empty strings"): + DescriptorConfig(output={"channel_pooling": {"": ["ch1"]}}) + + # Empty members + with pytest.raises(ValidationError, match="at least one channel"): + DescriptorConfig(output={"channel_pooling": {"G1": []}}) + + # Duplicate members + with pytest.raises(ValidationError, match="not contain duplicates"): + DescriptorConfig(output={"channel_pooling": {"G1": ["ch1", "ch1"]}}) + + # Coerce None/{} to none + assert ( + DescriptorConfig(output={"channel_pooling": None}).output.channel_pooling + == "none" + ) + assert ( + DescriptorConfig(output={"channel_pooling": {}}).output.channel_pooling + == "none" + ) + + +def test_coercion_logic_smoke(): + # Coerce bands + config = DescriptorConfig(families={"bands": {"bands": {"delta": [1, 4]}}}) + assert config.families.bands.bands["delta"] == (1.0, 4.0) + + # Coerce ratio_pairs + config = DescriptorConfig( + families={"bands": {"outputs": ["ratios"], "ratio_pairs": [["theta", "beta"]]}} + ) + assert config.families.bands.ratio_pairs == [("theta", "beta")] + + # Coerce None bands + from coco_pipe.descriptors.configs import BandDescriptorConfig + + assert BandDescriptorConfig(bands=None).bands["alpha"] == (8.0, 13.0) + + # Coerce None ratio_pairs + assert BandDescriptorConfig(ratio_pairs=None).ratio_pairs == [] diff --git a/tests/test_descriptors_core.py b/tests/test_descriptors_core.py new file mode 100644 index 0000000..a597531 --- /dev/null +++ b/tests/test_descriptors_core.py @@ -0,0 +1,881 @@ +import numpy as np +import pytest + +import coco_pipe.descriptors.core as descriptors_core +from coco_pipe.descriptors import DescriptorPipeline +from coco_pipe.descriptors.extractors.complexity import ComplexityDescriptorExtractor + + +def test_empty_pipeline_returns_explicit_result_structure(): + X = np.random.default_rng(0).normal(size=(5, 2, 64)) + pipe = DescriptorPipeline({}) + result = pipe.extract(X=X, sfreq=128.0) + + assert set(result) == {"X", "descriptor_names", "failures"} + assert result["X"].shape == (5, 0) + assert result["descriptor_names"] == [] + + +def test_band_pipeline_smoke(): + rng = np.random.default_rng(1) + X = rng.normal(size=(6, 3, 128)) + pipe = DescriptorPipeline( + { + "output": {"channel_pooling": "all"}, + "families": {"bands": {"enabled": True, "outputs": ["absolute_power"]}}, + } + ) + result = pipe.extract(X=X, sfreq=128.0, channel_names=["Fz", "Cz", "Pz"]) + + assert result["X"].shape[0] == 6 + assert result["X"].shape[1] == len(result["descriptor_names"]) + assert result["descriptor_names"][0].startswith("band_abs_") + + +def test_complexity_can_omit_sfreq_when_config_disables_it(): + X = np.random.default_rng(18).normal(size=(4, 2, 64)) + pipe = DescriptorPipeline( + { + "input": {"require_sfreq": False}, + "families": { + "complexity": { + "enabled": True, + "measures": ["sample_entropy"], + } + }, + "output": {"channel_pooling": "all"}, + } + ) + result = pipe.extract(X=X, channel_names=["Fz", "Cz"]) + + assert result["X"].shape == (4, 1) + + +def test_output_precision_is_respected(): + X = np.zeros((2, 2, 128), dtype=float) + pipe = DescriptorPipeline( + { + "output": { + "channel_pooling": "all", + "precision": "float64", + }, + "families": { + "parametric": { + "enabled": True, + "outputs": ["aperiodic"], + } + }, + "runtime": {"on_error": "collect"}, + } + ) + result = pipe.extract(X=X, sfreq=128.0, channel_names=["Fz", "Cz"]) + + assert result["X"].dtype == np.float64 + + +def test_missing_sfreq_is_explicit_error(): + X = np.random.default_rng(2).normal(size=(4, 2, 64)) + pipe = DescriptorPipeline( + {"families": {"bands": {"enabled": True, "outputs": ["absolute_power"]}}} + ) + with pytest.raises(ValueError, match="`sfreq`"): + pipe.extract(X=X) + + +def test_wrong_ndim_is_rejected(): + X = np.random.default_rng(21).normal(size=(4, 64)) + pipe = DescriptorPipeline({}) + + with pytest.raises(ValueError, match="Descriptors expect 3D input"): + pipe.extract(X=X, sfreq=128.0) + + +def test_wrong_channel_names_length_is_rejected(): + X = np.random.default_rng(22).normal(size=(4, 2, 64)) + pipe = DescriptorPipeline( + { + "families": { + "bands": { + "enabled": True, + "outputs": ["absolute_power"], + } + }, + } + ) + + with pytest.raises(ValueError, match="channel_names"): + pipe.extract(X=X, sfreq=128.0, channel_names=["C3"]) + + +def test_channel_pooling_groups_reject_unknown_channel_names(): + X = np.random.default_rng(22).normal(size=(4, 2, 64)) + pipe = DescriptorPipeline( + { + "output": {"channel_pooling": {"Frontal": ["C3", "C4"]}}, + "families": { + "bands": { + "enabled": True, + "outputs": ["absolute_power"], + } + }, + } + ) + + with pytest.raises(ValueError, match="unknown channel"): + pipe.extract(X=X, sfreq=128.0, channel_names=["Fz", "Cz"]) + + +def test_channel_pooling_groups_reject_overlapping_assignments(): + X = np.random.default_rng(22).normal(size=(4, 3, 64)) + pipe = DescriptorPipeline( + { + "output": { + "channel_pooling": { + "Frontal": ["Fz", "Cz"], + "Central": ["Cz", "Pz"], + } + }, + "families": { + "bands": { + "enabled": True, + "outputs": ["absolute_power"], + } + }, + } + ) + + with pytest.raises(ValueError, match="multiple channel_pooling groups"): + pipe.extract(X=X, sfreq=128.0, channel_names=["Fz", "Cz", "Pz"]) + + +def test_require_channel_names_flag_is_enforced(): + X = np.random.default_rng(25).normal(size=(4, 2, 64)) + pipe = DescriptorPipeline( + { + "input": {"require_channel_names": True}, + "output": {"channel_pooling": "all"}, + "families": { + "complexity": { + "enabled": True, + "measures": ["sample_entropy"], + } + }, + } + ) + + with pytest.raises(ValueError, match="channel_names"): + pipe.extract(X=X, sfreq=128.0) + + +def test_complexity_collects_short_segment_failures(): + X = np.ones((4, 2, 3), dtype=float) + pipe = DescriptorPipeline( + { + "output": {"channel_pooling": "all"}, + "families": { + "complexity": { + "enabled": True, + "measures": ["sample_entropy"], + } + }, + "runtime": {"on_error": "collect"}, + } + ) + result = pipe.extract(X=X, sfreq=128.0) + + assert result["X"].shape == (4, 1) + assert np.isnan(result["X"]).all() + assert result["failures"] + + +def test_bands_collect_short_window_resolution_failures(): + X = np.random.default_rng(7).normal(size=(3, 2, 8)) + pipe = DescriptorPipeline( + { + "output": {"channel_pooling": "all"}, + "families": { + "bands": { + "enabled": True, + "outputs": ["absolute_power", "relative_power"], + } + }, + "runtime": {"on_error": "collect"}, + } + ) + + result = pipe.extract(X=X, sfreq=160.0, channel_names=["C3", "C4"]) + + assert result["X"].shape == (3, 10) + assert any( + failure["exception_type"] == "BandResolutionError" + for failure in result["failures"] + ) + + +def test_warn_policy_emits_aggregate_warning(): + X = np.random.default_rng(23).normal(size=(3, 2, 8)) + pipe = DescriptorPipeline( + { + "output": {"channel_pooling": "all"}, + "families": { + "bands": { + "enabled": True, + "outputs": ["absolute_power"], + } + }, + "runtime": {"on_error": "warn"}, + } + ) + + with pytest.warns(UserWarning, match="Collected"): + result = pipe.extract(X=X, sfreq=160.0, channel_names=["C3", "C4"]) + + assert result["failures"] + + +def test_raise_policy_reraises_runtime_failure(): + X = np.random.default_rng(24).normal(size=(3, 2, 8)) + pipe = DescriptorPipeline( + { + "output": {"channel_pooling": "all"}, + "families": { + "bands": { + "enabled": True, + "outputs": ["absolute_power"], + } + }, + "runtime": {"on_error": "raise"}, + } + ) + + with pytest.raises(ValueError, match="does not overlap"): + pipe.extract(X=X, sfreq=160.0, channel_names=["C3", "C4"]) + + +def test_complexity_collects_nonfinite_output_as_nan(): + """Verify that real non-finite results are collected as NaNs.""" + X = np.ones((2, 2, 16), dtype=float) + pipe = DescriptorPipeline( + { + "output": {"channel_pooling": "all"}, + "families": { + "complexity": { + "enabled": True, + "measures": ["sample_entropy"], + } + }, + "runtime": {"on_error": "collect"}, + } + ) + result = pipe.extract(X=X, sfreq=128.0) + + assert np.isnan(result["X"]).all() + assert result["failures"] + + +def test_complexity_raise_policy_reraises_nonfinite_output(): + """Verify that real non-finite results reraise when policy is set to raise.""" + X = np.ones((2, 2, 16), dtype=float) + pipe = DescriptorPipeline( + { + "output": {"channel_pooling": "all"}, + "families": { + "complexity": { + "enabled": True, + "measures": ["sample_entropy"], + } + }, + "runtime": {"on_error": "raise"}, + } + ) + + with pytest.raises(ValueError, match="non-finite"): + pipe.extract(X=X, sfreq=128.0) + + +def test_constant_signal_parametric_skip_collects_failures(): + X = np.zeros((3, 2, 128), dtype=float) + pipe = DescriptorPipeline( + { + "output": {"channel_pooling": "all"}, + "families": { + "parametric": { + "enabled": True, + "outputs": ["aperiodic"], + } + }, + "runtime": {"on_error": "collect"}, + } + ) + result = pipe.extract(X=X, sfreq=128.0, channel_names=["Fz", "Cz"]) + + assert result["failures"] + assert np.isnan(result["X"]).all() + + +def test_missing_antropy_dependency_has_clear_install_hint(monkeypatch): + def _raise_import_error(self): + raise ImportError( + "antropy is required for complexity descriptor extraction. " + "Install it with 'pip install coco-pipe[descriptors]'." + ) + + monkeypatch.setattr( + ComplexityDescriptorExtractor, + "_load_antropy", + _raise_import_error, + ) + pipe = DescriptorPipeline( + { + "output": {"channel_pooling": "all"}, + "families": { + "complexity": { + "enabled": True, + "backend": "antropy", + "measures": ["sample_entropy"], + } + }, + } + ) + + with pytest.raises(ImportError, match=r"coco-pipe\[descriptors\]"): + pipe.extract(X=[[[1.0] * 32]], sfreq=128.0, channel_names=["Cz"]) + + +def test_multi_family_scale_smoke(): + rng = np.random.default_rng(4) + X = rng.normal(size=(24, 4, 256)) + pipe = DescriptorPipeline( + { + "output": {"channel_pooling": "all"}, + "families": { + "bands": { + "enabled": True, + "outputs": ["absolute_power", "relative_power"], + }, + "parametric": { + "enabled": True, + "outputs": ["aperiodic", "peak_summary"], + }, + "complexity": { + "enabled": True, + "measures": ["sample_entropy", "perm_entropy"], + }, + }, + "runtime": {"obs_chunk": 8}, + } + ) + result = pipe.extract( + X=X, + sfreq=256.0, + channel_names=["Fz", "Cz", "Pz", "Oz"], + ) + + assert result["X"].shape[0] == 24 + assert result["X"].shape[1] == len(result["descriptor_names"]) + + +def test_multi_family_parallel_matches_sequential(): + pytest.importorskip("joblib") + rng = np.random.default_rng(12) + X = rng.normal(size=(12, 3, 128)) + channel_names = ["Fz", "Cz", "Pz"] + base_config = { + "output": {"channel_pooling": "all"}, + "families": { + "bands": { + "enabled": True, + "outputs": ["absolute_power", "relative_power"], + }, + "complexity": { + "enabled": True, + "measures": ["sample_entropy", "hjorth_mobility"], + }, + }, + } + sequential = DescriptorPipeline( + { + **base_config, + "runtime": {"execution_backend": "sequential", "n_jobs": 1}, + } + ).extract(X=X, sfreq=128.0, channel_names=channel_names) + parallel = DescriptorPipeline( + { + **base_config, + "runtime": {"execution_backend": "joblib", "n_jobs": 2}, + } + ).extract(X=X, sfreq=128.0, channel_names=channel_names) + + assert sequential["descriptor_names"] == parallel["descriptor_names"] + assert np.allclose(sequential["X"], parallel["X"], equal_nan=True) + + +def test_parametric_parallel_matches_sequential(): + pytest.importorskip("joblib") + rng = np.random.default_rng(13) + t = np.linspace(0, 1, 128, endpoint=False) + X = rng.normal(scale=0.05, size=(6, 3, 128)) + X[:, 0, :] += np.sin(2 * np.pi * 10 * t) + X[:, 1, :] += np.sin(2 * np.pi * 18 * t) + X[:, 2, :] += np.sin(2 * np.pi * 6 * t) + + sequential = DescriptorPipeline( + { + "families": { + "parametric": { + "enabled": True, + "outputs": ["aperiodic", "fit_quality"], + } + }, + "output": {"channel_pooling": "all"}, + "runtime": {"execution_backend": "sequential", "n_jobs": 1}, + } + ).extract(X=X, sfreq=128.0, channel_names=["Fz", "Cz", "Pz"]) + parallel = DescriptorPipeline( + { + "families": { + "parametric": { + "enabled": True, + "outputs": ["aperiodic", "fit_quality"], + } + }, + "output": {"channel_pooling": "all"}, + "runtime": {"execution_backend": "joblib", "n_jobs": 2}, + } + ).extract(X=X, sfreq=128.0, channel_names=["Fz", "Cz", "Pz"]) + + assert sequential["descriptor_names"] == parallel["descriptor_names"] + assert np.allclose(sequential["X"], parallel["X"], equal_nan=True) + + +def test_multi_chunk_row_order_matches_unchunked(): + rng = np.random.default_rng(15) + X = rng.normal(size=(18, 3, 128)) + config = { + "output": {"channel_pooling": "all"}, + "families": {"bands": {"enabled": True, "outputs": ["absolute_power"]}}, + } + unchunked = DescriptorPipeline(config).extract( + X=X, + sfreq=128.0, + channel_names=["Fz", "Cz", "Pz"], + ) + chunked = DescriptorPipeline( + { + **config, + "runtime": {"obs_chunk": 5}, + } + ).extract( + X=X, + sfreq=128.0, + channel_names=["Fz", "Cz", "Pz"], + ) + + assert unchunked["descriptor_names"] == chunked["descriptor_names"] + assert np.allclose(unchunked["X"], chunked["X"], equal_nan=True) + + +def test_n_jobs_one_skips_joblib_loading(monkeypatch): + import builtins + import inspect + + rng = np.random.default_rng(16) + X = rng.normal(size=(4, 2, 64)) + real_import = builtins.__import__ + joblib_imports = 0 + + def _count_joblib_imports(name, *args, **kwargs): + nonlocal joblib_imports + caller = inspect.currentframe().f_back + caller_name = None if caller is None else caller.f_globals.get("__name__") + if name == "joblib" and caller_name in { + "coco_pipe.descriptors.core", + "coco_pipe.descriptors.extractors._parametric_fit", + }: + joblib_imports += 1 + return real_import(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", _count_joblib_imports) + result = DescriptorPipeline( + { + "families": { + "bands": { + "enabled": True, + "outputs": ["absolute_power"], + } + }, + "output": {"channel_pooling": "all"}, + "runtime": {"execution_backend": "joblib", "n_jobs": 1}, + } + ).extract(X=X, sfreq=128.0, channel_names=["Fz", "Cz"]) + + assert result["X"].shape == (4, 5) + assert joblib_imports == 0 + + +def test_parametric_parallel_n_jobs_all_cores_smoke(): + pytest.importorskip("joblib") + rng = np.random.default_rng(17) + # 4 seconds at 128Hz = 512 samples + t = np.linspace(0, 4, 512, endpoint=False) + X = rng.normal(scale=0.05, size=(4, 3, 512)) + # Add 1/f slope + freqs = np.fft.rfftfreq(512, 1 / 128.0) + weights = 1 / (freqs + 1.0) + for o in range(4): + for c in range(3): + X[o, c, :] = np.fft.irfft(np.fft.rfft(X[o, c, :]) * weights, n=512) + + X[:, 0, :] += 2.0 * np.sin(2 * np.pi * 10 * t) + X[:, 1, :] += 1.5 * np.sin(2 * np.pi * 16 * t) + X[:, 2, :] += 1.0 * np.sin(2 * np.pi * 6 * t) + + result = DescriptorPipeline( + { + "families": { + "parametric": { + "enabled": True, + "outputs": ["aperiodic"], + } + }, + "output": {"channel_pooling": "all"}, + "runtime": {"execution_backend": "joblib", "n_jobs": -1}, + } + ).extract(X=X, sfreq=128.0, channel_names=["Fz", "Cz", "Pz"]) + + # 2 features (offset, exponent) per observation + assert result["X"].shape == (4, 2) + + +def test_shared_psd_reuses_one_compute_per_batch_for_same_method(monkeypatch): + rng = np.random.default_rng(19) + t = np.linspace(0, 1, 128, endpoint=False) + X = rng.normal(scale=0.05, size=(8, 3, 128)) + X[:, 0, :] += np.sin(2 * np.pi * 10 * t) + X[:, 1, :] += np.sin(2 * np.pi * 18 * t) + X[:, 2, :] += np.sin(2 * np.pi * 6 * t) + + real_compute_psd = descriptors_core.compute_psd + calls: list[tuple[str, float, float]] = [] + + def _counted_compute_psd(*args, **kwargs): + calls.append((kwargs["method"], kwargs["fmin"], kwargs["fmax"])) + return real_compute_psd(*args, **kwargs) + + monkeypatch.setattr(descriptors_core, "compute_psd", _counted_compute_psd) + + DescriptorPipeline( + { + "output": {"channel_pooling": "all"}, + "families": { + "bands": { + "enabled": True, + "psd_method": "welch", + "outputs": ["absolute_power"], + }, + "parametric": { + "enabled": True, + "psd_method": "welch", + "outputs": ["aperiodic"], + }, + }, + "runtime": {"obs_chunk": 4}, + } + ).extract(X=X, sfreq=128.0, channel_names=["Fz", "Cz", "Pz"]) + + assert len(calls) == 2 + assert all(method == "welch" for method, _, _ in calls) + + +def test_corrected_bands_and_parametric_share_one_fit_batch_per_psd_group( + monkeypatch, +): + rng = np.random.default_rng(191) + t = np.linspace(0, 1, 128, endpoint=False) + X = rng.normal(scale=0.05, size=(8, 3, 128)) + X[:, 0, :] += np.sin(2 * np.pi * 10 * t) + X[:, 1, :] += np.sin(2 * np.pi * 18 * t) + X[:, 2, :] += np.sin(2 * np.pi * 6 * t) + + real_fit_batch = descriptors_core.fit_parametric_batch + calls = 0 + + def _counted_fit_batch(*args, **kwargs): + nonlocal calls + calls += 1 + return real_fit_batch(*args, **kwargs) + + monkeypatch.setattr(descriptors_core, "fit_parametric_batch", _counted_fit_batch) + + result = DescriptorPipeline( + { + "output": {"channel_pooling": "all"}, + "families": { + "bands": { + "enabled": True, + "psd_method": "welch", + "outputs": ["corrected_absolute_power"], + }, + "parametric": { + "enabled": True, + "psd_method": "welch", + "outputs": ["aperiodic"], + }, + }, + "runtime": {"obs_chunk": 4}, + } + ).extract(X=X, sfreq=128.0, channel_names=["Fz", "Cz", "Pz"]) + + assert calls == 2 + assert "band_corr_abs_alpha_ch-all" in result["descriptor_names"] + + +def test_shared_psd_splits_groups_by_method(monkeypatch): + rng = np.random.default_rng(20) + t = np.linspace(0, 1, 128, endpoint=False) + X = rng.normal(scale=0.05, size=(8, 3, 128)) + X[:, 0, :] += np.sin(2 * np.pi * 10 * t) + X[:, 1, :] += np.sin(2 * np.pi * 18 * t) + X[:, 2, :] += np.sin(2 * np.pi * 6 * t) + + real_compute_psd = descriptors_core.compute_psd + calls: list[str] = [] + + def _counted_compute_psd(*args, **kwargs): + calls.append(kwargs["method"]) + return real_compute_psd(*args, **kwargs) + + monkeypatch.setattr(descriptors_core, "compute_psd", _counted_compute_psd) + + DescriptorPipeline( + { + "output": {"channel_pooling": "all"}, + "families": { + "bands": { + "enabled": True, + "psd_method": "welch", + "outputs": ["absolute_power"], + }, + "parametric": { + "enabled": True, + "psd_method": "multitaper", + "outputs": ["aperiodic"], + }, + }, + "runtime": {"obs_chunk": 4}, + } + ).extract(X=X, sfreq=128.0, channel_names=["Fz", "Cz", "Pz"]) + + assert calls == ["welch", "multitaper", "welch", "multitaper"] + + +def test_shared_union_psd_matches_separate_family_outputs(): + rng = np.random.default_rng(21) + t = np.linspace(0, 1, 128, endpoint=False) + X = rng.normal(scale=0.05, size=(6, 3, 128)) + X[:, 0, :] += np.sin(2 * np.pi * 10 * t) + X[:, 1, :] += np.sin(2 * np.pi * 18 * t) + X[:, 2, :] += np.sin(2 * np.pi * 6 * t) + channel_names = ["Fz", "Cz", "Pz"] + + bands_cfg = { + "output": {"channel_pooling": "all"}, + "families": { + "bands": { + "enabled": True, + "psd_method": "welch", + "fmin": 1.0, + "fmax": 30.0, + "bands": { + "delta": [1.0, 4.0], + "theta": [4.0, 8.0], + "alpha": [8.0, 13.0], + "beta": [13.0, 30.0], + }, + "outputs": ["absolute_power", "relative_power"], + } + }, + } + param_cfg = { + "output": {"channel_pooling": "all"}, + "families": { + "parametric": { + "enabled": True, + "psd_method": "welch", + "freq_range": [1.0, 45.0], + "outputs": ["aperiodic", "fit_quality"], + } + }, + } + combined_cfg = { + "output": {"channel_pooling": "all"}, + "families": { + **bands_cfg["families"], + **param_cfg["families"], + }, + } + + bands_only = DescriptorPipeline(bands_cfg).extract( + X=X, + sfreq=128.0, + channel_names=channel_names, + ) + param_only = DescriptorPipeline(param_cfg).extract( + X=X, + sfreq=128.0, + channel_names=channel_names, + ) + combined = DescriptorPipeline(combined_cfg).extract( + X=X, + sfreq=128.0, + channel_names=channel_names, + ) + + band_names = [ + name for name in combined["descriptor_names"] if name.startswith("band_") + ] + param_names = [ + name for name in combined["descriptor_names"] if name.startswith("param_") + ] + band_indices = [combined["descriptor_names"].index(name) for name in band_names] + param_indices = [combined["descriptor_names"].index(name) for name in param_names] + + assert band_names == bands_only["descriptor_names"] + assert param_names == param_only["descriptor_names"] + assert np.allclose( + combined["X"][:, band_indices], + bands_only["X"], + equal_nan=True, + ) + assert np.allclose( + combined["X"][:, param_indices], + param_only["X"], + equal_nan=True, + ) + + +def test_obs_batch_parallel_disables_parametric_inner_joblib(monkeypatch): + import builtins + import inspect + + pytest.importorskip("joblib") + rng = np.random.default_rng(22) + t = np.linspace(0, 4, 512, endpoint=False) + X = rng.normal(scale=0.05, size=(6, 3, 512)) + freqs = np.fft.rfftfreq(512, 1 / 128.0) + weights = 1 / (freqs + 1.0) + for o in range(6): + for c in range(3): + X[o, c, :] = np.fft.irfft(np.fft.rfft(X[o, c, :]) * weights, n=512) + + X[:, 0, :] += 2.0 * np.sin(2 * np.pi * 10 * t) + X[:, 1, :] += 1.5 * np.sin(2 * np.pi * 18 * t) + X[:, 2, :] += 1.0 * np.sin(2 * np.pi * 6 * t) + real_import = builtins.__import__ + joblib_imports = 0 + + def _count_joblib_imports(name, *args, **kwargs): + nonlocal joblib_imports + caller = inspect.currentframe().f_back + caller_name = None if caller is None else caller.f_globals.get("__name__") + if name == "joblib" and caller_name in { + "coco_pipe.descriptors.core", + "coco_pipe.descriptors.extractors._parametric_fit", + }: + joblib_imports += 1 + return real_import(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", _count_joblib_imports) + result = DescriptorPipeline( + { + "families": { + "parametric": { + "enabled": True, + "outputs": ["aperiodic"], + } + }, + "output": {"channel_pooling": "all"}, + "runtime": { + "execution_backend": "joblib", + "n_jobs": 2, + "obs_chunk": 2, + }, + } + ).extract(X=X, sfreq=128.0, channel_names=["Fz", "Cz", "Pz"]) + + # 2 features (offset, exponent) for aperiodic 'fixed' mode + assert result["X"].shape == (X.shape[0], 2) + # Check that exponent is reasonable (> 0) and offset is finite + assert np.all(result["X"][:, 1] > 0) + assert np.all(np.isfinite(result["X"][:, 0])) + + +def test_single_psd_group_uses_psd_level_n_jobs(monkeypatch): + pytest.importorskip("joblib") + rng = np.random.default_rng(23) + X = rng.normal(size=(4, 3, 128)) + calls: list[int | None] = [] + real_compute_psd = descriptors_core.compute_psd + + def _counted_compute_psd(*args, **kwargs): + calls.append(kwargs["n_jobs"]) + return real_compute_psd(*args, **kwargs) + + monkeypatch.setattr(descriptors_core, "compute_psd", _counted_compute_psd) + DescriptorPipeline( + { + "families": { + "bands": { + "enabled": True, + "outputs": ["absolute_power"], + } + }, + "output": {"channel_pooling": "all"}, + "runtime": {"execution_backend": "joblib", "n_jobs": 2}, + } + ).extract(X=X, sfreq=128.0, channel_names=["Fz", "Cz", "Pz"]) + + assert calls == [2] + + +def test_validation_edge_cases_runtime(): + from coco_pipe.descriptors.configs import DescriptorConfig + from coco_pipe.descriptors.validation import ( + _normalize_channel_pooling, + validate_runtime_inputs, + ) + + config = DescriptorConfig(families={"bands": {"enabled": True}}) + X = np.zeros((2, 2, 64)) + + # sfreq <= 0 + with pytest.raises(ValueError, match="`sfreq` must be positive"): + validate_runtime_inputs(config, X=X, sfreq=0, channel_names=["ch1", "ch2"]) + + # ids alignment failure + with pytest.raises(ValueError, match="`ids` must align with n_obs=2"): + validate_runtime_inputs( + config, X=X, sfreq=100.0, ids=[1, 2, 3], channel_names=["ch1", "ch2"] + ) + + # channel_names alignment failure + with pytest.raises( + ValueError, match="`channel_names` must align with n_channels=2" + ): + validate_runtime_inputs(config, X=X, sfreq=100.0, channel_names=["ch1"]) + + # channel_names required but missing (when pooling is not 'all') + config_none = DescriptorConfig( + families={"bands": {"enabled": True}}, output={"channel_pooling": "none"} + ) + with pytest.raises(ValueError, match="`channel_names` must be passed explicitly"): + validate_runtime_inputs(config_none, X=X, sfreq=100.0, channel_names=None) + + # _normalize_channel_pooling edge cases + # missing channel_names when groups are used + with pytest.raises(ValueError, match="`channel_names` must be passed explicitly"): + _normalize_channel_pooling({"G1": ["ch1"]}, None) + + # duplicate channel_names when groups are used + with pytest.raises(ValueError, match="`channel_names` must be unique"): + _normalize_channel_pooling({"G1": ["ch1"]}, ["ch1", "ch1"]) diff --git a/tests/test_descriptors_extractors.py b/tests/test_descriptors_extractors.py new file mode 100644 index 0000000..fc7b44d --- /dev/null +++ b/tests/test_descriptors_extractors.py @@ -0,0 +1,492 @@ +""" +Comprehensive Test Suite for Descriptor Extractors +================================================== + +Unified tests for Spectral, Parametric, and Complexity descriptor extractors. +""" + +import sys +from unittest.mock import MagicMock + +import numpy as np +import pytest + +from coco_pipe.descriptors.configs import ( + BandDescriptorConfig, + ComplexityDescriptorConfig, + ParametricDescriptorConfig, +) +from coco_pipe.descriptors.extractors._parametric_fit import _ParametricFitBatch +from coco_pipe.descriptors.extractors._psd import compute_psd +from coco_pipe.descriptors.extractors.base import _DescriptorBlock +from coco_pipe.descriptors.extractors.complexity import ComplexityDescriptorExtractor +from coco_pipe.descriptors.extractors.parametric import ParametricDescriptorExtractor +from coco_pipe.descriptors.extractors.spectral import BandDescriptorExtractor +from coco_pipe.descriptors.extractors.utils import ( + average_channel_matrix, + make_failure_record, + pool_channel_descriptor_matrix, +) + +# --- Fixtures --- + + +@pytest.fixture +def signal_data(): + """Standard signal data: (n_obs, n_channels, n_times).""" + rng = np.random.default_rng(42) + sfreq = 250.0 + # Increase to 2 seconds for better Welch/entropy estimation + t = np.arange(0, 2, 1 / sfreq) + n_obs, n_chans = 5, 3 + + # Create 1/f-like noise + freqs = np.fft.rfftfreq(len(t), 1 / sfreq) + weights = 1 / (freqs + 1.0) + + X = np.zeros((n_obs, n_chans, len(t))) + for o in range(n_obs): + for c in range(n_chans): + white = rng.standard_normal(len(t)) + # Quick and dirty 1/f approximation + X[o, c, :] = np.fft.irfft(np.fft.rfft(white) * weights, n=len(t)) + + # Add strong oscillations to ensure fitting works + X[:, 0, :] += 2.0 * np.sin(2 * np.pi * 10 * t) + X[:, 1, :] += 1.5 * np.sin(2 * np.pi * 20 * t) + + return X, sfreq, ["Fz", "Cz", "Pz"] + + +@pytest.fixture +def psd_data(signal_data): + """Standard PSD data: (n_obs, n_channels, n_freqs).""" + X, sfreq, _ = signal_data + psds, freqs = compute_psd( + X, sfreq=sfreq, method="welch", fmin=1.0, fmax=45.0, n_jobs=1 + ) + return psds, freqs + + +@pytest.fixture +def mock_fit_batch(psd_data): + """A mock ParametricFitBatch for testing corrected bands.""" + psds, freqs = psd_data + n_obs, n_chans, n_freqs = psds.shape + + periodic_psds = np.zeros_like(psds) + # Add a "peak" at 10Hz (approx index) + f_idx = np.argmin(np.abs(freqs - 10.0)) + periodic_psds[:, :, f_idx] = 1.0 + + metrics = { + "offset": np.zeros((n_obs, n_chans)), + "exponent": np.ones((n_obs, n_chans)) * 1.5, + } + + return _ParametricFitBatch( + freqs=freqs, + metrics=metrics, + errors=[], + periodic_psds=periodic_psds, + ) + + +# --- 1. Base Interfaces and Utilities --- + + +def test_descriptor_block_structure(): + """Verify _DescriptorBlock simple data container.""" + X = np.zeros((5, 10)) + names = [f"desc_{i}" for i in range(10)] + + block = _DescriptorBlock(family="test", X=X, descriptor_names=names) + assert block.X.shape == (5, 10) + assert block.descriptor_names == names + + +def test_pool_channel_descriptor_matrix_logic(): + """Test all variants of channel pooling names and values.""" + # (n_obs, n_channels) + values = np.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]) + ch_names = ["Fz", "Cz"] + + # None + pooled, scopes = pool_channel_descriptor_matrix(values, ch_names, "none") + assert np.array_equal(pooled, values) + assert scopes == ["ch-Fz", "ch-Cz"] + + # All (mean across channels) + pooled, scopes = pool_channel_descriptor_matrix(values, ch_names, "all") + assert pooled.shape == (3, 1) + assert np.allclose(pooled[:, 0], [1.5, 3.5, 5.5]) + assert scopes == ["ch-all"] + + # Dict grouping + values3 = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) + ch_names3 = ["Fz", "Cz", "Pz"] + pooling = {"Frontal": ["Fz", "Cz"]} + pooled, scopes = pool_channel_descriptor_matrix(values3, ch_names3, pooling) + # Frontal (mean of 0,1) + Pz (remains) + assert pooled.shape == (2, 2) + assert np.allclose(pooled[:, 0], [1.5, 4.5]) # mean([1,2], [4,5]) + assert np.allclose(pooled[:, 1], [3.0, 6.0]) # Pz + assert scopes == ["chgrp-Frontal", "ch-Pz"] + + +def test_average_channel_matrix_robustness(): + """Verify averaging handles NaNs correctly.""" + X = np.array([[1.0, 2.0], [np.nan, 4.0], [5.0, np.nan], [np.nan, np.nan]]) + res = average_channel_matrix(X) + assert np.allclose(res[:3], [1.5, 4.0, 5.0]) + assert np.isnan(res[3]) + + +def test_make_failure_record_schema(): + """Check fixed schema for failure records.""" + rec = make_failure_record( + family="spectral", + obs_index=5, + exception_type="ValueError", + message="test error", + channel_index=2, + channel_name="Cz", + ) + assert rec["obs_index"] == 5 + assert rec["family"] == "spectral" + assert rec["exception_type"] == "ValueError" + assert rec["message"] == "test error" + assert rec["channel_index"] == 2 + assert rec["channel_name"] == "Cz" + + +# --- 2. Spectral (Band) Extractor --- + + +class TestBandExtractor: + def test_basic_extraction(self, psd_data, signal_data): + psds, freqs = psd_data + _, _, ch_names = signal_data + config = BandDescriptorConfig( + enabled=True, + outputs=["absolute_power", "relative_power"], + bands={"alpha": (8, 12), "beta": (15, 30)}, + ) + extractor = BandDescriptorExtractor(config) + + block = extractor.extract_psd( + psds, + freqs, + channel_names=ch_names, + channel_pooling="none", + ids=None, + runtime=MagicMock(), + ) + assert block.family == "bands" + # 2 bands * 2 outputs * 3 channels = 12 columns + assert block.X.shape == (psds.shape[0], 12) + assert "band_abs_alpha_ch-Fz" in block.descriptor_names + assert "band_rel_beta_ch-Pz" in block.descriptor_names + + def test_corrected_outputs(self, psd_data, signal_data, mock_fit_batch): + psds, freqs = psd_data + _, _, ch_names = signal_data + config = BandDescriptorConfig( + enabled=True, + outputs=["corrected_absolute_power", "corrected_ratios"], + bands={"alpha": (8, 12), "beta": (13, 30)}, + ratio_pairs=[("alpha", "beta")], + ) + extractor = BandDescriptorExtractor(config) + + block = extractor.extract_psd( + psds, + freqs, + channel_names=ch_names, + channel_pooling="all", + ids=None, + fit_batch=mock_fit_batch, + runtime=MagicMock(), + ) + # alpha, beta corrected + 1 ratio = 3 columns + assert block.X.shape == (psds.shape[0], 3) + assert "band_corr_abs_alpha_ch-all" in block.descriptor_names + assert "band_corr_ratio_alpha_beta_ch-all" in block.descriptor_names + + def test_missing_fit_batch_raises(self, psd_data, signal_data): + psds, freqs = psd_data + _, _, ch_names = signal_data + config = BandDescriptorConfig( + enabled=True, + outputs=["corrected_absolute_power"], + ) + extractor = BandDescriptorExtractor(config) + + with pytest.raises(ValueError, match="require a supplied parametric fit_batch"): + extractor.extract_psd( + psds, + freqs, + channel_names=ch_names, + channel_pooling="none", + ids=None, + runtime=MagicMock(), + ) + + def test_band_resolution_error(self, psd_data, signal_data): + psds, freqs = psd_data + _, _, ch_names = signal_data + config = BandDescriptorConfig( + enabled=True, + fmax=250.0, + bands={"ultra": (100, 200)}, + outputs=["absolute_power"], + ) + extractor = BandDescriptorExtractor(config) + + runtime = MagicMock() + runtime.on_error = "collect" + + block = extractor.extract_psd( + psds, + freqs, + channel_names=ch_names, + channel_pooling="all", + ids=None, + runtime=runtime, + ) + assert len(block.failures) > 0 + assert block.failures[0]["exception_type"] == "BandResolutionError" + + def test_spectral_capabilities_and_requests(self): + config = BandDescriptorConfig(enabled=True) + extractor = BandDescriptorExtractor(config) + assert extractor.capabilities["requires_sfreq"] is True + + # corrected without fit_config raises + config_corr = BandDescriptorConfig( + enabled=True, outputs=["corrected_absolute_power"] + ) + extractor_corr = BandDescriptorExtractor(config_corr, fit_config=None) + with pytest.raises(ValueError, match="Corrected band outputs require"): + extractor_corr.psd_request() + + def test_spectral_extract_psd_edge_cases(self, signal_data): + from unittest.mock import MagicMock + + X, sfreq, ch_names = signal_data + + # log_power coverage + config = BandDescriptorConfig( + enabled=True, outputs=["absolute_power"], log_power=True + ) + extractor = BandDescriptorExtractor(config) + psds = np.ones((1, 3, 10)) + freqs = np.linspace(1, 45, 10) + block = extractor.extract_psd(psds, freqs, ch_names, "all", None, MagicMock()) + assert "band_log_abs_delta_ch-all" in block.descriptor_names + + config_rel = BandDescriptorConfig( + enabled=True, + outputs=["relative_power"], + fmin=100.0, + fmax=200.0, + bands={"high": (120, 150)}, + ) + extractor_rel = BandDescriptorExtractor(config_rel) + block_rel = extractor_rel.extract_psd( + psds, freqs, ch_names, "all", None, MagicMock() + ) + assert np.isnan(block_rel.X).all() + + # fit_batch errors and corrected relative power empty freq + from coco_pipe.descriptors.extractors._parametric_fit import _ParametricFitBatch + + fit_batch = _ParametricFitBatch( + freqs=np.array([1, 10, 20]), # Must be non-empty and non-None + metrics={}, + periodic_psds=np.zeros((1, 3, 3)), + errors=[(0, 0, "FakeError", "Fit failed")], + meta={}, + ) + config_corr = BandDescriptorConfig( + enabled=True, outputs=["corrected_relative_power"] + ) + extractor_corr = BandDescriptorExtractor(config_corr) + block_corr = extractor_corr.extract_psd( + psds, freqs, ch_names, "all", None, MagicMock(), fit_batch=fit_batch + ) + assert len(block_corr.failures) > 0 + assert "FakeError" in block_corr.failures[0]["exception_type"] + + def test_spectral_standalone_extract(self, signal_data): + from unittest.mock import MagicMock + + X, sfreq, ch_names = signal_data + config = BandDescriptorConfig(enabled=True, outputs=["absolute_power"]) + extractor = BandDescriptorExtractor(config) + block = extractor.extract(X, sfreq, ch_names, "all", None, MagicMock()) + # 5 bands (delta, theta, alpha, beta, gamma) pooled to 'all' -> 5 columns + assert block.X.shape == (5, 5) + + def test_spectral_empty_output_block(self): + # Empty output when no families enabled or no outputs requested + config = BandDescriptorConfig(enabled=True, outputs=[]) + extractor = BandDescriptorExtractor(config) + psds = np.ones((2, 2, 10)) + freqs = np.linspace(1, 45, 10) + block = extractor.extract_psd( + psds, freqs, ["ch1", "ch2"], "all", None, MagicMock() + ) + assert block.X.shape == (2, 0) + + def test_spectral_standalone_extract_raises_for_corrected(self, signal_data): + from unittest.mock import MagicMock + + X, sfreq, ch_names = signal_data + config = BandDescriptorConfig( + enabled=True, outputs=["corrected_absolute_power"] + ) + extractor = BandDescriptorExtractor(config) + with pytest.raises( + ValueError, match="Corrected band outputs are only available" + ): + extractor.extract(X, sfreq, ch_names, "all", None, MagicMock()) + + +# --- 3. Parametric Extractor --- + + +class TestParametricExtractor: + def test_standalone_extract(self, signal_data): + """Verify real specparam-based extraction.""" + X, sfreq, ch_names = signal_data + config = ParametricDescriptorConfig( + enabled=True, + outputs=["aperiodic"], + psd_method="welch", + freq_range=(1.0, 45.0), + ) + extractor = ParametricDescriptorExtractor(config) + + block = extractor.extract( + X, + sfreq=sfreq, + channel_names=ch_names, + channel_pooling="all", + ids=None, + runtime=MagicMock(), + obs_offset=0, + ) + assert "param_exponent_ch-all" in block.descriptor_names + # 2 features (offset, exponent) for aperiodic 'fixed' mode + assert block.X.shape == (X.shape[0], 2) + # Check that exponent is reasonable (> 0) and offset is finite + assert np.all(block.X[:, 1] > 0) + assert np.all(np.isfinite(block.X[:, 0])) + + def test_extract_psd_requires_fit_batch(self, psd_data, signal_data): + psds, freqs = psd_data + _, _, ch_names = signal_data + extractor = ParametricDescriptorExtractor(ParametricDescriptorConfig()) + + with pytest.raises(ValueError, match="requires a supplied fit_batch"): + extractor.extract_psd( + psds, + freqs, + channel_names=ch_names, + channel_pooling="all", + ids=None, + runtime=MagicMock(), + obs_offset=0, + ) + + +# --- 4. Complexity Extractor --- + + +class TestComplexityExtractor: + def test_backend_dispatch_antropy(self, signal_data): + X, sfreq, ch_names = signal_data + config = ComplexityDescriptorConfig( + enabled=True, backend="antropy", measures=["spectral_entropy"] + ) + extractor = ComplexityDescriptorExtractor(config) + + block = extractor.extract( + X, + sfreq=sfreq, + channel_names=ch_names, + channel_pooling="all", + ids=None, + runtime=MagicMock(), + obs_offset=0, + ) + assert "complexity_spectral_entropy_ch-all" in block.descriptor_names + assert not np.isnan(block.X).any() + + def test_backend_dispatch_neurokit2(self, signal_data): + X, sfreq, ch_names = signal_data + config = ComplexityDescriptorConfig( + enabled=True, backend="neurokit2", measures=["perm_entropy"] + ) + extractor = ComplexityDescriptorExtractor(config) + + block = extractor.extract( + X, + sfreq=sfreq, + channel_names=ch_names, + channel_pooling="all", + ids=None, + runtime=MagicMock(), + obs_offset=0, + ) + assert "complexity_perm_entropy_ch-all" in block.descriptor_names + # Check that it's finite + assert not np.isnan(block.X).any() + + def test_mixed_execution_strategy(self, signal_data): + """Verify execution paths for combined complexity measures.""" + X, sfreq, ch_names = signal_data + config = ComplexityDescriptorConfig( + enabled=True, + backend="antropy", + measures=["spectral_entropy", "sample_entropy"], + ) + extractor = ComplexityDescriptorExtractor(config) + + block = extractor.extract( + X, + sfreq=sfreq, + channel_names=ch_names, + channel_pooling="none", + ids=None, + runtime=MagicMock(), + obs_offset=0, + ) + + # 2 measures * 3 channels = 6 columns + assert block.X.shape == (X.shape[0], 6) + assert "complexity_spectral_entropy_ch-Fz" in block.descriptor_names + assert "complexity_sample_entropy_ch-Pz" in block.descriptor_names + + +# --- 5. Lazy Loading and Dependency Guards --- + + +def test_lazy_loading_failure_antropy(monkeypatch): + """Verify informative error when antropy is missing.""" + monkeypatch.setitem(sys.modules, "antropy", None) + config = ComplexityDescriptorConfig(enabled=True, backend="antropy") + extractor = ComplexityDescriptorExtractor(config) + + with pytest.raises(ImportError, match="antropy"): + extractor._load_antropy() + + +def test_lazy_loading_failure_neurokit2(monkeypatch): + monkeypatch.setitem(sys.modules, "neurokit2", None) + config = ComplexityDescriptorConfig(enabled=True, backend="neurokit2") + extractor = ComplexityDescriptorExtractor(config) + + with pytest.raises(ImportError, match="neurokit"): + extractor._load_neurokit() diff --git a/tests/test_io_dataset.py b/tests/test_io_dataset.py index e99e76c..3e1eb05 100644 --- a/tests/test_io_dataset.py +++ b/tests/test_io_dataset.py @@ -236,7 +236,7 @@ def test_bids_dataset_mismatches(monkeypatch, tmp_path): monkeypatch.setattr( dataset, "_get_bids_path", - lambda: (lambda **k: types.SimpleNamespace(match=lambda: [], **k)), + lambda: lambda **k: types.SimpleNamespace(match=lambda: [], **k), ) # One subject, two sessions diff --git a/tests/test_io_structures.py b/tests/test_io_structures.py index f51f1e2..20d0fa3 100644 --- a/tests/test_io_structures.py +++ b/tests/test_io_structures.py @@ -1,6 +1,7 @@ import numpy as np import pytest +from coco_pipe.descriptors import DescriptorPipeline from coco_pipe.io.structures import DataContainer @@ -328,31 +329,175 @@ def test_aggregate(): dc.aggregate(by=[1, 2]) # 4. Aggregation Mean (Standard) - agg = dc.aggregate(by="Study ID", method="mean") + agg = dc.aggregate(by="Study ID", stats="mean") assert agg.shape == (2, 2) assert np.array_equal(agg.coords["obs"], ["A", "B"]) assert np.array_equal(agg.coords["Study ID"], ["A", "B"]) assert np.array_equal(agg.coords["site"], ["north", "south"]) assert "mixed" not in agg.coords + assert np.array_equal(agg.coords["epoch_count"], [2, 1]) # Group A: (1+2)/2 = 1.5. Group B: 3. assert agg.X[0, 0] == 1.5 assert agg.y is not None assert np.array_equal(agg.y, [0, 1]) # 0 is consistent for A # 5. Method variants - agg_std = dc.aggregate(by="Study ID", method="std") - assert agg_std.ids is None # Std voids IDs + agg_std = dc.aggregate(by="Study ID", stats="std") + assert np.array_equal(agg_std.ids, ["A", "B"]) - with pytest.raises(ValueError, match="Unknown method"): - dc.aggregate(by="Study ID", method="invalid") + with pytest.raises(ValueError, match="Unknown stats"): + dc.aggregate(by="Study ID", stats="invalid") def test_aggregate_unknown_method(): """Test unknown aggregation method error.""" X = np.zeros((2, 2)) dc = DataContainer(X, dims=("obs", "f")) - with pytest.raises(ValueError, match="Unknown method"): - dc.aggregate(by=[1, 2], method="magic") + with pytest.raises(ValueError, match="Unknown stats"): + dc.aggregate(by=[1, 2], stats="magic") + + +def _make_descriptor_container(X, *, descriptor_names=None): + return DataContainer( + X=np.asarray(X, dtype=np.float32), + dims=("obs", "feature"), + coords={ + "feature": descriptor_names or ["alpha_ch-all", "beta_ch-all"], + }, + ) + + +def _make_grouped_descriptor_container(): + return _make_descriptor_container( + [ + [np.nan, 1.0], + [np.nan, 2.0], + [3.0, 4.0], + [np.nan, np.nan], + ], + ) + + +def _make_signal_data(): + rng = np.random.default_rng(31) + t = np.linspace(0, 1, 256, endpoint=False) + X = rng.normal(scale=0.1, size=(6, 2, 256)) + X[:, 0, :] += np.sin(2 * np.pi * 10 * t) + X[:, 1, :] += np.sin(2 * np.pi * 6 * t) + return X + + +def _descriptor_result_container(result): + return DataContainer( + X=result["X"], + dims=("obs", "feature"), + coords={"feature": result["descriptor_names"]}, + ) + + +def test_aggregate_multiple_stats_insert_stat_dimension_in_requested_order(): + agg = _make_grouped_descriptor_container().aggregate( + by=["g1", "g1", "g2", "g2"], + stats=["mean", "std"], + ) + + assert agg.dims == ("obs", "stat", "feature") + assert agg.coords["stat"].tolist() == ["mean", "std"] + assert list(agg.coords["feature"]) == ["alpha_ch-all", "beta_ch-all"] + assert agg.X.shape == (2, 2, 2) + + +def test_aggregate_group_ids_follow_first_appearance_order(): + agg = _make_descriptor_container( + [[1.0], [2.0], [3.0], [4.0]], + descriptor_names=["alpha_ch-all"], + ).aggregate(by=["g2", "g1", "g2", "g3"]) + + assert agg.coords["obs"].tolist() == ["g2", "g1", "g3"] + assert agg.ids.tolist() == ["g2", "g1", "g3"] + + +def test_aggregate_count_sem_and_epoch_count_match_expected_values(): + count_agg = _make_grouped_descriptor_container().aggregate( + by=["g1", "g1", "g2", "g2"], + stats="count", + ) + sem_agg = _make_grouped_descriptor_container().aggregate( + by=["g1", "g1", "g2", "g2"], + stats="sem", + ) + + assert count_agg.X[0].tolist() == [0.0, 2.0] + assert count_agg.coords["epoch_count"].tolist() == [2, 2] + expected_sem = np.nanstd([1.0, 2.0]) / np.sqrt(2) + assert np.isclose(sem_agg.X[0, 1], expected_sem) + + +def test_aggregate_min_count_collect_policy_records_failure(): + agg = _make_grouped_descriptor_container().aggregate( + by=["g1", "g1", "g2", "g2"], + min_count=2, + on_insufficient="collect", + ) + + assert len(agg.meta["aggregate_failures"]) == 1 + assert agg.meta["aggregate_failures"][0]["exception_type"] == ( + "InsufficientObservations" + ) + assert agg.meta["aggregate_failures"][0]["valid_row_count"] == 1 + assert agg.meta["aggregate_failures"][0]["row_count"] == 2 + assert np.isnan(agg.X[1]).all() + + +def test_aggregate_min_count_warn_policy_emits_warning(): + with pytest.warns(UserWarning, match="requires at least 2"): + agg = _make_grouped_descriptor_container().aggregate( + by=["g1", "g1", "g2", "g2"], + min_count=2, + on_insufficient="warn", + ) + + assert np.isnan(agg.X[1]).all() + assert agg.meta["aggregate_failures"][0]["exception_type"] == ( + "InsufficientObservations" + ) + + +def test_aggregate_descriptor_pipeline_output_can_be_grouped(): + X = _make_signal_data() + result = DescriptorPipeline( + { + "output": {"channel_pooling": "all"}, + "families": {"bands": {"enabled": True, "outputs": ["absolute_power"]}}, + } + ).extract(X=X, sfreq=256.0, channel_names=["Fz", "Cz"]) + agg = _descriptor_result_container(result).aggregate( + by=["s1", "s1", "s1", "s2", "s2", "s2"], + stats="mean", + ) + + assert all("_global" not in name for name in result["descriptor_names"]) + assert any(name.endswith("_ch-all") for name in result["descriptor_names"]) + assert agg.X.shape == (2, result["X"].shape[1]) + assert agg.dims == ("obs", "feature") + + +def test_aggregate_descriptor_pipeline_preserves_channel_group_tokens(): + X = _make_signal_data() + result = DescriptorPipeline( + { + "output": {"channel_pooling": {"Frontal": ["Fz", "Cz"]}}, + "families": {"bands": {"enabled": True, "outputs": ["absolute_power"]}}, + } + ).extract(X=X, sfreq=256.0, channel_names=["Fz", "Cz"]) + agg = _descriptor_result_container(result).aggregate( + by=["s1", "s1", "s1", "s2", "s2", "s2"], + stats=["mean", "std"], + ) + + assert any(name.endswith("_chgrp-Frontal") for name in result["descriptor_names"]) + assert agg.dims == ("obs", "stat", "feature") + assert agg.coords["stat"].tolist() == ["mean", "std"] def test_unstack_basic(): @@ -432,3 +577,245 @@ def test_unstack_error_dim_not_found(): container = DataContainer(X=X, dims=("a", "b")) with pytest.raises(ValueError, match="Dimension 'c' not found"): container.unstack("c") + + +def test_aggregate_validation_errors(sample_container): + """Test validation of min_count, on_insufficient, and empty stats.""" + with pytest.raises(ValueError, match="`min_count` must be at least 1"): + sample_container.aggregate(by="group", min_count=0) + with pytest.raises(ValueError, match="`on_insufficient` must be one of"): + sample_container.aggregate(by="group", on_insufficient="invalid") + with pytest.raises(ValueError, match="`stats` must not be empty"): + sample_container.aggregate(by="group", stats=[]) + + +def test_aggregate_by_y(sample_container): + """Verify grouping using the target vector 'y'.""" + # sample_container.y is [0, 1] + agg = sample_container.aggregate(by="y", stats="mean") + assert agg.shape[0] == 2 + assert np.array_equal(agg.coords["obs"], [0, 1]) + + +def test_aggregate_obs_idx_not_zero(): + """Verify aggregation when 'obs' is not at the first axis.""" + X = np.zeros((3, 5, 10)) + # obs at index 1 + dc = DataContainer(X, dims=("channel", "obs", "time")) + # 5 observations grouped into 3: [G0, G0, G1, G1, G2] + agg = dc.aggregate(by=[0, 0, 1, 1, 2], stats="mean") + assert agg.dims == ("channel", "obs", "time") + assert agg.X.shape == (3, 3, 10) + + +def test_aggregate_all_stats(sample_container): + """Verify all supported statistical measures and aliases.""" + # Using legacy aliases to ensure normalization + stats = [ + "obs-mean", + "median", + "std", + "var", + "sem", + "min", + "max", + "first", + "count", + ] + # Reduce all observations to 1 group + agg = sample_container.aggregate(by=[0, 0], stats=stats) + assert agg.dims == ("obs", "stat", "channel", "time") + assert agg.coords["stat"].tolist() == [ + "mean", + "median", + "std", + "var", + "sem", + "min", + "max", + "first", + "count", + ] + + +def test_aggregate_1d_data(): + """Test aggregation of 1D data (only obs dimension).""" + X = np.array([1.0, 2.0, 3.0, 4.0]) + dc = DataContainer(X, dims=("obs",)) + agg = dc.aggregate(by=["A", "A", "B", "B"], stats="mean") + assert agg.dims == ("obs",) + assert agg.X.shape == (2,) + assert np.allclose(agg.X, [1.5, 3.5]) + + +def test_aggregate_insufficient_policy_raise(): + """Verify 'raise' policy for insufficient observations.""" + X = np.ones((1, 1)) + dc = DataContainer(X, dims=("obs", "f")) + with pytest.raises(ValueError, match="has 1 valid rows, requires at least 2"): + dc.aggregate(by=["G"], min_count=2, on_insufficient="raise") + + +def test_aggregate_y_inconsistency(): + """Verify 'y' is dropped if it varies within a group.""" + X = np.ones((2, 1)) + y = np.array([0, 1]) + dc = DataContainer(X, dims=("obs", "f"), y=y) + # Group [0, 1] has inconsistent y + agg = dc.aggregate(by=["Group", "Group"]) + assert agg.y is None + + +def test_normalization_inplace(sample_container): + """Verify in-place variants for center, zscore, and rms_scale.""" + # sample_container uses int X, must cast to float for subtract/div + sample_container.X = sample_container.X.astype(float) + + # center + dc_c = sample_container.center(dim="time") + assert not np.array_equal(dc_c.X, sample_container.X) + + import copy + + dc_c_in = copy.deepcopy(sample_container) + dc_c_in.center(dim="time", inplace=True) + assert np.allclose(np.nanmean(dc_c_in.X, axis=2), 0) + + # zscore + dc_z_in = copy.deepcopy(sample_container) + dc_z_in.zscore(dim="time", inplace=True) + assert np.allclose(np.nanstd(dc_z_in.X, axis=2), 1) + + # rms_scale + dc_rms_in = copy.deepcopy(sample_container) + dc_rms_in.rms_scale(dim="time", inplace=True) + # RMS should be 1 + rms = np.sqrt(np.mean(dc_rms_in.X**2, axis=2)) + assert np.allclose(rms, 1) + + +def test_isel_edge_cases(sample_container, caplog): + """Cover untested lines in isel (warnings and errors).""" + # 1. Unknown dim warning + import logging + + with caplog.at_level(logging.WARNING): + subset = sample_container.isel(unknown=[0]) + assert "Dimension unknown not in" in caplog.text + assert subset.shape == sample_container.shape + + # 2. Slicing failure (e.g. out of bounds) + with pytest.raises(IndexError): + sample_container.isel(obs=[10]) + + +def test_balance_complex_edge_cases(data_container_cls): + """Cover untested lines in balance (strata fallback and cleaning).""" + # 1. Target not found + X = np.zeros((2, 1)) + dc = data_container_cls(X, dims=("obs", "f")) + with pytest.raises(ValueError, match="Target 'missing' not found"): + dc.balance(target="missing") + + # 2. Covariate not found + y = np.array([0, 1]) + dc2 = data_container_cls(X, dims=("obs", "f"), y=y) + with pytest.raises(ValueError, match="Covariate 'missing' not found"): + dc2.balance(covariates=["missing"], target="y") + + # 3. Single-class stratum in oversample (fallback path) + y3 = np.array([0, 0, 1]) + s3 = np.array(["A", "A", "B"]) + dc3 = data_container_cls( + X=np.zeros((3, 1)), dims=("obs", "f"), y=y3, coords={"s": s3} + ) + # Group 'B' has only class 1. Undersample would fail, so we oversample. + balanced = dc3.balance(target="y", covariates=["s"], strategy="oversample") + assert balanced.shape[0] > 0 + + +def test_select_conflicting_selections(sample_container): + """Verify that conflicting selections on same axis raise ValueError.""" + # time=1 AND time=2 -> empty set + with pytest.raises(ValueError, match="resulted in empty set"): + sample_container.select(time=1).select(time=2) + + +def test_select_aux_coord_no_match(data_container_cls, caplog): + """Test selection on auxiliary coordinate that matches no dimension.""" + X = np.zeros((5, 10)) + # Aux coord with len 7 (matches neither 5 nor 10) + coords = {"aux": np.arange(7)} + dc = data_container_cls(X, dims=("obs", "feat"), coords=coords) + + import logging + + with caplog.at_level(logging.WARNING): + subset = dc.select(aux=1) + assert "matches no dimension" in caplog.text + assert subset.shape == (5, 10) + + +def test_select_fuzzy_no_match(sample_container, caplog): + """Verify warning when fuzzy matching fails to find candidates.""" + import logging + + with caplog.at_level(logging.WARNING): + # xyz is very far from Fz, Cz, Pz. + with pytest.raises(ValueError, match="resulted in empty set"): + sample_container.select(channel=["xyz"], fuzzy=True) + + assert "No fuzzy match found" in caplog.text + + +def test_select_no_coords(data_container_cls, caplog): + """Test selection on a dimension that has no defined coordinates.""" + X = np.zeros((2, 2)) + dc = data_container_cls(X, dims=("obs", "feat"), coords={}) + + import logging + + with caplog.at_level(logging.WARNING): + subset = dc.select(feat=[0]) + assert "is empty" in caplog.text.lower() + assert subset.shape == (2, 2) + + +def test_select_conflicting_axes(sample_container): + """Verify intersection logic for multiple selections on the same axis (y + obs).""" + # y=0 matches obs index 0. ids='s1' matches obs index 1. Intersection is empty. + with pytest.raises(ValueError, match="Conflicting selections"): + sample_container.select(y=[0], ids=["s1"]) + + +def test_unstack_shape_mismatch(sample_container): + """Verify error when unstacking with corrupted shape metadata.""" + stacked = sample_container.stack(dims=("obs", "time"), new_dim="obs") + # Manually corrupt shapes metadata: 10*10 = 100, but actual obs length is 8 (2*4) + stacked.meta["stacked_shapes"] = (10, 10) + with pytest.raises(ValueError, match="Shape mismatch"): + stacked.unstack("obs") + + +def test_normalization_invalid_dims(sample_container): + """Verify errors for invalid dimensions in zscore and rms_scale.""" + with pytest.raises(ValueError, match="not found"): + sample_container.zscore(dim="invalid") + with pytest.raises(ValueError, match="not found"): + sample_container.rms_scale(dim="invalid") + + +def test_baseline_correction_alias(sample_container): + """Verify that baseline_correction is a functional alias for center.""" + sample_container.X = sample_container.X.astype(float) + dc = sample_container.baseline_correction(dim="time") + assert np.allclose(np.nanmean(dc.X, axis=2), 0) + + +def test_aggregate_empty_feature_dim(): + """Verify valid_row_count calculation for data with no features.""" + X = np.zeros((2, 0)) + dc = DataContainer(X, dims=("obs", "feature")) + # Should not raise even if min_count=1 because valid_row_count will match row_count + agg = dc.aggregate(by=["A", "B"], min_count=1) + assert agg.shape == (2, 0) From 5bbe36200b9b1693cc9ed67989a07d8997e6fe9b Mon Sep 17 00:00:00 2001 From: Hamza Abdelhedi Date: Mon, 23 Mar 2026 16:43:17 -0600 Subject: [PATCH 2/7] Refactor: separate channel pooling from descriptor extraction and add DataContainer.aggregate_groups --- coco_pipe/descriptors/configs.py | 92 +---- coco_pipe/descriptors/core.py | 152 +++++++- coco_pipe/descriptors/extractors/base.py | 90 ++--- .../descriptors/extractors/complexity.py | 22 +- .../descriptors/extractors/parametric.py | 16 +- coco_pipe/descriptors/extractors/spectral.py | 71 ++-- coco_pipe/descriptors/extractors/utils.py | 189 ---------- coco_pipe/descriptors/validation.py | 56 +-- coco_pipe/io/structures.py | 347 ++++++++++++++++- tests/test_descriptors_configs.py | 73 ++-- tests/test_descriptors_core.py | 354 +++++++++++++----- tests/test_descriptors_extractors.py | 213 ++++++----- tests/test_io_structures.py | 277 +++++++++++++- 13 files changed, 1271 insertions(+), 681 deletions(-) delete mode 100644 coco_pipe/descriptors/extractors/utils.py diff --git a/coco_pipe/descriptors/configs.py b/coco_pipe/descriptors/configs.py index 6d0a52e..03943e1 100644 --- a/coco_pipe/descriptors/configs.py +++ b/coco_pipe/descriptors/configs.py @@ -9,7 +9,7 @@ - explicit runtime input requirements - family-specific configs for bands, parametric fitting, and complexity -- output formatting controls +- final output precision control - runtime execution controls These models validate local field structure and family-local constraints. The @@ -32,7 +32,6 @@ "ParametricDescriptorConfig", "ComplexityDescriptorConfig", "DescriptorFamiliesConfig", - "DescriptorOutputConfig", "DescriptorRuntimeConfig", "DescriptorConfig", ] @@ -48,9 +47,11 @@ _BAND_OUTPUTS = ( "absolute_power", + "log_absolute_power", "relative_power", "ratios", "corrected_absolute_power", + "corrected_log_absolute_power", "corrected_relative_power", "corrected_ratios", ) @@ -108,14 +109,17 @@ class BandDescriptorConfig(_StrictConfigModel): Global frequency window within which PSDs and bands are evaluated. bands : dict of str to tuple of float, default=canonical EEG bands Mapping from band name to ``(low, high)`` boundaries. - outputs : list of {"absolute_power", "relative_power", "ratios", \ -"corrected_absolute_power", "corrected_relative_power", "corrected_ratios"} + outputs : list of {"absolute_power", "log_absolute_power", \ +"relative_power", "ratios", "corrected_absolute_power", \ +"corrected_log_absolute_power", "corrected_relative_power", \ +"corrected_ratios"} Band descriptors to emit. ratio_pairs : list of tuple of str, default=[] Explicit numerator and denominator band names for ratio outputs. - log_power : bool, default=False - Whether to emit log-transformed absolute band power in addition to - absolute power when that output is enabled. + min_denominator_power : float, default=0.0 + Minimum denominator power required for relative-power and ratio + outputs. Any denominator at or below this threshold is treated as + undefined and yields ``NaN`` instead of an unstable division result. Notes ----- @@ -135,15 +139,17 @@ class BandDescriptorConfig(_StrictConfigModel): outputs: list[ Literal[ "absolute_power", + "log_absolute_power", "relative_power", "ratios", "corrected_absolute_power", + "corrected_log_absolute_power", "corrected_relative_power", "corrected_ratios", ] ] = Field(default_factory=lambda: ["absolute_power"]) ratio_pairs: list[tuple[str, str]] = Field(default_factory=list) - log_power: bool = False + min_denominator_power: float = Field(0.0, ge=0.0) @field_validator("bands", mode="before") @classmethod @@ -309,70 +315,6 @@ class DescriptorFamiliesConfig(_StrictConfigModel): ) -class DescriptorOutputConfig(_StrictConfigModel): - """ - Controls output precision and descriptor-level channel pooling. - - Parameters - ---------- - precision : {"float32", "float64"}, default="float32" - Output dtype used for the final descriptor matrix. - channel_pooling : {"none", "all"} or dict of str to list of str, default="none" - Descriptor-level channel pooling policy applied after per-channel - descriptors are computed. ``"none"`` keeps one descriptor per sensor, - ``"all"`` averages descriptor values across all sensors, and a mapping - averages descriptor values within each named group while leaving - ungrouped sensors unchanged. - - Notes - ----- - Output config is intentionally small. The descriptors module now returns a - minimal result object, so output controls are limited to matrix precision - and descriptor-level channel pooling. - """ - - precision: Literal["float32", "float64"] = "float32" - channel_pooling: Literal["none", "all"] | dict[str, list[str]] = "none" - - @field_validator("channel_pooling", mode="before") - @classmethod - def _coerce_channel_pooling( - cls, value: Any - ) -> Literal["none", "all"] | dict[str, list[str]]: - if value in (None, {}): - return "none" - if isinstance(value, str): - return value - return { - str(group_name): [str(member) for member in members] - for group_name, members in dict(value).items() - } - - @field_validator("channel_pooling") - @classmethod - def _validate_channel_pooling( - cls, value: Literal["none", "all"] | dict[str, list[str]] - ) -> Literal["none", "all"] | dict[str, list[str]]: - if isinstance(value, str): - if value not in {"none", "all"}: - raise ValueError("channel_pooling must be 'none', 'all', or a mapping.") - return value - for group_name, members in value.items(): - if not group_name: - raise ValueError( - "channel_pooling mapping keys must be non-empty strings." - ) - if not members: - raise ValueError( - f"channel_pooling['{group_name}'] must define at least one channel." - ) - if len(set(members)) != len(members): - raise ValueError( - f"channel_pooling['{group_name}'] must not contain duplicates." - ) - return value - - class DescriptorRuntimeConfig(_StrictConfigModel): """ Runtime execution controls for descriptor extraction. @@ -427,8 +369,8 @@ class DescriptorConfig(_StrictConfigModel): Runtime input requirements for explicit array extraction. families : DescriptorFamiliesConfig Enabled descriptor families and their typed configs. - output : DescriptorOutputConfig - Output precision and formatting settings. + precision : {"float32", "float64"} + Output dtype used for the final descriptor matrix. runtime : DescriptorRuntimeConfig Runtime execution and error-handling settings. @@ -442,5 +384,5 @@ class DescriptorConfig(_StrictConfigModel): input: DescriptorInputConfig = Field(default_factory=DescriptorInputConfig) families: DescriptorFamiliesConfig = Field(default_factory=DescriptorFamiliesConfig) - output: DescriptorOutputConfig = Field(default_factory=DescriptorOutputConfig) + precision: Literal["float32", "float64"] = "float32" runtime: DescriptorRuntimeConfig = Field(default_factory=DescriptorRuntimeConfig) diff --git a/coco_pipe/descriptors/core.py b/coco_pipe/descriptors/core.py index 8ed408b..c839f5f 100644 --- a/coco_pipe/descriptors/core.py +++ b/coco_pipe/descriptors/core.py @@ -310,7 +310,6 @@ def _process_psd_group( X_batch: np.ndarray, sfreq: float, channel_names: list[str] | None, - channel_pooling: str | dict[str, list[str]], ids_batch: np.ndarray | None, runtime, obs_offset: int, @@ -330,8 +329,6 @@ def _process_psd_group( Sampling frequency in Hertz. channel_names : list of str or None Runtime channel labels. - channel_pooling : {"none", "all"} or dict - Descriptor-level channel pooling policy. ids_batch : np.ndarray or None Observation identifiers aligned with ``X_batch``. runtime : DescriptorRuntimeConfig @@ -393,7 +390,6 @@ def _process_psd_group( psds, freqs, channel_names=channel_names, - channel_pooling=channel_pooling, ids=ids_batch, runtime=consumer_runtime, obs_offset=obs_offset, @@ -407,7 +403,6 @@ def _process_psd_group( psds, freqs, channel_names=channel_names, - channel_pooling=channel_pooling, ids=ids_batch, runtime=consumer_runtime, obs_offset=obs_offset, @@ -423,7 +418,6 @@ def _process_batch( X: np.ndarray, sfreq: float | None, channel_names: list[str] | None, - channel_pooling: str | dict[str, list[str]], ids: np.ndarray | None, signal_extractors: list[BaseDescriptorExtractor], psd_groups: list[_PSDGroup], @@ -443,8 +437,6 @@ def _process_batch( Sampling frequency in Hertz. channel_names : list of str or None Runtime channel labels. - channel_pooling : {"none", "all"} or dict - Descriptor-level channel pooling policy. ids : np.ndarray or None Observation identifiers aligned with ``X``. signal_extractors : list of BaseDescriptorExtractor @@ -475,7 +467,6 @@ def _signal_unit(extractor): X_batch, sfreq=sfreq, channel_names=channel_names, - channel_pooling=channel_pooling, ids=ids_batch, runtime=_sequential_runtime(runtime), obs_offset=obs_offset, @@ -488,7 +479,6 @@ def _psd_unit(group): X_batch, sfreq=sfreq, channel_names=channel_names, - channel_pooling=channel_pooling, ids_batch=ids_batch, runtime=_sequential_runtime(runtime), obs_offset=obs_offset, @@ -514,7 +504,6 @@ def _psd_unit(group): X_batch, sfreq=sfreq, channel_names=channel_names, - channel_pooling=channel_pooling, ids=ids_batch, runtime=signal_runtime, obs_offset=obs_offset, @@ -536,7 +525,6 @@ def _psd_unit(group): X_batch, sfreq=sfreq, channel_names=channel_names, - channel_pooling=channel_pooling, ids_batch=ids_batch, runtime=runtime if strategy == "parametric-inner" and group.needs_parametric_fit @@ -606,6 +594,7 @@ def __init__(self, config: DescriptorConfig | Mapping[str, Any]): ) corrected_outputs = { "corrected_absolute_power", + "corrected_log_absolute_power", "corrected_relative_power", "corrected_ratios", } @@ -740,7 +729,6 @@ def extract( X=inputs["X"], sfreq=inputs["sfreq"], channel_names=inputs["channel_names"], - channel_pooling=inputs["channel_pooling"], ids=inputs["ids"], signal_extractors=self.signal_extractors, psd_groups=self.psd_groups, @@ -761,7 +749,6 @@ def extract( X=inputs["X"], sfreq=inputs["sfreq"], channel_names=inputs["channel_names"], - channel_pooling=inputs["channel_pooling"], ids=inputs["ids"], signal_extractors=self.signal_extractors, psd_groups=self.psd_groups, @@ -780,7 +767,7 @@ def extract( X_desc, descriptor_names, failures = _merge_descriptor_blocks( blocks, n_obs=inputs["X"].shape[0], - precision=self.config.output.precision, + precision=self.config.precision, ) if self.config.runtime.on_error == "warn" and failures: @@ -794,3 +781,138 @@ def extract( "descriptor_names": descriptor_names, "failures": failures, } + + def pool_channels( + self, + result: Mapping[str, Any], + channel_groups: Mapping[str, Sequence[str]], + ) -> dict[str, Any]: + """Pool sensor-level descriptor columns into grouped channel outputs. + + Parameters + ---------- + result : mapping + Standard descriptor result produced by :meth:`extract`. + channel_groups : mapping of str to sequence of str + Channel groups used to replace sensor-level descriptor columns with + grouped ``"chgrp-..."`` outputs. + + Returns + ------- + dict[str, Any] + Descriptor result with grouped channel features and unchanged + failures. + + Raises + ------ + ValueError + If the provided result is malformed or if any requested group + cannot be formed from the sensor-level descriptor columns. + """ + if ( + "X" not in result + or "descriptor_names" not in result + or "failures" not in result + ): + raise ValueError( + "pool_channels() expects a result mapping with keys " + "'X', 'descriptor_names', and 'failures'." + ) + + X_desc = np.asarray(result["X"], dtype=float) + descriptor_names = [str(name) for name in result["descriptor_names"]] + if X_desc.ndim != 2: + raise ValueError("pool_channels() expects result['X'] to be 2D.") + if X_desc.shape[1] != len(descriptor_names): + raise ValueError( + "pool_channels() requires result['descriptor_names'] to align with " + "result['X'] columns." + ) + if not channel_groups: + raise ValueError("channel_groups must define at least one group.") + + base_to_channel_cols: dict[str, dict[str, int]] = {} + known_channels: set[str] = set() + for col_idx, descriptor_name in enumerate(descriptor_names): + if "_ch-" not in descriptor_name: + continue + base_name, channel_name = descriptor_name.rsplit("_ch-", 1) + base_to_channel_cols.setdefault(base_name, {})[channel_name] = col_idx + known_channels.add(channel_name) + + if not base_to_channel_cols: + raise ValueError( + "pool_channels() requires sensor-level descriptor names " + "ending in '_ch-'." + ) + + normalized_groups: dict[str, list[str]] = {} + assigned: dict[str, str] = {} + for raw_group_name, raw_members in channel_groups.items(): + group_name = str(raw_group_name) + if not group_name: + raise ValueError("channel_groups keys must be non-empty strings.") + members = [str(member) for member in raw_members] + if not members: + raise ValueError( + f"channel_groups['{group_name}'] must define at least one channel." + ) + if len(set(members)) != len(members): + raise ValueError( + f"channel_groups['{group_name}'] must not contain duplicates." + ) + for member in members: + if member not in known_channels: + raise ValueError( + f"channel_groups['{group_name}'] references unknown channel " + f"'{member}'." + ) + if member in assigned: + raise ValueError( + f"Channel '{member}' is assigned to multiple channel_groups: " + f"'{assigned[member]}' and '{group_name}'." + ) + assigned[member] = group_name + normalized_groups[group_name] = members + + output_columns: list[np.ndarray] = [] + pooled_names: list[str] = [] + seen_bases: set[str] = set() + for col_idx, descriptor_name in enumerate(descriptor_names): + if "_ch-" not in descriptor_name: + output_columns.append(X_desc[:, col_idx][:, None]) + pooled_names.append(descriptor_name) + continue + + base_name, _ = descriptor_name.rsplit("_ch-", 1) + if base_name in seen_bases: + continue + seen_bases.add(base_name) + + channel_to_col = base_to_channel_cols[base_name] + for group_name, members in normalized_groups.items(): + missing = [member for member in members if member not in channel_to_col] + if missing: + raise ValueError( + "pool_channels() could not form group " + f"'{group_name}' for descriptor base '{base_name}'. " + f"Missing channels: {missing}." + ) + member_indices = [channel_to_col[member] for member in members] + with warnings.catch_warnings(): + warnings.simplefilter("ignore", category=RuntimeWarning) + grouped = np.nanmean(X_desc[:, member_indices], axis=1) + output_columns.append(grouped[:, None]) + pooled_names.append(f"{base_name}_chgrp-{group_name}") + + X_pooled = _cast_precision( + np.concatenate(output_columns, axis=1) + if output_columns + else np.empty((X_desc.shape[0], 0), dtype=float), + self.config.precision, + ) + return { + "X": X_pooled, + "descriptor_names": pooled_names, + "failures": list(result["failures"]), + } diff --git a/coco_pipe/descriptors/extractors/base.py b/coco_pipe/descriptors/extractors/base.py index 6ccefa6..0b12df7 100644 --- a/coco_pipe/descriptors/extractors/base.py +++ b/coco_pipe/descriptors/extractors/base.py @@ -8,11 +8,12 @@ batches - `BasePSDDescriptorExtractor` for families that consume shared PSD batches - `_DescriptorBlock` as the private family output payload +- `make_failure_record` as the shared normalized failure-record helper The surrounding descriptors stack uses these interfaces to provide: - explicit runtime dispatch from `DescriptorPipeline` -- deterministic descriptor naming and channel reduction helpers +- deterministic sensor-level descriptor naming - family-wise metadata and failure collection - safe merging of family outputs into one stable result dictionary @@ -22,19 +23,6 @@ families. Unlike dim-reduction reducers, descriptor extractors are stateless at runtime and do not expose `fit`, persistence, or model objects. -Examples --------- -The shared finalization helper converts per-channel descriptor values into the -public column naming convention based on ``output.channel_pooling``: - -- ``channel_pooling="none"``: - ``band_abs_alpha_ch-Fz``, ``band_abs_alpha_ch-Cz`` -- ``channel_pooling="all"``: - ``band_abs_alpha_ch-all`` -- ``channel_pooling={"Frontal": ["Fz", "Cz"]}``: - ``band_abs_alpha_chgrp-Frontal`` plus any ungrouped channels such as - ``band_abs_alpha_ch-Pz`` - Author: Hamza Abdelhedi (hamza.abdelhedi@umontreal.ca) """ @@ -47,9 +35,12 @@ import numpy as np from ..configs import DescriptorRuntimeConfig -from .utils import pool_channel_descriptor_matrix -__all__ = ["BaseDescriptorExtractor", "BasePSDDescriptorExtractor"] +__all__ = [ + "BaseDescriptorExtractor", + "BasePSDDescriptorExtractor", + "make_failure_record", +] @dataclass(slots=True) @@ -84,13 +75,34 @@ class _DescriptorBlock: failures: list[dict[str, Any]] = field(default_factory=list) +def make_failure_record( + family: str, + obs_index: int, + obs_id: Any = None, + channel_index: int | None = None, + channel_name: str | None = None, + exception_type: str | None = None, + message: str | None = None, +) -> dict[str, Any]: + """Create one normalized extractor failure record.""" + return { + "family": family, + "obs_index": obs_index, + "obs_id": obs_id, + "channel_index": channel_index, + "channel_name": channel_name, + "exception_type": exception_type, + "message": message, + } + + class BaseDescriptorExtractor(ABC): """ Abstract base class for descriptor extraction families. Subclasses receive already validated NumPy inputs and must return one `_DescriptorBlock` aligned on the observation axis. The base class keeps - the extractor API narrow and provides a shared helper for channel + the extractor API narrow and provides a shared helper for sensor-level finalization and deterministic descriptor naming. Parameters @@ -129,7 +141,6 @@ class BaseDescriptorExtractor(ABC): ... X, ... sfreq, ... channel_names, - ... channel_pooling, ... ids, ... runtime, ... ): @@ -139,7 +150,6 @@ class BaseDescriptorExtractor(ABC): ... family_prefix="toy", ... metric_name="mean", ... channel_names=channel_names, - ... channel_pooling=channel_pooling, ... ) ... return _DescriptorBlock( ... family=self.family_name, @@ -184,7 +194,6 @@ def extract( X: np.ndarray, sfreq: float | None, channel_names: list[str] | None, - channel_pooling: str | dict[str, list[str]], ids: np.ndarray | None, runtime: DescriptorRuntimeConfig, obs_offset: int = 0, @@ -199,9 +208,6 @@ def extract( Sampling frequency in Hertz. channel_names : list of str, optional Explicit channel labels aligned with axis 1 of ``X``. - channel_pooling : {"none", "all"} or dict - Descriptor-level channel pooling policy applied after per-channel - descriptors are computed. ids : np.ndarray, optional Observation identifiers aligned with axis 0 of ``X``. runtime : DescriptorRuntimeConfig @@ -225,8 +231,8 @@ def extract( Notes ----- The recommended pattern is to keep family-specific computation local to - the extractor and delegate all public channel naming and channel pooling - behavior to :meth:`_finalize_descriptor`. + the extractor and delegate sensor-level naming behavior to + :meth:`_finalize_descriptor`. """ def _finalize_descriptor( @@ -235,9 +241,8 @@ def _finalize_descriptor( family_prefix: str, metric_name: str, channel_names: list[str] | None, - channel_pooling: str | dict[str, list[str]] = "none", ) -> tuple[np.ndarray, list[str]]: - """Pool channels and build deterministic descriptor names. + """Build deterministic sensor-level descriptor names. Parameters ---------- @@ -250,8 +255,6 @@ def _finalize_descriptor( Family-local metric identifier used in the descriptor name. channel_names : list of str, optional Channel labels used when building channel-resolved descriptor names. - channel_pooling : {"none", "all"} or dict, default="none" - Descriptor-level channel pooling policy. Returns ------- @@ -262,34 +265,26 @@ def _finalize_descriptor( Notes ----- This helper assumes ``values`` already represents descriptor values, not - raw signals. Pooling therefore always happens at the descriptor level: - - - ``"none"`` keeps one column per sensor - - ``"all"`` averages descriptor values across all sensors - - a mapping averages descriptor values within each named group and keeps - ungrouped sensors as individual columns + raw signals. It therefore only handles the stable sensor-level naming + convention used by the public extract result. Examples -------- Given ``channel_names=["Fz", "Cz", "Pz"]`` and ``metric_name="abs_alpha"``: - - ``channel_pooling="none"`` yields + - yields ``["band_abs_alpha_ch-Fz", "band_abs_alpha_ch-Cz", "band_abs_alpha_ch-Pz"]`` - - ``channel_pooling="all"`` yields - ``["band_abs_alpha_ch-all"]`` - - ``channel_pooling={"Frontal": ["Fz", "Cz"]}`` yields - ``["band_abs_alpha_chgrp-Frontal", "band_abs_alpha_ch-Pz"]`` """ if values.ndim == 1: values = values[:, None] - pooled_values, scopes = pool_channel_descriptor_matrix( - values, - channel_names=channel_names or [], - channel_pooling=channel_pooling, - ) + channel_names = channel_names or [f"ch-{idx}" for idx in range(values.shape[1])] + scopes = [ + channel_name if channel_name.startswith("ch-") else f"ch-{channel_name}" + for channel_name in channel_names + ] names = ["_".join((family_prefix, metric_name, scope)) for scope in scopes] - return pooled_values, names + return values, names class BasePSDDescriptorExtractor(BaseDescriptorExtractor): @@ -350,7 +345,6 @@ def extract_psd( psds: np.ndarray, freqs: np.ndarray, channel_names: list[str] | None, - channel_pooling: str | dict[str, list[str]], ids: np.ndarray | None, runtime: DescriptorRuntimeConfig, obs_offset: int = 0, @@ -366,8 +360,6 @@ def extract_psd( Frequency grid aligned with the last axis of ``psds``. channel_names : list of str, optional Explicit channel labels aligned with the channel axis. - channel_pooling : {"none", "all"} or dict - Descriptor-level channel pooling policy. ids : np.ndarray, optional Observation identifiers aligned with the observation axis. runtime : DescriptorRuntimeConfig diff --git a/coco_pipe/descriptors/extractors/complexity.py b/coco_pipe/descriptors/extractors/complexity.py index c0a7fd7..6267ebc 100644 --- a/coco_pipe/descriptors/extractors/complexity.py +++ b/coco_pipe/descriptors/extractors/complexity.py @@ -27,8 +27,7 @@ from ...utils import import_optional_dependency from ..configs import ComplexityDescriptorConfig -from .base import BaseDescriptorExtractor, _DescriptorBlock -from .utils import make_failure_record +from .base import BaseDescriptorExtractor, _DescriptorBlock, make_failure_record class ComplexityDescriptorExtractor(BaseDescriptorExtractor): @@ -56,8 +55,8 @@ class ComplexityDescriptorExtractor(BaseDescriptorExtractor): Notes ----- The extractor always computes descriptor values per sensor first. Public - output pooling, such as `channel_pooling="all"` or grouped channel pooling, - is applied afterward through :meth:`BaseDescriptorExtractor._finalize_descriptor`. + deterministic sensor-level naming is applied afterward through + :meth:`BaseDescriptorExtractor._finalize_descriptor`. When `antropy` is selected, the extractor uses batched calls where the backend supports them and falls back to scalar loops for measures that are @@ -136,7 +135,6 @@ def extract( X: np.ndarray, sfreq: float | None, channel_names: list[str] | None, - channel_pooling: str | dict[str, list[str]], ids: np.ndarray | None, runtime, obs_offset: int = 0, @@ -154,9 +152,6 @@ def extract( channel_names : list of str, optional Explicit channel labels aligned with axis 1 of ``X``. If omitted, fallback names ``"ch-0"``, ``"ch-1"``, ... are used internally. - channel_pooling : {"none", "all"} or dict - Descriptor-level channel pooling policy applied after per-sensor - complexity values are computed. ids : np.ndarray, optional Observation identifiers aligned with axis 0 of ``X``. runtime : DescriptorRuntimeConfig @@ -193,16 +188,12 @@ def extract( ``failures`` unless `runtime.on_error == "raise"`, in which case the extractor fails immediately. - Examples - -------- - With ``channel_pooling="none"`` and - ``channel_names=["Fz", "Cz"]``, a requested measure such as + Example + ------- + With ``channel_names=["Fz", "Cz"]``, a requested measure such as ``perm_entropy`` yields channel-resolved names like ``complexity_perm_entropy_ch-Fz`` and ``complexity_perm_entropy_ch-Cz``. - - With ``channel_pooling="all"``, the same metric yields one pooled - column named ``complexity_perm_entropy_ch-all``. """ channel_names = channel_names or [f"ch-{idx}" for idx in range(X.shape[1])] @@ -382,7 +373,6 @@ def extract( family_prefix="complexity", metric_name=measure, channel_names=channel_names, - channel_pooling=channel_pooling, ) chunk_features.append(feature) chunk_names.extend(names) diff --git a/coco_pipe/descriptors/extractors/parametric.py b/coco_pipe/descriptors/extractors/parametric.py index 8bfaa60..25369c1 100644 --- a/coco_pipe/descriptors/extractors/parametric.py +++ b/coco_pipe/descriptors/extractors/parametric.py @@ -29,8 +29,7 @@ from ..configs import ParametricDescriptorConfig from ._parametric_fit import _ParametricFitBatch, fit_parametric_batch from ._psd import compute_psd -from .base import BasePSDDescriptorExtractor, _DescriptorBlock -from .utils import make_failure_record +from .base import BasePSDDescriptorExtractor, _DescriptorBlock, make_failure_record class ParametricDescriptorExtractor(BasePSDDescriptorExtractor): @@ -57,8 +56,7 @@ class ParametricDescriptorExtractor(BasePSDDescriptorExtractor): Notes ----- The extractor always computes descriptor values per sensor first. Public - output pooling, such as `channel_pooling="all"` or grouped channel pooling, - is applied afterward through + sensor-level naming is applied afterward through :meth:`BaseDescriptorExtractor._finalize_descriptor`. When the pipeline provides a precomputed PSD batch through @@ -111,7 +109,6 @@ def extract_psd( psds: np.ndarray, freqs: np.ndarray, channel_names: list[str] | None, - channel_pooling: str | dict[str, list[str]], ids: np.ndarray | None, runtime, obs_offset: int = 0, @@ -130,9 +127,6 @@ def extract_psd( Explicit channel labels aligned with axis 1 of ``psds``. If omitted, fallback names ``"ch-0"``, ``"ch-1"``, ... are used internally. - channel_pooling : {"none", "all"} or dict - Descriptor-level channel pooling policy applied after per-sensor - parametric values are computed. ids : np.ndarray, optional Observation identifiers aligned with axis 0 of ``psds``. runtime : DescriptorRuntimeConfig @@ -195,7 +189,6 @@ def extract_psd( family_prefix="param", metric_name=metric_name, channel_names=channel_names, - channel_pooling=channel_pooling, ) chunk_features.append(feature) descriptor_names.extend(names) @@ -218,7 +211,6 @@ def extract( X: np.ndarray, sfreq: float | None, channel_names: list[str] | None, - channel_pooling: str | dict[str, list[str]], ids: np.ndarray | None, runtime, obs_offset: int = 0, @@ -234,9 +226,6 @@ def extract( Sampling frequency in Hertz. channel_names : list of str, optional Explicit channel labels aligned with axis 1 of ``X``. - channel_pooling : {"none", "all"} or dict - Descriptor-level channel pooling policy applied after per-sensor - parametric values are computed. ids : np.ndarray, optional Observation identifiers aligned with axis 0 of ``X``. runtime : DescriptorRuntimeConfig @@ -288,7 +277,6 @@ def extract( psds, freqs, channel_names=channel_names, - channel_pooling=channel_pooling, ids=ids, runtime=runtime, obs_offset=obs_offset, diff --git a/coco_pipe/descriptors/extractors/spectral.py b/coco_pipe/descriptors/extractors/spectral.py index 7a15130..10e67c2 100644 --- a/coco_pipe/descriptors/extractors/spectral.py +++ b/coco_pipe/descriptors/extractors/spectral.py @@ -17,10 +17,11 @@ reuses them for all requested outputs: - absolute power -- optional log absolute power +- log absolute power - relative power - band ratios - corrected absolute power +- corrected log absolute power - corrected relative power - corrected band ratios @@ -40,8 +41,7 @@ from ..configs import BandDescriptorConfig from ._parametric_fit import _ParametricFitBatch from ._psd import compute_psd -from .base import BasePSDDescriptorExtractor, _DescriptorBlock -from .utils import make_failure_record +from .base import BasePSDDescriptorExtractor, _DescriptorBlock, make_failure_record class BandDescriptorExtractor(BasePSDDescriptorExtractor): @@ -69,8 +69,7 @@ class BandDescriptorExtractor(BasePSDDescriptorExtractor): Notes ----- The extractor always computes descriptor values per sensor first. Public - output pooling, such as `channel_pooling="all"` or grouped channel pooling, - is applied afterward through + sensor-level naming is applied afterward through :meth:`BaseDescriptorExtractor._finalize_descriptor`. When the pipeline provides a precomputed PSD batch through @@ -146,6 +145,7 @@ def needs_parametric_fit(self) -> bool: output in { "corrected_absolute_power", + "corrected_log_absolute_power", "corrected_relative_power", "corrected_ratios", } @@ -166,7 +166,6 @@ def extract_psd( psds: np.ndarray, freqs: np.ndarray, channel_names: list[str] | None, - channel_pooling: str | dict[str, list[str]], ids: np.ndarray | None, runtime, obs_offset: int = 0, @@ -185,9 +184,6 @@ def extract_psd( Explicit channel labels aligned with axis 1 of ``psds``. If omitted, fallback names ``"ch-0"``, ``"ch-1"``, ... are used internally. - channel_pooling : {"none", "all"} or dict - Descriptor-level channel pooling policy applied after per-sensor - band values are computed. ids : np.ndarray, optional Observation identifiers aligned with axis 0 of ``psds``. runtime : DescriptorRuntimeConfig @@ -215,16 +211,15 @@ def extract_psd( The extractor first restricts the incoming PSD to the configured frequency window, then integrates one power value per configured band and sensor. Those band integrals are reused for all enabled outputs, - such as absolute power, relative power, log power, and ratios. + such as absolute power, log absolute power, relative power, and + ratios. Ratios are always derived from absolute band powers, not from + relative or log-transformed outputs. - Examples - -------- - With ``channel_pooling="none"`` and - ``channel_names=["Fz", "Cz"]``, an absolute alpha-band request yields - names such as ``band_abs_alpha_ch-Fz`` and ``band_abs_alpha_ch-Cz``. - - With ``channel_pooling="all"``, the same metric yields one pooled - column named ``band_abs_alpha_ch-all``. + Example + ------- + With ``channel_names=["Fz", "Cz"]``, an absolute alpha-band request + yields names such as ``band_abs_alpha_ch-Fz`` and + ``band_abs_alpha_ch-Cz``. """ channel_names = channel_names or [f"ch-{idx}" for idx in range(psds.shape[1])] eps = np.finfo(float).eps @@ -302,6 +297,7 @@ def append_band_outputs( missing_band_names: set[str], output_prefix: str | None, enabled_absolute_output: str, + enabled_log_output: str, enabled_relative_output: str, enabled_ratio_output: str, relative_message_prefix: str, @@ -309,6 +305,7 @@ def append_band_outputs( failed_pairs_to_skip: set[tuple[int, int]] | None = None, ) -> None: metric_prefix = [] if output_prefix is None else [output_prefix] + denominator_floor = self.config.min_denominator_power if enabled_absolute_output in self.config.outputs: for band_name, values in band_power_dict.items(): @@ -317,24 +314,21 @@ def append_band_outputs( family_prefix="band", metric_name="_".join(metric_prefix + ["abs", band_name]), channel_names=channel_names, - channel_pooling=channel_pooling, ) chunk_features.append(feature) descriptor_names.extend(names) - if self.config.log_power: - log_values = np.log10(np.clip(values, eps, None)) - feature, names = self._finalize_descriptor( - log_values, - family_prefix="band", - metric_name="_".join( - metric_prefix + ["log", "abs", band_name] - ), - channel_names=channel_names, - channel_pooling=channel_pooling, - ) - chunk_features.append(feature) - descriptor_names.extend(names) + if enabled_log_output in self.config.outputs: + for band_name, values in band_power_dict.items(): + log_values = np.log10(np.clip(values, eps, None)) + feature, names = self._finalize_descriptor( + log_values, + family_prefix="band", + metric_name="_".join(metric_prefix + ["log", "abs", band_name]), + channel_names=channel_names, + ) + chunk_features.append(feature) + descriptor_names.extend(names) if enabled_relative_output in self.config.outputs: for band_name, values in band_power_dict.items(): @@ -342,7 +336,7 @@ def append_band_outputs( values, total_power_array, out=np.full_like(values, np.nan, dtype=float), - where=total_power_array > 0, + where=total_power_array > denominator_floor, ) if band_name not in missing_band_names: for obs_rel, ch_idx in np.argwhere(~np.isfinite(relative)): @@ -374,7 +368,6 @@ def append_band_outputs( family_prefix="band", metric_name="_".join(metric_prefix + ["rel", band_name]), channel_names=channel_names, - channel_pooling=channel_pooling, ) chunk_features.append(feature) descriptor_names.extend(names) @@ -389,7 +382,7 @@ def append_band_outputs( np.nan, dtype=float, ), - where=band_power_dict[denominator] > 0, + where=band_power_dict[denominator] > denominator_floor, ) if ( numerator not in missing_band_names @@ -427,7 +420,6 @@ def append_band_outputs( metric_prefix + ["ratio", numerator, denominator] ), channel_names=channel_names, - channel_pooling=channel_pooling, ) chunk_features.append(feature) descriptor_names.extend(names) @@ -491,6 +483,7 @@ def append_band_outputs( missing_bands, None, "absolute_power", + "log_absolute_power", "relative_power", "ratios", "Relative power", @@ -502,6 +495,7 @@ def append_band_outputs( corrected_missing_bands, "corr", "corrected_absolute_power", + "corrected_log_absolute_power", "corrected_relative_power", "corrected_ratios", "Corrected relative power", @@ -537,7 +531,6 @@ def extract( X: np.ndarray, sfreq: float | None, channel_names: list[str] | None, - channel_pooling: str | dict[str, list[str]], ids: np.ndarray | None, runtime, obs_offset: int = 0, @@ -553,9 +546,6 @@ def extract( Sampling frequency in Hertz. channel_names : list of str, optional Explicit channel labels aligned with axis 1 of ``X``. - channel_pooling : {"none", "all"} or dict - Descriptor-level channel pooling policy applied after per-sensor - band values are computed. ids : np.ndarray, optional Observation identifiers aligned with axis 0 of ``X``. runtime : DescriptorRuntimeConfig @@ -595,7 +585,6 @@ def extract( psds, freqs, channel_names=channel_names, - channel_pooling=channel_pooling, ids=ids, runtime=runtime, obs_offset=obs_offset, diff --git a/coco_pipe/descriptors/extractors/utils.py b/coco_pipe/descriptors/extractors/utils.py deleted file mode 100644 index b36b76c..0000000 --- a/coco_pipe/descriptors/extractors/utils.py +++ /dev/null @@ -1,189 +0,0 @@ -""" -Utilities for descriptor extractors. - -This module contains small pure helpers shared across descriptor extractors. -They cover: - -- normalized failure-record creation -- descriptor-level channel pooling after per-channel values are computed - -Notes ------ -Pooling helpers in this module are descriptor-level only: - -- ``"none"`` keeps one descriptor column per input channel -- ``"all"`` averages descriptor values across all channels -- a mapping pools descriptor values within named groups and leaves ungrouped - channels unchanged - -Author: Hamza Abdelhedi (hamza.abdelhedi@umontreal.ca) -""" - -from __future__ import annotations - -from typing import Any - -import numpy as np - - -def make_failure_record( - family: str, - obs_index: int, - obs_id: Any = None, - channel_index: int | None = None, - channel_name: str | None = None, - exception_type: str | None = None, - message: str | None = None, -) -> dict[str, Any]: - """Create one normalized extractor failure record. - - Parameters - ---------- - family : str - Canonical family name that raised or collected the failure. - obs_index : int - Global observation index in the original input array. - obs_id : Any, optional - Optional user-provided observation identifier aligned with - ``obs_index``. - channel_index : int, optional - Channel index associated with the failure. - channel_name : str, optional - Explicit channel label associated with the failure. - exception_type : str - Exception class name or normalized failure type. - message : str - Stable human-readable failure description. - - Returns - ------- - dict[str, Any] - Failure record compatible with ``result["failures"]``. - """ - return { - "family": family, - "obs_index": obs_index, - "obs_id": obs_id, - "channel_index": channel_index, - "channel_name": channel_name, - "exception_type": exception_type, - "message": message, - } - - -def average_channel_matrix(values: np.ndarray) -> np.ndarray: - """Average a ``(n_obs, n_channels)`` descriptor matrix across channels. - - Parameters - ---------- - values : np.ndarray - Descriptor matrix with shape ``(n_obs, n_channels)``. A 1D vector is - returned unchanged. - - Returns - ------- - np.ndarray - Vector with shape ``(n_obs,)`` containing the NaN-aware mean across the - channel axis for each observation. - - Notes - ----- - Rows containing no finite values yield ``NaN``. - """ - if values.ndim == 1: - return values - out = np.empty(values.shape[0], dtype=float) - for idx, row in enumerate(values): - finite = row[np.isfinite(row)] - out[idx] = np.nan if finite.size == 0 else float(finite.mean()) - return out - - -def pool_channel_descriptor_matrix( - values: np.ndarray, - channel_names: list[str], - channel_pooling: str | dict[str, list[str]], -) -> tuple[np.ndarray, list[str]]: - """Pool per-channel descriptor values into the public output layout. - - Parameters - ---------- - values : np.ndarray - Descriptor matrix with shape ``(n_obs, n_channels)``. - channel_names : list of str - Channel labels aligned with the columns of ``values``. - channel_pooling : {"none", "all"} or dict of str to list of str - Pooling specification coming from ``output.channel_pooling``. - - Returns - ------- - tuple - ``(X_pooled, scopes)`` where ``X_pooled`` is the pooled descriptor - matrix and ``scopes`` is the aligned list of channel-scope tokens used - by the extractor base class to build final descriptor names. - - Raises - ------ - ValueError - If ``values`` is not 2D. - - Notes - ----- - Grouped outputs preserve first-sensor order: the first channel belonging to - a group determines where that group appears in the output column order. - Channels not assigned to a group remain as standalone outputs. - Channel labels that already carry the public ``"ch-"`` scope prefix are - preserved as-is instead of receiving a second prefix. - - Examples - -------- - Given ``channel_names=["Fz", "Cz", "Pz"]``: - - - ``channel_pooling="none"`` returns scopes - ``["ch-Fz", "ch-Cz", "ch-Pz"]`` - - ``channel_pooling="all"`` returns scopes ``["ch-all"]`` - - ``channel_pooling={"Frontal": ["Fz", "Cz"]}`` returns scopes - ``["chgrp-Frontal", "ch-Pz"]`` - """ - if values.ndim != 2: - raise ValueError("pool_channel_descriptor_matrix expects a 2D matrix.") - - def channel_scope(channel_name: str) -> str: - return channel_name if channel_name.startswith("ch-") else f"ch-{channel_name}" - - if channel_pooling == "none": - return values, [channel_scope(channel_name) for channel_name in channel_names] - if channel_pooling == "all": - return average_channel_matrix(values)[:, None], ["ch-all"] - - channel_to_index = {name: idx for idx, name in enumerate(channel_names)} - member_to_group = { - member: group_name - for group_name, members in channel_pooling.items() - for member in members - } - - grouped_columns: list[np.ndarray] = [] - scopes: list[str] = [] - emitted_groups: set[str] = set() - - for channel_name in channel_names: - group_name = member_to_group.get(channel_name) - if group_name is None: - grouped_columns.append(values[:, channel_to_index[channel_name]][:, None]) - scopes.append(channel_scope(channel_name)) - continue - if group_name in emitted_groups: - continue - member_indices = [ - channel_to_index[member] for member in channel_pooling[group_name] - ] - grouped_columns.append( - average_channel_matrix(values[:, member_indices])[:, None] - ) - scopes.append(f"chgrp-{group_name}") - emitted_groups.add(group_name) - - if not grouped_columns: - return np.empty((values.shape[0], 0), dtype=float), [] - return np.concatenate(grouped_columns, axis=1), scopes diff --git a/coco_pipe/descriptors/validation.py b/coco_pipe/descriptors/validation.py index 468850c..4324cc9 100644 --- a/coco_pipe/descriptors/validation.py +++ b/coco_pipe/descriptors/validation.py @@ -2,7 +2,7 @@ from __future__ import annotations -from collections.abc import Mapping, Sequence +from collections.abc import Sequence from typing import Any import numpy as np @@ -10,47 +10,6 @@ from .configs import DescriptorConfig -def _normalize_channel_pooling( - channel_pooling: str | Mapping[str, Sequence[str]], - channel_names: list[str] | None, -) -> str | dict[str, list[str]]: - if channel_pooling == "none": - return "none" - if channel_pooling == "all": - return "all" - if channel_names is None: - raise ValueError( - "`channel_names` must be passed explicitly when " - "output.channel_pooling uses named groups." - ) - if len(set(channel_names)) != len(channel_names): - raise ValueError( - "`channel_names` must be unique when output.channel_pooling uses " - "named groups." - ) - - known_channels = set(channel_names) - assigned: dict[str, str] = {} - normalized: dict[str, list[str]] = {} - for group_name, members in channel_pooling.items(): - normalized_members = [str(member) for member in members] - for member in normalized_members: - if member not in known_channels: - raise ValueError( - f"output.channel_pooling['{group_name}'] references unknown " - f"channel '{member}'." - ) - if member in assigned: - raise ValueError( - f"Channel '{member}' is assigned to multiple channel_pooling " - "groups: " - f"'{assigned[member]}' and '{group_name}'." - ) - assigned[member] = group_name - normalized[str(group_name)] = normalized_members - return normalized - - def validate_runtime_inputs( config: DescriptorConfig, *, @@ -109,12 +68,9 @@ def validate_runtime_inputs( raise ValueError("`sfreq` must be positive.") channel_names_out = None - channel_names_required = config.input.require_channel_names or ( - any( - getattr(config.families, family_name).enabled - for family_name in ("bands", "parametric", "complexity") - ) - and config.output.channel_pooling != "all" + channel_names_required = config.input.require_channel_names or any( + getattr(config.families, family_name).enabled + for family_name in ("bands", "parametric", "complexity") ) if channel_names is not None: channel_names_out = [str(name) for name in np.asarray(channel_names).tolist()] @@ -140,9 +96,5 @@ def validate_runtime_inputs( "X": X_arr, "ids": ids_out, "channel_names": channel_names_out, - "channel_pooling": _normalize_channel_pooling( - config.output.channel_pooling, - channel_names_out, - ), "sfreq": sfreq, } diff --git a/coco_pipe/io/structures.py b/coco_pipe/io/structures.py index 6445ced..104b521 100644 --- a/coco_pipe/io/structures.py +++ b/coco_pipe/io/structures.py @@ -177,6 +177,82 @@ def __repr__(self) -> str: f"coords={list(self.coords.keys())}>" ) + def obs_table( + self, + include_ids: bool = False, + id_col: str = "obs_id", + include_y: bool = False, + y_col: str = "y", + include_obs_coord: bool = False, + ) -> pd.DataFrame: + """ + Return one-dimensional coordinates aligned to the observation axis. + + This helper is useful when exporting a row-wise table from a container. + It only materializes metadata that can map cleanly to one row per + observation, skipping coordinates that belong to other axes such as + ``channel``, ``time``, ``feature``, or ``stat``. + + Parameters + ---------- + include_ids : bool, default=False + If True, include ``self.ids`` as the first column. + id_col : str, default="obs_id" + Column name used when exporting ``self.ids``. + include_y : bool, default=False + If True, include ``self.y`` as a column when present. + y_col : str, default="y" + Column name used when exporting ``self.y``. + include_obs_coord : bool, default=False + If True, include ``coords["obs"]`` when present. + + Returns + ------- + pandas.DataFrame + DataFrame containing only one-dimensional observation-aligned + metadata columns. + + Raises + ------ + ValueError + If the container has no ``obs`` dimension, or if ``include_ids`` is + requested when ``self.ids`` is missing. + """ + if "obs" not in self.dims: + raise ValueError("Observation metadata export requires an 'obs' dimension.") + + obs_len = self.X.shape[self.dims.index("obs")] + data: Dict[str, np.ndarray] = {} + + if include_ids: + if self.ids is None: + raise ValueError("`include_ids=True` requires `DataContainer.ids`.") + ids = np.asarray(self.ids, dtype=object) + if ids.ndim != 1 or len(ids) != obs_len: + raise ValueError("`DataContainer.ids` must be 1D and aligned to 'obs'.") + data[id_col] = ids + + if include_obs_coord and "obs" in self.coords: + obs_coord = np.asarray(self.coords["obs"], dtype=object) + if obs_coord.ndim == 1 and len(obs_coord) == obs_len: + data["obs"] = obs_coord + + for key, values in self.coords.items(): + if key == "obs": + continue + arr = np.asarray(values, dtype=object) + if arr.ndim == 1 and len(arr) == obs_len: + data[key] = arr + + if include_y and self.y is not None: + y = np.asarray(self.y, dtype=object) + if y.ndim != 1 or len(y) != obs_len: + raise ValueError("`DataContainer.y` must be 1D and aligned to 'obs'.") + if y_col not in data: + data[y_col] = y + + return pd.DataFrame(data) + def isel(self, **indexers) -> "DataContainer": """ Select data by integer indices on specified dimensions. @@ -1241,8 +1317,9 @@ def aggregate( stats : str or sequence of str, default="mean" Aggregation statistic or ordered list of statistics. Supported tokens are ``"mean"``, ``"median"``, ``"std"``, ``"var"``, - ``"sem"``, ``"min"``, ``"max"``, ``"count"``, and ``"first"``. - Legacy ``"obs-*"`` aliases are accepted and normalized. + ``"sem"``, ``"mad"``, ``"iqr"``, ``"min"``, ``"max"``, + ``"count"``, and ``"first"``. Legacy ``"obs-*"`` aliases are + accepted and normalized. min_count : int, default=1 Minimum number of valid observations required per group. A valid observation is one with at least one finite value across the @@ -1282,6 +1359,8 @@ def aggregate( "obs-std": "std", "obs-var": "var", "obs-sem": "sem", + "obs-mad": "mad", + "obs-iqr": "iqr", "obs-min": "min", "obs-max": "max", "obs-count": "count", @@ -1292,6 +1371,8 @@ def aggregate( "std", "var", "sem", + "mad", + "iqr", "min", "max", "count", @@ -1373,6 +1454,22 @@ def _reduce_group( values_flat = np.nanstd(group_X_flat, axis=0) / np.sqrt( counts_flat.astype(np.float64) ) + elif stat == "mad": + medians_flat = np.nanmedian(group_X_flat, axis=0) + values_flat = np.nanmedian( + np.abs(group_X_flat - medians_flat), + axis=0, + ) + elif stat == "iqr": + values_flat = np.nanpercentile( + group_X_flat, + 75, + axis=0, + ) - np.nanpercentile( + group_X_flat, + 25, + axis=0, + ) elif stat == "min": values_flat = np.nanmin(group_X_flat, axis=0) elif stat == "max": @@ -1534,3 +1631,249 @@ def _failure_record( coords=new_coords, meta=meta, ) + + def aggregate_groups( + self, + by: Union[str, np.ndarray, List[Any]], + groups: Sequence[Dict[str, Any]], + min_count: int = 1, + on_insufficient: str = "raise", + skip_empty: bool = True, + ) -> "DataContainer": + """ + Aggregate selected feature groups with different statistics. + + This is a thin wrapper around :meth:`aggregate` for tabular feature + containers. Each group spec selects a subset of feature columns and + applies one or more stats to that subset. The outputs are concatenated + along the ``feature`` dimension, and each resulting feature name is + prefixed with its stat (for example ``"mean_band_log_abs_alpha"``). + + Parameters + ---------- + by : str or array-like + Group definition for the observation axis. Passed through to + :meth:`aggregate`. + groups : sequence of dict + Ordered group specifications. Each group must provide ``"stats"`` + and may optionally provide include/exclude selectors: + + - ``names`` / ``exclude_names`` + - ``prefixes`` / ``exclude_prefixes`` + - ``suffixes`` / ``exclude_suffixes`` + - ``contains`` / ``exclude_contains`` + - ``regex`` / ``exclude_regex`` + + If a group provides no include selectors, it starts from all + features and then applies exclusions. + min_count : int, default=1 + Minimum number of valid observations required per group. Passed + through to :meth:`aggregate`. + on_insufficient : {"raise", "warn", "collect"}, default="raise" + Policy applied when a group has fewer than ``min_count`` valid + observations. Passed through to :meth:`aggregate`. + skip_empty : bool, default=True + If True, silently skip group specs that match no features. If + False, raise a ``ValueError`` when a group matches nothing. + + Returns + ------- + DataContainer + Aggregated container with dims ``("obs", "feature")`` and + stat-prefixed feature names. + + Raises + ------ + ValueError + If the container lacks a ``feature`` dimension or coord, no groups + are provided, a group spec is invalid, multiple groups would emit + the same output feature name, or no non-empty grouped outputs are + produced. + """ + if "feature" not in self.dims: + raise ValueError("aggregate_groups requires a 'feature' dimension.") + if "feature" not in self.coords: + raise ValueError("aggregate_groups requires a 'feature' coordinate.") + if not groups: + raise ValueError("`groups` must not be empty.") + + feature_names = np.asarray(self.coords["feature"], dtype=object) + feature_axis = self.dims.index("feature") + include_keys = ("names", "prefixes", "suffixes", "contains", "regex") + exclude_keys = tuple(f"exclude_{key}" for key in include_keys) + allowed_group_keys = {"name", "stats", *include_keys, *exclude_keys} + + def _normalize_patterns(value: Any) -> Tuple[str, ...]: + if value is None: + return () + if isinstance(value, str): + return (value,) + return tuple(str(item) for item in value) + + def _selector_mask(spec: Dict[str, Any], *, exclude: bool) -> np.ndarray: + selector_keys = exclude_keys if exclude else include_keys + mask = np.zeros(feature_names.size, dtype=bool) + for key in selector_keys: + patterns = _normalize_patterns(spec.get(key)) + if not patterns: + continue + base_key = key.removeprefix("exclude_") + if base_key == "names": + mask |= np.isin(feature_names.astype(str), patterns) + elif base_key == "prefixes": + mask |= np.array( + [ + any(str(name).startswith(pattern) for pattern in patterns) + for name in feature_names + ], + dtype=bool, + ) + elif base_key == "suffixes": + mask |= np.array( + [ + any(str(name).endswith(pattern) for pattern in patterns) + for name in feature_names + ], + dtype=bool, + ) + elif base_key == "contains": + mask |= np.array( + [ + any(pattern in str(name) for pattern in patterns) + for name in feature_names + ], + dtype=bool, + ) + elif base_key == "regex": + compiled = [re.compile(pattern) for pattern in patterns] + mask |= np.array( + [ + any(pattern.search(str(name)) for pattern in compiled) + for name in feature_names + ], + dtype=bool, + ) + return mask + + combined_parts: List[np.ndarray] = [] + combined_feature_names: List[str] = [] + aggregate_failures: List[Dict[str, Any]] = [] + base_agg: Optional["DataContainer"] = None + + for group_index, group in enumerate(groups): + if not isinstance(group, dict): + raise ValueError("Each entry in `groups` must be a dict.") + + unknown_keys = sorted(set(group) - allowed_group_keys) + if unknown_keys: + raise ValueError( + f"Unknown aggregate_groups keys: {unknown_keys}. " + f"Supported keys are: {sorted(allowed_group_keys)}" + ) + if "stats" not in group: + raise ValueError("Each aggregate_groups spec must include `stats`.") + + stats_spec = group["stats"] + if isinstance(stats_spec, str): + stats_out = [stats_spec] + else: + stats_out = [str(stat) for stat in stats_spec] + if not stats_out: + raise ValueError( + "Each aggregate_groups spec must include at least one stat." + ) + + include_mask = _selector_mask(group, exclude=False) + has_include_selectors = any( + group.get(key) is not None for key in include_keys + ) + if not has_include_selectors: + include_mask = np.ones(feature_names.size, dtype=bool) + exclude_mask = _selector_mask(group, exclude=True) + selected_mask = include_mask & ~exclude_mask + feature_indices = np.flatnonzero(selected_mask) + if feature_indices.size == 0: + if skip_empty: + continue + group_name = group.get("name", f"index {group_index}") + raise ValueError( + f"aggregate_groups spec {group_name!r} matched no features." + ) + + subset = self.isel(feature=feature_indices.tolist()) + for stat in stats_out: + grouped = subset.aggregate( + by=by, + stats=stat, + min_count=min_count, + on_insufficient=on_insufficient, + ) + prefixed_names = [ + f"{stat}_{name}" + for name in np.asarray(grouped.coords["feature"], dtype=object) + ] + duplicate_names = sorted( + set(prefixed_names).intersection(combined_feature_names) + ) + if duplicate_names: + raise ValueError( + "aggregate_groups would emit duplicate feature names: " + f"{duplicate_names}" + ) + + if base_agg is None: + base_agg = grouped + else: + if grouped.dims != base_agg.dims: + raise ValueError( + "aggregate_groups requires all grouped outputs to " + "share the same dimensions." + ) + if not np.array_equal(grouped.ids, base_agg.ids): + raise ValueError( + "aggregate_groups requires all grouped outputs to " + "share the same grouped observation ids." + ) + + combined_parts.append(np.asarray(grouped.X, dtype=np.float64)) + combined_feature_names.extend(prefixed_names) + + failures = grouped.meta.get("aggregate_failures", []) + for failure in failures: + failure_out = deepcopy(failure) + failure_out["aggregate_group_index"] = group_index + if "name" in group: + failure_out["aggregate_group_name"] = group["name"] + failure_out["aggregate_stat"] = stat + aggregate_failures.append(failure_out) + + if base_agg is None: # pragma: no cover - guarded above + raise ValueError("aggregate_groups produced no aggregated outputs.") + + new_coords = deepcopy(base_agg.coords) + new_coords["feature"] = np.asarray(combined_feature_names, dtype=object) + meta = deepcopy(self.meta) + unique_stats = list( + dict.fromkeys(name.split("_", 1)[0] for name in combined_feature_names) + ) + meta.update( + { + "aggregated": True, + "agg_by": by if isinstance(by, str) else None, + "agg_stats": unique_stats, + "agg_groups": deepcopy(list(groups)), + "min_count": int(min_count), + } + ) + if aggregate_failures: + meta["aggregate_failures"] = aggregate_failures + elif "aggregate_failures" in meta: + del meta["aggregate_failures"] + + X_out = np.concatenate(combined_parts, axis=feature_axis) + return replace( + base_agg, + X=X_out, + coords=new_coords, + meta=meta, + ) diff --git a/tests/test_descriptors_configs.py b/tests/test_descriptors_configs.py index 93d413d..249f061 100644 --- a/tests/test_descriptors_configs.py +++ b/tests/test_descriptors_configs.py @@ -53,18 +53,15 @@ def test_corrected_bands_require_parametric_fit_range_to_cover_band_window(): ) -def test_channel_pooling_accepts_none_all_or_mapping(): - assert ( - DescriptorConfig(output={"channel_pooling": "none"}).output.channel_pooling - == "none" - ) - assert ( - DescriptorConfig(output={"channel_pooling": "all"}).output.channel_pooling - == "all" - ) - assert DescriptorConfig( - output={"channel_pooling": {"Frontal": ["Fz", "Cz"]}} - ).output.channel_pooling == {"Frontal": ["Fz", "Cz"]} +def test_channel_pooling_config_field_is_rejected(): + with pytest.raises(ValidationError): + DescriptorConfig(output={"channel_pooling": "none"}) + + with pytest.raises(ValidationError): + DescriptorConfig(output={"channel_pooling": "all"}) + + with pytest.raises(ValidationError): + DescriptorConfig(output={"channel_pooling": {"Frontal": ["Fz", "Cz"]}}) def test_runtime_and_output_flags_parse_strictly(): @@ -74,10 +71,7 @@ def test_runtime_and_output_flags_parse_strictly(): "require_sfreq": False, "require_channel_names": True, }, - "output": { - "precision": "float64", - "channel_pooling": "all", - }, + "precision": "float64", "runtime": { "execution_backend": "joblib", "n_jobs": -1, @@ -89,8 +83,7 @@ def test_runtime_and_output_flags_parse_strictly(): assert config.input.require_sfreq is False assert config.input.require_channel_names is True - assert config.output.precision == "float64" - assert config.output.channel_pooling == "all" + assert config.precision == "float64" assert config.runtime.execution_backend == "joblib" assert config.runtime.n_jobs == -1 assert config.runtime.obs_chunk == 16 @@ -147,6 +140,13 @@ def test_removed_ceremonial_fields_are_rejected(): } ) + with pytest.raises(ValidationError): + DescriptorConfig.model_validate( + { + "families": {"bands": {"log_power": True}}, + } + ) + with pytest.raises(ValidationError): DescriptorConfig.model_validate( { @@ -154,6 +154,13 @@ def test_removed_ceremonial_fields_are_rejected(): } ) + with pytest.raises(ValidationError): + DescriptorConfig.model_validate( + { + "output": {"channel_pooling": "none"}, + } + ) + with pytest.raises(ValidationError): DescriptorConfig.model_validate( { @@ -230,6 +237,11 @@ def test_band_validation_edge_cases(): DescriptorConfig(families={"bands": {"outputs": ["non_existent"]}}) +def test_band_min_denominator_power_must_be_non_negative(): + with pytest.raises(ValidationError, match="min_denominator_power"): + DescriptorConfig(families={"bands": {"min_denominator_power": -1.0}}) + + def test_parametric_validation_edge_cases(): # Duplicate outputs with pytest.raises(ValidationError, match="duplicates"): @@ -263,35 +275,18 @@ def test_complexity_validation_edge_cases(): def test_channel_pooling_validation_edge_cases(): - # Invalid string - with pytest.raises( - ValidationError, - match="Input should be 'none' or 'all'|Input should be a valid dictionary", - ): + with pytest.raises(ValidationError): DescriptorConfig(output={"channel_pooling": "some_string"}) - # Empty group name - with pytest.raises(ValidationError, match="non-empty strings"): + with pytest.raises(ValidationError): DescriptorConfig(output={"channel_pooling": {"": ["ch1"]}}) - # Empty members - with pytest.raises(ValidationError, match="at least one channel"): + with pytest.raises(ValidationError): DescriptorConfig(output={"channel_pooling": {"G1": []}}) - # Duplicate members - with pytest.raises(ValidationError, match="not contain duplicates"): + with pytest.raises(ValidationError): DescriptorConfig(output={"channel_pooling": {"G1": ["ch1", "ch1"]}}) - # Coerce None/{} to none - assert ( - DescriptorConfig(output={"channel_pooling": None}).output.channel_pooling - == "none" - ) - assert ( - DescriptorConfig(output={"channel_pooling": {}}).output.channel_pooling - == "none" - ) - def test_coercion_logic_smoke(): # Coerce bands diff --git a/tests/test_descriptors_core.py b/tests/test_descriptors_core.py index a597531..7ab5eea 100644 --- a/tests/test_descriptors_core.py +++ b/tests/test_descriptors_core.py @@ -1,3 +1,6 @@ +import builtins +import inspect + import numpy as np import pytest @@ -16,12 +19,11 @@ def test_empty_pipeline_returns_explicit_result_structure(): assert result["descriptor_names"] == [] -def test_band_pipeline_smoke(): +def test_band_pipeline_smoke_sensor_level(): rng = np.random.default_rng(1) X = rng.normal(size=(6, 3, 128)) pipe = DescriptorPipeline( { - "output": {"channel_pooling": "all"}, "families": {"bands": {"enabled": True, "outputs": ["absolute_power"]}}, } ) @@ -29,7 +31,51 @@ def test_band_pipeline_smoke(): assert result["X"].shape[0] == 6 assert result["X"].shape[1] == len(result["descriptor_names"]) - assert result["descriptor_names"][0].startswith("band_abs_") + assert "band_abs_alpha_ch-Fz" in result["descriptor_names"] + assert "band_abs_alpha_ch-Cz" in result["descriptor_names"] + assert not any( + name.endswith("chgrp-Frontal") for name in result["descriptor_names"] + ) + + +def test_pool_channels_replaces_sensor_columns_with_grouped_columns(): + rng = np.random.default_rng(11) + X = rng.normal(size=(4, 3, 128)) + pipe = DescriptorPipeline( + { + "families": {"bands": {"enabled": True, "outputs": ["absolute_power"]}}, + } + ) + result = pipe.extract( + X=X, + sfreq=128.0, + channel_names=["Fz", "Cz", "Pz"], + ) + pooled = pipe.pool_channels(result, {"Frontal": ["Fz", "Cz"]}) + + assert "band_abs_alpha_chgrp-Frontal" in pooled["descriptor_names"] + assert "band_abs_alpha_ch-Fz" not in pooled["descriptor_names"] + fz_idx = result["descriptor_names"].index("band_abs_alpha_ch-Fz") + cz_idx = result["descriptor_names"].index("band_abs_alpha_ch-Cz") + grp_idx = pooled["descriptor_names"].index("band_abs_alpha_chgrp-Frontal") + expected = np.nanmean(result["X"][:, [fz_idx, cz_idx]], axis=1) + + assert np.allclose(pooled["X"][:, grp_idx], expected, equal_nan=True) + + +def test_pool_channels_preserves_non_channel_features(): + pipe = DescriptorPipeline({}) + result = { + "X": np.array([[1.0, 2.0, 3.0], [4.0, np.nan, 6.0]], dtype=float), + "descriptor_names": ["global_metric", "toy_mean_ch-Fz", "toy_mean_ch-Cz"], + "failures": [], + } + + pooled = pipe.pool_channels(result, {"Frontal": ["Fz", "Cz"]}) + + assert pooled["descriptor_names"] == ["global_metric", "toy_mean_chgrp-Frontal"] + assert np.allclose(pooled["X"][:, 0], result["X"][:, 0], equal_nan=True) + assert np.allclose(pooled["X"][:, 1], [2.5, 6.0], equal_nan=True) def test_complexity_can_omit_sfreq_when_config_disables_it(): @@ -43,22 +89,18 @@ def test_complexity_can_omit_sfreq_when_config_disables_it(): "measures": ["sample_entropy"], } }, - "output": {"channel_pooling": "all"}, } ) result = pipe.extract(X=X, channel_names=["Fz", "Cz"]) - assert result["X"].shape == (4, 1) + assert result["X"].shape == (4, 2) def test_output_precision_is_respected(): X = np.zeros((2, 2, 128), dtype=float) pipe = DescriptorPipeline( { - "output": { - "channel_pooling": "all", - "precision": "float64", - }, + "precision": "float64", "families": { "parametric": { "enabled": True, @@ -79,7 +121,7 @@ def test_missing_sfreq_is_explicit_error(): {"families": {"bands": {"enabled": True, "outputs": ["absolute_power"]}}} ) with pytest.raises(ValueError, match="`sfreq`"): - pipe.extract(X=X) + pipe.extract(X=X, channel_names=["Fz", "Cz"]) def test_wrong_ndim_is_rejected(): @@ -107,11 +149,10 @@ def test_wrong_channel_names_length_is_rejected(): pipe.extract(X=X, sfreq=128.0, channel_names=["C3"]) -def test_channel_pooling_groups_reject_unknown_channel_names(): +def test_pool_channels_reject_unknown_channel_names(): X = np.random.default_rng(22).normal(size=(4, 2, 64)) pipe = DescriptorPipeline( { - "output": {"channel_pooling": {"Frontal": ["C3", "C4"]}}, "families": { "bands": { "enabled": True, @@ -120,21 +161,16 @@ def test_channel_pooling_groups_reject_unknown_channel_names(): }, } ) + result = pipe.extract(X=X, sfreq=128.0, channel_names=["Fz", "Cz"]) with pytest.raises(ValueError, match="unknown channel"): - pipe.extract(X=X, sfreq=128.0, channel_names=["Fz", "Cz"]) + pipe.pool_channels(result, {"Frontal": ["C3", "C4"]}) -def test_channel_pooling_groups_reject_overlapping_assignments(): +def test_pool_channels_reject_overlapping_assignments(): X = np.random.default_rng(22).normal(size=(4, 3, 64)) pipe = DescriptorPipeline( { - "output": { - "channel_pooling": { - "Frontal": ["Fz", "Cz"], - "Central": ["Cz", "Pz"], - } - }, "families": { "bands": { "enabled": True, @@ -143,9 +179,74 @@ def test_channel_pooling_groups_reject_overlapping_assignments(): }, } ) + result = pipe.extract(X=X, sfreq=128.0, channel_names=["Fz", "Cz", "Pz"]) + + with pytest.raises(ValueError, match="multiple channel_groups"): + pipe.pool_channels( + result, + { + "Frontal": ["Fz", "Cz"], + "Central": ["Cz", "Pz"], + }, + ) + + +def test_pool_channels_reject_non_2d_x(): + pipe = DescriptorPipeline({}) + result = { + "X": np.zeros((2, 2, 2)), + "descriptor_names": ["a", "b"], + "failures": [], + } + with pytest.raises(ValueError, match="2D"): + pipe.pool_channels(result, {"G": ["ch1"]}) - with pytest.raises(ValueError, match="multiple channel_pooling groups"): - pipe.extract(X=X, sfreq=128.0, channel_names=["Fz", "Cz", "Pz"]) + +def test_pool_channels_reject_mismatched_names_and_columns(): + pipe = DescriptorPipeline({}) + result = { + "X": np.zeros((2, 1)), + "descriptor_names": ["a", "b"], + "failures": [], + } + with pytest.raises(ValueError, match=r"align with result\['X'\]"): + pipe.pool_channels(result, {"G": ["ch1"]}) + + +def test_pool_channels_reject_empty_group_definitions(): + pipe = DescriptorPipeline({}) + result = { + "X": np.zeros((2, 2)), + "descriptor_names": ["a_ch-Fz", "b_ch-Cz"], + "failures": [], + } + with pytest.raises(ValueError, match="at least one group"): + pipe.pool_channels(result, {}) + + with pytest.raises(ValueError, match="non-empty strings"): + pipe.pool_channels(result, {"": ["Fz"]}) + + with pytest.raises(ValueError, match="at least one channel"): + pipe.pool_channels(result, {"G": []}) + + with pytest.raises(ValueError, match="not contain duplicates"): + pipe.pool_channels(result, {"G": ["Fz", "Fz"]}) + + +def test_pool_channels_reject_incomplete_grouped_feature_base(): + pipe = DescriptorPipeline({}) + result = { + "X": np.array([[1.0, 3.0, 5.0], [2.0, 4.0, 6.0]], dtype=float), + "descriptor_names": [ + "toy_mean_ch-Fz", + "other_mean_ch-Fz", + "other_mean_ch-Cz", + ], + "failures": [], + } + + with pytest.raises(ValueError, match="could not form group"): + pipe.pool_channels(result, {"Frontal": ["Fz", "Cz"]}) def test_require_channel_names_flag_is_enforced(): @@ -153,7 +254,6 @@ def test_require_channel_names_flag_is_enforced(): pipe = DescriptorPipeline( { "input": {"require_channel_names": True}, - "output": {"channel_pooling": "all"}, "families": { "complexity": { "enabled": True, @@ -171,7 +271,6 @@ def test_complexity_collects_short_segment_failures(): X = np.ones((4, 2, 3), dtype=float) pipe = DescriptorPipeline( { - "output": {"channel_pooling": "all"}, "families": { "complexity": { "enabled": True, @@ -181,9 +280,9 @@ def test_complexity_collects_short_segment_failures(): "runtime": {"on_error": "collect"}, } ) - result = pipe.extract(X=X, sfreq=128.0) + result = pipe.extract(X=X, sfreq=128.0, channel_names=["Fz", "Cz"]) - assert result["X"].shape == (4, 1) + assert result["X"].shape == (4, 2) assert np.isnan(result["X"]).all() assert result["failures"] @@ -192,7 +291,6 @@ def test_bands_collect_short_window_resolution_failures(): X = np.random.default_rng(7).normal(size=(3, 2, 8)) pipe = DescriptorPipeline( { - "output": {"channel_pooling": "all"}, "families": { "bands": { "enabled": True, @@ -205,7 +303,7 @@ def test_bands_collect_short_window_resolution_failures(): result = pipe.extract(X=X, sfreq=160.0, channel_names=["C3", "C4"]) - assert result["X"].shape == (3, 10) + assert result["X"].shape == (3, 20) assert any( failure["exception_type"] == "BandResolutionError" for failure in result["failures"] @@ -216,7 +314,6 @@ def test_warn_policy_emits_aggregate_warning(): X = np.random.default_rng(23).normal(size=(3, 2, 8)) pipe = DescriptorPipeline( { - "output": {"channel_pooling": "all"}, "families": { "bands": { "enabled": True, @@ -237,7 +334,6 @@ def test_raise_policy_reraises_runtime_failure(): X = np.random.default_rng(24).normal(size=(3, 2, 8)) pipe = DescriptorPipeline( { - "output": {"channel_pooling": "all"}, "families": { "bands": { "enabled": True, @@ -253,11 +349,9 @@ def test_raise_policy_reraises_runtime_failure(): def test_complexity_collects_nonfinite_output_as_nan(): - """Verify that real non-finite results are collected as NaNs.""" X = np.ones((2, 2, 16), dtype=float) pipe = DescriptorPipeline( { - "output": {"channel_pooling": "all"}, "families": { "complexity": { "enabled": True, @@ -267,18 +361,16 @@ def test_complexity_collects_nonfinite_output_as_nan(): "runtime": {"on_error": "collect"}, } ) - result = pipe.extract(X=X, sfreq=128.0) + result = pipe.extract(X=X, sfreq=128.0, channel_names=["Fz", "Cz"]) assert np.isnan(result["X"]).all() assert result["failures"] def test_complexity_raise_policy_reraises_nonfinite_output(): - """Verify that real non-finite results reraise when policy is set to raise.""" X = np.ones((2, 2, 16), dtype=float) pipe = DescriptorPipeline( { - "output": {"channel_pooling": "all"}, "families": { "complexity": { "enabled": True, @@ -290,14 +382,13 @@ def test_complexity_raise_policy_reraises_nonfinite_output(): ) with pytest.raises(ValueError, match="non-finite"): - pipe.extract(X=X, sfreq=128.0) + pipe.extract(X=X, sfreq=128.0, channel_names=["Fz", "Cz"]) def test_constant_signal_parametric_skip_collects_failures(): X = np.zeros((3, 2, 128), dtype=float) pipe = DescriptorPipeline( { - "output": {"channel_pooling": "all"}, "families": { "parametric": { "enabled": True, @@ -327,7 +418,6 @@ def _raise_import_error(self): ) pipe = DescriptorPipeline( { - "output": {"channel_pooling": "all"}, "families": { "complexity": { "enabled": True, @@ -347,7 +437,6 @@ def test_multi_family_scale_smoke(): X = rng.normal(size=(24, 4, 256)) pipe = DescriptorPipeline( { - "output": {"channel_pooling": "all"}, "families": { "bands": { "enabled": True, @@ -381,7 +470,6 @@ def test_multi_family_parallel_matches_sequential(): X = rng.normal(size=(12, 3, 128)) channel_names = ["Fz", "Cz", "Pz"] base_config = { - "output": {"channel_pooling": "all"}, "families": { "bands": { "enabled": True, @@ -427,7 +515,6 @@ def test_parametric_parallel_matches_sequential(): "outputs": ["aperiodic", "fit_quality"], } }, - "output": {"channel_pooling": "all"}, "runtime": {"execution_backend": "sequential", "n_jobs": 1}, } ).extract(X=X, sfreq=128.0, channel_names=["Fz", "Cz", "Pz"]) @@ -439,7 +526,6 @@ def test_parametric_parallel_matches_sequential(): "outputs": ["aperiodic", "fit_quality"], } }, - "output": {"channel_pooling": "all"}, "runtime": {"execution_backend": "joblib", "n_jobs": 2}, } ).extract(X=X, sfreq=128.0, channel_names=["Fz", "Cz", "Pz"]) @@ -452,7 +538,6 @@ def test_multi_chunk_row_order_matches_unchunked(): rng = np.random.default_rng(15) X = rng.normal(size=(18, 3, 128)) config = { - "output": {"channel_pooling": "all"}, "families": {"bands": {"enabled": True, "outputs": ["absolute_power"]}}, } unchunked = DescriptorPipeline(config).extract( @@ -476,9 +561,6 @@ def test_multi_chunk_row_order_matches_unchunked(): def test_n_jobs_one_skips_joblib_loading(monkeypatch): - import builtins - import inspect - rng = np.random.default_rng(16) X = rng.normal(size=(4, 2, 64)) real_import = builtins.__import__ @@ -504,27 +586,27 @@ def _count_joblib_imports(name, *args, **kwargs): "outputs": ["absolute_power"], } }, - "output": {"channel_pooling": "all"}, "runtime": {"execution_backend": "joblib", "n_jobs": 1}, } ).extract(X=X, sfreq=128.0, channel_names=["Fz", "Cz"]) - assert result["X"].shape == (4, 5) + assert result["X"].shape == (4, len(result["descriptor_names"])) assert joblib_imports == 0 def test_parametric_parallel_n_jobs_all_cores_smoke(): pytest.importorskip("joblib") rng = np.random.default_rng(17) - # 4 seconds at 128Hz = 512 samples t = np.linspace(0, 4, 512, endpoint=False) X = rng.normal(scale=0.05, size=(4, 3, 512)) - # Add 1/f slope freqs = np.fft.rfftfreq(512, 1 / 128.0) weights = 1 / (freqs + 1.0) - for o in range(4): - for c in range(3): - X[o, c, :] = np.fft.irfft(np.fft.rfft(X[o, c, :]) * weights, n=512) + for obs_idx in range(4): + for ch_idx in range(3): + X[obs_idx, ch_idx, :] = np.fft.irfft( + np.fft.rfft(X[obs_idx, ch_idx, :]) * weights, + n=512, + ) X[:, 0, :] += 2.0 * np.sin(2 * np.pi * 10 * t) X[:, 1, :] += 1.5 * np.sin(2 * np.pi * 16 * t) @@ -538,13 +620,11 @@ def test_parametric_parallel_n_jobs_all_cores_smoke(): "outputs": ["aperiodic"], } }, - "output": {"channel_pooling": "all"}, "runtime": {"execution_backend": "joblib", "n_jobs": -1}, } ).extract(X=X, sfreq=128.0, channel_names=["Fz", "Cz", "Pz"]) - # 2 features (offset, exponent) per observation - assert result["X"].shape == (4, 2) + assert result["X"].shape == (4, 6) def test_shared_psd_reuses_one_compute_per_batch_for_same_method(monkeypatch): @@ -566,7 +646,6 @@ def _counted_compute_psd(*args, **kwargs): DescriptorPipeline( { - "output": {"channel_pooling": "all"}, "families": { "bands": { "enabled": True, @@ -587,9 +666,7 @@ def _counted_compute_psd(*args, **kwargs): assert all(method == "welch" for method, _, _ in calls) -def test_corrected_bands_and_parametric_share_one_fit_batch_per_psd_group( - monkeypatch, -): +def test_corrected_bands_and_parametric_share_one_fit_batch_per_psd_group(monkeypatch): rng = np.random.default_rng(191) t = np.linspace(0, 1, 128, endpoint=False) X = rng.normal(scale=0.05, size=(8, 3, 128)) @@ -609,7 +686,6 @@ def _counted_fit_batch(*args, **kwargs): result = DescriptorPipeline( { - "output": {"channel_pooling": "all"}, "families": { "bands": { "enabled": True, @@ -627,7 +703,7 @@ def _counted_fit_batch(*args, **kwargs): ).extract(X=X, sfreq=128.0, channel_names=["Fz", "Cz", "Pz"]) assert calls == 2 - assert "band_corr_abs_alpha_ch-all" in result["descriptor_names"] + assert "band_corr_abs_alpha_ch-Fz" in result["descriptor_names"] def test_shared_psd_splits_groups_by_method(monkeypatch): @@ -649,7 +725,6 @@ def _counted_compute_psd(*args, **kwargs): DescriptorPipeline( { - "output": {"channel_pooling": "all"}, "families": { "bands": { "enabled": True, @@ -679,7 +754,6 @@ def test_shared_union_psd_matches_separate_family_outputs(): channel_names = ["Fz", "Cz", "Pz"] bands_cfg = { - "output": {"channel_pooling": "all"}, "families": { "bands": { "enabled": True, @@ -697,7 +771,6 @@ def test_shared_union_psd_matches_separate_family_outputs(): }, } param_cfg = { - "output": {"channel_pooling": "all"}, "families": { "parametric": { "enabled": True, @@ -708,7 +781,6 @@ def test_shared_union_psd_matches_separate_family_outputs(): }, } combined_cfg = { - "output": {"channel_pooling": "all"}, "families": { **bands_cfg["families"], **param_cfg["families"], @@ -755,18 +827,18 @@ def test_shared_union_psd_matches_separate_family_outputs(): def test_obs_batch_parallel_disables_parametric_inner_joblib(monkeypatch): - import builtins - import inspect - pytest.importorskip("joblib") rng = np.random.default_rng(22) t = np.linspace(0, 4, 512, endpoint=False) X = rng.normal(scale=0.05, size=(6, 3, 512)) freqs = np.fft.rfftfreq(512, 1 / 128.0) weights = 1 / (freqs + 1.0) - for o in range(6): - for c in range(3): - X[o, c, :] = np.fft.irfft(np.fft.rfft(X[o, c, :]) * weights, n=512) + for obs_idx in range(6): + for ch_idx in range(3): + X[obs_idx, ch_idx, :] = np.fft.irfft( + np.fft.rfft(X[obs_idx, ch_idx, :]) * weights, + n=512, + ) X[:, 0, :] += 2.0 * np.sin(2 * np.pi * 10 * t) X[:, 1, :] += 1.5 * np.sin(2 * np.pi * 18 * t) @@ -794,7 +866,6 @@ def _count_joblib_imports(name, *args, **kwargs): "outputs": ["aperiodic"], } }, - "output": {"channel_pooling": "all"}, "runtime": { "execution_backend": "joblib", "n_jobs": 2, @@ -803,11 +874,18 @@ def _count_joblib_imports(name, *args, **kwargs): } ).extract(X=X, sfreq=128.0, channel_names=["Fz", "Cz", "Pz"]) - # 2 features (offset, exponent) for aperiodic 'fixed' mode - assert result["X"].shape == (X.shape[0], 2) - # Check that exponent is reasonable (> 0) and offset is finite - assert np.all(result["X"][:, 1] > 0) - assert np.all(np.isfinite(result["X"][:, 0])) + exponent_indices = [ + idx for idx, name in enumerate(result["descriptor_names"]) if "exponent" in name + ] + offset_indices = [ + idx for idx, name in enumerate(result["descriptor_names"]) if "offset" in name + ] + + assert exponent_indices + assert offset_indices + assert np.all(result["X"][:, exponent_indices] > 0) + assert np.all(np.isfinite(result["X"][:, offset_indices])) + assert joblib_imports >= 1 def test_single_psd_group_uses_psd_level_n_jobs(monkeypatch): @@ -830,7 +908,6 @@ def _counted_compute_psd(*args, **kwargs): "outputs": ["absolute_power"], } }, - "output": {"channel_pooling": "all"}, "runtime": {"execution_backend": "joblib", "n_jobs": 2}, } ).extract(X=X, sfreq=128.0, channel_names=["Fz", "Cz", "Pz"]) @@ -840,42 +917,123 @@ def _counted_compute_psd(*args, **kwargs): def test_validation_edge_cases_runtime(): from coco_pipe.descriptors.configs import DescriptorConfig - from coco_pipe.descriptors.validation import ( - _normalize_channel_pooling, - validate_runtime_inputs, - ) + from coco_pipe.descriptors.validation import validate_runtime_inputs config = DescriptorConfig(families={"bands": {"enabled": True}}) X = np.zeros((2, 2, 64)) - # sfreq <= 0 with pytest.raises(ValueError, match="`sfreq` must be positive"): validate_runtime_inputs(config, X=X, sfreq=0, channel_names=["ch1", "ch2"]) - # ids alignment failure with pytest.raises(ValueError, match="`ids` must align with n_obs=2"): validate_runtime_inputs( config, X=X, sfreq=100.0, ids=[1, 2, 3], channel_names=["ch1", "ch2"] ) - # channel_names alignment failure with pytest.raises( ValueError, match="`channel_names` must align with n_channels=2" ): validate_runtime_inputs(config, X=X, sfreq=100.0, channel_names=["ch1"]) - # channel_names required but missing (when pooling is not 'all') - config_none = DescriptorConfig( - families={"bands": {"enabled": True}}, output={"channel_pooling": "none"} - ) with pytest.raises(ValueError, match="`channel_names` must be passed explicitly"): - validate_runtime_inputs(config_none, X=X, sfreq=100.0, channel_names=None) + validate_runtime_inputs(config, X=X, sfreq=100.0, channel_names=None) - # _normalize_channel_pooling edge cases - # missing channel_names when groups are used - with pytest.raises(ValueError, match="`channel_names` must be passed explicitly"): - _normalize_channel_pooling({"G1": ["ch1"]}, None) - # duplicate channel_names when groups are used - with pytest.raises(ValueError, match="`channel_names` must be unique"): - _normalize_channel_pooling({"G1": ["ch1"]}, ["ch1", "ch1"]) +def test_pool_channels_requires_standard_result_structure(): + pipe = DescriptorPipeline({}) + + with pytest.raises(ValueError, match="'X', 'descriptor_names', and 'failures'"): + pipe.pool_channels({"X": np.ones((2, 2))}, {"G1": ["Fz"]}) + + +def test_pool_channels_requires_sensor_level_descriptor_names(): + pipe = DescriptorPipeline({}) + result = { + "X": np.ones((2, 1)), + "descriptor_names": ["global_metric"], + "failures": [], + } + + with pytest.raises(ValueError, match="sensor-level descriptor names"): + pipe.pool_channels(result, {"G1": ["Fz"]}) + + +def test_pool_channels_handles_mixture_of_sensor_and_global_features(): + pipe = DescriptorPipeline({}) + result = { + "X": np.array([[1.0, 10.0, 20.0], [2.0, 30.0, 40.0]], dtype=float), + "descriptor_names": ["global", "val_ch-Fz", "val_ch-Cz"], + "failures": [], + } + # Pool Fz, Cz -> 15, 35 + # Result: global=1,2, grouped=15,35 + pooled = pipe.pool_channels(result, {"Group": ["Fz", "Cz"]}) + assert pooled["descriptor_names"] == ["global", "val_chgrp-Group"] + assert np.allclose(pooled["X"][:, 0], [1.0, 2.0]) + assert np.allclose(pooled["X"][:, 1], [15.0, 35.0]) + + +def test_pool_channels_preserves_failures(): + pipe = DescriptorPipeline({}) + result = { + "X": np.zeros((2, 2)), + "descriptor_names": ["a_ch-Fz", "a_ch-Cz"], + "failures": [{"family": "toy", "message": "boom"}], + } + pooled = pipe.pool_channels(result, {"G": ["Fz"]}) + assert pooled["failures"] == result["failures"] + + +def test_pipeline_instantiation_validates_fit_range_coverage(): + config = { + "families": { + "bands": { + "enabled": True, + "fmin": 1.0, + "fmax": 45.0, + "outputs": ["corrected_absolute_power"], + }, + "parametric": { + "enabled": True, + "freq_range": [2.0, 50.0], # 2.0 > 1.0, bad + }, + } + } + with pytest.raises(ValueError, match="cover the band PSD window"): + DescriptorPipeline(config) + + +def test_pool_channels_reject_mismatched_columns(): + pipe = DescriptorPipeline({}) + result = { + "X": np.zeros((2, 2)), + "descriptor_names": ["a_ch-Fz"], # 2 columns vs 1 name + "failures": [], + } + with pytest.raises(ValueError, match=r"align with result\['X'\] columns"): + pipe.pool_channels(result, {"G": ["Fz"]}) + + +def test_pipeline_precision_is_propagated_to_pooled_output(): + pipe = DescriptorPipeline({"precision": "float32"}) + result = { + "X": np.array([[1.0, 2.0]], dtype=np.float64), + "descriptor_names": ["a_ch-Fz", "a_ch-Cz"], + "failures": [], + } + pooled = pipe.pool_channels(result, {"G": ["Fz", "Cz"]}) + assert pooled["X"].dtype == np.float32 + + +def test_empty_work_unit_parallel_smoke(): + pytest.importorskip("joblib") + # 0 signal extractors, 1 PSD group (1 consumer) -> sequential + pipe = DescriptorPipeline( + { + "families": {"bands": {"enabled": True}}, + "runtime": {"n_jobs": 2, "execution_backend": "joblib"}, + } + ) + X = np.zeros((2, 1, 64)) + result = pipe.extract(X, sfreq=100.0, channel_names=["ch1"]) + assert result["X"].shape[0] == 2 diff --git a/tests/test_descriptors_extractors.py b/tests/test_descriptors_extractors.py index fc7b44d..d72a2cb 100644 --- a/tests/test_descriptors_extractors.py +++ b/tests/test_descriptors_extractors.py @@ -18,15 +18,13 @@ ) from coco_pipe.descriptors.extractors._parametric_fit import _ParametricFitBatch from coco_pipe.descriptors.extractors._psd import compute_psd -from coco_pipe.descriptors.extractors.base import _DescriptorBlock +from coco_pipe.descriptors.extractors.base import ( + _DescriptorBlock, + make_failure_record, +) from coco_pipe.descriptors.extractors.complexity import ComplexityDescriptorExtractor from coco_pipe.descriptors.extractors.parametric import ParametricDescriptorExtractor from coco_pipe.descriptors.extractors.spectral import BandDescriptorExtractor -from coco_pipe.descriptors.extractors.utils import ( - average_channel_matrix, - make_failure_record, - pool_channel_descriptor_matrix, -) # --- Fixtures --- @@ -36,11 +34,9 @@ def signal_data(): """Standard signal data: (n_obs, n_channels, n_times).""" rng = np.random.default_rng(42) sfreq = 250.0 - # Increase to 2 seconds for better Welch/entropy estimation t = np.arange(0, 2, 1 / sfreq) n_obs, n_chans = 5, 3 - # Create 1/f-like noise freqs = np.fft.rfftfreq(len(t), 1 / sfreq) weights = 1 / (freqs + 1.0) @@ -48,10 +44,8 @@ def signal_data(): for o in range(n_obs): for c in range(n_chans): white = rng.standard_normal(len(t)) - # Quick and dirty 1/f approximation X[o, c, :] = np.fft.irfft(np.fft.rfft(white) * weights, n=len(t)) - # Add strong oscillations to ensure fitting works X[:, 0, :] += 2.0 * np.sin(2 * np.pi * 10 * t) X[:, 1, :] += 1.5 * np.sin(2 * np.pi * 20 * t) @@ -105,43 +99,6 @@ def test_descriptor_block_structure(): assert block.descriptor_names == names -def test_pool_channel_descriptor_matrix_logic(): - """Test all variants of channel pooling names and values.""" - # (n_obs, n_channels) - values = np.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]) - ch_names = ["Fz", "Cz"] - - # None - pooled, scopes = pool_channel_descriptor_matrix(values, ch_names, "none") - assert np.array_equal(pooled, values) - assert scopes == ["ch-Fz", "ch-Cz"] - - # All (mean across channels) - pooled, scopes = pool_channel_descriptor_matrix(values, ch_names, "all") - assert pooled.shape == (3, 1) - assert np.allclose(pooled[:, 0], [1.5, 3.5, 5.5]) - assert scopes == ["ch-all"] - - # Dict grouping - values3 = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) - ch_names3 = ["Fz", "Cz", "Pz"] - pooling = {"Frontal": ["Fz", "Cz"]} - pooled, scopes = pool_channel_descriptor_matrix(values3, ch_names3, pooling) - # Frontal (mean of 0,1) + Pz (remains) - assert pooled.shape == (2, 2) - assert np.allclose(pooled[:, 0], [1.5, 4.5]) # mean([1,2], [4,5]) - assert np.allclose(pooled[:, 1], [3.0, 6.0]) # Pz - assert scopes == ["chgrp-Frontal", "ch-Pz"] - - -def test_average_channel_matrix_robustness(): - """Verify averaging handles NaNs correctly.""" - X = np.array([[1.0, 2.0], [np.nan, 4.0], [5.0, np.nan], [np.nan, np.nan]]) - res = average_channel_matrix(X) - assert np.allclose(res[:3], [1.5, 4.0, 5.0]) - assert np.isnan(res[3]) - - def test_make_failure_record_schema(): """Check fixed schema for failure records.""" rec = make_failure_record( @@ -178,7 +135,6 @@ def test_basic_extraction(self, psd_data, signal_data): psds, freqs, channel_names=ch_names, - channel_pooling="none", ids=None, runtime=MagicMock(), ) @@ -203,15 +159,13 @@ def test_corrected_outputs(self, psd_data, signal_data, mock_fit_batch): psds, freqs, channel_names=ch_names, - channel_pooling="all", ids=None, fit_batch=mock_fit_batch, runtime=MagicMock(), ) - # alpha, beta corrected + 1 ratio = 3 columns - assert block.X.shape == (psds.shape[0], 3) - assert "band_corr_abs_alpha_ch-all" in block.descriptor_names - assert "band_corr_ratio_alpha_beta_ch-all" in block.descriptor_names + assert block.X.shape == (psds.shape[0], 9) + assert "band_corr_abs_alpha_ch-Fz" in block.descriptor_names + assert "band_corr_ratio_alpha_beta_ch-Pz" in block.descriptor_names def test_missing_fit_batch_raises(self, psd_data, signal_data): psds, freqs = psd_data @@ -227,7 +181,6 @@ def test_missing_fit_batch_raises(self, psd_data, signal_data): psds, freqs, channel_names=ch_names, - channel_pooling="none", ids=None, runtime=MagicMock(), ) @@ -250,7 +203,6 @@ def test_band_resolution_error(self, psd_data, signal_data): psds, freqs, channel_names=ch_names, - channel_pooling="all", ids=None, runtime=runtime, ) @@ -275,15 +227,31 @@ def test_spectral_extract_psd_edge_cases(self, signal_data): X, sfreq, ch_names = signal_data - # log_power coverage + # explicit log output coverage config = BandDescriptorConfig( - enabled=True, outputs=["absolute_power"], log_power=True + enabled=True, + outputs=["log_absolute_power", "corrected_log_absolute_power"], ) extractor = BandDescriptorExtractor(config) psds = np.ones((1, 3, 10)) freqs = np.linspace(1, 45, 10) - block = extractor.extract_psd(psds, freqs, ch_names, "all", None, MagicMock()) - assert "band_log_abs_delta_ch-all" in block.descriptor_names + fit_batch = _ParametricFitBatch( + freqs=freqs, + metrics={}, + periodic_psds=np.ones((1, 3, 10)), + errors=[], + meta={}, + ) + block = extractor.extract_psd( + psds, + freqs, + ch_names, + None, + MagicMock(), + fit_batch=fit_batch, + ) + assert "band_log_abs_delta_ch-Fz" in block.descriptor_names + assert "band_corr_log_abs_delta_ch-Fz" in block.descriptor_names config_rel = BandDescriptorConfig( enabled=True, @@ -293,14 +261,9 @@ def test_spectral_extract_psd_edge_cases(self, signal_data): bands={"high": (120, 150)}, ) extractor_rel = BandDescriptorExtractor(config_rel) - block_rel = extractor_rel.extract_psd( - psds, freqs, ch_names, "all", None, MagicMock() - ) + block_rel = extractor_rel.extract_psd(psds, freqs, ch_names, None, MagicMock()) assert np.isnan(block_rel.X).all() - # fit_batch errors and corrected relative power empty freq - from coco_pipe.descriptors.extractors._parametric_fit import _ParametricFitBatch - fit_batch = _ParametricFitBatch( freqs=np.array([1, 10, 20]), # Must be non-empty and non-None metrics={}, @@ -313,7 +276,7 @@ def test_spectral_extract_psd_edge_cases(self, signal_data): ) extractor_corr = BandDescriptorExtractor(config_corr) block_corr = extractor_corr.extract_psd( - psds, freqs, ch_names, "all", None, MagicMock(), fit_batch=fit_batch + psds, freqs, ch_names, None, MagicMock(), fit_batch=fit_batch ) assert len(block_corr.failures) > 0 assert "FakeError" in block_corr.failures[0]["exception_type"] @@ -324,9 +287,8 @@ def test_spectral_standalone_extract(self, signal_data): X, sfreq, ch_names = signal_data config = BandDescriptorConfig(enabled=True, outputs=["absolute_power"]) extractor = BandDescriptorExtractor(config) - block = extractor.extract(X, sfreq, ch_names, "all", None, MagicMock()) - # 5 bands (delta, theta, alpha, beta, gamma) pooled to 'all' -> 5 columns - assert block.X.shape == (5, 5) + block = extractor.extract(X, sfreq, ch_names, None, MagicMock()) + assert block.X.shape == (5, 15) def test_spectral_empty_output_block(self): # Empty output when no families enabled or no outputs requested @@ -334,9 +296,7 @@ def test_spectral_empty_output_block(self): extractor = BandDescriptorExtractor(config) psds = np.ones((2, 2, 10)) freqs = np.linspace(1, 45, 10) - block = extractor.extract_psd( - psds, freqs, ["ch1", "ch2"], "all", None, MagicMock() - ) + block = extractor.extract_psd(psds, freqs, ["ch1", "ch2"], None, MagicMock()) assert block.X.shape == (2, 0) def test_spectral_standalone_extract_raises_for_corrected(self, signal_data): @@ -350,7 +310,37 @@ def test_spectral_standalone_extract_raises_for_corrected(self, signal_data): with pytest.raises( ValueError, match="Corrected band outputs are only available" ): - extractor.extract(X, sfreq, ch_names, "all", None, MagicMock()) + extractor.extract(X, sfreq, ch_names, None, MagicMock()) + + def test_ratio_denominator_floor_turns_near_zero_ratios_into_nan(self): + from unittest.mock import MagicMock + + freqs = np.array([9.0, 10.0, 20.0, 21.0], dtype=float) + psds = np.zeros((1, 1, freqs.size), dtype=float) + psds[0, 0, :2] = 1.0 + psds[0, 0, 2:] = 1e-14 + + config = BandDescriptorConfig( + enabled=True, + outputs=["ratios"], + bands={"alpha": (8.0, 12.0), "beta": (19.0, 22.0)}, + ratio_pairs=[("alpha", "beta")], + min_denominator_power=1e-12, + ) + extractor = BandDescriptorExtractor(config) + + block = extractor.extract_psd( + psds, + freqs, + channel_names=["Fz"], + ids=np.array(["obs-0"], dtype=object), + runtime=MagicMock(), + ) + + assert block.descriptor_names == ["band_ratio_alpha_beta_ch-Fz"] + assert np.isnan(block.X[0, 0]) + assert len(block.failures) == 1 + assert "NumericalIssue" in block.failures[0]["exception_type"] # --- 3. Parametric Extractor --- @@ -372,17 +362,16 @@ def test_standalone_extract(self, signal_data): X, sfreq=sfreq, channel_names=ch_names, - channel_pooling="all", ids=None, runtime=MagicMock(), obs_offset=0, ) - assert "param_exponent_ch-all" in block.descriptor_names - # 2 features (offset, exponent) for aperiodic 'fixed' mode - assert block.X.shape == (X.shape[0], 2) - # Check that exponent is reasonable (> 0) and offset is finite - assert np.all(block.X[:, 1] > 0) - assert np.all(np.isfinite(block.X[:, 0])) + assert "param_exponent_ch-Fz" in block.descriptor_names + assert block.X.shape == (X.shape[0], 6) + exponent_idx = block.descriptor_names.index("param_exponent_ch-Fz") + offset_idx = block.descriptor_names.index("param_offset_ch-Fz") + assert np.all(block.X[:, exponent_idx] > 0) + assert np.all(np.isfinite(block.X[:, offset_idx])) def test_extract_psd_requires_fit_batch(self, psd_data, signal_data): psds, freqs = psd_data @@ -394,7 +383,6 @@ def test_extract_psd_requires_fit_batch(self, psd_data, signal_data): psds, freqs, channel_names=ch_names, - channel_pooling="all", ids=None, runtime=MagicMock(), obs_offset=0, @@ -416,12 +404,11 @@ def test_backend_dispatch_antropy(self, signal_data): X, sfreq=sfreq, channel_names=ch_names, - channel_pooling="all", ids=None, runtime=MagicMock(), obs_offset=0, ) - assert "complexity_spectral_entropy_ch-all" in block.descriptor_names + assert "complexity_spectral_entropy_ch-Fz" in block.descriptor_names assert not np.isnan(block.X).any() def test_backend_dispatch_neurokit2(self, signal_data): @@ -435,12 +422,11 @@ def test_backend_dispatch_neurokit2(self, signal_data): X, sfreq=sfreq, channel_names=ch_names, - channel_pooling="all", ids=None, runtime=MagicMock(), obs_offset=0, ) - assert "complexity_perm_entropy_ch-all" in block.descriptor_names + assert "complexity_perm_entropy_ch-Fz" in block.descriptor_names # Check that it's finite assert not np.isnan(block.X).any() @@ -458,7 +444,6 @@ def test_mixed_execution_strategy(self, signal_data): X, sfreq=sfreq, channel_names=ch_names, - channel_pooling="none", ids=None, runtime=MagicMock(), obs_offset=0, @@ -469,6 +454,62 @@ def test_mixed_execution_strategy(self, signal_data): assert "complexity_spectral_entropy_ch-Fz" in block.descriptor_names assert "complexity_sample_entropy_ch-Pz" in block.descriptor_names + def test_complexity_collect_nonfinite_numerical_issue(self, signal_data): + from unittest.mock import patch + + from coco_pipe.descriptors.configs import DescriptorRuntimeConfig + + X, sfreq, ch_names = signal_data + config = ComplexityDescriptorConfig( + enabled=True, + measures=["sample_entropy"], + ) + extractor = ComplexityDescriptorExtractor(config) + # Mock backend result to include inf + with patch.object( + extractor, + "_load_antropy", + return_value=MagicMock( + sample_entropy=lambda x, **kwargs: np.array([np.inf] * x.shape[0]) + ), + ): + block = extractor.extract( + X, + sfreq=sfreq, + channel_names=ch_names, + ids=None, + runtime=DescriptorRuntimeConfig(on_error="collect"), + ) + assert np.isnan(block.X).all() + assert any(f["exception_type"] == "NumericalIssue" for f in block.failures) + + def test_complexity_raise_on_error(self, signal_data): + from unittest.mock import patch + + from coco_pipe.descriptors.configs import DescriptorRuntimeConfig + + X, sfreq, ch_names = signal_data + config = ComplexityDescriptorConfig( + enabled=True, + measures=["sample_entropy"], + ) + extractor = ComplexityDescriptorExtractor(config) + with patch.object( + extractor, + "_load_antropy", + return_value=MagicMock( + sample_entropy=lambda x, **kwargs: np.array([[np.inf]]) + ), + ): + with pytest.raises(ValueError, match="became non-finite"): + extractor.extract( + X, + sfreq=sfreq, + channel_names=ch_names, + ids=None, + runtime=DescriptorRuntimeConfig(on_error="raise"), + ) + # --- 5. Lazy Loading and Dependency Guards --- diff --git a/tests/test_io_structures.py b/tests/test_io_structures.py index 20d0fa3..3af14c5 100644 --- a/tests/test_io_structures.py +++ b/tests/test_io_structures.py @@ -228,6 +228,64 @@ def test_repr(): assert "obs=5" in r +def test_obs_table_exports_only_obs_aligned_vectors(sample_container): + sample_container.coords["bad_matrix"] = np.array([[1, 2], [3, 4]]) + + metadata = sample_container.obs_table( + include_ids=True, include_y=True, y_col="target" + ) + + assert metadata.columns.tolist() == [ + "obs_id", + "Study ID", + "group", + "target", + ] + assert metadata["obs_id"].tolist() == ["s0", "s1"] + assert metadata["Study ID"].tolist() == ["S0", "S1"] + assert metadata["group"].tolist() == ["control", "patient"] + assert metadata["target"].tolist() == [0, 1] + + +def test_obs_table_requires_ids_when_requested(sample_container): + sample_container.ids = None + + with pytest.raises(ValueError, match="include_ids=True"): + sample_container.obs_table(include_ids=True) + + +def test_obs_table_validation_errors(sample_container): + """Test alignment and dimensionality validation in obs_table.""" + # 1. No obs dimension + X = np.zeros((10, 10)) + dc_no_obs = DataContainer(X, dims=("a", "b")) + with pytest.raises(ValueError, match="requires an 'obs' dimension"): + dc_no_obs.obs_table() + + # 2. Corrupted ids (2D) + sample_container.ids = np.zeros((2, 2)) + with pytest.raises(ValueError, match="must be 1D and aligned"): + sample_container.obs_table(include_ids=True) + + # 3. Corrupted y (aligned but 2D) + sample_container.ids = ["s0", "s1"] + sample_container.y = np.zeros((2, 2)) + with pytest.raises(ValueError, match="must be 1D and aligned"): + sample_container.obs_table(include_y=True) + + +def test_obs_table_include_obs_coord(sample_container): + """Verify include_obs_coord parameter.""" + sample_container.coords["obs"] = ["O1", "O2"] + table = sample_container.obs_table(include_obs_coord=True) + assert "obs" in table.columns + assert table["obs"].tolist() == ["O1", "O2"] + + # Also check coordinate loop skip logic for 'obs' + table2 = sample_container.obs_table(include_obs_coord=False) + assert "obs" not in table2.columns + + def test_save_load_errors(tmp_path): """Test save/load failure modes.""" DataContainer(np.zeros((2, 2)), dims=("a", "b")) @@ -433,6 +491,207 @@ def test_aggregate_count_sem_and_epoch_count_match_expected_values(): assert np.isclose(sem_agg.X[0, 1], expected_sem) +def test_aggregate_mad_and_iqr_ignore_nans_per_feature(): + agg = _make_grouped_descriptor_container().aggregate( + by=["g1", "g1", "g2", "g2"], + stats=["mad", "iqr"], + ) + + assert agg.dims == ("obs", "stat", "feature") + assert agg.coords["stat"].tolist() == ["mad", "iqr"] + assert np.isnan(agg.X[0, 0, 0]) + assert np.isnan(agg.X[0, 1, 0]) + assert np.isclose(agg.X[0, 0, 1], 0.5) + assert np.isclose(agg.X[0, 1, 1], 0.5) + assert np.isclose(agg.X[1, 0, 0], 0.0) + assert np.isclose(agg.X[1, 1, 0], 0.0) + assert np.isclose(agg.X[1, 0, 1], 0.0) + assert np.isclose(agg.X[1, 1, 1], 0.0) + + +def test_aggregate_groups_applies_selected_stats_in_requested_order(): + container = _make_descriptor_container( + [ + [1.0, 10.0, 100.0, 1000.0], + [3.0, 14.0, 120.0, 1100.0], + [5.0, 18.0, 130.0, 1300.0], + [7.0, 22.0, 150.0, 1500.0], + ], + descriptor_names=[ + "band_abs_alpha_ch-all", + "band_log_abs_alpha_ch-all", + "complexity_entropy_ch-all", + "param_offset_ch-all", + ], + ) + + agg = container.aggregate_groups( + by=["s1", "s1", "s2", "s2"], + groups=[ + { + "stats": "mean", + "exclude_prefixes": ["band_abs_"], + }, + { + "prefixes": ["band_log_abs_"], + "stats": ["median", "iqr"], + }, + { + "prefixes": ["complexity_"], + "stats": ["median", "mad"], + }, + { + "prefixes": ["param_"], + "stats": ["median", "iqr"], + }, + ], + ) + + assert agg.dims == ("obs", "feature") + assert agg.coords["obs"].tolist() == ["s1", "s2"] + assert agg.ids.tolist() == ["s1", "s2"] + assert agg.coords["feature"].tolist() == [ + "mean_band_log_abs_alpha_ch-all", + "mean_complexity_entropy_ch-all", + "mean_param_offset_ch-all", + "median_band_log_abs_alpha_ch-all", + "iqr_band_log_abs_alpha_ch-all", + "median_complexity_entropy_ch-all", + "mad_complexity_entropy_ch-all", + "median_param_offset_ch-all", + "iqr_param_offset_ch-all", + ] + assert "mean_band_abs_alpha_ch-all" not in agg.coords["feature"].tolist() + assert np.allclose( + agg.X[0], + [12.0, 110.0, 1050.0, 12.0, 2.0, 110.0, 10.0, 1050.0, 50.0], + ) + assert np.allclose( + agg.X[1], + [20.0, 140.0, 1400.0, 20.0, 2.0, 140.0, 10.0, 1400.0, 100.0], + ) + assert agg.meta["agg_stats"] == ["mean", "median", "iqr", "mad"] + + +def test_aggregate_groups_rejects_duplicate_output_feature_names(): + container = _make_descriptor_container( + [[1.0], [2.0], [3.0], [4.0]], + descriptor_names=["band_log_abs_alpha_ch-all"], + ) + + with pytest.raises(ValueError, match="duplicate feature names"): + container.aggregate_groups( + by=["s1", "s1", "s2", "s2"], + groups=[ + {"stats": "mean", "names": ["band_log_abs_alpha_ch-all"]}, + {"stats": "mean", "prefixes": ["band_log_abs_"]}, + ], + ) + + +def test_aggregate_groups_selectors_and_validation(sample_container): + """Test selector helpers and input validation in aggregate_groups.""" + sample_container.dims = ("obs", "feature") + sample_container.coords["feature"] = ["a", "b"] + sample_container.X = np.zeros((2, 2)) + + # 1. Unknown keys + with pytest.raises(ValueError, match="Unknown aggregate_groups keys"): + groups = [{"stats": "mean", "invalid": 1}] + sample_container.aggregate_groups(by="group", groups=groups) + + # 2. Missing stats + with pytest.raises(ValueError, match="must include `stats`"): + sample_container.aggregate_groups(by="group", groups=[{"names": ["a"]}]) + + # 3. No match when skipping is disabled + with pytest.raises(ValueError, match="matched no features"): + sample_container.aggregate_groups( + by="group", + groups=[{"names": ["missing"], "stats": "mean"}], + skip_empty=False, + ) + + agg_empty_skipped = sample_container.aggregate_groups( + by=[0, 1], + groups=[ + {"names": ["missing"], "stats": "mean"}, + {"names": ["a"], "stats": "mean"}, + ], + ) + assert agg_empty_skipped.coords["feature"].tolist() == ["mean_a"] + + # Selector variety + container = _make_descriptor_container( + np.zeros((4, 4)), descriptor_names=["aa", "ab", "ba", "bb"] + ) + # Check regex and contains in one go + agg = container.aggregate_groups( + by=[0, 0, 1, 1], + groups=[ + {"contains": ["a"], "stats": "mean"}, + {"regex": ["^b"], "stats": "median"}, + {"suffixes": ["b"], "stats": "max"}, + ], + ) + assert "mean_aa" in agg.coords["feature"] + assert "median_ba" in agg.coords["feature"] + assert "max_bb" in agg.coords["feature"] + agg2 = container.aggregate_groups( + by=[0, 0, 1, 1], + groups=[ + {"contains": ["a"], "stats": "mean"}, + {"suffixes": ["b"], "stats": "median"}, + ], + ) + assert "mean_aa" in agg2.coords["feature"] + assert "median_ab" in agg2.coords["feature"] + + +def test_aggregate_groups_meta_and_consistency(sample_container): + """Test meta propagation and consistency checks in aggregate_groups.""" + sample_container.dims = ("obs", "feature") + sample_container.coords["feature"] = ["f1", "f2"] + sample_container.X = np.zeros((2, 2)) + + # 2. Failures collection + sample_container.meta["aggregate_failures"] = [ + { + "family": "bands", + "obs_index": 0, + "obs_id": "s0", + "channel_index": 0, + "channel_name": "ch0", + "exception_type": "Error", + "message": "msg", + } + ] + # isel will keep meta. aggregate will keep meta. + agg = sample_container.aggregate_groups( + by=[0, 1], groups=[{"names": ["f1"], "stats": "mean", "name": "G1"}] + ) + assert len(agg.meta.get("aggregate_failures", [])) > 0 + assert agg.meta["aggregate_failures"][0]["aggregate_group_name"] == "G1" + + +def test_aggregate_groups_consistency_checks(sample_container): + """Verify mixed per-group stats are flattened into one feature axis.""" + sample_container.dims = ("obs", "feature") + sample_container.coords["feature"] = ["f1", "f2"] + sample_container.X = np.zeros((2, 2)) + + agg = sample_container.aggregate_groups( + by=[0, 1], + groups=[ + {"names": ["f1"], "stats": ["mean", "median"]}, + {"names": ["f2"], "stats": "mean"}, + ], + ) + + assert agg.dims == ("obs", "feature") + assert agg.coords["feature"].tolist() == ["mean_f1", "median_f1", "mean_f2"] + + def test_aggregate_min_count_collect_policy_records_failure(): agg = _make_grouped_descriptor_container().aggregate( by=["g1", "g1", "g2", "g2"], @@ -467,7 +726,6 @@ def test_aggregate_descriptor_pipeline_output_can_be_grouped(): X = _make_signal_data() result = DescriptorPipeline( { - "output": {"channel_pooling": "all"}, "families": {"bands": {"enabled": True, "outputs": ["absolute_power"]}}, } ).extract(X=X, sfreq=256.0, channel_names=["Fz", "Cz"]) @@ -477,19 +735,24 @@ def test_aggregate_descriptor_pipeline_output_can_be_grouped(): ) assert all("_global" not in name for name in result["descriptor_names"]) - assert any(name.endswith("_ch-all") for name in result["descriptor_names"]) + assert any(name.endswith("_ch-Fz") for name in result["descriptor_names"]) assert agg.X.shape == (2, result["X"].shape[1]) assert agg.dims == ("obs", "feature") def test_aggregate_descriptor_pipeline_preserves_channel_group_tokens(): X = _make_signal_data() - result = DescriptorPipeline( + pipe = DescriptorPipeline( { - "output": {"channel_pooling": {"Frontal": ["Fz", "Cz"]}}, "families": {"bands": {"enabled": True, "outputs": ["absolute_power"]}}, } - ).extract(X=X, sfreq=256.0, channel_names=["Fz", "Cz"]) + ) + result = pipe.extract( + X=X, + sfreq=256.0, + channel_names=["Fz", "Cz"], + ) + result = pipe.pool_channels(result, {"Frontal": ["Fz", "Cz"]}) agg = _descriptor_result_container(result).aggregate( by=["s1", "s1", "s1", "s2", "s2", "s2"], stats=["mean", "std"], @@ -617,6 +880,8 @@ def test_aggregate_all_stats(sample_container): "std", "var", "sem", + "obs-mad", + "obs-iqr", "min", "max", "first", @@ -631,6 +896,8 @@ def test_aggregate_all_stats(sample_container): "std", "var", "sem", + "mad", + "iqr", "min", "max", "first", From 68c46c27277cf3ccdf0f152e80d0edc536087e8a Mon Sep 17 00:00:00 2001 From: Hamza Abdelhedi Date: Mon, 23 Mar 2026 17:51:15 -0600 Subject: [PATCH 3/7] Enhance: add extended peak metrics and advanced complexity measures --- README.md | 7 + coco_pipe/descriptors/configs.py | 12 + .../descriptors/extractors/_parametric_fit.py | 51 ++- .../descriptors/extractors/complexity.py | 217 +++++++-- .../descriptors/extractors/parametric.py | 11 +- tests/test_descriptors_configs.py | 35 ++ tests/test_descriptors_extractors.py | 417 +++++++++++++++++- 7 files changed, 699 insertions(+), 51 deletions(-) diff --git a/README.md b/README.md index 8191e3c..3330f3d 100644 --- a/README.md +++ b/README.md @@ -193,6 +193,13 @@ Contributions are welcome! If you have suggestions or find any bugs, please open - Plan and implement enhancements for visualization features. - Integrate new visual components and testing. +#### Descriptors Module +- Add a future connectivity descriptor family built on `mne-connectivity`. +- Start that connectivity family with phase-based measures such as `PLV`, with room for later extensions like `ciPLV`, `PLI`, and `wPLI`. +- Add a future wavelet-based descriptor batch built on `PyWavelets`. +- Start that wavelet batch with `sure_entropy`. +- Keep `log_energy_entropy` on the roadmap, but finalize its scientific definition before implementation. + #### Dim reduction: - Add parallelism diff --git a/coco_pipe/descriptors/configs.py b/coco_pipe/descriptors/configs.py index 03943e1..34779c6 100644 --- a/coco_pipe/descriptors/configs.py +++ b/coco_pipe/descriptors/configs.py @@ -60,6 +60,18 @@ "sample_entropy", "perm_entropy", "spectral_entropy", + "approx_entropy", + "svd_entropy", + "petrosian_fd", + "katz_fd", + "higuchi_fd", + "shannon_entropy", + "fuzzy_entropy", + "dispersion_entropy", + "hurst_exponent", + "zero_crossings", + "kurtosis", + "rms", "hjorth_mobility", "hjorth_complexity", "lziv_complexity", diff --git a/coco_pipe/descriptors/extractors/_parametric_fit.py b/coco_pipe/descriptors/extractors/_parametric_fit.py index 0d46755..8b97324 100644 --- a/coco_pipe/descriptors/extractors/_parametric_fit.py +++ b/coco_pipe/descriptors/extractors/_parametric_fit.py @@ -23,6 +23,8 @@ from ...utils import import_optional_dependency from ..configs import ParametricDescriptorConfig +_ALPHA_BAND = (8.0, 13.0) + @dataclass(slots=True) class _ParametricFitBatch: @@ -124,12 +126,41 @@ def fit_single_spectrum( peak_count = 0.0 dominant_freq = np.nan dominant_power = np.nan + dominant_bandwidth = np.nan + alpha_peak_freq = np.nan + alpha_peak_power = np.nan else: periodic = np.atleast_2d(periodic) peak_count = float(periodic.shape[0]) - dominant_idx = int(np.nanargmax(periodic[:, 1])) - dominant_freq = float(periodic[dominant_idx, 0]) - dominant_power = float(periodic[dominant_idx, 1]) + power = np.asarray(periodic[:, 1], dtype=float) + valid_power = np.isfinite(power) + if np.any(valid_power): + valid_indices = np.flatnonzero(valid_power) + dominant_idx = int(valid_indices[np.nanargmax(power[valid_power])]) + dominant_freq = float(periodic[dominant_idx, 0]) + dominant_power = float(periodic[dominant_idx, 1]) + dominant_bandwidth = ( + float(periodic[dominant_idx, 2]) if periodic.shape[1] >= 3 else np.nan + ) + else: + dominant_freq = np.nan + dominant_power = np.nan + dominant_bandwidth = np.nan + + alpha_mask = ( + valid_power + & np.isfinite(periodic[:, 0]) + & (periodic[:, 0] >= _ALPHA_BAND[0]) + & (periodic[:, 0] <= _ALPHA_BAND[1]) + ) + if np.any(alpha_mask): + alpha_indices = np.flatnonzero(alpha_mask) + alpha_idx = int(alpha_indices[np.nanargmax(power[alpha_mask])]) + alpha_peak_freq = float(periodic[alpha_idx, 0]) + alpha_peak_power = float(periodic[alpha_idx, 1]) + else: + alpha_peak_freq = np.nan + alpha_peak_power = np.nan offset = float(aperiodic[0]) if aperiodic.size >= 1 else np.nan knee = float(aperiodic[1]) if aperiodic.size == 3 else np.nan @@ -161,6 +192,9 @@ def fit_single_spectrum( "peak_count": peak_count, "peak_freq_dom": dominant_freq, "peak_power_dom": dominant_power, + "peak_bandwidth_dom": dominant_bandwidth, + "alpha_peak_freq": alpha_peak_freq, + "alpha_peak_power": alpha_peak_power, }, periodic_psd @@ -211,7 +245,16 @@ def fit_parametric_batch( if "fit_quality" in config.outputs: metric_names.extend(["fit_error", "r_squared"]) if "peak_summary" in config.outputs: - metric_names.extend(["peak_count", "peak_freq_dom", "peak_power_dom"]) + metric_names.extend( + [ + "peak_count", + "peak_freq_dom", + "peak_power_dom", + "peak_bandwidth_dom", + "alpha_peak_freq", + "alpha_peak_power", + ] + ) metric_arrays = { metric_name: np.full( (local_psds.shape[0], local_psds.shape[1]), diff --git a/coco_pipe/descriptors/extractors/complexity.py b/coco_pipe/descriptors/extractors/complexity.py index 6267ebc..3d04707 100644 --- a/coco_pipe/descriptors/extractors/complexity.py +++ b/coco_pipe/descriptors/extractors/complexity.py @@ -13,8 +13,13 @@ - `spectral_entropy`, `hjorth_mobility`, and `hjorth_complexity` use batched `antropy` calls over flattened observation-channel units -- `sample_entropy`, `perm_entropy`, and `lziv_complexity` are still evaluated - one 1D signal at a time +- `sample_entropy`, `perm_entropy`, `approx_entropy`, `svd_entropy`, + `petrosian_fd`, `katz_fd`, `higuchi_fd`, and `lziv_complexity` are still + evaluated one 1D signal at a time +- `shannon_entropy`, `fuzzy_entropy`, `dispersion_entropy`, and + `hurst_exponent` use scalar `neurokit2` calls +- `zero_crossings`, `kurtosis`, and `rms` are computed as simple scalar + channelwise signal descriptors Author: Hamza Abdelhedi (hamza.abdelhedi@umontreal.ca) """ @@ -24,11 +29,50 @@ from typing import Any import numpy as np +from scipy.stats import kurtosis as scipy_kurtosis from ...utils import import_optional_dependency from ..configs import ComplexityDescriptorConfig from .base import BaseDescriptorExtractor, _DescriptorBlock, make_failure_record +_ANTROPY_BATCHED_MEASURES = frozenset( + {"spectral_entropy", "hjorth_mobility", "hjorth_complexity"} +) +_ANTROPY_SCALAR_MEASURES = frozenset( + { + "sample_entropy", + "perm_entropy", + "approx_entropy", + "svd_entropy", + "petrosian_fd", + "katz_fd", + "higuchi_fd", + "lziv_complexity", + } +) +_NEUROKIT_SCALAR_MEASURES = frozenset( + { + "sample_entropy", + "perm_entropy", + "spectral_entropy", + "shannon_entropy", + "fuzzy_entropy", + "dispersion_entropy", + "hurst_exponent", + } +) +_CUSTOM_SCALAR_MEASURES = frozenset({"zero_crossings", "kurtosis", "rms"}) + + +def _normalize_scalar_output(value: Any) -> float: + """Normalize backend scalar outputs to one plain float.""" + if isinstance(value, tuple): + value = value[0] + array = np.asarray(value, dtype=float) + if array.size != 1: + raise ValueError("Complexity backend returned a non-scalar result.") + return float(array.reshape(-1)[0]) + class ComplexityDescriptorExtractor(BaseDescriptorExtractor): """ @@ -58,9 +102,12 @@ class ComplexityDescriptorExtractor(BaseDescriptorExtractor): deterministic sensor-level naming is applied afterward through :meth:`BaseDescriptorExtractor._finalize_descriptor`. - When `antropy` is selected, the extractor uses batched calls where the - backend supports them and falls back to scalar loops for measures that are - inherently one-signal-at-a-time in the current backend API. + When `backend="auto"` is selected, the extractor resolves each measure to + the preferred available implementation: + + - `antropy` for the existing antropy-backed measures + - `neurokit2` for measures that are only supported there + - built-in NumPy/SciPy implementations for simple scalar signal summaries """ family_name = "complexity" @@ -181,8 +228,9 @@ def extract( - batched `antropy` calls for `spectral_entropy`, `hjorth_mobility`, and `hjorth_complexity` - - scalar calls for `sample_entropy`, `perm_entropy`, and - `lziv_complexity` + - scalar `antropy` calls for the remaining antropy-backed measures + - scalar `neurokit2` calls for `shannon_entropy`, `fuzzy_entropy`, + `dispersion_entropy`, and `hurst_exponent` Non-finite outputs are converted to `NaN` and recorded under ``failures`` unless `runtime.on_error == "raise"`, in which case the @@ -210,11 +258,64 @@ def extract( flat_signals = X.reshape(-1, X.shape[-1]) batched_outputs: dict[str, np.ndarray] = {} scalar_dispatch: dict[str, Any] = {} + measure_backends: dict[str, str] = {} + custom_scalar_dispatch = { + "zero_crossings": lambda signal, kwargs, sfreq: float( + np.count_nonzero(np.diff(np.signbit(np.asarray(signal, dtype=float)))) + ), + "kurtosis": lambda signal, kwargs, sfreq: float( + scipy_kurtosis( + np.asarray(signal, dtype=float), + fisher=kwargs.get("fisher", True), + bias=kwargs.get("bias", False), + ) + ), + "rms": lambda signal, kwargs, sfreq: float( + np.sqrt(np.mean(np.square(np.asarray(signal, dtype=float)))) + ), + } + + unsupported: list[str] = [] + for measure in self.config.measures: + if measure in _CUSTOM_SCALAR_MEASURES: + measure_backends[measure] = "custom" + continue + if self.config.backend == "antropy": + if ( + measure in _ANTROPY_BATCHED_MEASURES + or measure in _ANTROPY_SCALAR_MEASURES + ): + measure_backends[measure] = "antropy" + else: + unsupported.append(measure) + continue + if self.config.backend == "neurokit2": + if measure in _NEUROKIT_SCALAR_MEASURES: + measure_backends[measure] = "neurokit2" + else: + unsupported.append(measure) + continue + if ( + measure in _ANTROPY_BATCHED_MEASURES + or measure in _ANTROPY_SCALAR_MEASURES + ): + measure_backends[measure] = "antropy" + elif measure in _NEUROKIT_SCALAR_MEASURES: + measure_backends[measure] = "neurokit2" + else: + unsupported.append(measure) - if self.config.backend in {"antropy", "auto"}: + if unsupported: + raise ValueError( + f"Measures {sorted(unsupported)} are not supported by backend " + f"'{self.config.backend}'." + ) + + ant = None + if "antropy" in measure_backends.values(): ant = self._load_antropy() - if "spectral_entropy" in self.config.measures: + if measure_backends.get("spectral_entropy") == "antropy": batched_outputs["spectral_entropy"] = np.asarray( ant.spectral_entropy( flat_signals, @@ -226,55 +327,98 @@ def extract( ) if ( - "hjorth_mobility" in self.config.measures - or "hjorth_complexity" in self.config.measures + measure_backends.get("hjorth_mobility") == "antropy" + or measure_backends.get("hjorth_complexity") == "antropy" ): mobility, complexity = ant.hjorth_params( flat_signals, axis=-1, ) - if "hjorth_mobility" in self.config.measures: + if measure_backends.get("hjorth_mobility") == "antropy": batched_outputs["hjorth_mobility"] = np.asarray( mobility, dtype=float, ) - if "hjorth_complexity" in self.config.measures: + if measure_backends.get("hjorth_complexity") == "antropy": batched_outputs["hjorth_complexity"] = np.asarray( complexity, dtype=float, ) - scalar_dispatch = { - "sample_entropy": lambda signal, kwargs, sfreq: float( - ant.sample_entropy(signal, **kwargs) + antropy_scalar_dispatch = { + "sample_entropy": lambda signal, kwargs, sfreq: ( + _normalize_scalar_output(ant.sample_entropy(signal, **kwargs)) ), - "perm_entropy": lambda signal, kwargs, sfreq: float( + "perm_entropy": lambda signal, kwargs, sfreq: _normalize_scalar_output( ant.perm_entropy(signal, **kwargs) ), - "lziv_complexity": lambda signal, kwargs, sfreq: float( - ant.lziv_complexity( - (signal > np.median(signal)).astype(int), - **kwargs, + "approx_entropy": lambda signal, kwargs, sfreq: ( + _normalize_scalar_output(ant.app_entropy(signal, **kwargs)) + ), + "svd_entropy": lambda signal, kwargs, sfreq: _normalize_scalar_output( + ant.svd_entropy(signal, **kwargs) + ), + "petrosian_fd": lambda signal, kwargs, sfreq: _normalize_scalar_output( + ant.petrosian_fd(signal, **kwargs) + ), + "katz_fd": lambda signal, kwargs, sfreq: _normalize_scalar_output( + ant.katz_fd(signal, **kwargs) + ), + "higuchi_fd": lambda signal, kwargs, sfreq: _normalize_scalar_output( + ant.higuchi_fd(signal, **kwargs) + ), + "lziv_complexity": lambda signal, kwargs, sfreq: ( + _normalize_scalar_output( + ant.lziv_complexity( + (signal > np.median(signal)).astype(int), + **kwargs, + ) ) ), } - else: + for measure, func in antropy_scalar_dispatch.items(): + if measure_backends.get(measure) == "antropy": + scalar_dispatch[measure] = func + + nk = None + if "neurokit2" in measure_backends.values(): nk = self._load_neurokit() - scalar_dispatch = { - "sample_entropy": lambda signal, kwargs, sfreq: float( - nk.entropy_sample(signal, **kwargs)[0] + neurokit_scalar_dispatch = { + "sample_entropy": lambda signal, kwargs, sfreq: ( + _normalize_scalar_output(nk.entropy_sample(signal, **kwargs)) ), - "perm_entropy": lambda signal, kwargs, sfreq: float( - nk.entropy_permutation(signal, **kwargs)[0] + "perm_entropy": lambda signal, kwargs, sfreq: _normalize_scalar_output( + nk.entropy_permutation(signal, **kwargs) + ), + "spectral_entropy": lambda signal, kwargs, sfreq: ( + _normalize_scalar_output( + nk.entropy_spectral( + signal, + sampling_rate=sfreq, + **kwargs, + ) + ) ), - "spectral_entropy": lambda signal, kwargs, sfreq: float( - nk.entropy_spectral( - signal, - sampling_rate=sfreq, - **kwargs, - )[0] + "shannon_entropy": lambda signal, kwargs, sfreq: ( + _normalize_scalar_output(nk.entropy_shannon(signal, **kwargs)) + ), + "fuzzy_entropy": lambda signal, kwargs, sfreq: _normalize_scalar_output( + nk.entropy_fuzzy(signal, **kwargs) + ), + "dispersion_entropy": lambda signal, kwargs, sfreq: ( + _normalize_scalar_output(nk.entropy_dispersion(signal, **kwargs)) + ), + "hurst_exponent": lambda signal, kwargs, sfreq: ( + _normalize_scalar_output(nk.fractal_hurst(signal, **kwargs)) ), } + for measure, func in neurokit_scalar_dispatch.items(): + if measure_backends.get(measure) == "neurokit2": + scalar_dispatch[measure] = func + + for measure, func in custom_scalar_dispatch.items(): + if measure_backends.get(measure) == "custom": + scalar_dispatch[measure] = func for measure, flat_values in batched_outputs.items(): values = np.asarray(flat_values, dtype=float).reshape( @@ -308,12 +452,6 @@ def extract( for measure in self.config.measures if measure not in batched_outputs ] - unsupported = sorted(set(scalar_measures) - set(scalar_dispatch)) - if unsupported: - raise ValueError( - f"Measures {unsupported} are not supported by backend " - f"'{self.config.backend}'." - ) for obs_rel in range(X.shape[0]): unit_signals = X[obs_rel] @@ -389,6 +527,7 @@ def extract( "backend": self.config.backend, "measures": list(self.config.measures), "batched_measures": sorted(batched_outputs), + "measure_backends": dict(measure_backends), }, failures=failures, ) diff --git a/coco_pipe/descriptors/extractors/parametric.py b/coco_pipe/descriptors/extractors/parametric.py index 25369c1..ade30e4 100644 --- a/coco_pipe/descriptors/extractors/parametric.py +++ b/coco_pipe/descriptors/extractors/parametric.py @@ -164,7 +164,16 @@ def extract_psd( if "fit_quality" in self.config.outputs: metrics.extend(["fit_error", "r_squared"]) if "peak_summary" in self.config.outputs: - metrics.extend(["peak_count", "peak_freq_dom", "peak_power_dom"]) + metrics.extend( + [ + "peak_count", + "peak_freq_dom", + "peak_power_dom", + "peak_bandwidth_dom", + "alpha_peak_freq", + "alpha_peak_power", + ] + ) failures: list[dict[str, Any]] = [] for obs_rel, unit_idx, exception_type, message in fit_batch.errors: if runtime.on_error == "raise": diff --git a/tests/test_descriptors_configs.py b/tests/test_descriptors_configs.py index 249f061..d886688 100644 --- a/tests/test_descriptors_configs.py +++ b/tests/test_descriptors_configs.py @@ -274,6 +274,41 @@ def test_complexity_validation_edge_cases(): DescriptorConfig(families={"complexity": {"measures": ["non_existent"]}}) +def test_new_complexity_measures_are_accepted(): + measures = [ + "approx_entropy", + "svd_entropy", + "petrosian_fd", + "katz_fd", + "higuchi_fd", + "shannon_entropy", + "fuzzy_entropy", + "dispersion_entropy", + "hurst_exponent", + "zero_crossings", + "kurtosis", + "rms", + ] + config = DescriptorConfig.model_validate( + {"families": {"complexity": {"enabled": True, "measures": measures}}} + ) + assert config.families.complexity.measures == measures + + +def test_unknown_complexity_measure_is_rejected(): + with pytest.raises(ValidationError, match="Unknown complexity measures"): + DescriptorConfig.model_validate( + { + "families": { + "complexity": { + "enabled": True, + "measures": ["approx_entropy", "non_existent"], + } + } + } + ) + + def test_channel_pooling_validation_edge_cases(): with pytest.raises(ValidationError): DescriptorConfig(output={"channel_pooling": "some_string"}) diff --git a/tests/test_descriptors_extractors.py b/tests/test_descriptors_extractors.py index d72a2cb..28f21d0 100644 --- a/tests/test_descriptors_extractors.py +++ b/tests/test_descriptors_extractors.py @@ -10,7 +10,9 @@ import numpy as np import pytest +from scipy.stats import kurtosis as scipy_kurtosis +import coco_pipe.descriptors.extractors._parametric_fit as param_fit_module from coco_pipe.descriptors.configs import ( BandDescriptorConfig, ComplexityDescriptorConfig, @@ -388,6 +390,150 @@ def test_extract_psd_requires_fit_batch(self, psd_data, signal_data): obs_offset=0, ) + def test_peak_summary_emits_extended_peak_metrics(self, psd_data, signal_data): + psds, freqs = psd_data + _, _, ch_names = signal_data + n_obs, n_chans, _ = psds.shape + extractor = ParametricDescriptorExtractor( + ParametricDescriptorConfig(enabled=True, outputs=["peak_summary"]) + ) + fit_batch = _ParametricFitBatch( + freqs=freqs, + metrics={ + "peak_count": np.full((n_obs, n_chans), 2.0), + "peak_freq_dom": np.full((n_obs, n_chans), 10.0), + "peak_power_dom": np.full((n_obs, n_chans), 0.5), + "peak_bandwidth_dom": np.full((n_obs, n_chans), 2.0), + "alpha_peak_freq": np.full((n_obs, n_chans), 10.0), + "alpha_peak_power": np.full((n_obs, n_chans), 0.5), + }, + errors=[], + periodic_psds=None, + ) + + block = extractor.extract_psd( + psds, + freqs, + channel_names=ch_names, + ids=None, + runtime=MagicMock(), + fit_batch=fit_batch, + obs_offset=0, + ) + + assert block.X.shape == (n_obs, 18) + assert "param_peak_bandwidth_dom_ch-Fz" in block.descriptor_names + assert "param_alpha_peak_freq_ch-Fz" in block.descriptor_names + assert "param_alpha_peak_power_ch-Fz" in block.descriptor_names + + def test_fit_single_spectrum_reads_dominant_and_alpha_peak_metrics( + self, monkeypatch + ): + class _FakeResultsModel: + def get_component(self, name): + raise AssertionError("Periodic PSD reconstruction is not used here.") + + class _FakeResults: + has_model = True + model = _FakeResultsModel() + + def get_params(self, name): + if name == "aperiodic": + return np.array([0.1, 1.5], dtype=float) + if name == "periodic": + return np.array( + [ + [6.0, 0.2, 1.0], + [10.0, 0.7, 2.5], + [20.0, 0.5, 3.0], + ], + dtype=float, + ) + raise KeyError(name) + + def get_metrics(self, *names): + if names == ("error",): + return np.array([0.05], dtype=float) + if names == ("gof", "rsquared"): + return np.array([0.98], dtype=float) + raise KeyError(names) + + class _FakeSpectralModel: + def __init__(self, **kwargs): + self.results = _FakeResults() + + def fit(self, freqs, spectrum, freq_range): + return None + + monkeypatch.setattr( + param_fit_module, + "import_optional_dependency", + lambda *args, **kwargs: _FakeSpectralModel, + ) + + metrics, periodic_psd = param_fit_module.fit_single_spectrum( + freqs=np.array([1.0, 10.0, 20.0], dtype=float), + spectrum=np.array([1.0, 2.0, 1.5], dtype=float), + config=ParametricDescriptorConfig(enabled=True, outputs=["peak_summary"]), + need_periodic_psd=False, + ) + + assert periodic_psd is None + assert metrics["peak_count"] == 3.0 + assert metrics["peak_freq_dom"] == pytest.approx(10.0) + assert metrics["peak_power_dom"] == pytest.approx(0.7) + assert metrics["peak_bandwidth_dom"] == pytest.approx(2.5) + assert metrics["alpha_peak_freq"] == pytest.approx(10.0) + assert metrics["alpha_peak_power"] == pytest.approx(0.7) + + def test_fit_single_spectrum_returns_nan_for_missing_alpha_peak(self, monkeypatch): + class _FakeResultsModel: + def get_component(self, name): + raise AssertionError("Periodic PSD reconstruction is not used here.") + + class _FakeResults: + has_model = True + model = _FakeResultsModel() + + def get_params(self, name): + if name == "aperiodic": + return np.array([0.2, 1.1], dtype=float) + if name == "periodic": + return np.array([[20.0, 0.5, 3.0]], dtype=float) + raise KeyError(name) + + def get_metrics(self, *names): + if names == ("error",): + return np.array([0.02], dtype=float) + if names == ("gof", "rsquared"): + return np.array([0.99], dtype=float) + raise KeyError(names) + + class _FakeSpectralModel: + def __init__(self, **kwargs): + self.results = _FakeResults() + + def fit(self, freqs, spectrum, freq_range): + return None + + monkeypatch.setattr( + param_fit_module, + "import_optional_dependency", + lambda *args, **kwargs: _FakeSpectralModel, + ) + + metrics, periodic_psd = param_fit_module.fit_single_spectrum( + freqs=np.array([1.0, 10.0, 20.0], dtype=float), + spectrum=np.array([1.0, 2.0, 1.5], dtype=float), + config=ParametricDescriptorConfig(enabled=True, outputs=["peak_summary"]), + need_periodic_psd=False, + ) + + assert periodic_psd is None + assert metrics["peak_count"] == 1.0 + assert np.isnan(metrics["alpha_peak_freq"]) + assert np.isnan(metrics["alpha_peak_power"]) + # --- 4. Complexity Extractor --- @@ -469,9 +615,7 @@ def test_complexity_collect_nonfinite_numerical_issue(self, signal_data): with patch.object( extractor, "_load_antropy", - return_value=MagicMock( - sample_entropy=lambda x, **kwargs: np.array([np.inf] * x.shape[0]) - ), + return_value=MagicMock(sample_entropy=lambda x, **kwargs: np.inf), ): block = extractor.extract( X, @@ -497,11 +641,9 @@ def test_complexity_raise_on_error(self, signal_data): with patch.object( extractor, "_load_antropy", - return_value=MagicMock( - sample_entropy=lambda x, **kwargs: np.array([[np.inf]]) - ), + return_value=MagicMock(sample_entropy=lambda x, **kwargs: np.inf), ): - with pytest.raises(ValueError, match="became non-finite"): + with pytest.raises(ValueError, match="produced a non-finite result"): extractor.extract( X, sfreq=sfreq, @@ -510,6 +652,267 @@ def test_complexity_raise_on_error(self, signal_data): runtime=DescriptorRuntimeConfig(on_error="raise"), ) + def test_easy_batch_complexity_measures_emit_expected_columns(self, signal_data): + X, sfreq, ch_names = signal_data + measures = [ + "approx_entropy", + "svd_entropy", + "petrosian_fd", + "katz_fd", + "higuchi_fd", + "zero_crossings", + "kurtosis", + "rms", + ] + extractor = ComplexityDescriptorExtractor( + ComplexityDescriptorConfig( + enabled=True, + backend="antropy", + measures=measures, + ) + ) + + block = extractor.extract( + X, + sfreq=sfreq, + channel_names=ch_names, + ids=None, + runtime=MagicMock(), + obs_offset=0, + ) + + assert block.X.shape == (X.shape[0], len(measures) * len(ch_names)) + for measure in measures: + assert f"complexity_{measure}_ch-Fz" in block.descriptor_names + + def test_zero_crossings_matches_manual_count(self): + X = np.array([[[-1.0, 1.0, -2.0, 3.0, -4.0]]], dtype=float) + extractor = ComplexityDescriptorExtractor( + ComplexityDescriptorConfig( + enabled=True, + backend="neurokit2", + measures=["zero_crossings"], + ) + ) + + block = extractor.extract( + X, + sfreq=250.0, + channel_names=["Fz"], + ids=None, + runtime=MagicMock(), + obs_offset=0, + ) + + assert block.descriptor_names == ["complexity_zero_crossings_ch-Fz"] + assert block.X[0, 0] == pytest.approx(4.0) + + def test_rms_matches_manual_calculation(self): + X = np.array([[[1.0, -1.0, 1.0, -1.0]]], dtype=float) + extractor = ComplexityDescriptorExtractor( + ComplexityDescriptorConfig( + enabled=True, + backend="neurokit2", + measures=["rms"], + ) + ) + + block = extractor.extract( + X, + sfreq=250.0, + channel_names=["Fz"], + ids=None, + runtime=MagicMock(), + obs_offset=0, + ) + + assert block.descriptor_names == ["complexity_rms_ch-Fz"] + assert block.X[0, 0] == pytest.approx(1.0) + + def test_kurtosis_matches_scipy_definition(self): + signal = np.array([1.0, 2.0, 2.5, 4.0, 8.0], dtype=float) + X = signal.reshape(1, 1, -1) + extractor = ComplexityDescriptorExtractor( + ComplexityDescriptorConfig( + enabled=True, + backend="neurokit2", + measures=["kurtosis"], + ) + ) + + block = extractor.extract( + X, + sfreq=250.0, + channel_names=["Fz"], + ids=None, + runtime=MagicMock(), + obs_offset=0, + ) + + assert block.descriptor_names == ["complexity_kurtosis_ch-Fz"] + assert block.X[0, 0] == pytest.approx( + scipy_kurtosis(signal, fisher=True, bias=False) + ) + + @pytest.mark.parametrize( + ("measure", "backend_attr", "backend_result", "expected"), + [ + ("shannon_entropy", "entropy_shannon", (1.25, {"base": 2}), 1.25), + ("fuzzy_entropy", "entropy_fuzzy", (0.75, {"Tolerance": 0.2}), 0.75), + ( + "dispersion_entropy", + "entropy_dispersion", + (0.5, {"dimension": 3}), + 0.5, + ), + ("hurst_exponent", "fractal_hurst", np.array([0.62]), 0.62), + ], + ) + def test_medium_batch_neurokit_measures_are_wired( + self, + measure, + backend_attr, + backend_result, + expected, + ): + from unittest.mock import patch + + X = np.array([[[0.1, 0.2, 0.4, 0.8, 0.3, 0.1]]], dtype=float) + extractor = ComplexityDescriptorExtractor( + ComplexityDescriptorConfig( + enabled=True, + backend="neurokit2", + measures=[measure], + ) + ) + fake_nk = MagicMock() + setattr( + fake_nk, + backend_attr, + lambda signal, _result=backend_result, **kwargs: _result, + ) + + with patch.object(extractor, "_load_neurokit", return_value=fake_nk): + block = extractor.extract( + X, + sfreq=250.0, + channel_names=["Fz"], + ids=None, + runtime=MagicMock(), + obs_offset=0, + ) + + assert block.descriptor_names == [f"complexity_{measure}_ch-Fz"] + assert block.X[0, 0] == pytest.approx(expected) + + def test_antropy_backend_rejects_medium_batch_measures(self): + X = np.array([[[0.1, 0.2, 0.4, 0.8, 0.3, 0.1]]], dtype=float) + extractor = ComplexityDescriptorExtractor( + ComplexityDescriptorConfig( + enabled=True, + backend="antropy", + measures=["fuzzy_entropy"], + ) + ) + + with pytest.raises(ValueError, match="not supported by backend 'antropy'"): + extractor.extract( + X, + sfreq=250.0, + channel_names=["Fz"], + ids=None, + runtime=MagicMock(), + obs_offset=0, + ) + + def test_auto_backend_mixes_antropy_and_neurokit_measures(self): + from unittest.mock import patch + + X = np.array([[[0.1, 0.2, 0.4, 0.8, 0.3, 0.1]]], dtype=float) + extractor = ComplexityDescriptorExtractor( + ComplexityDescriptorConfig( + enabled=True, + backend="auto", + measures=["sample_entropy", "fuzzy_entropy"], + ) + ) + fake_ant = MagicMock(sample_entropy=lambda signal, **kwargs: 0.25) + fake_nk = MagicMock(entropy_fuzzy=lambda signal, **kwargs: (0.75, {})) + + with ( + patch.object(extractor, "_load_antropy", return_value=fake_ant), + patch.object(extractor, "_load_neurokit", return_value=fake_nk), + ): + block = extractor.extract( + X, + sfreq=250.0, + channel_names=["Fz"], + ids=None, + runtime=MagicMock(), + obs_offset=0, + ) + + assert "complexity_sample_entropy_ch-Fz" in block.descriptor_names + assert "complexity_fuzzy_entropy_ch-Fz" in block.descriptor_names + assert block.meta["measure_backends"] == { + "sample_entropy": "antropy", + "fuzzy_entropy": "neurokit2", + } + + def test_medium_batch_neurokit_collects_nonfinite_values(self): + from unittest.mock import patch + + from coco_pipe.descriptors.configs import DescriptorRuntimeConfig + + X = np.array([[[0.1, 0.2, 0.4, 0.8, 0.3, 0.1]]], dtype=float) + extractor = ComplexityDescriptorExtractor( + ComplexityDescriptorConfig( + enabled=True, + backend="neurokit2", + measures=["fuzzy_entropy"], + ) + ) + fake_nk = MagicMock(entropy_fuzzy=lambda signal, **kwargs: (np.inf, {})) + + with patch.object(extractor, "_load_neurokit", return_value=fake_nk): + block = extractor.extract( + X, + sfreq=250.0, + channel_names=["Fz"], + ids=None, + runtime=DescriptorRuntimeConfig(on_error="collect"), + obs_offset=0, + ) + + assert np.isnan(block.X[0, 0]) + assert any(f["exception_type"] == "NumericalIssue" for f in block.failures) + + def test_medium_batch_neurokit_raises_on_nonfinite_values(self): + from unittest.mock import patch + + from coco_pipe.descriptors.configs import DescriptorRuntimeConfig + + X = np.array([[[0.1, 0.2, 0.4, 0.8, 0.3, 0.1]]], dtype=float) + extractor = ComplexityDescriptorExtractor( + ComplexityDescriptorConfig( + enabled=True, + backend="neurokit2", + measures=["fuzzy_entropy"], + ) + ) + fake_nk = MagicMock(entropy_fuzzy=lambda signal, **kwargs: (np.inf, {})) + + with patch.object(extractor, "_load_neurokit", return_value=fake_nk): + with pytest.raises(ValueError, match="produced a non-finite result"): + extractor.extract( + X, + sfreq=250.0, + channel_names=["Fz"], + ids=None, + runtime=DescriptorRuntimeConfig(on_error="raise"), + obs_offset=0, + ) + # --- 5. Lazy Loading and Dependency Guards --- From da5e2a9e1228864f1546dcff918a54e77e7322e1 Mon Sep 17 00:00:00 2001 From: Hamza Abdelhedi Date: Mon, 23 Mar 2026 18:09:40 -0600 Subject: [PATCH 4/7] Fix: update examples to new DescriptorConfig schema and pool_channels API --- configs/run_descriptors_eeg.yml | 3 +-- examples/descriptors_example.py | 7 +++++-- scripts/run_descriptors.py | 7 +++++++ 3 files changed, 13 insertions(+), 4 deletions(-) diff --git a/configs/run_descriptors_eeg.yml b/configs/run_descriptors_eeg.yml index 723d664..6afc5f7 100644 --- a/configs/run_descriptors_eeg.yml +++ b/configs/run_descriptors_eeg.yml @@ -26,9 +26,7 @@ descriptors: input: require_sfreq: true require_channel_names: false - output: precision: float32 - channel_pooling: all include_failure_summary: true include_runtime_meta: true families: @@ -68,3 +66,4 @@ descriptors: on_error: collect save_path: outputs/descriptors_eeg.npz +channel_groups: all diff --git a/examples/descriptors_example.py b/examples/descriptors_example.py index a0257ab..45160b7 100644 --- a/examples/descriptors_example.py +++ b/examples/descriptors_example.py @@ -19,7 +19,6 @@ def main() -> None: ids = np.asarray([f"obs-{idx:02d}" for idx in range(12)]) config = { - "output": {"channel_pooling": "all"}, "families": { "bands": { "enabled": True, @@ -36,14 +35,18 @@ def main() -> None: }, } + channels = ["Fz", "Cz", "Pz"] pipe = DescriptorPipeline(config) result = pipe.extract( X=X, ids=ids, sfreq=256.0, - channel_names=["Fz", "Cz", "Pz"], + channel_names=channels, ) + # Apply pooling after extraction + result = pipe.pool_channels(result, {"all": channels}) + print("Descriptor matrix shape:", result["X"].shape) print("First five names:", result["descriptor_names"][:5]) print("Failure count:", len(result["failures"])) diff --git a/scripts/run_descriptors.py b/scripts/run_descriptors.py index ce36897..2965f63 100644 --- a/scripts/run_descriptors.py +++ b/scripts/run_descriptors.py @@ -117,6 +117,13 @@ def main() -> None: pipe = DescriptorPipeline(descriptor_cfg) result = pipe.extract(**explicit_inputs) + channel_groups = payload.get("channel_groups") + if channel_groups == "all": + channel_groups = {"all": explicit_inputs["channel_names"]} + + if channel_groups: + result = pipe.pool_channels(result, channel_groups) + if save_path: _save_result(Path(save_path), result) From 55c31f976198b3fabe889fe2e6657201b71fc654 Mon Sep 17 00:00:00 2001 From: Hamza Abdelhedi Date: Thu, 16 Apr 2026 09:20:24 -0400 Subject: [PATCH 5/7] loosen up bids file detection --- coco_pipe/io/dataset.py | 2 +- coco_pipe/io/utils.py | 7 ++++++- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/coco_pipe/io/dataset.py b/coco_pipe/io/dataset.py index 8496182..9e6f9d5 100644 --- a/coco_pipe/io/dataset.py +++ b/coco_pipe/io/dataset.py @@ -689,7 +689,7 @@ def load(self) -> DataContainer: stem_parts.append(f"run-{run}") pre_epoched_suffix = self.suffix or "epo" - stem = "_".join(stem_parts) + stem = "*_".join(stem_parts) matches = sorted( pre_epoched_dir.glob(f"{stem}*_{pre_epoched_suffix}.fif") ) diff --git a/coco_pipe/io/utils.py b/coco_pipe/io/utils.py index 6904b14..ac50ba4 100644 --- a/coco_pipe/io/utils.py +++ b/coco_pipe/io/utils.py @@ -359,7 +359,12 @@ def detect_runs( Detect available runs for a given subject/session/task. """ bp = _get_bids_path()( - root=root, subject=subject, session=session, task=task, datatype=datatype + root=root, + subject=subject, + session=session, + task=task, + datatype=datatype, + check=False, ) matches = bp.match() runs = set() From 5170264f1bca32ea3713c7b64346561984239f3e Mon Sep 17 00:00:00 2001 From: Hamza Abdelhedi Date: Fri, 17 Apr 2026 12:35:54 -0400 Subject: [PATCH 6/7] add interactive tables for reports and LR x Balanced Accuracy as a seperation metric for dim_reduction --- coco_pipe/decoding/__init__.py | 2 + coco_pipe/decoding/utils.py | 79 ++++++ coco_pipe/dim_reduction/core.py | 9 +- coco_pipe/dim_reduction/evaluation/core.py | 68 ++++- coco_pipe/report/__init__.py | 11 +- coco_pipe/report/core.py | 93 ++++++- coco_pipe/report/templates/base.html | 275 +++++++++++++++++++-- tests/test_dimred_evaluation.py | 154 ++++++++++++ tests/test_report_core.py | 42 ++++ 9 files changed, 690 insertions(+), 43 deletions(-) diff --git a/coco_pipe/decoding/__init__.py b/coco_pipe/decoding/__init__.py index c8d5cb5..3215082 100644 --- a/coco_pipe/decoding/__init__.py +++ b/coco_pipe/decoding/__init__.py @@ -1,10 +1,12 @@ from .configs import ExperimentConfig from .core import Experiment from .registry import get_estimator_cls, register_estimator +from .utils import cross_validate_score __all__ = [ "ExperimentConfig", "register_estimator", "get_estimator_cls", "Experiment", + "cross_validate_score", ] diff --git a/coco_pipe/decoding/utils.py b/coco_pipe/decoding/utils.py index a1f433c..b660042 100644 --- a/coco_pipe/decoding/utils.py +++ b/coco_pipe/decoding/utils.py @@ -17,6 +17,7 @@ import numpy as np import pandas as pd +from sklearn.base import BaseEstimator, clone from sklearn.metrics import ( accuracy_score, balanced_accuracy_score, @@ -39,6 +40,8 @@ StratifiedKFold, train_test_split, ) +from sklearn.pipeline import Pipeline +from sklearn.preprocessing import StandardScaler from .configs import CVConfig @@ -262,3 +265,79 @@ def get_scorer(name: str) -> Callable: f"Unknown metric '{name}'. Available: {sorted(list(metrics.keys()))}" ) return metrics[name] + + +def cross_validate_score( + estimator: BaseEstimator, + X: np.ndarray, + y: Sequence, + *, + groups: Optional[Sequence] = None, + cv_config: Optional[CVConfig] = None, + metric: str = "balanced_accuracy", + use_scaler: bool = False, +) -> float: + """ + Compute one mean cross-validated score for an estimator. + + Parameters + ---------- + estimator : BaseEstimator + Estimator to fit inside each fold. + X : np.ndarray + Input features with shape ``(n_samples, n_features)``. + y : sequence + Target labels aligned with ``X``. + groups : sequence, optional + Group labels aligned with ``X``. + cv_config : CVConfig, optional + Cross-validation configuration. Defaults to a 5-fold stratified + strategy, or 5-fold stratified-group strategy when groups are + provided. + metric : str, default="balanced_accuracy" + Metric name resolved through :func:`get_scorer`. + use_scaler : bool, default=False + When ``True``, wraps the estimator in a ``StandardScaler`` pipeline. + + Returns + ------- + float + Mean cross-validated score. + """ + X = np.asarray(X) + y = np.asarray(y).reshape(-1) + if len(X) != len(y): + raise ValueError("X and y must have matching sample counts.") + + group_values = None + if groups is not None: + group_values = np.asarray(groups).reshape(-1) + if len(group_values) != len(y): + raise ValueError("groups must align with X and y.") + + if cv_config is None: + cv_config = CVConfig( + strategy="stratified_group_kfold" + if group_values is not None + else "stratified", + n_splits=5, + shuffle=True, + random_state=42, + ) + + scorer = get_scorer(metric) + cv = get_cv_splitter(cv_config, groups=group_values) + base_estimator = estimator + if use_scaler: + base_estimator = Pipeline( + [("scaler", StandardScaler()), ("clf", clone(estimator))] + ) + + scores = [] + for train_idx, test_idx in cv.split(X, y, group_values): + model = clone(base_estimator) + model.fit(X[train_idx], y[train_idx]) + y_pred = model.predict(X[test_idx]) + scores.append(float(scorer(y[test_idx], y_pred))) + + return float(np.nanmean(scores)) if scores else float("nan") diff --git a/coco_pipe/dim_reduction/core.py b/coco_pipe/dim_reduction/core.py index af00781..c355567 100644 --- a/coco_pipe/dim_reduction/core.py +++ b/coco_pipe/dim_reduction/core.py @@ -278,6 +278,7 @@ def score( metrics: Optional[List[str]] = None, k_values: Optional[List[int]] = None, labels: Optional[np.ndarray] = None, + groups: Optional[np.ndarray] = None, times: Optional[np.ndarray] = None, separation_method: str = "centroid", ) -> Dict[str, Dict[str, Any]]: @@ -301,7 +302,12 @@ def score( Neighborhood sizes used for multi-scale standard metric evaluation. labels : np.ndarray, optional Optional labels aligned with the embedding. Used for trajectory - separation when ``X_emb`` is 3D. + separation when ``X_emb`` is 3D and for explicit supervised 2D + metrics when requested. + groups : np.ndarray, optional + Optional grouping variable aligned with the embedding. Required by + grouped supervised evaluation metrics such as + ``separation_logreg_balanced_accuracy``. times : np.ndarray, optional Optional trajectory time coordinates aligned with the trajectory length axis. @@ -327,6 +333,7 @@ def score( method_name=self.method, metrics=metrics, labels=labels, + groups=groups, times=times, quality_metadata=self.get_quality_metadata(), diagnostics=self.get_diagnostics(), diff --git a/coco_pipe/dim_reduction/evaluation/core.py b/coco_pipe/dim_reduction/evaluation/core.py index d1b503e..a97bd1b 100644 --- a/coco_pipe/dim_reduction/evaluation/core.py +++ b/coco_pipe/dim_reduction/evaluation/core.py @@ -36,10 +36,13 @@ import numpy as np import pandas as pd +from sklearn.linear_model import LogisticRegression if TYPE_CHECKING: from ..core import DimReduction +from ...decoding.configs import CVConfig +from ...decoding.utils import cross_validate_score from .geometry import ( trajectory_acceleration, trajectory_curvature, @@ -63,6 +66,7 @@ __all__ = ["evaluate_embedding", "MethodSelector"] METRIC_COLUMNS = ("method", "metric", "value", "scope", "scope_value") +SEPARATION_LOGREG_BALANCED_ACCURACY = "separation_logreg_balanced_accuracy" SWEEP_METRICS = ( "trustworthiness", "continuity", @@ -90,6 +94,7 @@ "continuity": "desc", "lcmc": "desc", "shepard_correlation": "desc", + SEPARATION_LOGREG_BALANCED_ACCURACY: "desc", "mrre_intrusion": "asc", "mrre_extrusion": "asc", "mrre_total": "asc", @@ -441,6 +446,7 @@ def evaluate_embedding( method_name: str = "embedding", metrics: Optional[Sequence[str]] = None, labels: Optional[np.ndarray] = None, + groups: Optional[np.ndarray] = None, times: Optional[np.ndarray] = None, quality_metadata: Optional[Dict[str, Any]] = None, diagnostics: Optional[Dict[str, Any]] = None, @@ -469,8 +475,13 @@ def evaluate_embedding( Metric selectors to compute. ``None`` computes all metrics available for the provided inputs. labels : np.ndarray, optional - Optional trajectory labels used by ``trajectory_separation`` for native - 3D embeddings. + Optional labels aligned with the embedding. Used by + ``trajectory_separation`` for native 3D embeddings and by explicit + supervised 2D metrics such as + ``separation_logreg_balanced_accuracy`` when requested. + groups : np.ndarray, optional + Optional grouping variable aligned with ``X_emb``. Required by + ``separation_logreg_balanced_accuracy``. times : np.ndarray, optional Optional trajectory time coordinates used for separation AUC integration when trajectory metrics are evaluated. @@ -552,6 +563,7 @@ def evaluate_embedding( metric_selection = None if metrics is None else set(metrics) standard_metric_names = set(SWEEP_METRICS) | {"shepard_correlation"} + supervised_metric_names = {SEPARATION_LOGREG_BALANCED_ACCURACY} trajectory_metric_names = set(DEFAULT_SCORE_METRICS) - standard_metric_names metrics_payload: Dict[str, Any] = {} @@ -572,11 +584,13 @@ def evaluate_embedding( if X_emb.ndim == 2: if metric_selection is None: - metric_selection = standard_metric_names + standard_selection = standard_metric_names + supervised_selection = set() else: - metric_selection = metric_selection & standard_metric_names + standard_selection = metric_selection & standard_metric_names + supervised_selection = metric_selection & supervised_metric_names - if metric_selection: + if standard_selection: if X is None: raise ValueError( "Original data `X` is required to evaluate standard metrics " @@ -592,7 +606,7 @@ def evaluate_embedding( method_name=method_name, X_eval=X, X_emb_eval=X_emb, - metric_selection=metric_selection, + metric_selection=standard_selection, n_neighbors=n_neighbors, k_values=k_values, random_state=random_state, @@ -600,6 +614,36 @@ def evaluate_embedding( metrics_payload.update(std_metrics) diagnostics_payload.update(std_diagnostics) records.extend(std_records) + if SEPARATION_LOGREG_BALANCED_ACCURACY in supervised_selection: + if labels is None or groups is None: + raise ValueError( + f"`labels` and `groups` are required for " + f"'{SEPARATION_LOGREG_BALANCED_ACCURACY}'." + ) + separation_score = cross_validate_score( + LogisticRegression(max_iter=1000, class_weight="balanced"), + X_emb, + labels, + groups=groups, + cv_config=CVConfig( + strategy="stratified_group_kfold", + n_splits=5, + shuffle=True, + random_state=42, + ), + metric="balanced_accuracy", + use_scaler=True, + ) + metrics_payload[SEPARATION_LOGREG_BALANCED_ACCURACY] = separation_score + records.append( + { + "method": method_name, + "metric": SEPARATION_LOGREG_BALANCED_ACCURACY, + "value": separation_score, + "scope": "global", + "scope_value": "global", + } + ) elif X_emb.ndim == 3: if metric_selection is None: metric_selection = trajectory_metric_names @@ -719,6 +763,18 @@ def __init__( self.metric_records_ = [] + @classmethod + def from_records(cls, records: List[Dict[str, Any]]) -> "MethodSelector": + """Create a selector directly from long-form metric records.""" + selector = cls({}) + selector.metric_records_ = [dict(record) for record in records] + return selector + + @classmethod + def from_frame(cls, frame: pd.DataFrame) -> "MethodSelector": + """Create a selector directly from a metric-record DataFrame.""" + return cls.from_records(frame.to_dict(orient="records")) + def collect(self) -> "MethodSelector": """ Collect cached metric records from already-scored reducers. diff --git a/coco_pipe/report/__init__.py b/coco_pipe/report/__init__.py index 9b184fb..e2cfbe9 100644 --- a/coco_pipe/report/__init__.py +++ b/coco_pipe/report/__init__.py @@ -7,9 +7,17 @@ def __getattr__(name): - if name in ["Report", "Section", "PlotlyElement", "TableElement", "ImageElement"]: + if name in [ + "Report", + "Section", + "PlotlyElement", + "TableElement", + "InteractiveTableElement", + "ImageElement", + ]: from .core import ( # noqa: F401 ImageElement, + InteractiveTableElement, PlotlyElement, Report, Section, @@ -41,6 +49,7 @@ def __getattr__(name): "Section", "PlotlyElement", "TableElement", + "InteractiveTableElement", "ImageElement", "from_container", "from_bids", diff --git a/coco_pipe/report/core.py b/coco_pipe/report/core.py index 9ff295d..a67e8d9 100644 --- a/coco_pipe/report/core.py +++ b/coco_pipe/report/core.py @@ -8,6 +8,7 @@ import base64 import gzip +import html import io import json import re @@ -342,22 +343,22 @@ def __init__(self, data: Any, title: Optional[str] = None): self.title = title self.table_id = f"table-{uuid.uuid4().hex[:8]}" - def render(self) -> str: - # Convert to DataFrame - if isinstance(self.data, pd.DataFrame): - df = self.data - elif isinstance(self.data, dict): - # Check if all values are scalars - # (to avoid "If using all scalar values, you must pass an index") + @staticmethod + def _to_frame(data: Any) -> pd.DataFrame: + """Normalize supported table-like inputs to a DataFrame.""" + if isinstance(data, pd.DataFrame): + return data + if isinstance(data, dict): if all( isinstance(v, (int, float, str, np.number)) or v is None - for v in self.data.values() + for v in data.values() ): - df = pd.DataFrame([self.data]) - else: - df = pd.DataFrame(self.data) - else: - df = pd.DataFrame(self.data) + return pd.DataFrame([data]) + return pd.DataFrame(data) + return pd.DataFrame(data) + + def render(self) -> str: + df = self._to_frame(self.data) # Basic Tailwind Styling html = '
' @@ -419,6 +420,72 @@ def _render_row(self, row, idx) -> str: return html +class InteractiveTableElement(Element): + """Render a payload-backed interactive data table.""" + + def __init__( + self, + data: Any, + title: Optional[str] = None, + selector_columns: Optional[List[str]] = None, + default_sort: Optional[Dict[str, str]] = None, + page_size: int = 50, + ): + self.data = data + self.title = title + self.selector_columns = list(selector_columns or []) + self.default_sort = dict(default_sort) if default_sort else None + self.page_size = int(page_size) + self.registry_id: Optional[str] = None + + def collect_payload(self, registry: Dict[str, Any]) -> None: + if self.registry_id is None: + self.registry_id = str(uuid.uuid4()) + + df = TableElement._to_frame(self.data) + payload = { + "columns": [str(column) for column in df.columns], + "rows": json.loads(df.to_json(orient="records", date_format="iso")), + } + registry[self.registry_id] = payload + + def render(self) -> str: + if self.registry_id is None: + self.registry_id = str(uuid.uuid4()) + + config = { + "title": self.title, + "selector_columns": self.selector_columns, + "default_sort": self.default_sort, + "page_size": self.page_size, + } + config_json = html.escape(json.dumps(config), quote=True) + title_html = "" + if self.title: + title_html = f""" +
+

+ {self.title} +

+
+ """ + + return f""" +
+ {title_html} +
+
+ Loading interactive table... +
+
+
+ """ + + class MetricsTableElement(TableElement): """ Comparison table that highlights best values. diff --git a/coco_pipe/report/templates/base.html b/coco_pipe/report/templates/base.html index 53600da..c0b6301 100644 --- a/coco_pipe/report/templates/base.html +++ b/coco_pipe/report/templates/base.html @@ -387,28 +387,28 @@

Execut document.querySelectorAll('.lazy-plot').forEach(p => plotObserver.observe(p)); - // --- 5. Export Table to CSV --- - function exportTableToCSV(tableId, title) { - const table = document.getElementById(tableId); - if (!table) return; - - let csv = []; - - // Get Headers - const headerRow = table.querySelector('thead tr'); - if (headerRow) { - const headers = Array.from(headerRow.querySelectorAll('th')).map(th => `"${th.innerText}"`); - csv.push(headers.join(",")); - } + // --- 5. Interactive Tables + CSV Export --- + function escapeHtml(value) { + return String(value ?? '') + .replace(/&/g, '&') + .replace(//g, '>') + .replace(/"/g, '"') + .replace(/'/g, '''); + } - // Get Rows - const rows = table.querySelectorAll('tbody tr'); + function exportRowsToCSV(columns, rows, title) { + const csv = []; + csv.push(columns.map(col => `"${String(col).replace(/"/g, '""')}"`).join(",")); rows.forEach(row => { - const cols = Array.from(row.querySelectorAll('td')).map(td => `"${td.innerText.replace(/"/g, '""')}"`); // Escape quotes - csv.push(cols.join(",")); + const values = columns.map(col => { + const value = row[col]; + const text = value == null ? '' : String(value); + return `"${text.replace(/"/g, '""')}"`; + }); + csv.push(values.join(",")); }); - // Download const csvFile = new Blob([csv.join("\n")], { type: "text/csv;charset=utf-8;" }); const link = document.createElement("a"); if (link.download !== undefined) { @@ -422,10 +422,241 @@

Execut } } - // --- 6. Linked Brushing (Zoom Sync) --- - // Basic implementation: Sync Zoom across plots of same type/dimensionality? - // Or just basic log for now. - // For MVP Phase E, we just stop here. + function exportTableToCSV(tableId, title) { + const table = document.getElementById(tableId); + if (!table) return; + + const columns = []; + const headerRow = table.querySelector('thead tr'); + if (headerRow) { + Array.from(headerRow.querySelectorAll('th')).forEach(th => { + columns.push(th.innerText); + }); + } + + const rows = []; + table.querySelectorAll('tbody tr').forEach(row => { + const values = Array.from(row.querySelectorAll('td')).map(td => td.innerText); + const record = {}; + columns.forEach((column, idx) => { + record[column] = values[idx] ?? ''; + }); + rows.push(record); + }); + exportRowsToCSV(columns, rows, title); + } + + function initInteractiveTables() { + document.querySelectorAll('.interactive-table').forEach(container => { + const dataId = container.getAttribute('data-id'); + const payload = REPORT_DATA[dataId]; + if (!payload) { + container.innerHTML = '
Error: interactive table payload missing.
'; + return; + } + + const config = JSON.parse(container.getAttribute('data-config') || '{}'); + const columns = Array.isArray(payload.columns) ? payload.columns : []; + const allRows = Array.isArray(payload.rows) ? payload.rows : []; + const selectorColumns = Array.isArray(config.selector_columns) ? config.selector_columns.filter(column => columns.includes(column)) : []; + const state = { + search: '', + sortColumn: config.default_sort && columns.includes(config.default_sort.column) ? config.default_sort.column : null, + sortDirection: config.default_sort && config.default_sort.direction === 'desc' ? 'desc' : 'asc', + pageSize: Math.max(1, Number(config.page_size || 50)), + page: 1, + selectorFilters: {}, + }; + + function applyFilters() { + let rows = allRows.slice(); + + selectorColumns.forEach(column => { + const value = state.selectorFilters[column]; + if (value) { + rows = rows.filter(row => String(row[column] ?? '') === value); + } + }); + + if (state.search) { + const query = state.search.toLowerCase(); + rows = rows.filter(row => + columns.some(column => String(row[column] ?? '').toLowerCase().includes(query)) + ); + } + + if (state.sortColumn) { + rows.sort((left, right) => { + const a = left[state.sortColumn]; + const b = right[state.sortColumn]; + const aNum = Number(a); + const bNum = Number(b); + let cmp = 0; + if (!Number.isNaN(aNum) && !Number.isNaN(bNum) && a !== '' && b !== '') { + cmp = aNum - bNum; + } else { + cmp = String(a ?? '').localeCompare(String(b ?? ''), undefined, { numeric: true, sensitivity: 'base' }); + } + return state.sortDirection === 'asc' ? cmp : -cmp; + }); + } + return rows; + } + + function render() { + const filteredRows = applyFilters(); + const totalPages = Math.max(1, Math.ceil(filteredRows.length / state.pageSize)); + state.page = Math.min(state.page, totalPages); + const start = (state.page - 1) * state.pageSize; + const visibleRows = filteredRows.slice(start, start + state.pageSize); + + const selectorHtml = selectorColumns.map(column => { + const options = Array.from(new Set(allRows.map(row => String(row[column] ?? '')).filter(value => value !== ''))).sort((a, b) => a.localeCompare(b, undefined, { numeric: true, sensitivity: 'base' })); + const optionsHtml = ['', ...options.map(option => ``)].join(''); + return ` + + `; + }).join(''); + + const headerHtml = columns.map(column => { + const active = state.sortColumn === column; + const indicator = active ? (state.sortDirection === 'asc' ? ' ▲' : ' ▼') : ''; + return ` + + ${escapeHtml(column)}${indicator} + + `; + }).join(''); + + const bodyHtml = visibleRows.length > 0 + ? visibleRows.map(row => ` + + ${columns.map(column => `${escapeHtml(row[column] ?? '')}`).join('')} + + `).join('') + : `No rows match the current filters.`; + + container.innerHTML = ` +
+
+
+ +
+ ${selectorHtml} + + +
+
+
+ + ${headerHtml} + ${bodyHtml} +
+
+
Showing ${filteredRows.length === 0 ? 0 : start + 1}-${Math.min(start + state.pageSize, filteredRows.length)} of ${filteredRows.length} rows
+
+ + Page ${state.page} / ${totalPages} + +
+
+
+ `; + + const searchEl = container.querySelector('[data-table-search]'); + if (searchEl) { + searchEl.addEventListener('input', event => { + state.search = event.target.value; + state.page = 1; + render(); + }); + } + + container.querySelectorAll('[data-selector-column]').forEach(selectEl => { + selectEl.addEventListener('change', event => { + state.selectorFilters[event.target.getAttribute('data-selector-column')] = event.target.value; + state.page = 1; + render(); + }); + }); + + const pageSizeEl = container.querySelector('[data-page-size]'); + if (pageSizeEl) { + pageSizeEl.addEventListener('change', event => { + state.pageSize = Math.max(1, Number(event.target.value)); + state.page = 1; + render(); + }); + } + + container.querySelectorAll('[data-sort-column]').forEach(headerEl => { + headerEl.addEventListener('click', () => { + const column = headerEl.getAttribute('data-sort-column'); + if (state.sortColumn === column) { + state.sortDirection = state.sortDirection === 'asc' ? 'desc' : 'asc'; + } else { + state.sortColumn = column; + state.sortDirection = 'asc'; + } + render(); + }); + }); + + const exportEl = container.querySelector('[data-export-table]'); + if (exportEl) { + exportEl.addEventListener('click', () => { + exportRowsToCSV(columns, filteredRows, config.title || 'interactive-table'); + }); + } + + container.querySelectorAll('[data-page-action]').forEach(button => { + button.addEventListener('click', () => { + const action = button.getAttribute('data-page-action'); + if (action === 'prev' && state.page > 1) { + state.page -= 1; + } else if (action === 'next' && state.page < totalPages) { + state.page += 1; + } + render(); + }); + }); + } + + render(); + }); + } + + initInteractiveTables(); diff --git a/tests/test_dimred_evaluation.py b/tests/test_dimred_evaluation.py index 2573dc5..062562e 100644 --- a/tests/test_dimred_evaluation.py +++ b/tests/test_dimred_evaluation.py @@ -29,6 +29,10 @@ trajectory_turning_angle, trustworthiness, ) +from coco_pipe.dim_reduction.evaluation.core import ( + SEPARATION_LOGREG_BALANCED_ACCURACY, + evaluate_embedding, +) from coco_pipe.viz.dim_reduction import plot_metrics @@ -113,6 +117,116 @@ def test_method_selector_to_frame(data): assert set(metrics_df["method"]) == {"PCA"} +def test_evaluate_embedding_supervised_metric_records(): + rng = np.random.RandomState(42) + group_labels = np.array([0] * 10 + [1] * 10) + group_features = rng.normal( + loc=group_labels[:, None] * 3.0, scale=0.4, size=(20, 6) + ) + X = np.repeat(group_features, 2, axis=0) + rng.normal(scale=0.1, size=(40, 6)) + y = np.repeat(group_labels, 2) + groups = np.repeat(np.arange(20), 2) + X_emb = X[:, :2] + + payload = evaluate_embedding( + X_emb, + metrics=[SEPARATION_LOGREG_BALANCED_ACCURACY], + labels=y, + groups=groups, + ) + + assert SEPARATION_LOGREG_BALANCED_ACCURACY in payload["metrics"] + records = pd.DataFrame.from_records(payload["records"]) + assert not records.empty + assert records.iloc[0]["metric"] == SEPARATION_LOGREG_BALANCED_ACCURACY + assert records.iloc[0]["scope"] == "global" + assert records.iloc[0]["scope_value"] == "global" + + +def test_evaluate_embedding_supervised_metric_requires_labels_and_groups(): + X_emb = np.random.rand(12, 2) + y = np.array([0, 1] * 6) + groups = np.repeat(np.arange(6), 2) + + with pytest.raises(ValueError, match="labels` and `groups` are required"): + evaluate_embedding( + X_emb, + metrics=[SEPARATION_LOGREG_BALANCED_ACCURACY], + labels=y, + ) + + with pytest.raises(ValueError, match="labels` and `groups` are required"): + evaluate_embedding( + X_emb, + metrics=[SEPARATION_LOGREG_BALANCED_ACCURACY], + groups=groups, + ) + + +def test_dim_reduction_score_supports_grouped_supervised_metric(): + rng = np.random.RandomState(7) + group_labels = np.array([0] * 8 + [1] * 8) + group_features = rng.normal( + loc=group_labels[:, None] * 2.5, scale=0.5, size=(16, 5) + ) + X = np.repeat(group_features, 2, axis=0) + rng.normal(scale=0.05, size=(32, 5)) + y = np.repeat(group_labels, 2) + groups = np.repeat(np.arange(16), 2) + + reducer = DimReduction("PCA", n_components=2) + X_emb = reducer.fit_transform(X) + scores = reducer.score( + X_emb, + metrics=[SEPARATION_LOGREG_BALANCED_ACCURACY], + labels=y, + groups=groups, + ) + + assert SEPARATION_LOGREG_BALANCED_ACCURACY in scores["metrics"] + assert any( + record["metric"] == SEPARATION_LOGREG_BALANCED_ACCURACY + for record in reducer.metric_records_ + ) + + +def test_method_selector_from_records_and_frame_preserve_extra_columns(): + records = [ + { + "method": "PCA", + "metric": SEPARATION_LOGREG_BALANCED_ACCURACY, + "value": 0.72, + "scope": "global", + "scope_value": "global", + "fit_id": "fit-a", + "eval_name": "epilepsy", + }, + { + "method": "UMAP", + "metric": SEPARATION_LOGREG_BALANCED_ACCURACY, + "value": 0.84, + "scope": "global", + "scope_value": "global", + "fit_id": "fit-a", + "eval_name": "epilepsy", + }, + ] + + selector = MethodSelector.from_records(records) + frame = selector.to_frame() + assert {"fit_id", "eval_name"} <= set(frame.columns) + + ranked = selector.rank_methods(SEPARATION_LOGREG_BALANCED_ACCURACY) + assert ranked.iloc[0]["method"] == "UMAP" + + selector_from_frame = MethodSelector.from_frame(frame) + frame_roundtrip = selector_from_frame.to_frame() + assert {"fit_id", "eval_name"} <= set(frame_roundtrip.columns) + ranked_roundtrip = selector_from_frame.rank_methods( + SEPARATION_LOGREG_BALANCED_ACCURACY + ) + assert ranked_roundtrip.iloc[0]["method"] == "UMAP" + + def test_velocity_fields(linear_data): X = linear_data X_emb = X @@ -1268,3 +1382,43 @@ def test_method_selector_rank_methods_missing_records(): selector = MethodSelector([]) with pytest.raises(ValueError, match="No evaluation metrics available"): selector.rank_methods(selection_metric="trustworthiness") + + +def test_evaluate_embedding_k_normalization_skip(): + """Test skipping of k-metrics when normalizer is non-positive.""" + X = np.random.rand(10, 5) + X_emb = X[:, :2] + # For n_samples=10, 2*n_samples - 3*k - 1 <= 0 means 20 - 3*k - 1 <= 0 + # => 19 <= 3*k => k >= 7 + # trustworthiness requires (2*n_samples - 3*k - 1) > 0 + result = evaluate_embedding( + X_emb, + X=X, + metrics=["trustworthiness"], + k_values=[7], + ) + # The metric should be skipped, so it shouldn't be in result['metrics'] + assert "trustworthiness" not in result["metrics"] + + +def test_evaluate_embedding_default_metrics_selection(): + """Test default metric selection.""" + X = np.random.rand(20, 5) + X_emb = X[:, :2] + # Calling with metrics=None (default) should trigger the standard selection + result = evaluate_embedding(X_emb, X=X, metrics=None) + assert "trustworthiness" in result["metrics"] + assert SEPARATION_LOGREG_BALANCED_ACCURACY not in result["metrics"] + + +def test_method_selector_init_dict_invalid_type(): + """Test MethodSelector initialization with dict and invalid reducer type.""" + with pytest.raises(TypeError, match="only accepts scored DimReduction objects"): + MethodSelector({"pca": "not a reducer"}) + + +def test_evaluate_embedding_invalid_dim(): + """Test evaluate_embedding with invalid X_emb dimensionality.""" + X_emb = np.random.rand(10, 2, 2, 2) + with pytest.raises(ValueError, match="must be either 2D or 3D"): + evaluate_embedding(X_emb) diff --git a/tests/test_report_core.py b/tests/test_report_core.py index 5de9b5d..4f0d553 100644 --- a/tests/test_report_core.py +++ b/tests/test_report_core.py @@ -8,6 +8,7 @@ from coco_pipe.report.core import ( HtmlElement, ImageElement, + InteractiveTableElement, MetricsTableElement, PlotlyElement, Report, @@ -128,6 +129,47 @@ def test_table_element_dict_inputs(): assert "2" in el_non_scalar.render() +def test_interactive_table_element_payload_and_render(tmp_report_file): + df = pd.DataFrame( + { + "eval_name": ["epilepsy", "adhd"], + "reducer": ["PCA", "UMAP"], + "score": [0.71, 0.82], + } + ) + element = InteractiveTableElement( + df, + title="Interactive Metrics", + selector_columns=["eval_name", "reducer"], + default_sort={"column": "score", "direction": "desc"}, + page_size=25, + ) + + registry = {} + element.collect_payload(registry) + assert len(registry) == 1 + payload = next(iter(registry.values())) + assert payload["columns"] == ["eval_name", "reducer", "score"] + assert len(payload["rows"]) == 2 + + html = element.render() + assert 'class="interactive-table"' in html + assert 'data-id="' in html + assert 'data-config="' in html + + report = Report(title="Interactive Table Report") + report.add_element(element) + report.save(str(tmp_report_file)) + content = tmp_report_file.read_text(encoding="utf-8") + assert "interactive-table" in content + assert "initInteractiveTables" in content + assert "data-table-search" in content + assert "data-sort-column" in content + assert "data-selector-column" in content + assert "data-export-table" in content + assert "data-page-size" in content + + def test_metrics_table_highlighting(): df = pd.DataFrame( {"method": ["A", "B"], "acc": [0.8, 0.9], "loss": [0.2, 0.1]} From 6dce0f247353aef5b9252342a80dcadb6699637e Mon Sep 17 00:00:00 2001 From: Hamza Abdelhedi Date: Sun, 19 Apr 2026 14:03:34 -0400 Subject: [PATCH 7/7] minor fixes to report sections --- coco_pipe/report/templates/section.html | 11 ++++++----- coco_pipe/viz/plotly_utils.py | 4 +++- 2 files changed, 9 insertions(+), 6 deletions(-) diff --git a/coco_pipe/report/templates/section.html b/coco_pipe/report/templates/section.html index 78e3466..f991feb 100644 --- a/coco_pipe/report/templates/section.html +++ b/coco_pipe/report/templates/section.html @@ -1,7 +1,7 @@
+ onclick="const content=this.nextElementSibling; if(content){content.classList.toggle('hidden');} const chevron=this.querySelector('[data-chevron]'); if(chevron){chevron.classList.toggle('rotate-180');}">

{% if icon %} {{ icon | safe }} @@ -23,20 +23,21 @@

{% if status != "OK" %} {{ status }} - {% else %} - - {% endif %} +

-
+
{% if findings %}
diff --git a/coco_pipe/viz/plotly_utils.py b/coco_pipe/viz/plotly_utils.py index c6df79d..4e2a6c5 100644 --- a/coco_pipe/viz/plotly_utils.py +++ b/coco_pipe/viz/plotly_utils.py @@ -267,12 +267,14 @@ def _marker_payload( if is_categorical(values): if hasattr(values, "cat"): categories = values.cat.categories.tolist() + lookup_values = values.astype(str) else: categories = sorted( pd.Series(values).dropna().astype(str).unique().tolist() ) + lookup_values = pd.Series(values).astype(str) cat_map = {cat: i for i, cat in enumerate(categories)} - mapped = [cat_map.get(v, np.nan) for v in values] + mapped = [cat_map.get(v, np.nan) for v in lookup_values] _, colorscale = _discrete_colorscale(categories, palette=palette) payload = { "color": mapped,