diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index e00ed75..a50793d 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -113,6 +113,28 @@ jobs: files: coverage.xml fail_ci_if_error: false + test-mne: + name: Test MNE wrapper (optional [mne] extra) + needs: [lint, typecheck] + runs-on: ubuntu-latest + env: + PYTORCH_ENABLE_MPS_FALLBACK: "1" + steps: + - uses: actions/checkout@v4 + - uses: astral-sh/setup-uv@v5 + with: + python-version: "3.12" + enable-cache: true + # mne IS pip-installable on Linux (unlike the Apple-only mlx extra), so the + # AMICAICA wrapper (issue #139) is exercised for real in CI here. Its tests + # live in pamica/tests/mne_tests and self-skip in the base `test` job + # (importorskip on mne); this job installs the [mne] extra so they run on + # the real sample EEG. pamica/mne_compat is omitted from coverage in + # pyproject (the base env has no mne), so no fail-under gate applies here. + - run: uv sync --extra mne + - name: pytest (mne wrapper, real sample EEG) + run: uv run pytest pamica/tests/mne_tests + build: name: Build + import (Python ${{ matrix.python-version }}) needs: [lint, typecheck] diff --git a/docs/api/index.md b/docs/api/index.md index e9b7e45..3be3169 100644 --- a/docs/api/index.md +++ b/docs/api/index.md @@ -27,3 +27,13 @@ from pamica.mlx_impl import AMICAMLXNG # requires the `mlx` extra - **[`AMICAMLXNG`](mlx-backend.md)** — the optional Apple-GPU (MLX) backend; the fastest option on Apple Silicon (float32). + +The optional MNE-Python wrapper is likewise imported explicitly: + +```python +from pamica.mne_compat import AMICAICA # requires the `mne` extra +``` + +- **[`AMICAICA`](mne-compat.md)** — fit AMICA from an MNE `Raw`/`Epochs` and + interoperate with `mne.preprocessing.ICA` (`get_sources`, `apply`, + `plot_components`, `to_mne_ica`). diff --git a/docs/api/mne-compat.md b/docs/api/mne-compat.md new file mode 100644 index 0000000..3ff177f --- /dev/null +++ b/docs/api/mne-compat.md @@ -0,0 +1,125 @@ +# MNE-Python compatibility (AMICAICA) + +`AMICAICA` fits AMICA directly from an [MNE-Python](https://mne.tools) +`Raw`/`Epochs` and hands the result back through the standard MNE ICA surface. +It is **additive**: the scikit-learn-style [`AMICA`](amica.md) interface and the +byte-identical [EEGLAB output](../guides/eeglab.md) are unchanged; this is a +second entry point for MNE users, not a replacement. + +MNE is an optional dependency, so `import pamica` never requires it. Install the +extra and import the wrapper explicitly: + +```bash +pip install pamica[mne] +``` + +```python +import mne +from pamica.mne_compat import AMICAICA + +raw = mne.io.read_raw_eeglab("subject.set", preload=True) + +ica = AMICAICA(n_mix=3, random_state=42).fit(raw, picks="eeg", max_iter=100) + +sources = ica.get_sources(raw) # an mne.io.RawArray of component activations +maps = ica.get_components() # scalp maps, shape (n_channels, n_components) +ica.plot_components() # native mne.viz topographies + +# Reconstruct with some components removed: +clean = ica.apply(raw.copy(), exclude=[0, 3]) +``` + +`fit` accepts a `Raw` or `Epochs` (epochs are concatenated along time, as MNE's +own ICA does), any MNE `picks` selector, and forwards remaining keywords +(`max_iter`, `lrate`, `do_newton`, ...) to [`AMICA.fit`](amica.md). It rejects +non-finite input and PCA reduction (`pcakeep`/`pcadb`, which leaves the sphere +rank-deficient so the full-rank export would be invalid), and a degenerate +(diverged) fit is refused by the consumer methods rather than emitting NaNs. + +## Interoperating with `mne.preprocessing.ICA` + +`to_mne_ica()` returns a fully-populated +[`mne.preprocessing.ICA`](https://mne.tools/stable/generated/mne.preprocessing.ICA.html), +so the entire MNE ICA ecosystem (plotting, `find_bads_eog`/`_ecg`, exclusion +workflows) works on an AMICA decomposition: + +```python +mne_ica = ica.to_mne_ica() +eog_idx, scores = mne_ica.find_bads_eog(raw) +mne_ica.plot_scores(scores) +``` + +The wrapper's `get_sources`, `apply`, `get_components`, `plot_components` and +`plot_sources` delegate to this object, so they reproduce `AMICA.transform` +exactly: MNE +computes sources as `unmixing_matrix_ @ pca_components_ @ (X - pca_mean_)`, and +the export maps pamica's mean, symmetric-ZCA sphere and unmixing into those +matrices (writing the sphere as `V diag(1/√e) Vᵀ` with `V` orthonormal so MNE's +scalp maps come out in channel space). The equivalence +`to_mne_ica().get_sources(raw) == AMICA.transform(X)` is pinned by the test +suite on real sample EEG. + +## Multi-model fits + +AMICA can learn a mixture of ICA models (`n_models > 1`). MNE's `ICA` represents +only one unmixing matrix, so each model is exported as its own single-model +`mne.preprocessing.ICA`, and the per-sample *model dominance* (which model best +explains each timepoint) is exposed directly, since MNE has no concept for it: + +```python +ica = AMICAICA(n_models=2, random_state=42).fit(raw, max_iter=100) + +# Per-model: model_idx selects the model on every consumer method. +sources_m1 = ica.get_sources(raw, model_idx=1) +ica.plot_components(model_idx=1) +model1 = ica.to_mne_ica(model_idx=1) # a standard ICA for model 1 + +# Model dominance over time (P(model | sample), columns sum to 1): +prob = ica.get_model_probability(raw) # (n_models, n_samples) +ica.plot_model_probability(raw) # per-model probability + best-model LL +``` + +`get_model_probability`/`plot_model_probability` build on the public +`AMICA.model_loglik`/`model_probability` accessors, which score arbitrary data +through the fitted sphere and mean. Each per-model export folds that model's +data-space center into `pca_mean_`, so `to_mne_ica(model_idx=h).get_sources(raw)` +reproduces `AMICA.transform(X, model_idx=h)` (with `X` the picked channel array) +for every model, not just the first. + +## Inspecting pamica-specific metadata + +An `mne.preprocessing.ICA` has no field for AMICA's adaptive source densities or +component sharing, so rather than drop them, the wrapper exposes them directly: + +```python +from pamica.mne_compat import AMICAICA, PDFTYPE_NAMES + +ica = AMICAICA(n_models=2, random_state=42).fit(raw, max_iter=100) + +families = ica.get_pdftype(model_idx=0) # (n_components,) codes 0-4 +names = [PDFTYPE_NAMES[c] for c in families] # e.g. "generalized_gaussian" +rho = ica.get_rho(model_idx=0) # (n_mix, n_components) GG shape +shared = ica.shared_components() # [(model, comp), ...] groups +``` + +`get_pdftype` returns each component's density family (0 generalized Gaussian, +1 super-Gaussian cosh, 2 Gaussian, 3 logistic, 4 sub-Gaussian cosh; they differ +per component only under the adaptive switcher `pdftype=1`). `get_rho` is the +generalized-Gaussian shape (meaningful for `pdftype=0`). `shared_components` +lists components merged across models by `share_comps` (empty otherwise). The +same three accessors exist on the scikit-learn-style [`AMICA`](amica.md). + +## Separation-quality metrics + +The [MIR/PMI metrics](metrics.md) are available directly on an MNE object, so +MNE-side users get the same separation-quality numbers as EEGLAB-side users: + +```python +mir_nats, variance = ica.mir(raw, model_idx=0) # mutual information reduced +mi_matrix = ica.pmi(raw, model_idx=0) # pairwise MI between sources +``` + +Both extract the fitted channels from the `Raw`/`Epochs` and delegate to +`AMICA.mir`/`pmi`, reproducing the array API exactly. + +::: pamica.mne_compat.AMICAICA diff --git a/docs/changelog.md b/docs/changelog.md index fc562a2..dacbced 100644 --- a/docs/changelog.md +++ b/docs/changelog.md @@ -3,6 +3,60 @@ Release notes are also published on the [GitHub releases page](https://github.com/sccn/pAMICA/releases). +## 0.3.0 + +MNE-Python compatibility layer (epic #139), additive: the scikit-learn-style +`AMICA` API and the byte-identical EEGLAB I/O are unchanged. + +- `pamica.mne_compat.AMICAICA`, an MNE-facing wrapper that fits AMICA directly + from an `mne.io.Raw`/`Epochs` (`picks=...`, epochs concatenated along time like + MNE's own ICA) and interoperates with the standard MNE ICA consumer surface: + `get_sources`, `apply`, `get_components`, `plot_components` and `plot_sources`. + `to_mne_ica()` + returns a fully-populated `mne.preprocessing.ICA` (including `reject_`/ + `n_samples_`, so `ICA.save` and `plot_properties` work), so the whole MNE ICA + ecosystem (component plotting, `find_bads_eog`/`_ecg`, exclusion workflows) + works on an AMICA decomposition. The export maps pamica's mean, symmetric-ZCA + sphere and unmixing into MNE's `pca_mean_`/`pca_components_`/`unmixing_matrix_`, + writing the sphere as `V diag(1/sqrt(e)) V^T` with `V` orthonormal so MNE's + scalp maps are in channel space; `to_mne_ica().get_sources(raw)` reproduces + `AMICA.transform(X)` to float64 precision, pinned on real sample EEG. `fit` + rejects PCA reduction (`pcakeep`/`pcadb`, which leaves the sphere rank-deficient + and the export invalid) and non-finite input, and a degenerate fit is refused + by the consumer methods rather than emitting NaNs. MNE is an + optional extra (`pip install pamica[mne]`); `import pamica` never requires it, + and a dedicated CI job runs the wrapper tests with the extra installed (phase 1, + single-model, #140). +- Multi-model exposure through the MNE wrapper: `AMICAICA(n_models=...)` fits a + mixture of ICA models, and since MNE's `ICA` represents only one unmixing, + each model is exported as its own single-model `mne.preprocessing.ICA` via + `to_mne_ica(model_idx=...)` (and the `model_idx` argument on `get_sources`/ + `apply`/`get_components`/`plot_components`/`plot_sources`). The per-sample model + dominance MNE cannot represent is exposed directly: `get_model_probability(inst)` + returns `P(model | sample)` (`(n_models, n_samples)`, columns sum to 1) and + `plot_model_probability(inst)` draws the per-model probability plus best-model + log-likelihood over time. These build on a new public live accessor, + `AMICA.model_loglik`/`model_probability` (and the `AMICATorchNG` equivalents), + which score arbitrary data through the stored sphere/mean; the training-data + path (without `do_reject`) is pinned bit-for-bit against the E-step's own `Lht`. The per-model export + folds each model's data-space center `c` into `pca_mean_`, so the round trip + holds for the multi-model case too. `pamica.viz.plot_model_probability` now also + accepts a live `lht` array, not only a written `AmicaOutput` (phase 2, #141). +- pamica-specific fitted metadata is inspectable through the MNE wrapper rather + than silently dropped by the `mne.preprocessing.ICA` export: `get_pdftype(model_idx=...)` + returns each component's source-density family code (0-4, named by + `pamica.mne_compat.PDFTYPE_NAMES`), `get_rho(model_idx=...)` the + generalized-Gaussian shape parameters, and `shared_components()` the components + merged across models by `share_comps`. The same accessors are added to + `AMICA`/`AMICATorchNG` (phase 3, #142). +- Separation-quality metrics are available directly on an MNE object: + `AMICAICA.mir(inst, model_idx=...)` (Mutual Information Reduction, in nats) and + `AMICAICA.pmi(inst, model_idx=...)` (pairwise mutual information between the + fitted sources), so MNE-side users get the same metrics as EEGLAB-side users. + Both extract the fitted channels from the `Raw`/`Epochs` and delegate to + `AMICA.mir`/`pmi` (#133); the results match the array API exactly (phase 4, + #143). + ## 0.2.2 GitHub repository rename to pAMICA and a `__version__` fix. diff --git a/mkdocs.yml b/mkdocs.yml index bce8084..65307d3 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -102,6 +102,7 @@ nav: - MLX backend: api/mlx-backend.md - NumPy backend: api/numpy-backend.md - Native backend: api/native-backend.md + - MNE compatibility: api/mne-compat.md - Separation metrics: api/metrics.md - Visualization: api/viz.md - Development: diff --git a/pamica/amica.py b/pamica/amica.py index 869c220..b264107 100644 --- a/pamica/amica.py +++ b/pamica/amica.py @@ -419,6 +419,105 @@ def pmi( return self.model_.pmi(X, model_idx=model_idx, nbins=nbins) + def model_loglik(self, X: np.ndarray) -> np.ndarray: + """Per-model, per-sample log-likelihood ``Lht`` on ``X`` (issue #141). + + Delegates to :meth:`AMICATorchNG.model_loglik`. For a multi-model fit + this is the joint log-likelihood of each model at each sample, from + which the per-sample model posterior (dominance) is + ``softmax(Lht, axis=0)``; see :meth:`model_probability`. + + Parameters + ---------- + X : np.ndarray of shape (n_channels, n_samples) + Raw (unpreprocessed) data. + + Returns + ------- + Lht : np.ndarray of shape (n_models, n_samples) + + Raises + ------ + ValueError + If the model is unfitted, or if ``X`` is non-finite. + RuntimeError + If the fit ended degenerate (issue #50). + """ + self._check_usable("compute the model log-likelihood") + assert self.model_ is not None + + return self.model_.model_loglik(X) + + def model_probability(self, X: np.ndarray) -> np.ndarray: + """Per-sample posterior probability of each model (issue #141). + + Delegates to :meth:`AMICATorchNG.model_probability`: the column-wise + ``softmax`` over models of :meth:`model_loglik`, i.e. ``P(model h | + x_t)``. Each column sums to 1; all ones for a single model. + + Parameters + ---------- + X : np.ndarray of shape (n_channels, n_samples) + Raw (unpreprocessed) data. + + Returns + ------- + prob : np.ndarray of shape (n_models, n_samples) + + Raises + ------ + ValueError + If the model is unfitted, if ``X`` is non-finite, or if every model + underflows to ``-inf`` log-likelihood at some sample. + RuntimeError + If the fit ended degenerate (issue #50). + """ + self._check_usable("compute the model probability") + assert self.model_ is not None + + return self.model_.model_probability(X) + + def get_pdftype(self, model_idx: int = 0) -> np.ndarray: + """Per-source density-family code for model ``model_idx`` (issue #142). + + Delegates to :meth:`AMICATorchNG.get_pdftype`. One integer per source + component (0-4; see :data:`pamica.torch_impl.PDFTYPE_NAMES`). + + Returns + ------- + np.ndarray of int, shape (n_sources,) + """ + self._check_usable("get the density family") + assert self.model_ is not None + + return self.model_.get_pdftype(model_idx=model_idx) + + def get_rho(self, model_idx: int = 0) -> np.ndarray: + """Generalized-Gaussian shape ``rho`` for model ``model_idx`` (issue #142). + + Delegates to :meth:`AMICATorchNG.get_rho`. + + Returns + ------- + np.ndarray of float, shape (n_mix, n_sources) + """ + self._check_usable("get rho") + assert self.model_ is not None + + return self.model_.get_rho(model_idx=model_idx) + + def shared_components(self) -> list: + """Components shared across models by ``share_comps`` (issue #142). + + Delegates to :meth:`AMICATorchNG.shared_components`: one group of + ``(model_idx, source_idx)`` pairs per shared column; empty when nothing + is shared. + """ + self._check_usable("get the shared components") + assert self.model_ is not None + + return self.model_.shared_components() + def variance_order( self, model_idx: int = 0, return_svar: bool = False ) -> Union[np.ndarray, tuple]: diff --git a/pamica/mne_compat/__init__.py b/pamica/mne_compat/__init__.py new file mode 100644 index 0000000..2bc1ef3 --- /dev/null +++ b/pamica/mne_compat/__init__.py @@ -0,0 +1,15 @@ +"""MNE-Python compatibility layer for pamica (issue #139). + +This subpackage is an *optional* entry point: it imports ``mne`` at module +import time, so it is intentionally NOT imported by :mod:`pamica.__init__`. +``import pamica`` therefore never requires MNE; reach the wrapper explicitly:: + + from pamica.mne_compat import AMICAICA + +Install the dependency with ``pip install pamica[mne]`` (same lazy-extra +pattern as the optional ``mlx`` backend). +""" + +from .core import AMICAICA, PDFTYPE_NAMES + +__all__ = ["AMICAICA", "PDFTYPE_NAMES"] diff --git a/pamica/mne_compat/core.py b/pamica/mne_compat/core.py new file mode 100644 index 0000000..b309452 --- /dev/null +++ b/pamica/mne_compat/core.py @@ -0,0 +1,602 @@ +"""MNE-Python-facing wrapper over pamica's AMICA backend (issue #139). + +:class:`AMICAICA` fits AMICA directly from an :class:`mne.io.Raw` / +:class:`mne.Epochs` and exposes the standard MNE ICA consumer surface +(``get_sources``, ``apply``, ``get_components``, ``plot_components``, +``plot_sources``). Its core value-add is :meth:`AMICAICA.to_mne_ica`, which maps +pamica's fitted mean/sphere/unmixing into a fully-populated +:class:`mne.preprocessing.ICA`; the other methods delegate to that object, so MNE +performs the export/plot/back-projection machinery unchanged. + +Multi-model AMICA (``n_models > 1``) is exposed per-model (each model is its own +single-model MNE ICA), plus a per-sample model-dominance accessor +(:meth:`AMICAICA.get_model_probability`) that MNE's ``ICA`` cannot represent +(issue #141). The pamica-specific fitted metadata MNE cannot hold -- source- +density family, GG shape, component sharing -- stays inspectable via +:meth:`AMICAICA.get_pdftype` / :meth:`get_rho` / :meth:`shared_components` +(issue #142), and the separation-quality metrics :meth:`AMICAICA.mir` / +:meth:`pmi` run directly on an MNE object (issue #143). PCA reduction is +unsupported (full-rank whitening only). +""" + +from typing import Optional, Union + +import numpy as np +import torch + +# MNE is an optional extra (`pip install pamica[mne]`); this subpackage is never +# imported by ``pamica.__init__``, so ``import pamica`` does not require it. The +# base type-check/CI env has no mne, hence the scoped ignore (issue #139). +import mne # ty: ignore[unresolved-import] +from mne.preprocessing import ICA as _MNEICA # ty: ignore[unresolved-import] + +from ..amica import AMICA +from ..torch_impl import PDFTYPE_NAMES + +__all__ = ["AMICAICA", "PDFTYPE_NAMES"] + + +class AMICAICA: + """Fit AMICA from MNE objects and interoperate with ``mne.preprocessing.ICA``. + + The wrapper fits pamica's natural-gradient AMICA backend on the data of an + MNE :class:`~mne.io.Raw` or :class:`~mne.Epochs` and lets MNE consume the + result: :meth:`get_sources`, :meth:`apply`, :meth:`get_components`, + :meth:`plot_components` and :meth:`plot_sources` all delegate to a real + :class:`mne.preprocessing.ICA` built by :meth:`to_mne_ica`. + + For a multi-model fit (``n_models > 1``) each model is exported as its own + single-model MNE ICA (``to_mne_ica(model_idx=...)`` / the ``model_idx`` + argument on the consumer methods), and the per-sample model dominance -- + which MNE's ``ICA`` cannot represent -- is exposed directly by + :meth:`get_model_probability` / :meth:`plot_model_probability`. + + Separation-quality metrics (issue #133) are available directly on an MNE + object: :meth:`mir` (Mutual Information Reduction) and :meth:`pmi` (pairwise + mutual information between sources). The pamica-specific fitted metadata MNE + cannot hold -- source-density family, GG shape, component sharing -- is + inspectable via :meth:`get_pdftype` / :meth:`get_rho` / :meth:`shared_components`. + + Parameters + ---------- + n_models : int, default=1 + Number of ICA models to learn (AMICA ``n_models``). + n_mix : int, default=3 + Number of mixture components per source (AMICA ``n_mix``). + random_state : int or None, default=None + Seed for the AMICA fit (passed through as the backend ``seed``) and + stored on the exported :class:`~mne.preprocessing.ICA`. + device : str or torch.device, optional + Torch device for the fit (``None`` = auto; the float64 parity backend + falls back to CPU when auto-selection lands on MPS). See :class:`AMICA`. + verbose : bool, default=True + Whether the underlying :class:`AMICA` prints fit progress. + + Attributes + ---------- + amica_ : AMICA + The fitted pamica model (holds all ``n_models`` models). + info_ : mne.Info + The picked measurement info the fit was run on (channel subset only). + ch_names_ : list of str + Names of the fitted channels, in order. + n_components_ : int + Number of ICA components (equals the number of fitted channels; AMICA + keeps ``n_sources == n_channels``). + converged_ : bool + Whether the last fit ended usable (not degenerate). A degenerate fit is + kept for inspection but refused by the consumer methods (issue #50). + stop_reason_ : str or None + Why the backend fit stopped (e.g. ``"max_iter"``, ``"nan_ll"``). + + Notes + ----- + Model ``h``'s AMICA transform is ``S = W_fort @ (sphere @ (X - mean) - c_h)``, + where ``c_h`` is that model's data-space center (identically zero for a + single model, since the ``c`` update is gated to ``n_models > 1``). MNE + computes sources as ``S = unmixing_matrix_ @ pca_components_ @ (X - pca_mean_)`` + (with a unit pre-whitener). Writing the symmetric-ZCA ``sphere`` as + ``V @ diag(1/sqrt(e)) @ V.T`` with ``V`` orthonormal, the exported ICA for + model ``h`` uses ``pca_components_ = V.T``, + ``unmixing_matrix_ = W_fort @ sphere @ V`` and + ``pca_mean_ = mean + inv(sphere) @ c_h`` (which reduces to ``mean`` when + ``c_h`` is zero). Keeping ``pca_components_`` orthonormal is what makes MNE's + ``get_components`` (scalp maps ``inv(sphere) @ inv(W_fort)``) come out right, + since MNE assumes orthonormal PCA rows. The mapping is pinned by a round-trip + test (``to_mne_ica(model_idx=h).get_sources(raw)`` equals + ``amica_.transform(X, model_idx=h)``). + + PCA reduction (``pcakeep``/``pcadb``) is unsupported: it leaves the sphere + rank-deficient, so the export (which assumes a full-rank whitening) would be + numerically invalid, and :meth:`fit` rejects it. Rank-deficient *data* (for + example average-referenced EEG) is the same hazard reached through data + conditioning rather than a keyword; such a fit typically diverges and is + refused as degenerate. + """ + + def __init__( + self, + n_models: int = 1, + n_mix: int = 3, + random_state: Optional[int] = None, + device: Optional[Union[str, torch.device]] = None, + verbose: bool = True, + ): + self.n_models = n_models + self.n_mix = n_mix + self.random_state = random_state + self.device = device + self.verbose = verbose + + self.amica_: Optional[AMICA] = None + self.info_ = None + self.ch_names_: Optional[list] = None + self.n_components_: Optional[int] = None + self.converged_: bool = False + self.stop_reason_: Optional[str] = None + self._n_samples: Optional[int] = None + self._fit_kind: Optional[str] = None + # Per-model export cache (one mne.preprocessing.ICA per model_idx). + self._mne_ica_cache: dict = {} + + # ------------------------------------------------------------------ + # Fit + # ------------------------------------------------------------------ + def fit( + self, + inst, + picks=None, + start: Optional[int] = None, + stop: Optional[int] = None, + **fit_kwargs, + ) -> "AMICAICA": + """Fit AMICA to the data of an MNE ``Raw`` or ``Epochs``. + + Parameters + ---------- + inst : mne.io.BaseRaw | mne.BaseEpochs + The data to decompose. ``Epochs`` are concatenated along time + (``np.hstack``), matching how MNE's own ICA fits epoched data. + picks : str | list | slice | None, default=None + Channels to fit, in any form MNE accepts (e.g. ``"eeg"``, + ``"data"``, a name list). ``None`` selects all good data channels + (bads excluded), matching MNE's ICA default. + start, stop : int | None + Sample range for ``Raw`` input, passed to ``get_data``; ``None`` uses + the full recording. Not supported for ``Epochs`` (raises). + **fit_kwargs + Forwarded to :meth:`AMICA.fit` (e.g. ``max_iter``, ``lrate``, + ``do_newton``) and the backend constructor (e.g. ``block_size``). + ``pcakeep``/``pcadb`` (PCA reduction) are rejected, since they leave + the sphere rank-deficient and the MNE export invalid. + + Returns + ------- + self : AMICAICA + + Raises + ------ + TypeError + If ``inst`` is not an MNE ``Raw``/``Epochs``. + ValueError + If ``start``/``stop`` are given for ``Epochs``, ``stop`` exceeds the + recording length, the selected data is non-finite, or the fit used + PCA reduction. + """ + if not isinstance(inst, (mne.io.BaseRaw, mne.BaseEpochs)): + raise TypeError( + "AMICAICA.fit expects an mne.io.Raw or mne.Epochs, got " + f"{type(inst).__name__}." + ) + is_raw = isinstance(inst, mne.io.BaseRaw) + if not is_raw and (start is not None or stop is not None): + raise ValueError("start/stop are only supported for Raw input, not Epochs.") + + picked = inst.copy().pick("data" if picks is None else picks, exclude="bads") + ch_names = list(picked.ch_names) + + if is_raw: + if stop is not None and stop > inst.n_times: + raise ValueError( + f"stop={stop} exceeds the recording length ({inst.n_times} " + "samples)." + ) + range_kw = {} + if start is not None: + range_kw["start"] = start + if stop is not None: + range_kw["stop"] = stop + X = picked.get_data(**range_kw) + fit_kind = "raw" + else: + # Concatenate epochs along time, as MNE's ICA does for fitting. + X = np.hstack(picked.get_data()) + fit_kind = "epochs" + + X = np.ascontiguousarray(X, dtype=np.float64) + if not np.isfinite(X).all(): + bad = [ch_names[i] for i in np.flatnonzero(~np.isfinite(X).all(axis=1))] + raise ValueError( + "AMICAICA.fit: input contains non-finite (NaN/Inf) samples in " + f"channel(s) {bad}; clean bad segments/interpolation before " + "fitting." + ) + + if isinstance(self.random_state, (int, np.integer)): + fit_kwargs.setdefault("seed", int(self.random_state)) + + amica = AMICA( + n_models=self.n_models, + n_mix=self.n_mix, + device=self.device, + verbose=self.verbose, + ) + amica.fit(X, **fit_kwargs) + + if amica.model_ is not None and amica.model_._pca_reduced(): + raise ValueError( + "AMICAICA does not support PCA reduction (pcakeep/pcadb): it " + "leaves the sphere rank-deficient, so the MNE ICA export (which " + "assumes a full-rank whitening) would be numerically invalid. " + "Refit without pcakeep/pcadb." + ) + + # Publish to self only after every fallible step above succeeds, so a + # failed (re)fit leaves the previously fitted state intact rather than a + # mix of the old model and the new attempt's metadata (mirrors the + # local-first pattern in AMICA.fit). + self.info_ = picked.info + self.ch_names_ = ch_names + self.n_components_ = X.shape[0] + self._n_samples = X.shape[1] + self._fit_kind = fit_kind + self.amica_ = amica + self.converged_ = amica.converged_ + self.stop_reason_ = amica.stop_reason_ + self._mne_ica_cache = {} # invalidate any cached exports + return self + + # ------------------------------------------------------------------ + # Export + # ------------------------------------------------------------------ + def to_mne_ica(self, model_idx: int = 0) -> _MNEICA: + """Build (and cache) a fully-populated :class:`mne.preprocessing.ICA`. + + The returned object is a genuine MNE ICA: ``get_sources``, ``apply``, + ``get_components``, ``save`` and the ``mne.viz`` plotters operate on it + natively. See the class :class:`Notes ` for the + mean/sphere/unmixing to ``pca_mean_``/``pca_components_``/ + ``unmixing_matrix_`` mapping. + + For a multi-model fit each model is exported as its own single-model + MNE ICA (MNE has no multi-model concept); pass ``model_idx`` to pick + one. The per-model exports are cached and returned by reference: + mutating one (for example setting ``.exclude``) persists across + subsequent :meth:`apply`/:meth:`get_sources` calls for that model until + the next :meth:`fit`. + + Parameters + ---------- + model_idx : int, default=0 + Which AMICA model to export (``0..n_models-1``). + + Returns + ------- + ica : mne.preprocessing.ICA + """ + self._check_fitted("build the MNE ICA") + model_idx = self._check_model_idx(model_idx) + if model_idx in self._mne_ica_cache: + return self._mne_ica_cache[model_idx] + + amica = self.amica_ + if amica is None or amica.model_ is None or self.ch_names_ is None: + raise RuntimeError( + "AMICAICA: internal state is inconsistent; refit before to_mne_ica()." + ) + backend = amica.model_ + if backend.mean is None or backend.sphere is None or backend.c is None: + raise RuntimeError( + "AMICAICA: the fitted backend is missing mean/sphere/c; refit " + "before to_mne_ica()." + ) + + mean = backend.mean.cpu().numpy().ravel() + sphere = backend.sphere.cpu().numpy() + w_fort = amica.get_unmixing_matrix(model_idx=model_idx) + c = backend.c.cpu().numpy()[:, model_idx] # per-model center (sphered space) + n_ch = sphere.shape[0] + + # Orthonormal eigenbasis of the symmetric-ZCA sphere + # (sphere = V diag(1/sqrt(cov_eval)) V.T). eigh gives ascending + # sphere-eigenvalues (= 1/sqrt(cov_eval)); reorder to descending + # explained variance so pca_components_ matches MNE's PCA convention. + sphere_evals, evecs = np.linalg.eigh(sphere) + cov_evals = 1.0 / sphere_evals**2 + order = np.argsort(cov_evals)[::-1] + v = evecs[:, order] + cov_evals = cov_evals[order] + + pca_components = v.T + unmixing = w_fort @ sphere @ v + # Fold the per-model center c (in sphered space) into pca_mean via the + # data-space offset inv(sphere) @ c, so MNE's (X - pca_mean) reproduces + # AMICA's W(sphere(X - mean) - c). c is identically zero for a single + # model, leaving pca_mean == mean bit-for-bit. + pca_mean = mean + np.linalg.solve(sphere, c) if np.any(c) else mean + + ica = _MNEICA( + n_components=n_ch, + method="infomax", + max_iter="auto", + random_state=self.random_state, + ) + # `method` is inert here: the fitted attributes are hand-populated and + # ICA.fit (the only place `method` branches) is never called, but + # ICA.__init__ still requires a valid method name. + ica.info = self.info_ + ica.ch_names = list(self.ch_names_) + ica.n_components_ = n_ch + ica.pca_mean_ = pca_mean + ica.pca_components_ = pca_components + ica.pca_explained_variance_ = cov_evals + ica.unmixing_matrix_ = unmixing + ica.pre_whitener_ = np.ones((n_ch, 1)) + ica.n_iter_ = max(int(getattr(backend, "iteration", 0)), 1) + # MNE's own fit sets these; read_ica_eeglab (the precedent for building + # an ICA from an external decomposition) sets reject_=None. Without them + # ICA.save()/plot_properties raise AttributeError. + ica.reject_ = None + ica.n_samples_ = int(self._n_samples) if self._n_samples is not None else 0 + ica._update_mixing_matrix() + ica._update_ica_names() + ica.current_fit = self._fit_kind + + self._mne_ica_cache[model_idx] = ica + return ica + + # ------------------------------------------------------------------ + # MNE consumer surface (delegates to the exported ICA) + # ------------------------------------------------------------------ + def get_sources(self, inst, *args, model_idx: int = 0, **kwargs): + """Sources for ``inst`` from model ``model_idx`` (see ``ICA.get_sources``). + + ``model_idx`` is keyword-only so positional arguments pass straight + through to MNE's ``ICA.get_sources`` (e.g. ``add_channels``, + ``start``/``stop``). + """ + return self.to_mne_ica(model_idx).get_sources(inst, *args, **kwargs) + + def apply(self, inst, *args, model_idx: int = 0, **kwargs): + """Remove selected components of model ``model_idx`` and back-project. + + ``model_idx`` is keyword-only so positional arguments pass straight + through to MNE's ``ICA.apply``. Pass ``exclude=[...]`` (or set it on the + exported ICA) to drop components; with no exclusions this reconstructs + the input. + """ + return self.to_mne_ica(model_idx).apply(inst, *args, **kwargs) + + def get_components(self, *, model_idx: int = 0) -> np.ndarray: + """Scalp maps of model ``model_idx``, ``(n_channels, n_components)``.""" + return self.to_mne_ica(model_idx).get_components() + + def plot_components(self, *args, model_idx: int = 0, **kwargs): + """Plot model ``model_idx`` component topographies (see ``ICA.plot_components``).""" + return self.to_mne_ica(model_idx).plot_components(*args, **kwargs) + + def plot_sources(self, inst, *args, model_idx: int = 0, **kwargs): + """Plot model ``model_idx`` component time courses (see ``ICA.plot_sources``).""" + return self.to_mne_ica(model_idx).plot_sources(inst, *args, **kwargs) + + # ------------------------------------------------------------------ + # Multi-model dominance (issue #141) + # ------------------------------------------------------------------ + def get_model_probability(self, inst) -> np.ndarray: + """Per-sample posterior probability of each model on ``inst`` (dominance). + + Returns ``P(model | sample)`` as ``(n_models, n_samples)`` via + :meth:`AMICA.model_probability` on ``inst``'s data (``Epochs`` are + concatenated along time). Each column sums to 1; all ones for a single + model. MNE's own ``ICA`` has no multi-model concept, so this is exposed + here rather than through the exported per-model ICA objects. + """ + self._check_fitted("compute the model probability") + if self.amica_ is None: + raise RuntimeError( + "AMICAICA: internal state is inconsistent; refit before scoring." + ) + return self.amica_.model_probability(self._data_for(inst)) + + def plot_model_probability(self, inst, *, srate: Optional[float] = None, **kwargs): + """Plot per-model probability + best-model log-likelihood over ``inst``. + + Delegates to :func:`pamica.viz.plot_model_probability` with the live + per-model log-likelihood (:meth:`AMICA.model_loglik`) on ``inst``'s + data. ``srate`` defaults to the fitted recording's sampling rate, so the + x-axis is in seconds; extra keywords (``smooth_sec``, ``window_sec``, + ``axes``) pass through. + """ + from ..viz import plot_model_probability as _plot_model_probability + + self._check_fitted("plot the model probability") + if self.amica_ is None or self.info_ is None: + raise RuntimeError( + "AMICAICA: internal state is inconsistent; refit before plotting." + ) + lht = self.amica_.model_loglik(self._data_for(inst)) + if srate is None: + srate = float(self.info_["sfreq"]) + return _plot_model_probability(lht=lht, srate=srate, **kwargs) + + # ------------------------------------------------------------------ + # Separation-quality metrics (issue #143, on top of #133) + # ------------------------------------------------------------------ + def mir(self, inst, *, model_idx: int = 0, nbins: Optional[int] = None) -> tuple: + """Mutual Information Reduction of model ``model_idx`` on ``inst``. + + How much mutual information the fitted unmixing removes from the data, + in nats (issue #133). Delegates to :meth:`AMICA.mir` on ``inst``'s + fitted-channel data; MIR is shift-invariant, so mean/``c`` centering is + irrelevant. + + Parameters + ---------- + inst : mne.io.BaseRaw | mne.BaseEpochs + Data to score (``Epochs`` concatenated along time). + model_idx : int, default=0 + Which model's unmixing to use. + nbins : int, optional + Histogram bin count; see :func:`pamica.metrics.mir`. + + Returns + ------- + mir_nats : float + variance : float + """ + self._check_fitted("compute MIR") + model_idx = self._check_model_idx(model_idx) + assert self.amica_ is not None + return self.amica_.mir(self._data_for(inst), model_idx=model_idx, nbins=nbins) + + def pmi( + self, inst, *, model_idx: int = 0, nbins: Optional[int] = None + ) -> np.ndarray: + """Pairwise Mutual Information between model ``model_idx``'s sources on ``inst``. + + The residual pairwise dependence between fitted sources, in nats + (issue #133). Delegates to :meth:`AMICA.pmi` on ``inst``'s fitted-channel + data. + + Parameters + ---------- + inst : mne.io.BaseRaw | mne.BaseEpochs + Data to score (``Epochs`` concatenated along time). + model_idx : int, default=0 + Which model's sources to use. + nbins : int, optional + Histogram bin count; see :func:`pamica.metrics.pairwise_mi`. + + Returns + ------- + mi_matrix : np.ndarray of shape (n_components, n_components) + Symmetric; the diagonal is each source's own entropy. + """ + self._check_fitted("compute PMI") + model_idx = self._check_model_idx(model_idx) + assert self.amica_ is not None + return self.amica_.pmi(self._data_for(inst), model_idx=model_idx, nbins=nbins) + + # ------------------------------------------------------------------ + # pamica-specific metadata (issue #142) + # + # MNE's ``ICA`` carries no source-density family, GG shape, or + # component-sharing state, so these are exposed here rather than silently + # dropped by the ``mne.preprocessing.ICA`` export. + # ------------------------------------------------------------------ + def get_pdftype(self, *, model_idx: int = 0) -> np.ndarray: + """Per-component source-density family code for model ``model_idx``. + + One integer per ICA component (0-4); map to names with + :data:`pamica.mne_compat.PDFTYPE_NAMES`. All components share one family + unless the adaptive switcher (``pdftype=1``) moved them (issue #26). + """ + self._check_fitted("get the density family") + model_idx = self._check_model_idx(model_idx) + assert self.amica_ is not None + return self.amica_.get_pdftype(model_idx=model_idx) + + def get_rho(self, *, model_idx: int = 0) -> np.ndarray: + """Generalized-Gaussian shape ``rho`` for model ``model_idx``. + + Shape ``(n_mix, n_components)``; ``rho == 2`` is Gaussian-shaped, + ``rho == 1`` Laplacian, ``rho < 1`` heavier-tailed. Meaningful only for + the generalized-Gaussian family (``pdftype=0``). + """ + self._check_fitted("get rho") + model_idx = self._check_model_idx(model_idx) + assert self.amica_ is not None + return self.amica_.get_rho(model_idx=model_idx) + + def shared_components(self) -> list: + """Components shared across models by ``share_comps`` (issue #60). + + One group of ``(model_idx, component_idx)`` pairs per shared column; + empty when nothing is shared (always so for a single model or a default + multi-model fit with ``share_comps`` off). + """ + self._check_fitted("get the shared components") + assert self.amica_ is not None + return self.amica_.shared_components() + + # ------------------------------------------------------------------ + # Helpers + # ------------------------------------------------------------------ + def _data_for(self, inst) -> np.ndarray: + """The fitted-channel data of ``inst`` as ``(n_channels, n_samples)``. + + Selects the exact channels the fit used (by name) so the array aligns + with the stored sphere/unmixing; ``Epochs`` are concatenated along time. + """ + if not isinstance(inst, (mne.io.BaseRaw, mne.BaseEpochs)): + raise TypeError( + f"expected an mne.io.Raw or mne.Epochs, got {type(inst).__name__}." + ) + picked = inst.copy().pick(self.ch_names_) + if isinstance(inst, mne.io.BaseRaw): + X = picked.get_data() + else: + X = np.hstack(picked.get_data()) + return np.ascontiguousarray(X, dtype=np.float64) + + def _check_model_idx(self, model_idx: int) -> int: + if not isinstance(model_idx, (int, np.integer)): + raise TypeError( + f"model_idx must be an int, got {type(model_idx).__name__}." + ) + # Bound against the fitted backend's model count (the source of truth), + # not the mutable constructor hyperparameter self.n_models. + n = ( + self.amica_.model_.n_models + if self.amica_ is not None and self.amica_.model_ is not None + else self.n_models + ) + if not (0 <= model_idx < n): + raise ValueError( + f"model_idx={model_idx} out of range for a {n}-model fit " + f"(valid: 0..{n - 1})." + ) + return int(model_idx) + + def _check_fitted(self, action: str = "this call") -> None: + """Raise if no usable model is available. + + ``ValueError`` when never fitted (matching :class:`AMICA`'s convention); + ``RuntimeError`` when the fit ended degenerate (non-finite parameters, + issue #50), so the failure surfaces here rather than as opaque NaNs + downstream. + """ + if self.amica_ is None: + raise ValueError(f"AMICAICA must be fitted before {action}; run fit().") + if not self.converged_: + raise RuntimeError( + f"Refusing to {action}: the AMICA fit ended degenerate " + f"(stop_reason={self.stop_reason_!r}), so it holds non-finite " + "parameters and would produce NaN output. Lower lrate, disable " + "Newton, or check data conditioning, then refit." + ) + + def __repr__(self) -> str: + if self.amica_ is None: + return ( + f"" + ) + if not self.converged_: + return ( + f"" + ) + return ( + f"" + ) diff --git a/pamica/tests/mne_tests/__init__.py b/pamica/tests/mne_tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/pamica/tests/mne_tests/test_mne_wrapper.py b/pamica/tests/mne_tests/test_mne_wrapper.py new file mode 100644 index 0000000..cdd2231 --- /dev/null +++ b/pamica/tests/mne_tests/test_mne_wrapper.py @@ -0,0 +1,535 @@ +"""AMICAICA MNE-wrapper tests (issue #139, single-model phase 1 + multi-model #141). + +Real sample EEG only (no synthetic/mock): the bundled EEGLAB ``eeglab_data.set`` +(32 channels, 30504 samples) is loaded through ``mne.io.read_raw_eeglab`` -- the +same entry point a real MNE user would use. The whole module self-skips when the +optional ``mne`` extra is absent, so the base CI env (no mne) skips it; the +dedicated ``test-mne`` CI job installs ``pamica[mne]`` and runs it for real. + +The load-bearing check is the round trip: the ICA that ``to_mne_ica`` builds must +reproduce ``AMICA.transform`` bit-for-bit (up to float64 eigh residual), which +pins the mean/sphere/unmixing -> pca_mean_/pca_components_/unmixing_matrix_ map. +""" + +from pathlib import Path + +import numpy as np +import pytest + +mne = pytest.importorskip("mne") + +from pamica.mne_compat import AMICAICA, PDFTYPE_NAMES # noqa: E402 (after importorskip) + +mne.set_log_level("ERROR") + +SAMPLE_DIR = Path(__file__).resolve().parents[2] / "sample_data" +SET_FILE = SAMPLE_DIR / "eeglab_data.set" +SEED = 42 +MAX_ITER = 12 # enough to move off the init; parity is orientation, not convergence + +pytestmark = pytest.mark.skipif( + not SET_FILE.exists(), reason="sample eeglab_data.set missing" +) + + +@pytest.fixture(scope="module") +def raw(): + """Real continuous EEG as an MNE Raw (32 EEG channels, 128 Hz).""" + return mne.io.read_raw_eeglab(str(SET_FILE), preload=True) + + +@pytest.fixture(scope="module") +def fitted(raw): + """A single AMICAICA fit reused across the read-only assertions.""" + return AMICAICA(n_mix=3, random_state=SEED, device="cpu", verbose=False).fit( + raw, max_iter=MAX_ITER + ) + + +def _picked_data(raw): + """The exact array AMICAICA fits on (all good data channels, float64).""" + return raw.copy().pick("data", exclude="bads").get_data().astype(np.float64) + + +# --- the mapping crux ------------------------------------------------------- +def test_get_sources_matches_amica_transform(raw, fitted): + """to_mne_ica().get_sources == AMICA.transform on the same data.""" + s_mne = fitted.to_mne_ica().get_sources(raw).get_data() + s_amica = fitted.amica_.transform(_picked_data(raw)) + assert s_mne.shape == s_amica.shape == (fitted.n_components_, raw.n_times) + np.testing.assert_allclose(s_mne, s_amica, rtol=1e-6, atol=1e-9) + + +def test_get_components_are_channel_space_maps(fitted): + """get_components() equals the channel-space mixing inv(sphere) @ inv(W).""" + comps = fitted.get_components() + assert comps.shape == (fitted.n_components_, fitted.n_components_) + sphere = fitted.amica_.model_.sphere.cpu().numpy() + w_fort = fitted.amica_.get_unmixing_matrix(0) + maps_ref = np.linalg.inv(sphere) @ np.linalg.inv(w_fort) + np.testing.assert_allclose(comps, maps_ref, rtol=1e-6, atol=1e-10) + + +def test_apply_without_exclude_reconstructs(raw, fitted): + """apply() with no excluded components returns the input unchanged.""" + recon = fitted.apply(raw.copy()) + np.testing.assert_allclose(recon.get_data(), raw.get_data(), rtol=1e-6, atol=1e-12) + + +def test_apply_excludes_component(raw, fitted): + """Excluding a component changes the data and zeroes that source.""" + ica = fitted.to_mne_ica() + cleaned = ica.apply(raw.copy(), exclude=[0]) + # The reconstruction must differ from the input (a component was removed). + assert not np.allclose(cleaned.get_data(), raw.get_data()) + # Re-deriving sources from the cleaned data: component 0 is gone. + src_after = ica.get_sources(cleaned).get_data() + assert np.abs(src_after[0]).max() < 1e-6 * np.abs(src_after).max() + + +# --- object validity -------------------------------------------------------- +def test_to_mne_ica_returns_valid_fitted_ica(raw, fitted): + ica = fitted.to_mne_ica() + assert isinstance(ica, mne.preprocessing.ICA) + assert ica.current_fit != "unfitted" + assert ica.n_components_ == fitted.n_components_ + assert ica.ch_names == fitted.ch_names_ + # A genuine MNE ICA exposes the full component interface. + assert ica.get_components().shape == (fitted.n_components_, fitted.n_components_) + + +def test_export_is_cached_and_invalidated_on_refit(raw): + ica = AMICAICA(random_state=SEED, device="cpu", verbose=False).fit( + raw, max_iter=MAX_ITER + ) + first = ica.to_mne_ica() + assert ica.to_mne_ica() is first # cached + ica.fit(raw, max_iter=MAX_ITER) # refit invalidates the cache + assert ica.to_mne_ica() is not first + + +# --- picks / epochs / plotting ---------------------------------------------- +def test_fit_with_channel_subset(raw): + picks = raw.ch_names[:10] + ica = AMICAICA(random_state=SEED, device="cpu", verbose=False).fit( + raw, picks=picks, max_iter=MAX_ITER + ) + assert ica.n_components_ == 10 + assert ica.ch_names_ == picks + s_mne = ica.to_mne_ica().get_sources(raw).get_data() + x = raw.copy().pick(picks).get_data().astype(np.float64) + assert ica.amica_ is not None + np.testing.assert_allclose(s_mne, ica.amica_.transform(x), rtol=1e-6, atol=1e-9) + + +def test_fit_from_epochs(raw): + epochs = mne.make_fixed_length_epochs(raw, duration=2.0, preload=True) + ica = AMICAICA(random_state=SEED, device="cpu", verbose=False).fit( + epochs, max_iter=MAX_ITER + ) + assert ica._fit_kind == "epochs" + src = ica.to_mne_ica().get_sources(epochs).get_data() # (n_ep, n_comp, n_time) + x = np.hstack(epochs.copy().pick("data", exclude="bads").get_data()) + assert ica.amica_ is not None + np.testing.assert_allclose( + np.hstack(src), ica.amica_.transform(x), rtol=1e-6, atol=1e-9 + ) + + +def test_plot_components_returns_figure(fitted): + import matplotlib + + matplotlib.use("Agg") + import matplotlib.figure + import matplotlib.pyplot as plt + + out = fitted.plot_components(picks=[0, 1], show=False) + figs = out if isinstance(out, list) else [out] + assert figs and all(isinstance(f, matplotlib.figure.Figure) for f in figs) + plt.close("all") + + +def test_pca_components_orthonormal_and_variance_ordered(raw, fitted): + """The eigenbasis ordering (invisible to the V-cancelling round trip) is real. + + get_sources/get_components reduce to W@sphere and inv(sphere)@inv(W) for ANY + orthonormal V, so they cannot see a wrong eigenvector order. Pin it here: + pca_components_ must be orthonormal and its rows ordered by descending data + variance, with pca_explained_variance_ equal to that per-row variance. + """ + ica = fitted.to_mne_ica() + p = ica.pca_components_ + np.testing.assert_allclose(p @ p.T, np.eye(p.shape[0]), atol=1e-10) + ev = ica.pca_explained_variance_ + assert np.all(np.diff(ev) <= 0), "explained variance must be descending" + # Independent oracle: project the mean-centered data onto pca_components_; + # each row's variance is that data-covariance eigenvalue, in the same order. + x = _picked_data(raw) + proj = p @ (x - ica.pca_mean_[:, None]) + np.testing.assert_allclose(proj.var(axis=1), ev, rtol=1e-5) + + +def test_fit_with_start_stop_subranges_raw(raw): + stop = 5000 + ica = AMICAICA(random_state=SEED, device="cpu", verbose=False).fit( + raw, start=100, stop=stop, max_iter=MAX_ITER + ) + assert ica.to_mne_ica().n_samples_ == stop - 100 + # A different range yields a different decomposition (the range is honored). + full = AMICAICA(random_state=SEED, device="cpu", verbose=False).fit( + raw, max_iter=MAX_ITER + ) + assert not np.allclose(ica.get_components(), full.get_components()) + + +def test_apply_on_epochs(raw): + epochs = mne.make_fixed_length_epochs(raw, duration=2.0, preload=True) + ica = AMICAICA(random_state=SEED, device="cpu", verbose=False).fit( + epochs, max_iter=MAX_ITER + ) + recon = ica.apply(epochs.copy()) + np.testing.assert_allclose( + recon.get_data(), epochs.get_data(), rtol=1e-6, atol=1e-12 + ) + cleaned = ica.apply(epochs.copy(), exclude=[0]) + assert not np.allclose(cleaned.get_data(), epochs.get_data()) + + +def test_exported_ica_saves_and_reloads(fitted, raw, tmp_path): + """reject_/n_samples_ must be set, else ICA.save() raises AttributeError.""" + ica = fitted.to_mne_ica() + fname = tmp_path / "amica-ica.fif" + ica.save(fname) + reloaded = mne.preprocessing.read_ica(fname) + np.testing.assert_allclose( + reloaded.get_sources(raw).get_data(), + ica.get_sources(raw).get_data(), + rtol=1e-6, + atol=1e-9, + ) + + +def test_plot_sources_returns_figure(fitted, raw): + import matplotlib + + matplotlib.use("Agg") + import matplotlib.figure + import matplotlib.pyplot as plt + + fig = fitted.plot_sources(raw, show=False) + assert isinstance(fig, matplotlib.figure.Figure) + plt.close("all") + + +def test_seed_determinism(raw): + a = AMICAICA(random_state=1, device="cpu", verbose=False).fit( + raw, max_iter=MAX_ITER + ) + b = AMICAICA(random_state=1, device="cpu", verbose=False).fit( + raw, max_iter=MAX_ITER + ) + np.testing.assert_allclose(a.get_components(), b.get_components()) + c = AMICAICA(random_state=2, device="cpu", verbose=False).fit( + raw, max_iter=MAX_ITER + ) + assert not np.allclose(a.get_components(), c.get_components()) + + +# --- guards ----------------------------------------------------------------- +def test_unfitted_calls_raise(): + ica = AMICAICA() + with pytest.raises(ValueError, match="must be fitted"): + ica.to_mne_ica() + + +def test_fit_rejects_non_mne_input(): + with pytest.raises(TypeError, match="mne.io.Raw or mne.Epochs"): + AMICAICA().fit(np.random.randn(8, 500)) + + +def test_fit_rejects_start_stop_on_epochs(raw): + epochs = mne.make_fixed_length_epochs(raw, duration=2.0, preload=True) + with pytest.raises(ValueError, match="only supported for Raw"): + AMICAICA().fit(epochs, stop=100) + + +def test_fit_rejects_out_of_range_stop(raw): + with pytest.raises(ValueError, match="exceeds the recording length"): + AMICAICA().fit(raw, stop=raw.n_times + 1) + + +def test_fit_rejects_non_finite_data(raw): + bad = raw.copy() + bad._data[3, 100:200] = np.nan + with pytest.raises(ValueError, match="non-finite"): + AMICAICA(device="cpu", verbose=False).fit(bad, max_iter=MAX_ITER) + + +def test_fit_rejects_pca_reduction(raw): + with pytest.raises(ValueError, match="PCA reduction"): + AMICAICA(device="cpu", verbose=False).fit(raw, max_iter=5, pcakeep=10) + + +def test_degenerate_fit_is_refused_and_labeled(raw): + """A degenerate (non-converged) fit is kept for inspection but refused. + + Real divergence is platform-dependent, so mark the wrapper's own + convergence signal (the documented contract) to exercise the guard + deterministically: `__repr__` labels it and the consumer methods refuse it. + """ + ica = AMICAICA(random_state=SEED, device="cpu", verbose=False).fit( + raw, max_iter=MAX_ITER + ) + ica.converged_ = False + ica.stop_reason_ = "nan_ll" + assert "degenerate" in repr(ica) + with pytest.raises(RuntimeError, match="degenerate"): + ica.to_mne_ica() + + +def test_failed_refit_leaves_prior_state_intact(raw): + """An exception mid-refit must not corrupt a previously fitted wrapper.""" + ica = AMICAICA(random_state=SEED, device="cpu", verbose=False).fit( + raw, max_iter=MAX_ITER + ) + before = ica.get_components() + with pytest.raises(ValueError): # stop past the end aborts before publishing + ica.fit(raw, picks=raw.ch_names[:10], stop=raw.n_times + 1, max_iter=MAX_ITER) + assert ica.n_components_ == before.shape[0] # still the 32-channel fit + np.testing.assert_array_equal(ica.get_components(), before) + + +# --- multi-model (issue #141) ---------------------------------------------- +@pytest.fixture(scope="module") +def fitted_2m(raw): + """A two-model AMICAICA fit reused across the multi-model assertions.""" + return AMICAICA( + n_models=2, n_mix=3, random_state=SEED, device="cpu", verbose=False + ).fit(raw, max_iter=MAX_ITER) + + +def test_per_model_export_roundtrips_with_nonzero_c(raw, fitted_2m): + """Each model exports its own ICA reproducing that model's transform. + + Unlike the single-model case, the per-model center c is nonzero, so this + exercises the inv(sphere)@c fold into pca_mean_ that phase 1 deferred. + """ + x = _picked_data(raw) + c = fitted_2m.amica_.model_.c.cpu().numpy() + for h in range(2): + assert np.any(c[:, h]), f"model {h} center c should be nonzero" + s_mne = fitted_2m.get_sources(raw, model_idx=h).get_data() + np.testing.assert_allclose( + s_mne, fitted_2m.amica_.transform(x, model_idx=h), rtol=1e-6, atol=1e-9 + ) + + +def test_per_model_export_is_distinct_and_cached(fitted_2m): + ica0, ica1 = fitted_2m.to_mne_ica(0), fitted_2m.to_mne_ica(1) + assert ica0 is not ica1 + assert fitted_2m.to_mne_ica(0) is ica0 # cached per model + assert not np.allclose(ica0.unmixing_matrix_, ica1.unmixing_matrix_) + + +def test_per_model_apply_reconstructs(raw, fitted_2m): + for h in range(2): + recon = fitted_2m.apply(raw.copy(), model_idx=h) + np.testing.assert_allclose( + recon.get_data(), raw.get_data(), rtol=1e-6, atol=1e-12 + ) + + +def test_per_model_apply_exclude_differs_by_model(raw, fitted_2m): + """apply must honor model_idx: excluding a component of different models + removes different signals (guards against model_idx being ignored).""" + c0 = fitted_2m.apply(raw.copy(), model_idx=0, exclude=[0]).get_data() + c1 = fitted_2m.apply(raw.copy(), model_idx=1, exclude=[0]).get_data() + assert not np.allclose(c0, c1) + + +def test_per_model_get_components_and_plot(fitted_2m): + import matplotlib + + matplotlib.use("Agg") + import matplotlib.figure + import matplotlib.pyplot as plt + + maps0 = fitted_2m.get_components(model_idx=0) + maps1 = fitted_2m.get_components(model_idx=1) + assert not np.allclose(maps0, maps1) # per-model scalp maps differ + out = fitted_2m.plot_components(model_idx=1, picks=[0, 1], show=False) + figs = out if isinstance(out, list) else [out] + assert all(isinstance(f, matplotlib.figure.Figure) for f in figs) + plt.close("all") + + +def test_model_probability_is_normalized(raw, fitted_2m): + prob = fitted_2m.get_model_probability(raw) + assert prob.shape == (2, raw.n_times) + np.testing.assert_allclose(prob.sum(axis=0), 1.0, atol=1e-10) + assert np.all(prob >= 0.0) and np.all(prob <= 1.0) + + +def test_single_model_probability_is_all_ones(raw, fitted): + np.testing.assert_allclose(fitted.get_model_probability(raw), 1.0, atol=1e-10) + + +def test_model_probability_and_plot_on_epochs(raw, fitted_2m): + import matplotlib + + matplotlib.use("Agg") + import matplotlib.pyplot as plt + + epochs = mne.make_fixed_length_epochs(raw, duration=2.0, preload=True) + n = len(epochs) * len(epochs.times) + prob = fitted_2m.get_model_probability(epochs) + assert prob.shape == (2, n) # concatenated along time + np.testing.assert_allclose(prob.sum(axis=0), 1.0, atol=1e-10) + fig = fitted_2m.plot_model_probability(epochs) + assert fig is not None + plt.close("all") + + +def test_plot_model_probability_default_srate_gives_seconds(raw, fitted_2m): + import matplotlib + + matplotlib.use("Agg") + import matplotlib.figure + import matplotlib.pyplot as plt + + fig = fitted_2m.plot_model_probability(raw) + assert isinstance(fig, matplotlib.figure.Figure) + ax = fig.axes[1] + assert ax.get_xlabel() == "Time (s)" + # srate default must be the fitted recording's rate, not just any non-None: + # the x-data must equal sample-index / sfreq (catches a wrong/stale binding). + expected_t = np.arange(raw.n_times) / raw.info["sfreq"] + np.testing.assert_allclose( + np.asarray(ax.lines[0].get_xdata(), dtype=float), expected_t + ) + plt.close("all") + + +def test_out_of_range_model_idx_raises(fitted_2m): + with pytest.raises(ValueError, match="out of range"): + fitted_2m.to_mne_ica(model_idx=2) + with pytest.raises(ValueError, match="out of range"): + fitted_2m.get_sources(None, model_idx=-1) + with pytest.raises(TypeError, match="model_idx must be an int"): + fitted_2m.to_mne_ica(model_idx=1.0) + + +def test_get_model_probability_rejects_non_mne_input(fitted_2m): + with pytest.raises(TypeError, match="Raw or mne.Epochs"): + fitted_2m.get_model_probability(np.zeros((32, 100))) + + +def test_degenerate_fit_refused_on_dominance_methods(raw): + ica = AMICAICA(n_models=2, random_state=SEED, device="cpu", verbose=False).fit( + raw, max_iter=MAX_ITER + ) + ica.converged_ = False + ica.stop_reason_ = "nan_ll" + with pytest.raises(RuntimeError, match="degenerate"): + ica.get_model_probability(raw) + with pytest.raises(RuntimeError, match="degenerate"): + ica.plot_model_probability(raw) + + +# --- pamica-specific metadata (issue #142) ---------------------------------- +def test_get_pdftype_and_rho_per_model(fitted_2m): + for h in range(2): + pdf = fitted_2m.get_pdftype(model_idx=h) + assert pdf.shape == (fitted_2m.n_components_,) + assert np.array_equal(np.unique(pdf), [0]) # default generalized Gaussian + assert PDFTYPE_NAMES[int(pdf[0])] == "generalized_gaussian" + rho = fitted_2m.get_rho(model_idx=h) + assert rho.shape == (fitted_2m.n_mix, fitted_2m.n_components_) + assert np.all(np.isfinite(rho)) + + +def test_shared_components_empty_by_default(fitted_2m): + assert fitted_2m.shared_components() == [] + + +def test_metadata_respects_model_idx_bounds(fitted_2m): + with pytest.raises(ValueError, match="out of range"): + fitted_2m.get_pdftype(model_idx=2) + with pytest.raises(ValueError, match="out of range"): + fitted_2m.get_rho(model_idx=5) + + +def test_metadata_requires_fit(): + ica = AMICAICA() + with pytest.raises(ValueError, match="must be fitted"): + ica.get_pdftype() + with pytest.raises(ValueError, match="must be fitted"): + ica.shared_components() + + +# --- separation-quality metrics (issue #143) -------------------------------- +def test_mir_matches_amica_per_model(raw, fitted_2m): + x = _picked_data(raw) + mirs = [] + for h in range(2): + m = fitted_2m.mir(raw, model_idx=h) + np.testing.assert_allclose(m, fitted_2m.amica_.mir(x, model_idx=h)) + mirs.append(m[0]) + assert not np.isclose(mirs[0], mirs[1]) # model_idx honored, not hardcoded to 0 + + +def test_pmi_matches_amica_per_model(raw, fitted_2m): + x = _picked_data(raw) + pmis = [] + for h in range(2): + pmi_w = fitted_2m.pmi(raw, model_idx=h) + assert pmi_w.shape == (fitted_2m.n_components_, fitted_2m.n_components_) + np.testing.assert_allclose(pmi_w, fitted_2m.amica_.pmi(x, model_idx=h)) + pmis.append(pmi_w) + assert not np.allclose(pmis[0], pmis[1]) # per-model, not hardcoded to 0 + + +def test_mir_pmi_on_epochs(raw, fitted_2m): + epochs = mne.make_fixed_length_epochs(raw, duration=2.0, preload=True) + x = np.hstack(epochs.copy().pick("data", exclude="bads").get_data()) + np.testing.assert_allclose(fitted_2m.mir(epochs), fitted_2m.amica_.mir(x)) + np.testing.assert_allclose(fitted_2m.pmi(epochs), fitted_2m.amica_.pmi(x)) + + +def test_mir_pmi_nbins_passthrough(raw, fitted_2m): + """nbins must reach the metric: a non-default nbins matches AMICA with the + same nbins (a dropped forward would silently fall back to the default).""" + x = _picked_data(raw) + np.testing.assert_allclose( + fitted_2m.mir(raw, nbins=50), fitted_2m.amica_.mir(x, nbins=50) + ) + np.testing.assert_allclose( + fitted_2m.pmi(raw, nbins=50), fitted_2m.amica_.pmi(x, nbins=50) + ) + + +def test_mir_pmi_respect_model_idx_bounds(fitted_2m, raw): + with pytest.raises(ValueError, match="out of range"): + fitted_2m.mir(raw, model_idx=2) + with pytest.raises(ValueError, match="out of range"): + fitted_2m.pmi(raw, model_idx=-1) + + +def test_mir_pmi_require_fit(raw): + ica = AMICAICA() + with pytest.raises(ValueError, match="must be fitted"): + ica.mir(raw) + with pytest.raises(ValueError, match="must be fitted"): + ica.pmi(raw) + + +def test_mir_pmi_refuse_degenerate(raw): + ica = AMICAICA(n_models=2, random_state=SEED, device="cpu", verbose=False).fit( + raw, max_iter=MAX_ITER + ) + ica.converged_ = False + ica.stop_reason_ = "nan_ll" + with pytest.raises(RuntimeError, match="degenerate"): + ica.mir(raw) + with pytest.raises(RuntimeError, match="degenerate"): + ica.pmi(raw) diff --git a/pamica/tests/test_viz.py b/pamica/tests/test_viz.py index 8854404..c2d9f76 100644 --- a/pamica/tests/test_viz.py +++ b/pamica/tests/test_viz.py @@ -387,3 +387,38 @@ def test_plot_model_probability_raises_without_lht(two_model_output): out_no_lht = dataclasses.replace(two_model_output, Lht=None) with pytest.raises(ValueError, match="Lht"): plot_model_probability(out_no_lht) + + +# --- live `lht` array path (issue #141) -------------------------------------- +def test_plot_model_probability_lht_array_matches_out(two_model_output): + """Passing a raw Lht array reproduces the AmicaOutput path exactly.""" + lht = np.asarray(two_model_output.Lht, dtype=np.float64) + fig_out = plot_model_probability(two_model_output) + fig_lht = plot_model_probability(lht=lht) + for line_out, line_lht in zip(fig_out.axes[0].lines, fig_lht.axes[0].lines): + np.testing.assert_array_equal(line_out.get_ydata(), line_lht.get_ydata()) + np.testing.assert_array_equal( + fig_out.axes[1].lines[0].get_ydata(), fig_lht.axes[1].lines[0].get_ydata() + ) + + +def test_plot_model_probability_requires_exactly_one_source(two_model_output): + lht = np.asarray(two_model_output.Lht, dtype=np.float64) + with pytest.raises(ValueError, match="exactly one"): + plot_model_probability(two_model_output, lht=lht) # both + with pytest.raises(ValueError, match="provide"): + plot_model_probability() # neither + + +def test_plot_model_probability_rejects_non_2d_lht(): + with pytest.raises(ValueError, match="2-D"): + plot_model_probability(lht=np.zeros(10)) + + +def test_plot_model_probability_lht_with_smoothing(two_model_output, eeglab_metadata): + """The srate/smoothing kwargs (passed through by AMICAICA) work on lht=.""" + lht = np.asarray(two_model_output.Lht, dtype=np.float64) + srate = eeglab_metadata["srate"] + fig = plot_model_probability(lht=lht, srate=srate, smooth_sec=1.0) + probs = np.array([line.get_ydata() for line in fig.axes[0].lines]) + np.testing.assert_allclose(probs.sum(axis=0), 1.0, atol=1e-8) diff --git a/pamica/tests/torch_tests/test_amica_ng_wrapper.py b/pamica/tests/torch_tests/test_amica_ng_wrapper.py index 4d4cb9f..0e30e30 100644 --- a/pamica/tests/torch_tests/test_amica_ng_wrapper.py +++ b/pamica/tests/torch_tests/test_amica_ng_wrapper.py @@ -643,6 +643,12 @@ def test_unfitted_output_raises_not_fitted(): model.mir(np.zeros((NW, 16))) with pytest.raises(ValueError, match="fitted"): model.pmi(np.zeros((NW, 16))) + with pytest.raises(ValueError, match="fitted"): + model.get_pdftype() + with pytest.raises(ValueError, match="fitted"): + model.get_rho() + with pytest.raises(ValueError, match="fitted"): + model.shared_components() def test_degenerate_fit_refuses_output(real_data, tmp_path, caplog): @@ -675,6 +681,9 @@ def test_degenerate_fit_refuses_output(real_data, tmp_path, caplog): lambda: model.variance_order(), lambda: model.mir(real_data[:, :512]), lambda: model.pmi(real_data[:, :512]), + lambda: model.get_pdftype(), + lambda: model.get_rho(), + lambda: model.shared_components(), ): with pytest.raises(RuntimeError, match="degenerate.*nan_ll"): action() diff --git a/pamica/tests/torch_tests/test_ng_metadata.py b/pamica/tests/torch_tests/test_ng_metadata.py new file mode 100644 index 0000000..968c14a --- /dev/null +++ b/pamica/tests/torch_tests/test_ng_metadata.py @@ -0,0 +1,145 @@ +"""Fitted-parameter metadata accessors (issue #142). + +``get_pdftype``/``get_rho``/``shared_components`` expose the source-density +family, GG shape, and component-sharing state that the EEGLAB/MNE exports do not +carry. Real sample EEG only (no synthetic data). +""" + +from pathlib import Path + +import numpy as np +import pytest + +from pamica import AMICATorchNG +from pamica.torch_impl import PDFTYPE_NAMES +from pamica.torch_impl.utils import load_eeglab_data + +SAMPLE_DIR = Path(__file__).resolve().parents[2] / "sample_data" +DATA_FILE = SAMPLE_DIR / "eeglab_data.fdt" +NW = 32 +FIELD = 30504 +NMIX = 3 +SEED = 42 + +pytestmark = pytest.mark.skipif(not DATA_FILE.exists(), reason="sample data missing") + + +@pytest.fixture(scope="module") +def real_data(): + return load_eeglab_data(str(DATA_FILE), data_dim=NW, field_dim=FIELD).astype( + np.float64 + ) + + +def _fit(X, n_models=1, max_iter=8, **kw): + m = AMICATorchNG( + n_channels=NW, n_models=n_models, n_mix=NMIX, seed=SEED, device="cpu", **kw + ) + m.fit(X, max_iter=max_iter, verbose=False) + return m + + +def test_pdftype_names_cover_all_codes(): + assert set(PDFTYPE_NAMES) == {0, 1, 2, 3, 4} + + +def test_get_pdftype_default_is_generalized_gaussian(real_data): + m = _fit(real_data) + pdf = m.get_pdftype() + assert pdf.shape == (NW,) + assert np.array_equal(np.unique(pdf), [0]) # pdftype=0 default + assert PDFTYPE_NAMES[int(pdf[0])] == "generalized_gaussian" + + +def test_get_pdftype_reflects_fixed_family(real_data): + m = _fit(real_data, pdftype=2) # Gaussian + assert np.array_equal(np.unique(m.get_pdftype()), [2]) + + +def test_get_rho_shape_and_bounds(real_data): + m = _fit(real_data) + rho = m.get_rho() + assert rho.shape == (NMIX, NW) + assert np.all(rho >= m.minrho) and np.all(rho <= m.maxrho) + + +def test_shared_components_empty_without_sharing(real_data): + assert _fit(real_data, n_models=1).shared_components() == [] + assert _fit(real_data, n_models=2).shared_components() == [] # share off default + + +def test_shared_components_reports_partial_merges(real_data): + """A PARTIAL merge (some columns single-model) exercises both the grouping + and the exclude-filter: shared groups appear, non-shared columns do not.""" + m = _fit( + real_data[:, :4096], + n_models=2, + max_iter=25, + block_size=1024, + share_comps=True, + share_start=8, + share_iter=10, + do_newton=True, + ) + groups = m.shared_components() + # Partial: not zero (something merged), and fewer than n_sources (some + # columns stay single-model and are correctly excluded from the groups). + assert 0 < len(groups) < NW + for group in groups: + models = {h for h, _ in group} + assert len(models) >= 2 # a shared column spans >= 2 models + assert all(0 <= h < 2 and 0 <= i < NW for h, i in group) + + +def test_get_rho_differs_per_model(real_data): + """get_rho indexes comp_list per model, so a 2-model fit's rho differs.""" + m = _fit(real_data, n_models=2, max_iter=10) + assert not np.allclose(m.get_rho(model_idx=0), m.get_rho(model_idx=1)) + + +def test_get_pdftype_adaptive_switcher_is_mixed(real_data): + """The adaptive switcher (pdftype=1) moves sources into families {1, 4}; + get_pdftype must surface those per-source codes, not a uniform array.""" + m = AMICATorchNG( + n_channels=NW, + n_models=1, + n_mix=1, # adaptive mode is single-component + seed=SEED, + device="cpu", + pdftype=1, + kurt_start=3, + num_kurt=5, + kurt_int=1, + ) + m.fit(real_data, max_iter=20, verbose=False) + codes = set(np.unique(m.get_pdftype()).tolist()) + assert codes.issubset({1, 4}) and codes # only the two cosh families + + +def test_metadata_accessors_reject_bad_model_idx(real_data): + m = _fit(real_data, n_models=2, max_iter=5) + for call in (m.get_pdftype, m.get_rho): + with pytest.raises(ValueError, match="out of range"): + call(model_idx=2) + with pytest.raises(ValueError, match="out of range"): + call(model_idx=-1) # negative would silently wrap without the guard + + +def test_metadata_accessors_require_fit(): + m = AMICATorchNG(n_channels=NW, n_models=1, n_mix=NMIX, seed=SEED, device="cpu") + for call in (m.get_pdftype, m.get_rho, m.shared_components): + with pytest.raises(RuntimeError, match="fitted"): + call() + + +def test_amica_wrapper_delegates_metadata(real_data): + """The scikit-learn-style AMICA wrapper exposes the same accessors.""" + from pamica import AMICA + + a = AMICA(n_models=2, n_mix=NMIX, device="cpu", verbose=False) + a.fit(real_data, max_iter=8, seed=SEED) + assert a.get_pdftype(model_idx=1).shape == (NW,) + assert a.get_rho(model_idx=1).shape == (NMIX, NW) + assert a.shared_components() == [] # share off + with pytest.raises(ValueError, match="out of range"): + a.get_pdftype(model_idx=-1) # guard propagates through the wrapper diff --git a/pamica/tests/torch_tests/test_ng_model_posterior.py b/pamica/tests/torch_tests/test_ng_model_posterior.py new file mode 100644 index 0000000..98700df --- /dev/null +++ b/pamica/tests/torch_tests/test_ng_model_posterior.py @@ -0,0 +1,101 @@ +"""Live per-model posterior accessor tests (issue #141). + +``model_loglik``/``model_probability`` expose, for arbitrary data, the per-model +per-sample log-likelihood the E-step already computes internally. Real sample +EEG only (no synthetic data): the internal training-data ``_llt_lht`` (Fortran's +LLt, issue #155) is an exact oracle for ``model_loglik`` evaluated on that same +data, so the live accessor is pinned bit-for-bit rather than by eyeball. The +equality holds when the fit did not use ``do_reject`` (the default here); +``_compute_full_posterior_ll`` zeroes rejected columns as Fortran's ``load_rej`` +sentinel, which ``model_loglik`` (rejection-unaware) does not reproduce. +""" + +from pathlib import Path + +import numpy as np +import pytest + +from pamica import AMICATorchNG +from pamica.torch_impl.utils import load_eeglab_data + +SAMPLE_DIR = Path(__file__).resolve().parents[2] / "sample_data" +DATA_FILE = SAMPLE_DIR / "eeglab_data.fdt" +NW = 32 +FIELD = 30504 +SEED = 42 + +pytestmark = pytest.mark.skipif(not DATA_FILE.exists(), reason="sample data missing") + + +@pytest.fixture(scope="module") +def real_data(): + return load_eeglab_data(str(DATA_FILE), data_dim=NW, field_dim=FIELD).astype( + np.float64 + ) + + +def _fit(X, n_models, max_iter=10): + m = AMICATorchNG(n_channels=NW, n_models=n_models, n_mix=3, seed=SEED, device="cpu") + m.fit(X, max_iter=max_iter, verbose=False) + return m + + +def test_model_loglik_matches_internal_lht(real_data): + """model_loglik on the training data equals the stored _llt_lht exactly.""" + m = _fit(real_data, n_models=2) + lht = m.model_loglik(real_data) + assert lht.shape == (2, real_data.shape[1]) + # _llt_lht is the E-step's own per-model per-sample LL, recomputed post-fit + # from the same parameters -- an exact oracle for the live accessor. + np.testing.assert_array_equal(lht, m._llt_lht) + + +def test_model_probability_is_normalized(real_data): + m = _fit(real_data, n_models=2) + prob = m.model_probability(real_data) + assert prob.shape == (2, real_data.shape[1]) + np.testing.assert_allclose(prob.sum(axis=0), 1.0, atol=1e-12) + assert np.all(prob >= 0.0) and np.all(prob <= 1.0) + + +def test_single_model_probability_is_all_ones(real_data): + m = _fit(real_data, n_models=1, max_iter=5) + np.testing.assert_allclose(m.model_probability(real_data), 1.0, atol=1e-12) + + +def test_model_loglik_uses_stored_sphere_not_reprocess(real_data): + """Scoring new data must not overwrite the fitted sphere/mean.""" + m = _fit(real_data, n_models=2) + sphere_before = m.sphere.clone() + mean_before = m.mean.clone() + _ = m.model_loglik(real_data[:, :1000]) # a different-length slice + assert np.array_equal(m.sphere.cpu().numpy(), sphere_before.cpu().numpy()) + assert np.array_equal(m.mean.cpu().numpy(), mean_before.cpu().numpy()) + + +def test_model_loglik_requires_fit(): + m = AMICATorchNG(n_channels=NW, n_models=2, n_mix=3, seed=SEED, device="cpu") + with pytest.raises(RuntimeError, match="fitted"): + m.model_loglik(np.zeros((NW, 10))) + + +def test_model_loglik_rejects_non_finite_input(real_data): + m = _fit(real_data, n_models=2, max_iter=5) + bad = real_data.copy() + bad[3, 100] = np.nan + with pytest.raises(ValueError, match="non-finite"): + m.model_loglik(bad) + with pytest.raises(ValueError, match="non-finite"): + m.model_probability(bad) + + +def test_amica_wrapper_delegates(real_data): + """The scikit-learn-style AMICA wrapper exposes the same accessors.""" + from pamica import AMICA + + a = AMICA(n_models=2, n_mix=3, device="cpu", verbose=False) + a.fit(real_data, max_iter=8, seed=SEED) + lht = a.model_loglik(real_data) + prob = a.model_probability(real_data) + assert lht.shape == prob.shape == (2, real_data.shape[1]) + np.testing.assert_allclose(prob.sum(axis=0), 1.0, atol=1e-12) diff --git a/pamica/torch_impl/__init__.py b/pamica/torch_impl/__init__.py index 065c84a..09f6512 100644 --- a/pamica/torch_impl/__init__.py +++ b/pamica/torch_impl/__init__.py @@ -5,11 +5,12 @@ with support for CUDA, ROCm, and Apple Silicon (MPS) backends. """ -from .core import AMICATorchNG +from .core import AMICATorchNG, PDFTYPE_NAMES from .utils import setup_device, check_numerical_stability __all__ = [ "AMICATorchNG", + "PDFTYPE_NAMES", "setup_device", "check_numerical_stability", ] diff --git a/pamica/torch_impl/core.py b/pamica/torch_impl/core.py index 1a17ca8..8f4f89d 100644 --- a/pamica/torch_impl/core.py +++ b/pamica/torch_impl/core.py @@ -58,6 +58,17 @@ logger = logging.getLogger(__name__) +# Human-readable names for the ``pdftype``/``pdtype`` source-density family codes +# (issue #26; amica15.f90). Exposed alongside the numeric codes so a fitted +# model's per-source density family is inspectable (issue #142). +PDFTYPE_NAMES = { + 0: "generalized_gaussian", + 1: "super_gaussian_cosh", + 2: "gaussian", + 3: "logistic", + 4: "sub_gaussian_cosh", +} + _LOG2 = math.log(2.0) _LOG4 = math.log(4.0) # logistic-family normalizer (amica15.f90:1328) _HALF_LOG_PI = 0.5 * math.log(math.pi) @@ -1953,6 +1964,23 @@ def _reject_outliers(self, ll_vec: torch.Tensor): int(self.good_idx.numel()), ) + def _check_model_idx(self, model_idx: int) -> None: + """Validate a model index against the fitted ``n_models``. + + Raises a clear ``ValueError`` (rejecting negatives, which torch's + negative indexing would otherwise turn into a silent wrong-model result) + instead of an opaque tensor ``IndexError``. + """ + if not isinstance(model_idx, (int, np.integer)): + raise TypeError( + f"model_idx must be an int, got {type(model_idx).__name__}." + ) + if not (0 <= model_idx < self.n_models): + raise ValueError( + f"model_idx={model_idx} out of range for a {self.n_models}-model " + f"fit (valid: 0..{self.n_models - 1})." + ) + def transform(self, X: np.ndarray, model_idx: int = 0) -> np.ndarray: """Apply the learned unmixing matrix to (new) data. @@ -1966,6 +1994,7 @@ def transform(self, X: np.ndarray, model_idx: int = 0) -> np.ndarray: raise RuntimeError( "AMICATorchNG.transform() requires a fitted model; call fit() first." ) + self._check_model_idx(model_idx) X_t = torch.from_numpy(np.ascontiguousarray(X)).to(self.device, self.dtype) X_t = self.sphere @ (X_t - self.mean) # c is the per-model data-space center: unmix as W(x - c) (issue #27). @@ -1979,6 +2008,7 @@ def get_mixing_matrix(self, model_idx: int = 0) -> np.ndarray: "AMICATorchNG.get_mixing_matrix() requires a fitted model; call " "fit() first." ) + self._check_model_idx(model_idx) return self.A[:, self.comp_list[:, model_idx]].T.cpu().numpy() def get_unmixing_matrix(self, model_idx: int = 0) -> np.ndarray: @@ -1988,6 +2018,7 @@ def get_unmixing_matrix(self, model_idx: int = 0) -> np.ndarray: "AMICATorchNG.get_unmixing_matrix() requires a fitted model; call " "fit() first." ) + self._check_model_idx(model_idx) return self.W[:, :, model_idx].T.cpu().numpy() def _pca_reduced(self) -> bool: @@ -2031,6 +2062,7 @@ def mir( raise RuntimeError( "AMICATorchNG.mir() requires a fitted model; call fit() first." ) + self._check_model_idx(model_idx) if self._pca_reduced(): raise ValueError( "mir() is incompatible with PCA reduction (pcakeep/pcadb): " @@ -2069,6 +2101,188 @@ def pmi( """ return pairwise_mi(self.transform(X, model_idx=model_idx), nbins) + # ------------------------------------------------------------------ + # Multi-model posterior (issue #141) + # ------------------------------------------------------------------ + def model_loglik(self, X: np.ndarray) -> np.ndarray: + """Per-model, per-sample log-likelihood ``Lht`` on (new) data. + + For each model ``h`` and sample ``t`` this is the joint log-likelihood + ``log(gm[h]) + log|det W_h| + sldet + sum_i log p_h(s_i)`` (Fortran's + ``Lht``/``modloglik``), evaluated on arbitrary raw data via the STORED + sphere/mean -- never re-preprocessing, which would overwrite them. The + per-sample posterior over models (model dominance) is + ``softmax(Lht, axis=0)``; see :meth:`model_probability`. + + This does not replicate a training-time ``do_reject`` mask: it scores + every sample of ``X``. On a ``do_reject`` fit's own training data it + therefore returns real values where the stored ``_llt_lht`` carries + Fortran's sentinel zeros for rejected samples (issue #155), so the two + agree bit-for-bit only when the fit did not use ``do_reject``. Like + :meth:`transform`, it assumes a usable (non-degenerate) fit; the + :class:`~pamica.AMICA` wrapper enforces that via ``_check_usable``. + + Parameters + ---------- + X : np.ndarray of shape (n_channels, n_samples) + Raw (unpreprocessed) data. + + Returns + ------- + Lht : np.ndarray of shape (n_models, n_samples) + + Raises + ------ + RuntimeError + If the model is unfitted. + ValueError + If ``X`` contains non-finite (NaN/Inf) values. + """ + if self.sphere is None or self.mean is None or self.W is None: + raise RuntimeError( + "AMICATorchNG.model_loglik() requires a fitted model; call fit() first." + ) + X = np.ascontiguousarray(X) + if not np.isfinite(X).all(): + bad = np.flatnonzero(~np.isfinite(X).all(axis=1)) + raise ValueError( + "AMICATorchNG.model_loglik(): input contains non-finite (NaN/Inf) " + f"values in {bad.size} channel(s) {bad.tolist()}; clean bad " + "segments before scoring." + ) + X_t = torch.from_numpy(X).to(self.device, self.dtype) + X_t = self.sphere @ (X_t - self.mean) + n_samples = X_t.shape[1] + Lht = np.zeros((self.n_models, n_samples)) + for start in range(0, n_samples, self.block_size): + end = min(start + self.block_size, n_samples) + logV, *_ = self._forward(X_t[:, start:end]) + Lht[:, start:end] = logV.T.detach().cpu().numpy() + return Lht + + def model_probability(self, X: np.ndarray) -> np.ndarray: + """Per-sample posterior probability of each model (model dominance). + + The column-wise ``softmax`` over models of :meth:`model_loglik`, i.e. + ``P(model h | x_t)``; each column sums to 1. For a single model this is + all ones. + + Parameters + ---------- + X : np.ndarray of shape (n_channels, n_samples) + Raw (unpreprocessed) data. + + Returns + ------- + prob : np.ndarray of shape (n_models, n_samples) + + Raises + ------ + RuntimeError + If the model is unfitted. + ValueError + If ``X`` is non-finite, or if every model underflows to ``-inf`` + log-likelihood at some sample (the posterior is undefined there). + """ + Lht = self.model_loglik(X) + col_max = Lht.max(axis=0, keepdims=True) + if not np.isfinite(col_max).all(): + n_bad = int((~np.isfinite(col_max)).sum()) + raise ValueError( + f"AMICATorchNG.model_probability(): every model has -inf " + f"log-likelihood at {n_bad} sample(s), so the posterior is " + "undefined there (an extreme outlier under a tight source " + "density)." + ) + ex = np.exp(Lht - col_max) + return ex / ex.sum(axis=0, keepdims=True) + + # ------------------------------------------------------------------ + # Fitted-parameter metadata (issue #142) + # ------------------------------------------------------------------ + def get_pdftype(self, model_idx: int = 0) -> np.ndarray: + """Per-source density-family code for model ``model_idx``. + + One integer per source component (0-4; see + :data:`pamica.torch_impl.PDFTYPE_NAMES`): 0 generalized Gaussian, 1 + super-Gaussian cosh, 2 Gaussian, 3 logistic, 4 sub-Gaussian cosh. All + sources share ``pdftype`` unless the adaptive switcher (``pdftype=1``) + moved them individually (issue #26). + + Returns + ------- + np.ndarray of int, shape (n_sources,) + """ + if self.pdtype is None: + raise RuntimeError( + "AMICATorchNG.get_pdftype() requires a fitted model; call fit() first." + ) + self._check_model_idx(model_idx) + return self.pdtype[:, model_idx].detach().cpu().numpy() + + def get_rho(self, model_idx: int = 0) -> np.ndarray: + """Generalized-Gaussian shape parameter ``rho`` for model ``model_idx``. + + One value per (mixture component, source): ``rho == 2`` is Gaussian- + shaped, ``rho == 1`` Laplacian, ``rho < 1`` heavier-tailed. Only the + generalized-Gaussian family (``pdftype=0``) updates ``rho``; for every + non-zero code (1-4, the fixed and adaptive cosh families) it stays frozen + at ``rho0`` and does not describe the fitted density. + + Returns + ------- + np.ndarray of float, shape (n_mix, n_sources) + """ + if self.rho is None or self.comp_list is None: + raise RuntimeError( + "AMICATorchNG.get_rho() requires a fitted model; call fit() first." + ) + self._check_model_idx(model_idx) + # Defense-in-depth, matching state_dict()'s isfinite sweep: a degenerate + # multi-model fit can leave one model's rho non-finite without the + # aggregate LL tripping nan_ll, and _check_usable only inspects + # stop_reason. Refuse rather than return a silent NaN. + if not torch.isfinite(self.rho).all(): + raise RuntimeError( + "AMICATorchNG.get_rho(): rho holds non-finite values (a " + "degenerate fit); inspect stop_reason and refit." + ) + idx = self.comp_list[:, model_idx] + return self.rho[:, idx].detach().cpu().numpy() + + def shared_components(self) -> list: + """Components shared across models by ``share_comps`` (issue #60). + + ``share_comps`` folds near-collinear components of different models onto + one shared mixing column + density, recorded as a repeated index in + ``comp_list``. Returns one group per shared column: a list of + ``(model_idx, source_idx)`` pairs that all reference it. Empty when no + component is shared across two or more models (always for one model, and + for a default multi-model fit with ``share_comps`` off). + + Note that a merge synchronizes only the mixture parameters routed through + ``comp_list`` (``mu``/``alpha``/``beta``/``rho``); the per-source density + *family* code ``pdtype`` is a separate tensor and is not synchronized, so + under the adaptive switcher (``pdftype=1``) a shared pair can still report + different :meth:`get_pdftype` codes. + + Returns + ------- + list of list of tuple(int, int) + """ + if self.comp_list is None: + raise RuntimeError( + "AMICATorchNG.shared_components() requires a fitted model; call " + "fit() first." + ) + cl = self.comp_list.detach().cpu().numpy() # (n_sources, n_models) + groups = [] + for col in np.unique(cl): + src, mdl = np.where(cl == col) + if np.unique(mdl).size >= 2: + groups.append([(int(h), int(i)) for i, h in zip(src, mdl)]) + return groups + # ------------------------------------------------------------------ # EEGLAB drop-in output (issue #92) # ------------------------------------------------------------------ @@ -2114,6 +2328,7 @@ def variance_order( "AMICATorchNG.variance_order() requires a fitted model; call " "fit() first." ) + self._check_model_idx(model_idx) cl = self.comp_list[:, model_idx].cpu().numpy() alpha = self.alpha[:, cl].cpu().numpy() mu = self.mu[:, cl].cpu().numpy() diff --git a/pamica/viz.py b/pamica/viz.py index dd42e6c..fcfa589 100644 --- a/pamica/viz.py +++ b/pamica/viz.py @@ -147,8 +147,9 @@ def plot_pmi_heatmap( def plot_model_probability( - out: AmicaOutput, + out: AmicaOutput | None = None, *, + lht: np.ndarray | None = None, srate: float | None = None, smooth_sec: float | None = None, window_sec: float | None = None, @@ -161,11 +162,19 @@ def plot_model_probability( single most probable model at each timepoint (``Lht.max(axis=0)``), not the total ``Lt``, matching the observed MATLAB behaviour. + Provide exactly one source of ``Lht``: a written ``out`` (an + :class:`AmicaOutput`), or a live ``lht`` array (for example from + :meth:`pamica.AMICA.model_loglik`, so an MNE-side fit can plot dominance + without writing an ``amicaout`` directory, issue #141). + Parameters ---------- - out : AmicaOutput + out : AmicaOutput, optional Must have `out.Lht` populated (per-model per-sample log-likelihood; - see `pamica.numpy_impl.load.loadmodout`). + see `pamica.numpy_impl.load.loadmodout`). Mutually exclusive with `lht`. + lht : np.ndarray of shape (n_models, n_samples), optional + A per-model per-sample log-likelihood array. Mutually exclusive with + `out`. srate : float, optional Sampling rate in Hz. pamica has no built-in notion of sample rate (`AmicaOutput`/`loadmodout`/`load_data_file` carry none), so a seconds @@ -192,15 +201,35 @@ def plot_model_probability( Raises ------ ValueError - If `out.Lht` is `None`, or if `smooth_sec`/`window_sec` is given - without `srate`. + If neither or both of `out`/`lht` are given, if `out.Lht` is `None`, if + `lht` is not 2-D, or if `smooth_sec`/`window_sec` is given without + `srate`. """ - if out.Lht is None: + if lht is not None and out is not None: + raise ValueError( + "plot_model_probability: provide exactly one of `out` (an " + "AmicaOutput with Lht) or `lht` (a (n_models, n_samples) array)." + ) + if lht is not None: + Lht = np.asarray(lht, dtype=np.float64) + elif out is not None: + if out.Lht is None: + raise ValueError( + "plot_model_probability: out.Lht is None (no LLt saved with this " + "fit); refit with a backend that writes LLt (see " + "AMICA.write_amica_output) or pass an AmicaOutput loaded from a " + "directory that has it." + ) + Lht = np.asarray(out.Lht, dtype=np.float64) + else: + raise ValueError( + "plot_model_probability: provide `out` (an AmicaOutput with Lht) or " + "`lht` (a (n_models, n_samples) array)." + ) + if Lht.ndim != 2: raise ValueError( - "plot_model_probability: out.Lht is None (no LLt saved with this " - "fit); refit with a backend that writes LLt (see " - "AMICA.write_amica_output) or pass an AmicaOutput loaded from a " - "directory that has it." + "plot_model_probability: lht must be 2-D (n_models, n_samples), got " + f"shape {Lht.shape}." ) if smooth_sec is not None and srate is None: raise ValueError( @@ -214,7 +243,6 @@ def plot_model_probability( "seconds-based x-axis width; pass srate explicitly)." ) - Lht = np.asarray(out.Lht, dtype=np.float64) n_models, n_samples = Lht.shape if smooth_sec is not None: diff --git a/pyproject.toml b/pyproject.toml index db086d0..528437d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -49,6 +49,10 @@ docs = [ # Visualization for the benchmark sweeps: MNE builds the scalp-map montage for # the IC topomaps and (issue #91) the distributed channel-subset positions. viz = ["mne>=1.6"] +# MNE-Python compatibility layer (AMICAICA, issue #139). MNE is not a core +# dependency, so install with `pip install pamica[mne]`; `import pamica` never +# requires it (pamica.mne_compat is imported lazily, same pattern as mlx). +mne = ["mne>=1.6"] [project.urls] Homepage = "https://github.com/sccn/pAMICA" @@ -57,7 +61,7 @@ Repository = "https://github.com/sccn/pAMICA" [tool.setuptools] # List subpackages explicitly so the wheel is deterministic across platforms # (a bare ["pamica"] dropped torch_impl/ on some setuptools versions). -packages = ["pamica", "pamica.numpy_impl", "pamica.torch_impl", "pamica.mlx_impl", "pamica.metrics", "pamica.native"] +packages = ["pamica", "pamica.numpy_impl", "pamica.torch_impl", "pamica.mlx_impl", "pamica.metrics", "pamica.native", "pamica.mne_compat"] # numpy_impl/params.json is the NumPy backend's default parameter file, loaded at # runtime via importlib/__file__, so it must ship in the wheel. package-data = {"pamica" = ["data/*"], "pamica.numpy_impl" = ["params.json"]} @@ -83,6 +87,10 @@ omit = [ # GPU, so CI (ubuntu, no mlx) cannot exercise it and always measures 0%. It is # covered locally by tests/mlx_tests/ on Apple hardware (issue #76). "pamica/mlx_impl/*", + # Optional MNE compatibility layer (issue #139): imports mne, absent from the + # base test env, so the base coverage run always measures 0%. Exercised by + # the dedicated `test-mne` CI job (uv sync --extra mne) on real sample EEG. + "pamica/mne_compat/*", ] [tool.coverage.report] diff --git a/uv.lock b/uv.lock index 7fd617f..ed80b15 100644 --- a/uv.lock +++ b/uv.lock @@ -1212,6 +1212,9 @@ docs = [ mlx = [ { name = "mlx" }, ] +mne = [ + { name = "mne" }, +] viz = [ { name = "mne" }, ] @@ -1234,6 +1237,7 @@ requires-dist = [ { name = "mkdocs-material", marker = "extra == 'docs'", specifier = ">=9.5" }, { name = "mkdocstrings", extras = ["python"], marker = "extra == 'docs'", specifier = ">=0.26" }, { name = "mlx", marker = "extra == 'mlx'", specifier = ">=0.32" }, + { name = "mne", marker = "extra == 'mne'", specifier = ">=1.6" }, { name = "mne", marker = "extra == 'viz'", specifier = ">=1.6" }, { name = "numpy" }, { name = "scipy" }, @@ -1241,7 +1245,7 @@ requires-dist = [ { name = "torch", specifier = ">=2.12.1" }, { name = "tqdm" }, ] -provides-extras = ["mlx", "docs", "viz"] +provides-extras = ["mlx", "docs", "viz", "mne"] [package.metadata.requires-dev] dev = [