Phase 2: multi-model exposure for MNE users

[neuromechanist]

Summary

Exposes multi-model AMICA through the MNE wrapper (AMICAICA(n_models=...)). MNE's ICA represents only one unmixing matrix, so we work with current MNE, not around it: each model is exported as its own single-model mne.preprocessing.ICA, and the per-sample model dominance MNE cannot represent is exposed directly. No MNE fork or extension needed.

Closes #141. Part of epic #139.

What's new

Backend (a public live posterior accessor): AMICATorchNG.model_loglik(X) returns the per-model per-sample log-likelihood Lht for arbitrary data, scored through the stored sphere/mean (never re-preprocessing), and model_probability(X) is its column-wise softmax over models (P(model | sample)). Both are delegated on the AMICA wrapper. The training-data path is pinned bit-for-bit against the E-step's own _llt_lht (Fortran Lht, #155).

Wrapper (AMICAICA):

• n_models constructor argument.
• to_mne_ica(model_idx=...) exports any model; the consumer methods (get_sources/apply/get_components/plot_components/plot_sources) all take model_idx. Per-model exports are cached separately.
• The per-model export folds each model's data-space center c into pca_mean_ (mean + inv(sphere) @ c), the piece phase 1 could skip because single-model c is zero. Verified: to_mne_ica(model_idx=h).get_sources(raw) reproduces AMICA.transform(raw, model_idx=h) for every model at ~6e-14, with c confirmed nonzero.
• get_model_probability(inst) returns (n_models, n_samples) dominance; plot_model_probability(inst) draws per-model probability + best-model LL (srate defaults to the fitted recording's rate).

viz: pamica.viz.plot_model_probability now also accepts a live lht array (not only a written AmicaOutput), reused by the wrapper's dominance plot. Backward compatible; all 34 existing viz tests pass.

Test plan

• Backend: pamica/tests/torch_tests/test_ng_model_posterior.py (base CI, real EEG) — model_loglik vs _llt_lht bit-for-bit, posterior normalization, single-model all-ones, stored-sphere invariance, unfit guard, AMICA-wrapper delegation.
• viz: live-lht path equals the AmicaOutput path, exactly-one-source and 2-D guards.
• MNE wrapper (test-mne job): per-model round trip with nonzero c, distinct per-model exports, per-model apply, dominance normalization, single-model all-ones, the dominance plot, out-of-range model_idx.
• Local gates green: ruff, ty check ., mkdocs build --strict, typos, and the full backend/viz/mne suites.

[neuromechanist]

Review (5 Sonnet reviewers: code, tests, comments, silent-failure, type-design)

All actionable findings addressed in commits 0044c98 / fd500f9 / 0743b44. The code-reviewer found nothing above its confidence bar and independently confirmed the c-fold math, softmax stability, model_idx bounds, cache invalidation, and single-model parity. No false positives.

Correctness / robustness

• CRITICAL (silent-failure, reproduced): scoring path had no finiteness validation. model_loglik/model_probability let NaN/Inf silently propagate — a regression against the mir/pmi pattern. Added an np.isfinite guard (channel-named) at the shared AMICATorchNG.model_loglik entry point, plus an all-models-underflow guard in model_probability (all -inf column -> clear ValueError instead of silent NaN). Pinned by test_model_loglik_rejects_non_finite_input.
• HIGH (type-design): model_idx collided with MNE's positional args. get_sources(inst, model_idx) / apply(inst, model_idx) sat where MNE uses add_channels/include. Made model_idx keyword-only on all consumer methods and forward *args/**kwargs, matching plot_components/plot_sources.
• MED (type-design): _check_model_idx now type-guards model_idx (clear TypeError, not a raw comparison error) and bounds against the fitted backend's n_models, not the mutable constructor attribute.
• MED (silent-failure): bare asserts in get_model_probability/plot_model_probability replaced with explicit RuntimeErrors (survive python -O, match to_mne_ica).

Docs / comments

• CRITICAL (comments): the class Notes mapping — the canonical reference to_mne_ica points at — was stale, stating pca_mean_ = mean unconditionally while this PR's to_mne_ica folds in inv(sphere)@c for multi-model. Rewritten to the general form pca_mean_ = mean + inv(sphere)@c_h (reduces to mean when c_h=0).
• Fixed a transform(raw) vs transform(X) slip in the API page, and scoped the _llt_lht bit-for-bit oracle claim to do_reject=False (the accessor is rejection-unaware) in the test docstring, model_loglik docstring, and changelog. Added Raises to the AMICA delegating docstrings.

Tests (the reviewers' gaps)

• test_per_model_apply_reconstructs didn't discriminate (no-exclude reconstructs regardless of model) -> added test_per_model_apply_exclude_differs_by_model.
• Added: Epochs get_model_probability/plot_model_probability; the dominance plot's default srate now asserted by x-data (not just the label); get_components/plot_components with model_idx=1; the TypeError guard on non-MNE input; degenerate-fit refusal on both dominance methods; and smoothing + live-lht in viz.

Kept as intended design (with rationale)

• model_loglik/model_probability return the whole (n_models, n_samples) matrix rather than taking model_idx — they are inherently about the ensemble (dominance), so a per-model model_idx would be meaningless; passing one raises cleanly. Documented, not silent.
• AMICAICA(n_models, n_mix, ...) puts n_models first to match the sibling AMICA signature; both phases 1-2 are unreleased (0.3.0), so there is no external break — the reviewers agreed this is the right call made before shipping.
• viz.plot_model_probability(out | lht) mutual exclusion stays runtime-enforced with distinct, tested error messages (a single small function, low complexity).

Local gates green after fixes: ruff, ty check ., mkdocs build --strict, typos, and the full backend (43) + mne (35) suites on real sample EEG.