diff --git a/pamica/numpy_impl/core.py b/pamica/numpy_impl/core.py index 839ffe7..ac34112 100644 --- a/pamica/numpy_impl/core.py +++ b/pamica/numpy_impl/core.py @@ -1222,14 +1222,49 @@ def _update_parameters(self, updates: Dict): self.iter, ) - # Update mixture weights - self.alpha = updates["dalpha_n"] / np.sum(updates["dalpha_n"], axis=0) + # A component merged away by share_comps is no longer referenced by + # comp_list, so no sufficient statistic accumulates into its column and + # the divisions below are 0/0 = NaN. Update only used columns and freeze + # the rest at their last finite value (rho is excluded: its own 1e-8 + # floor keeps it finite, so it drifts rather than freezing -- harmless, + # since no dead column is read downstream), as AMICATorchNG does (Fortran + # carries the NaN harmlessly behind its comp_used mask; keeping them + # finite means a fit cannot report success while holding NaN parameters, + # issue #240). All-True with the default comp_list, so the ordinary path + # is unchanged. + used = ( + self.comp_used + if self.comp_used is not None + else np.ones(self.num_comps, dtype=bool) + ) + + # Update mixture weights. errstate because np.where evaluates both + # branches: an unused column's 0/0 is computed and discarded, and would + # otherwise warn on every iteration after a merge. + with np.errstate(invalid="ignore", divide="ignore"): + alpha_next = updates["dalpha_n"] / np.sum(updates["dalpha_n"], axis=0) + self.alpha = np.where(used, alpha_next, self.alpha) + # errstate above silences the 0/0 that np.where computes for a dead + # column and discards. That also silenced numpy's warning for a genuine + # 0/0 in a LIVE column (a component whose responsibility mass collapses + # to exactly zero), which used to be the only signal it happened. Check + # explicitly instead, mirroring the mu/beta canary below, so the origin + # is not lost to a later unattributable nan-LL stop. + if not np.all(np.isfinite(self.alpha)): + self.logger.warning( + "Non-finite alpha at iter %d (component responsibility mass " + "collapsed).", + self.iter, + ) # Exact-EM mixture location/scale (Fortran :1978/:1993). These are # fixed-point updates -- mu += dmu_n/dmu_d, beta *= sqrt(dbeta_n/dbeta_d) # -- NOT first-order gradient steps, so they carry no lrate. - self.mu = self.mu + updates["dmu_n"] / updates["dmu_d"] - self.beta = self.beta * np.sqrt(updates["dbeta_n"] / updates["dbeta_d"]) + with np.errstate(invalid="ignore", divide="ignore"): + mu_next = self.mu + updates["dmu_n"] / updates["dmu_d"] + beta_next = self.beta * np.sqrt(updates["dbeta_n"] / updates["dbeta_d"]) + self.mu = np.where(used, mu_next, self.mu) + self.beta = np.where(used, beta_next, self.beta) self.beta = np.clip(self.beta, self.invsigmin, self.invsigmax) # Fortran keeps a live "NaN in sbeta!" canary here (amica17.f90:1996-2000); # the exact-EM mu/beta divisions are unguarded (matching Fortran), so @@ -1379,6 +1414,11 @@ def _update_parameters(self, updates: Dict): # (LEFT-multiply by the TRANSPOSED direction). Right-multiply by the # untransposed dir is invisible at the fixed point but sends the fit # downhill -- issue #24 root cause. + # Per-model loop. For a disjoint comp_list this equals Fortran's single + # weighted DAXPY, but a column shared across models takes one step per + # contributing model instead of one averaged step -- issue #242, which + # ships separately because no test written for it so far distinguishes + # the two. for h in range(self.num_models): idx = self.comp_list[:, h] self.A[:, idx] = self.A[:, idx] - self.lrate * np.dot( diff --git a/pamica/numpy_impl/utils.py b/pamica/numpy_impl/utils.py index 32a141c..7c9dede 100644 --- a/pamica/numpy_impl/utils.py +++ b/pamica/numpy_impl/utils.py @@ -131,8 +131,6 @@ def identify_shared_components(A, W, comp_list, comp_thresh=0.99): num_comps = A.shape[1] data_dim = A.shape[0] - comp_used = np.ones(num_comps, dtype=bool) - # Compare components between models for h1 in range(num_models): for h2 in range(h1 + 1, num_models): @@ -160,9 +158,17 @@ def identify_shared_components(A, W, comp_list, comp_thresh=0.99): if not shared: # Merge components - comp_used[k2] = False comp_list[comp_list == k2] = k1 + # Derive comp_used from the final comp_list rather than tracking it during + # the merge loop. A fresh np.ones() per call forgot every column merged away + # in an earlier call: once comp_list is fully merged the k1 == k2 guard skips + # every pair, and the mask came back all-True while half the columns were + # dead (issue #240). Matches AMICATorchNG.comp_used, which is a property + # derived the same way. + comp_used = np.zeros(num_comps, dtype=bool) + comp_used[np.unique(comp_list)] = True + return comp_list, comp_used diff --git a/pamica/tests/test_numpy_share_comps.py b/pamica/tests/test_numpy_share_comps.py new file mode 100644 index 0000000..670867e --- /dev/null +++ b/pamica/tests/test_numpy_share_comps.py @@ -0,0 +1,139 @@ +"""``share_comps`` on the NumPy backend (issues #240, #242). + +Two defects, both on the path a merge opens up and neither reachable with the +default disjoint ``comp_list``: + +* ``comp_used`` was rebuilt fresh on every ``identify_shared_components`` call, + so once ``comp_list`` was fully merged the ``k1 == k2`` guard skipped every + pair and the mask came back all-True while half the columns were dead. The + unmasked mixture update then divided 0/0 and left NaN in ``mu``/``beta`` + while the fit reported success (#240). +The A-update defect on the same path (#242) is not covered here and not fixed +here: a shared column still takes one step per contributing model. Three +attempts at a test for it passed equally with the loop and with Fortran's single +weighted application, so it ships with a test that actually distinguishes them. + +Real sample EEG throughout. Sharing is forced by construction where a short fit +would not reliably produce a merge. +""" + +from pathlib import Path + +import numpy as np +import pytest + +from pamica import AMICA_NumPy as AMICA +from pamica.numpy_impl.data import load_data_file +from pamica.numpy_impl.utils import identify_shared_components + +_FDT = Path(__file__).resolve().parent.parent / "sample_data" / "eeglab_data.fdt" + +pytestmark = pytest.mark.skipif(not _FDT.exists(), reason="sample data missing") + + +def _real_data(n_samples: int = 4096) -> np.ndarray: + data = load_data_file(str(_FDT), 32, 30504, dtype=np.float32) + return data[:, :n_samples].astype(np.float64) + + +def _shared_fit(max_iter: int = 5, share_comps: bool = True): + model = AMICA( + num_models=2, + num_mix=3, + max_iter=max_iter, + seed=7, + share_comps=share_comps, + share_start=1, + share_int=2, + ) + model.fit(_real_data()) + assert model.comp_list is not None + return model + + +# --- comp_used staleness (#240) --------------------------------------------- +def _force_collinear_pair(model): + """Make model 1's first column a copy of model 0's, so a merge must fire. + + Deterministic on purpose. Keying the test on whether a merge happened to + occur would make it skip under the very bug it guards: the stale mask is + all-True, which reads as "no merge fired". + """ + k0 = int(model.comp_list[0, 0]) + k1 = int(model.comp_list[0, 1]) + model.A[:, k1] = model.A[:, k0] + return k0, k1 + + +def test_comp_used_survives_a_second_identify_call(): + """The mask must not forget columns merged away by an earlier call. + + Calling twice is the crux: the second call sees an already-merged + ``comp_list``, so every pair hits the ``k1 == k2`` guard and no merge fires. + A mask built during that loop comes back all-True. + """ + model = _shared_fit(max_iter=3, share_comps=False) + _force_collinear_pair(model) + + comp_list_after, used_first = identify_shared_components( + model.A, model.W, model.comp_list.copy(), model.comp_thresh + ) + assert not used_first.all(), "setup failed: the forced collinear pair did not merge" + + _, used_second = identify_shared_components( + model.A, model.W, comp_list_after.copy(), model.comp_thresh + ) + assert used_second.sum() == used_first.sum(), ( + "comp_used was rebuilt from scratch and forgot the earlier merge" + ) + + +def test_comp_used_matches_the_columns_comp_list_references(): + """The mask is exactly the set of referenced columns, as in AMICATorchNG.""" + model = _shared_fit() + referenced = np.zeros(model.num_comps, dtype=bool) + referenced[np.unique(model.comp_list)] = True + np.testing.assert_array_equal(model.comp_used, referenced) + + +# --- NaN mixture parameters (#240) ------------------------------------------ +def test_sharing_leaves_finite_mixture_parameters(): + """A merged-away column receives no mass, so its update is 0/0. + + Before the fix this left NaN in half of ``mu`` and ``beta`` while the fit + returned normally. + """ + model = _shared_fit() + for name in ("A", "mu", "beta", "gm", "alpha"): + value = np.asarray(getattr(model, name)) + assert np.all(np.isfinite(value)), f"{name} holds non-finite values" + + +def test_unused_columns_keep_their_last_finite_value(): + """Frozen, not zeroed: an unused column keeps the value it last held. + + The merge is forced rather than hoped for, so a stale all-True mask fails + here instead of skipping. + """ + model = _shared_fit(max_iter=3, share_comps=False) + _force_collinear_pair(model) + model.comp_list, model.comp_used = identify_shared_components( + model.A, model.W, model.comp_list, model.comp_thresh + ) + unused = ~model.comp_used + assert unused.any(), "setup failed: no column was merged away" + + updates = model._get_updates_and_likelihood() + model._update_parameters(updates) + assert np.all(np.isfinite(model.mu[:, unused])) + assert np.all(model.beta[:, unused] > 0.0) + assert np.all(np.isfinite(model.alpha[:, unused])) + + +def test_default_comp_list_is_unaffected(): + """Every column has one contributor without sharing, so nothing changes.""" + model = AMICA(num_models=2, num_mix=3, max_iter=5, seed=42) + model.fit(_real_data()) + assert model.comp_used is None or model.comp_used.all() + for name in ("A", "mu", "beta"): + assert np.all(np.isfinite(np.asarray(getattr(model, name))))