Merge pull request #191 from sccn/feature/issue-139-epic-mne-compat

MNE-Python compatibility layer (epic #139)

Author: neuromechanist

.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]

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`).

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

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.

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:

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]:

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"]