Add tests: persistence, kurtosis guard, Newton, rejection

Author: neuromechanist

pyAMICA/tests/torch_tests/test_ng_pdf_families.py

@@ -24,7 +24,7 @@
 NW = 32
 FIELD = 30504
 
-# Fortran log-normalizer literals (amica15.f90:1315/1327/1341/1353).
+# Fortran log-normalizer literals (amica15.f90:1315/1328/1341/1353).
 _LOG4 = math.log(4.0)
 _LSQ2PI = math.log(2.506628274)
 _LNSUB = math.log(4.132731354)
@@ -135,10 +135,47 @@ def test_pdftype_validation():
     AMICATorchNG(n_channels=NW, n_mix=1, pdftype=4, device="cpu")
 
 
+def test_kurt_schedule_validation():
+    """The adaptive-switch schedule params are validated at construction so a
+    bad value fails loudly instead of crashing deep in fit()."""
+    for bad in dict(kurt_int=0, kurt_start=0, num_kurt=-1).items():
+        with pytest.raises(ValueError):
+            AMICATorchNG(n_channels=NW, n_mix=1, pdftype=1, device="cpu", **dict([bad]))
+    # The same params are inert (unvalidated) for non-adaptive pdftype.
+    AMICATorchNG(n_channels=NW, n_mix=3, pdftype=0, device="cpu", kurt_int=0)
+
+
+def test_pdtype_from_kurtosis_decision():
+    """The pure kurtosis->family decision: super-G(+)->1, sub-G(-)->4, and a
+    non-finite / dead-model kurtosis keeps the prior pdtype (the guard). This is
+    the sub-Gaussian (code 4) switch branch that real EEG rarely triggers."""
+    # Single model: cover +, -, NaN (keep prior), - again.
+    m = AMICATorchNG(n_channels=4, n_mix=1, pdftype=1, device="cpu", seed=0)
+    m._initialize_parameters()  # sets self.pdtype to all-1 (prior)
+    kurt = torch.tensor([[2.0], [-2.0], [float("nan")], [-0.5]], dtype=torch.float64)
+    nsub = torch.tensor([10.0], dtype=torch.float64)
+    out = m._pdtype_from_kurtosis(kurt, nsub)
+    assert out.flatten().tolist() == [1, 4, 1, 4]  # NaN -> kept prior (1)
+
+    # Two models, model 1 dead (nsub==0): its sources keep the prior (1) even
+    # though their (finite) negative kurtosis would otherwise pick code 4.
+    m2 = AMICATorchNG(
+        n_channels=2, n_models=2, n_mix=1, pdftype=1, device="cpu", seed=0
+    )
+    m2._initialize_parameters()
+    kurt2 = torch.full((2, 2), -2.0, dtype=torch.float64)
+    nsub2 = torch.tensor([10.0, 0.0], dtype=torch.float64)
+    out2 = m2._pdtype_from_kurtosis(kurt2, nsub2)
+    assert out2[:, 0].tolist() == [4, 4]  # live model switched to sub-Gaussian
+    assert out2[:, 1].tolist() == [1, 1]  # dead model kept prior
+
+
 @pytest.mark.skipif(not DATA_FILE.exists(), reason="sample data missing")
 @pytest.mark.parametrize("pdftype,n_mix", [(0, 3), (2, 3), (3, 3), (4, 1), (1, 1)])
 def test_family_fit_finite_and_monotone(pdftype: int, n_mix: int):
