diff --git a/docs/api/mne-compat.md b/docs/api/mne-compat.md index 7750a0b..c4afba5 100644 --- a/docs/api/mne-compat.md +++ b/docs/api/mne-compat.md @@ -59,9 +59,31 @@ 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. -!!! note "Single-model in this release" - `AMICAICA` covers the single-model case (`n_models == 1`). Multi-model - exposure for MNE users (per-model sources and model dominance) is tracked in - issue #141. +## 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. ::: pamica.mne_compat.AMICAICA diff --git a/docs/changelog.md b/docs/changelog.md index 55d35bd..1669e00 100644 --- a/docs/changelog.md +++ b/docs/changelog.md @@ -27,6 +27,21 @@ MNE-Python compatibility layer (epic #139), additive: the scikit-learn-style 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). ## 0.2.2 diff --git a/pamica/amica.py b/pamica/amica.py index 869c220..6a9f66f 100644 --- a/pamica/amica.py +++ b/pamica/amica.py @@ -419,6 +419,64 @@ 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 variance_order( self, model_idx: int = 0, return_svar: bool = False ) -> Union[np.ndarray, tuple]: diff --git a/pamica/mne_compat/core.py b/pamica/mne_compat/core.py index de0ad0f..05a7e67 100644 --- a/pamica/mne_compat/core.py +++ b/pamica/mne_compat/core.py @@ -8,8 +8,10 @@ :class:`mne.preprocessing.ICA`; the other methods delegate to that object, so MNE performs the export/plot/back-projection machinery unchanged. -This phase covers the single-model, full-rank case (``n_models == 1``, no PCA -reduction). Multi-model exposure is issue #141. +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). PCA reduction is unsupported (full-rank whitening only). """ from typing import Optional, Union @@ -35,8 +37,16 @@ class AMICAICA: :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`. + 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 @@ -51,7 +61,7 @@ class AMICAICA: Attributes ---------- amica_ : AMICA - The fitted single-model pamica model. + 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 @@ -67,18 +77,20 @@ class AMICAICA: Notes ----- - The single-model AMICA transform is ``S = W_fort @ sphere @ (X - mean)`` (the - per-model data-space center ``c`` is identically zero for one model, since - its 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 uses - ``pca_components_ = V.T``, ``unmixing_matrix_ = W_fort @ sphere @ V`` and - ``pca_mean_ = mean``. 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().get_sources(raw)`` equals - ``amica_.transform(X)``). + 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 @@ -90,11 +102,13 @@ class AMICAICA: 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 @@ -108,7 +122,8 @@ def __init__( self.stop_reason_: Optional[str] = None self._n_samples: Optional[int] = None self._fit_kind: Optional[str] = None - self._mne_ica: Optional[_MNEICA] = None + # Per-model export cache (one mne.preprocessing.ICA per model_idx). + self._mne_ica_cache: dict = {} # ------------------------------------------------------------------ # Fit @@ -197,7 +212,10 @@ def fit( fit_kwargs.setdefault("seed", int(self.random_state)) amica = AMICA( - n_models=1, n_mix=self.n_mix, device=self.device, verbose=self.verbose + n_models=self.n_models, + n_mix=self.n_mix, + device=self.device, + verbose=self.verbose, ) amica.fit(X, **fit_kwargs) @@ -221,13 +239,13 @@ def fit( self.amica_ = amica self.converged_ = amica.converged_ self.stop_reason_ = amica.stop_reason_ - self._mne_ica = None # invalidate any cached export + self._mne_ica_cache = {} # invalidate any cached exports return self # ------------------------------------------------------------------ # Export # ------------------------------------------------------------------ - def to_mne_ica(self) -> _MNEICA: + 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``, @@ -236,17 +254,26 @@ def to_mne_ica(self) -> _MNEICA: mean/sphere/unmixing to ``pca_mean_``/``pca_components_``/ ``unmixing_matrix_`` mapping. - The object is cached and returned by reference: mutating it (for example - setting ``.exclude``) persists across subsequent :meth:`apply`/ - :meth:`get_sources` calls until the next :meth:`fit`. + 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") - if self._mne_ica is not None: - return self._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: @@ -254,15 +281,16 @@ def to_mne_ica(self) -> _MNEICA: "AMICAICA: internal state is inconsistent; refit before to_mne_ica()." ) backend = amica.model_ - if backend.mean is None or backend.sphere is None: + 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; refit " + "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=0) + 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 @@ -277,9 +305,11 @@ def to_mne_ica(self) -> _MNEICA: pca_components = v.T unmixing = w_fort @ sphere @ v - # Single-model AMICA has the per-model center c identically zero (its - # update is gated to n_models>1), so no data-space offset enters pca_mean. - pca_mean = mean + # 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, @@ -308,39 +338,122 @@ def to_mne_ica(self) -> _MNEICA: ica._update_ica_names() ica.current_fit = self._fit_kind - self._mne_ica = ica + self._mne_ica_cache[model_idx] = ica return ica # ------------------------------------------------------------------ # MNE consumer surface (delegates to the exported ICA) # ------------------------------------------------------------------ - def get_sources(self, inst): - """Sources for ``inst`` as an MNE object (see ``ICA.get_sources``).""" - return self.to_mne_ica().get_sources(inst) + def get_sources(self, inst, *args, model_idx: int = 0, **kwargs): + """Sources for ``inst`` from model ``model_idx`` (see ``ICA.get_sources``). - def apply(self, inst, **kwargs): - """Remove selected components and back-project (see ``ICA.apply``). + ``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) - Pass ``exclude=[...]`` (or set it on the exported ICA) to drop - components; with no exclusions this reconstructs the input. + 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().apply(inst, **kwargs) + return self.to_mne_ica(model_idx).apply(inst, *args, **kwargs) - def get_components(self) -> np.ndarray: - """Scalp maps, shape ``(n_channels, n_components)`` (``ICA.get_components``).""" - return self.to_mne_ica().get_components() + 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, **kwargs): - """Plot component topographies (see ``ICA.plot_components``).""" - return self.to_mne_ica().plot_components(*args, **kwargs) + 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, **kwargs): - """Plot component time courses (see ``ICA.plot_sources``).""" - return self.to_mne_ica().plot_sources(inst, *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) # ------------------------------------------------------------------ # 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. @@ -361,13 +474,15 @@ def _check_fitted(self, action: str = "this call") -> None: def __repr__(self) -> str: if self.amica_ is None: - return f"" + return ( + f"" + ) if not self.converged_: return ( f"" + f"n_models={self.n_models}, n_mix={self.n_mix}, {self._fit_kind})>" ) return ( f"" + f"n_models={self.n_models}, n_mix={self.n_mix}, {self._fit_kind})>" ) diff --git a/pamica/tests/mne_tests/test_mne_wrapper.py b/pamica/tests/mne_tests/test_mne_wrapper.py index 7ea150a..da89323 100644 --- a/pamica/tests/mne_tests/test_mne_wrapper.py +++ b/pamica/tests/mne_tests/test_mne_wrapper.py @@ -1,4 +1,4 @@ -"""AMICAICA MNE-wrapper tests (issue #139 phase 1, single-model). +"""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 @@ -297,3 +297,140 @@ def test_failed_refit_leaves_prior_state_intact(raw): 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) 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_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/core.py b/pamica/torch_impl/core.py index 1a17ca8..ff80d3f 100644 --- a/pamica/torch_impl/core.py +++ b/pamica/torch_impl/core.py @@ -2069,6 +2069,102 @@ 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) + # ------------------------------------------------------------------ # EEGLAB drop-in output (issue #92) # ------------------------------------------------------------------ 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: