Address #51 review: degenerate final_ll_, n_kurt_done, tests

PR review (4 Sonnet reviewers) findings:
- final_ll_ was left at the last finite ll_history value after a degenerate
  stop while the model held diverged params; now set to NaN there
  (silent-failure review).
- Roll back n_kurt_done with pdtype in the best-iterate snapshot so a restored
  model's adaptive-PDF switch count stays consistent with its pdtype.
- Warn once when keep_best is requested together with do_reject (the safeguard
  is silently inactive there).
- Fix the restore-skip comment (a singular_ll stop leaves A/W finite, not
  non-finite; the skip is a scope decision deferred to #50) and qualify the
  _KEEP_BEST_TOL min_dll scale reference (Fortran's, not the legacy NumPy one).

Tests added/expanded:
- keep_best inactive under do_reject (final_ll_ == ll_history[-1]).
- keep_best False-vs-True differential on an overshooting multi-model run, with
  an explicit skip when the run is monotone so restore-branch coverage is
  visible, not vacuous.
- wrapper final_ll_ survives save/load.

69 -> 72 torch tests pass; ruff clean.

Author: neuromechanist

pyAMICA/tests/torch_tests/test_amica_ng_wrapper.py

@@ -63,6 +63,10 @@ def test_ng_save_load_roundtrip(fitted_ng, real_data, tmp_path):
     assert loaded.n_models == fitted_ng.n_models
     assert loaded.n_mix == fitted_ng.n_mix
     assert loaded.ll_history_ == fitted_ng.ll_history_
+    # final_ll_ (the fitted model's LL, issue #51) is populated and survives the
+    # round-trip -- use it, not ll_history_[-1], as the model's log-likelihood.
+    assert fitted_ng.final_ll_ is not None
+    assert loaded.final_ll_ == fitted_ng.final_ll_
 
     # torch.save/load restores tensors bit-exactly and CPU matmul is
     # deterministic, so transform() on the restored tensors reproduces the

pyAMICA/tests/torch_tests/test_ng_backend.py

@@ -1162,32 +1162,68 @@ def test_keep_best_snapshot_restore_roundtrip():
         assert torch.equal(getattr(m, name), snap[name]), name
 
 
+def _multimodel_keep_best(seed: int, keep_best: bool) -> AMICATorchNG:
+    m = AMICATorchNG(
+        n_channels=NW, n_models=2, n_mix=NMIX, seed=seed, device="cpu",
+        dtype=torch.float64, block_size=512, lrate=0.05, maxdecs=3,
+        do_newton=True, newt_start=50, newt_ramp=10, newtrate=1.0,
+        keep_best=keep_best,
+    )  # fmt: skip
+    m.fit(_load_real_data(), max_iter=100, verbose=False)
+    return m
+
+
 @pytest.mark.skipif(not DATA_FILE.exists(), reason="sample data missing")
 def test_keep_best_returns_within_tol_of_peak():
     """On a real multi-model fit the safeguard (issue #51) returns a model whose
     log-likelihood is within tolerance of the best iterate seen and never below
     the last iterate, and ``final_ll_`` reflects the *returned* parameters (not
-    the raw ``ll_history[-1]``, which stays the true trajectory). seed 8 reaches
-    the plateau where natural-gradient AMICA dips below its own peak, so the
-    restore branch runs; the assertions are contract invariants that also hold
-    if a platform's BLAS makes the run monotone."""
-    data = _load_real_data()
-    m = AMICATorchNG(
-        n_channels=NW, n_models=2, n_mix=NMIX, seed=8, device="cpu",
-        dtype=torch.float64, block_size=512, lrate=0.05, maxdecs=3,
-        do_newton=True, newt_start=50, newt_ramp=10, newtrate=1.0, keep_best=True,
-    )  # fmt: skip
-    m.fit(data, max_iter=100, verbose=False)
+    the raw ``ll_history[-1]``, which stays the true trajectory). keep_best does
+    not change the optimization path, only which iterate is returned, so the
+    ``keep_best=False`` run has the same trajectory but returns the (lower) last
+    iterate. seed 8 reaches the plateau where natural-gradient AMICA dips below
+    its own peak, so the restore branch runs; the invariants also hold if a
+    platform's BLAS makes the run monotone (see the explicit skip below)."""
+    on = _multimodel_keep_best(8, keep_best=True)
+    off = _multimodel_keep_best(8, keep_best=False)
+
+    # keep_best does not alter the trajectory, only the returned iterate.
+    assert on.ll_history == off.ll_history
+    # return-last returns exactly the last iterate; keep_best is never worse.
+    assert off.final_ll_ == off.ll_history[-1]
+    assert on.final_ll_ >= off.final_ll_
 
-    peak = max(m.ll_history)
+    peak = max(on.ll_history)
     # The returned LL is within tolerance of the peak and never below the last.
-    assert abs(m.final_ll_ - peak) <= _KEEP_BEST_TOL
-    assert m.final_ll_ >= m.ll_history[-1] - 1e-12
-    # ll_history stays the true (possibly-overshooting) trajectory: keep_best
-    # returns a value at least as good as the raw last iterate would give.
-    assert m.final_ll_ >= m.ll_history[-1]
+    assert abs(on.final_ll_ - peak) <= _KEEP_BEST_TOL
+    assert on.final_ll_ >= on.ll_history[-1]
     # The returned parameters really sit at final_ll_ (recompute the E-step LL).
-    X_t = m._preprocess(data)
-    acc = m._accumulate_blocks(X_t)
+    data = _load_real_data()
+    X_t = on._preprocess(data)
+    acc = on._accumulate_blocks(X_t)
     ll_model = float(acc["ll"] / (X_t.shape[1] * NW))
-    assert abs(ll_model - m.final_ll_) < 1e-9
+    assert abs(ll_model - on.final_ll_) < 1e-9
+
+    # Make branch coverage visible rather than silently vacuous: if this run did
+    # not overshoot on this platform, the restore branch was not exercised.
+    if peak - on.ll_history[-1] <= _KEEP_BEST_TOL:
+        pytest.skip("seed 8 did not overshoot here; restore branch not exercised")
+    # It did overshoot, so keep_best strictly beat return-last.
+    assert on.final_ll_ > off.final_ll_
+
+
+@pytest.mark.skipif(not DATA_FILE.exists(), reason="sample data missing")
+def test_keep_best_inactive_under_reject():
+    """The safeguard is disabled under ``do_reject`` (the good-sample set, hence
+    the LL normalization, changes across iterations, so per-iteration LLs are not
+    comparable): ``final_ll_`` is exactly the last trajectory value, no restore
+    fires (issue #51)."""
+    data = _load_real_data()
+    m = AMICATorchNG(
+        n_channels=NW, n_models=2, n_mix=NMIX, seed=SEED, device="cpu",
+        dtype=torch.float64, block_size=512, do_reject=True, rejsig=2.0,
+        rejstart=2, rejint=3, maxrej=2, keep_best=True,
+    )  # fmt: skip
+    m.fit(data, max_iter=12, verbose=False)
+    assert m.numrej >= 1  # rejection actually fired, so the good set changed
+    assert m.final_ll_ == m.ll_history[-1]  # no best-iterate restore under reject

pyAMICA/torch_impl/amica_torch_ng.py

@@ -75,8 +75,11 @@
 # multi-model LL variance -- one seed peaked at -3.357 then crashed to -3.545 in
 # its last iterations). fit() therefore tracks the highest-LL iterate and
 # restores it when the final LL falls more than this tolerance below that peak.
-# Units: mean log-likelihood per sample-channel (same scale as min_dll), so 1e-9
-# reads as "numerical noise, not a real overshoot". The threshold also keeps a
+# Units: mean log-likelihood per sample-channel (the same scale as Fortran's
+# min_dll, amica17.f90:1866, which normalizes LL(iter) by numgoodsum*nw before
+# comparing -- NOT the legacy NumPy pyAMICA.min_dll, which compares un-normalized
+# summed LL), so 1e-9 reads as "numerical noise, not a real overshoot". The
+# threshold also keeps a
 # monotone single-model run (issue #24 parity) a bit-exact no-op: its final
 # iterate already IS the best, the gap is 0 < tol, and no restore fires.
 _KEEP_BEST_TOL = 1e-9
@@ -1138,18 +1141,27 @@ def _pdtype_from_kurtosis(
             )
         return result
 
-    def _snapshot_params(self) -> Dict[str, torch.Tensor]:
-        """Clone the fitted parameter tensors for the best-iterate safeguard
-        (issue #51). Clones (not aliases) so the live in-place M-step updates do
-        not roll the snapshot forward. Covers exactly ``_PARAM_TENSORS``; the
-        constant preprocessing tensors (``mean``/``sphere``) are included so a
-        restore is a total, unambiguous rollback of model state."""
-        return {name: getattr(self, name).clone() for name in self._PARAM_TENSORS}
+    def _snapshot_params(self) -> Dict[str, object]:
+        """Snapshot the fitted state for the best-iterate safeguard (issue #51).
+
+        Clones each ``_PARAM_TENSORS`` tensor (not an alias) so the live in-place
+        M-step updates do not roll the snapshot forward; the constant
+        preprocessing tensors (``mean``/``sphere``) are included so a restore is a
+        total rollback. Also captures the scalar ``n_kurt_done`` (the adaptive-PDF
+        switch counter that gates ``pdtype``) so a restored model's switch count
+        stays consistent with its rolled-back ``pdtype`` -- otherwise a switch
+        applied after the peak iterate would leave the two out of sync in a saved
+        model (silent-failure review)."""
+        snap: Dict[str, object] = {
+            name: getattr(self, name).clone() for name in self._PARAM_TENSORS
+        }
+        snap["n_kurt_done"] = self.n_kurt_done
+        return snap
 
-    def _restore_params(self, snapshot: Dict[str, torch.Tensor]) -> None:
-        """Restore parameter tensors captured by :meth:`_snapshot_params`."""
-        for name, tensor in snapshot.items():
-            setattr(self, name, tensor)
+    def _restore_params(self, snapshot: Dict[str, object]) -> None:
+        """Restore the state captured by :meth:`_snapshot_params`."""
+        for name, value in snapshot.items():
+            setattr(self, name, value)
 
     # ------------------------------------------------------------------
     # Public API
@@ -1201,7 +1213,15 @@ def fit(
         # LLs are not comparable.
         track_best = self.keep_best and not self.do_reject
         best_ll = -math.inf
-        best_snapshot: Optional[Dict[str, torch.Tensor]] = None
+        best_snapshot: Optional[Dict[str, object]] = None
+        if self.keep_best and self.do_reject:
+            # keep_best defaults on, so a user enabling rejection would otherwise
+            # silently lose the safeguard; surface it once (silent-failure review).
+            logger.warning(
+                "keep_best is inactive under do_reject: the good-sample set (and "
+                "the per-iteration LL normalization) changes as samples are "
+                "rejected, so best-iterate selection by LL is not well-defined."
+            )
 
         iterator = tqdm(range(max_iter), desc="AMICA-NG", disable=not verbose)
         for it in iterator:
@@ -1309,17 +1329,26 @@ def fit(
 
             iterator.set_postfix({"LL": f"{ll:.4f}", "lrate": f"{self.lrate:.4g}"})
 
-        # Log-likelihood of the parameters fit() returns. Defaults to the last
-        # trajectory value; overwritten with the best iterate's LL below if the
-        # safeguard restores it.
-        self.final_ll_ = self.ll_history[-1] if self.ll_history else float("nan")
+        # Log-likelihood of the parameters fit() returns. A degenerate stop
+        # leaves the model on the diverged parameters, whose LL is NOT the last
+        # finite ll_history value (the guard breaks before appending), so report
+        # NaN there rather than a stale healthy-looking number (silent-failure
+        # review). Otherwise it is the last trajectory value, overwritten with the
+        # best iterate's LL below if the safeguard restores it.
+        if self.stop_reason in self._DEGENERATE_STOP_REASONS:
+            self.final_ll_ = float("nan")
+        else:
+            self.final_ll_ = self.ll_history[-1] if self.ll_history else float("nan")
 
         # Restore the best iterate if the run ended materially below it (issue
-        # #51). Skipped for a degenerate stop (params are already non-finite and
-        # state_dict()/the wrapper reject them, so there is nothing good to keep)
-        # and when the final LL is within _KEEP_BEST_TOL of the best -- a monotone
-        # single-model run has final == best, so no restore fires and issue #24
-        # parity stays bit-exact.
+        # #51). Skipped for a degenerate stop -- not because the parameters are
+        # necessarily non-finite (a singular_ll stop leaves A/W finite but
+        # singular) but because salvaging a diverged run here would pre-empt issue
+        # #50's degenerate-fit contract; state_dict() already refuses to persist
+        # any model whose stop_reason is degenerate. Also skipped when the final
+        # LL is within _KEEP_BEST_TOL of the best -- a monotone single-model run
+        # has final == best, so no restore fires and issue #24 parity stays
+        # bit-exact.
         if (
             track_best
             and best_snapshot is not None