-    """Every family fits real EEG to a finite, monotonically improving LL."""
+    """Every family fits real EEG to a finite LL that does not regress below its
+    starting value (natural-gradient AMICA can dip mid-run, so this checks net
+    non-decrease, last >= first, not strict monotonicity)."""
     data = _load_real_data()
     kw = dict(num_kurt=0) if pdftype == 1 else {}  # fixed super-G, no switching
     m = AMICATorchNG(
@@ -154,7 +191,8 @@ def test_family_fit_finite_and_monotone(pdftype: int, n_mix: int):
 @pytest.mark.skipif(not DATA_FILE.exists(), reason="sample data missing")
 def test_auto_switcher_runs_and_is_stable():
     """The extended-Infomax switcher runs the full schedule, keeps every source
-    in a valid family, and stays finite/monotone on real EEG."""
+    in a valid family, and stays finite with a net non-decreasing LL (last >=
+    first) on real EEG."""
     data = _load_real_data()
     m = AMICATorchNG(
         n_channels=NW,
@@ -188,6 +226,99 @@ def test_auto_switch_noop_when_num_kurt_zero():
     assert np.all(fixed.pdtype.cpu().numpy() == 1)
 
 
+@pytest.mark.skipif(not DATA_FILE.exists(), reason="sample data missing")
+def test_state_dict_roundtrips_pdftype_state():
+    """save/load must preserve the density family. A fixed non-GG model and the
+    adaptive switcher's per-source pdtype/n_kurt_done must survive a round-trip
+    (else a reloaded model silently reverts to GG)."""
+    data = _load_real_data()
+
+    # Fixed logistic family.
+    m = AMICATorchNG(n_channels=NW, n_mix=3, pdftype=3, device="cpu", seed=0)
+    m.fit(data, max_iter=8, verbose=False)
+    loaded = AMICATorchNG.from_state_dict(m.state_dict(), device="cpu")
+    assert loaded.pdftype == 3 and loaded.dorho is False
+    assert torch.equal(loaded.pdtype, m.pdtype)
+    assert np.allclose(loaded.transform(data), m.transform(data))
+
+    # Adaptive switcher: per-source pdtype and the switch counter must persist.
+    ad = AMICATorchNG(
+        n_channels=NW,
+        n_mix=1,
+        pdftype=1,
+        device="cpu",
+        seed=0,
+        kurt_start=3,
+        num_kurt=5,
+        kurt_int=1,
+    )
+    ad.fit(data, max_iter=12, verbose=False)
+    ad_loaded = AMICATorchNG.from_state_dict(ad.state_dict(), device="cpu")
+    assert ad_loaded.pdftype == 1 and ad_loaded.do_choose_pdfs is True
+    assert ad_loaded.n_kurt_done == ad.n_kurt_done == 5
+    assert torch.equal(ad_loaded.pdtype, ad.pdtype)
+
+
+@pytest.mark.skipif(not DATA_FILE.exists(), reason="sample data missing")
+@pytest.mark.parametrize("pdftype,n_mix", [(2, 3), (3, 3), (4, 1)])
+def test_family_fit_with_newton(pdftype: int, n_mix: int):
+    """Non-GG families run with the Newton preconditioner (as Fortran does) and
+    stay finite/monotone on real EEG; the cosh curvature may fall back to natural
+    gradient, which is expected and must not crash."""
+    data = _load_real_data()
+    m = AMICATorchNG(
+        n_channels=NW,
+        n_mix=n_mix,
+        pdftype=pdftype,
+        device="cpu",
+        seed=0,
+        do_newton=True,
+        newt_start=5,
+    )
+    m.fit(data, max_iter=15, verbose=False)
+    ll = np.asarray(m.ll_history)
+    assert np.all(np.isfinite(ll))
+    assert ll[-1] >= ll[0] - 1e-6
+
+
+@pytest.mark.skipif(not DATA_FILE.exists(), reason="sample data missing")
+def test_adaptive_switch_with_rejection():
+    """The adaptive switcher runs after outlier rejection has shrunk the sample
+    set (it consumes the post-rejection X_use) without crashing on real EEG."""
+    data = _load_real_data()
+    m = AMICATorchNG(
+        n_channels=NW,
+        n_mix=1,
+        pdftype=1,
+        device="cpu",
+        seed=0,
+        do_reject=True,
+        rejstart=2,
+        rejint=2,
+        maxrej=1,
+        kurt_start=3,
+        num_kurt=3,
+        kurt_int=1,
+    )
+    m.fit(data, max_iter=12, verbose=False)
+    assert np.all(np.isfinite(m.ll_history))
+    assert set(np.unique(m.pdtype.cpu().numpy())).issubset({1, 4})
+
+
+@pytest.mark.skipif(not DATA_FILE.exists(), reason="sample data missing")
+def test_multimodel_fixed_family():
+    """A fixed non-GG family works with n_models>1 (exercises the per-model
+    _pdtype_h / _choose_pdfs indexing path)."""
+    data = _load_real_data()
+    m = AMICATorchNG(
+        n_channels=NW, n_models=2, n_mix=3, pdftype=2, device="cpu", seed=0
+    )
+    m.fit(data, max_iter=10, verbose=False)
+    ll = np.asarray(m.ll_history)
+    assert np.all(np.isfinite(ll))
+    assert m.pdtype.shape == (NW, 2)
+
+
 @pytest.mark.skipif(
     os.environ.get("AMICA_RUN_FORTRAN") != "1",
     reason="opt-in Fortran-binary integration test (set AMICA_RUN_FORTRAN=1)",