Fix comp_used staleness and NaN under sharing

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(

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
 
 

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