Fix MLX rholrate to ratchet at maxdecs (#195) (#197)

* Fix MLX rholrate to ratchet at maxdecs (#195)

* Match newt_start default to torch/numpy; add reset test

Author: neuromechanist

pamica/mlx_impl/core.py

@@ -93,6 +93,7 @@ def __init__(
         lratefact: float = 0.5,
         maxdecs: int = 5,
         newt_ramp: int = 10,
+        newt_start: int = 20,
         do_newton: bool = False,
         rho0: float = 1.5,
         minrho: float = 1.0,
@@ -134,6 +135,11 @@ def __init__(
         self.lratefact = lratefact
         self.maxdecs = maxdecs
         self.newt_ramp = newt_ramp
+        # Schedule threshold for the rholrate ceiling ratchet (Fortran
+        # amica15.f90:1049 gates it on iter > newt_start, independent of
+        # do_newton). MLX does no Newton, but keeps newt_start as this gate so the
+        # rho-rate schedule matches Fortran/torch.
+        self.newt_start = newt_start
         self.do_newton = False
 
         self.rho0 = rho0
@@ -602,7 +608,18 @@ def fit(
             self.ll_history.append(ll)
 
             # Learning-rate control (Fortran amica17.f90:1062-1108): anneal on an
-            # LL decrease; ratchet the ceiling after maxdecs persistent decreases.
+            # LL decrease; ratchet the ceilings after maxdecs persistent decreases.
+            #
+            # rholrate is a maxdecs-ratcheted CEILING, not a per-decrease-annealed
+            # working rate. Fortran resets rholrate=rholrate0 every iteration before
+            # the rho update (amica15.f90:1788) and only tightens the rholrate0
+            # ceiling at maxdecs (amica15.f90:1050, gated on iter > newt_start), so
+            # its per-decrease rholrate*=rholratefact (:1045) never reaches the rho
+            # update. rho has no ramp, so self.rholrate carries that ceiling directly
+            # (reset to rholrate0 at fit start, nothing re-inflates it) and must
+            # ratchet ONLY at maxdecs. The previous per-decrease decay collapsed the
+            # rho rate to ~1e-5 within a few hundred iterations and froze rho at a
+            # stale shape (issue #195, mirroring the torch/numpy fix in #193/#194).
             if len(self.ll_history) > 1 and ll < self.ll_history[-2]:
                 if self.lrate <= self.minlrate:
                     logger.warning(
@@ -613,10 +630,11 @@ def fit(
                     self.stop_reason = "lrate_floor"
                     break
                 self.lrate *= self.lratefact
-                self.rholrate *= self.rholratefact
                 numdecs += 1
                 if numdecs >= self.maxdecs:
                     self.lrate_cap *= self.lratefact
+                    if it > self.newt_start:
+                        self.rholrate *= self.rholratefact
                     numdecs = 0
 
         if self.stop_reason in self._DEGENERATE_STOP_REASONS:

pamica/tests/mlx_tests/test_mlx_backend.py

@@ -310,3 +310,81 @@ def test_multimodel_dead_model_keeps_prior_c():
     m.c = mx.array(np.stack([np.zeros(NW, np.float32), marker], axis=1))
     m._update_parameters(acc, xt.shape[1])
     assert np.allclose(np.array(m.c)[:, 1], marker)  # dead model kept its prior c
+
+
+def test_rholrate_ratchets_at_maxdecs_not_per_decrease():
+    """Issue #195 (mirrors #193/#194 for torch/numpy): the MLX rho learning rate is
+    a maxdecs-ratcheted CEILING, not a per-LL-decrease monotone decay.
+
+    Fortran resets ``rholrate = rholrate0`` each iteration before the rho update
+    (amica15.f90:1788) and only tightens the ceiling at ``maxdecs``
+    (amica15.f90:1050, gated on ``iter > newt_start``). MLX previously decayed
+    ``rholrate`` on EVERY LL decrease with no reset, collapsing it toward ~1e-5
+    within a few hundred iterations and freezing rho at a stale shape.
+
+    An aggressive-lrate natural-gradient run on the real sample overshoots and
+    triggers several LL decreases; with ``newt_start=0`` (gate always open) the
+    ceiling ratchets once per ``maxdecs`` decreases, NOT once per decrease.
+    """
+    from pamica.mlx_impl import AMICAMLXNG
+
+    data = _load_real_data()
+    m = AMICAMLXNG(
+        n_channels=NW, n_mix=NMIX, seed=SEED, block_size=256,
+        lrate=0.5, lratefact=0.5, rholrate=0.05, rholratefact=0.5,
+        maxdecs=3, newt_start=0,
+    )  # fmt: skip
+    m.fit(data, max_iter=400, verbose=False)
+
+    ll = m.ll_history
+    n_dec = sum(1 for i in range(1, len(ll)) if ll[i] < ll[i - 1])
+    assert m.lrate > m.minlrate, "run hit the lrate floor; use a gentler config"
+    assert n_dec >= m.maxdecs, "too few LL decreases to exercise the ratchet"
+
+    # The ceiling ratcheted a whole number of times at the maxdecs cadence
+    # (rholrate0 * rholratefact**k), far fewer steps than one-per-decrease.
+    assert m.rholrate < m.rholrate0
+    k = round(np.log(m.rholrate / m.rholrate0) / np.log(m.rholratefact))
+    assert m.rholrate == pytest.approx(m.rholrate0 * m.rholratefact**k)
+    assert k <= n_dec // m.maxdecs + 1
+    # Guard against a regression to the old per-decrease decay (orders below).
+    buggy = m.rholrate0 * (m.rholratefact**n_dec)
+    assert m.rholrate > buggy * 10
+
+
+def test_rholrate_ceiling_ratchet_gated_on_newt_start():
+    """The rholrate ceiling ratchet is gated on ``iter > newt_start`` (Fortran
+    amica15.f90:1049, independent of do_newton). With ``newt_start`` past the whole
+    budget, LL decreases must NOT tighten the rho ceiling -- it stays at
+    ``rholrate0`` exactly (only ``lrate_cap`` ratchets, which is not gated)."""
+    from pamica.mlx_impl import AMICAMLXNG
+
+    data = _load_real_data()
+    m = AMICAMLXNG(
+        n_channels=NW, n_mix=NMIX, seed=SEED, block_size=256,
+        lrate=0.5, lratefact=0.5, rholrate=0.05, rholratefact=0.5,
+        maxdecs=3, newt_start=100_000,
+    )  # fmt: skip
+    m.fit(data, max_iter=400, verbose=False)
+
+    ll = m.ll_history
+    n_dec = sum(1 for i in range(1, len(ll)) if ll[i] < ll[i - 1])
+    assert n_dec >= m.maxdecs, "config did not exercise the decrease path"
+    # Gated off for the whole run: the rho ceiling never moved.
+    assert m.rholrate == m.rholrate0
+
+
+def test_rholrate_ceiling_resets_at_fit_start():
+    """Issue #195: the rho-rate ceiling is reset to rholrate0 at fit start
+    (``_initialize_parameters``, the same call ``fit`` makes), so a previously
+    ratcheted ``rholrate`` does not carry across a re-fit/restart -- parity with
+    the numpy backend's ``test_reinitialize_for_restart_resets_rho_ceiling``."""
+    from pamica.mlx_impl import AMICAMLXNG
+
+    m = AMICAMLXNG(n_channels=NW, n_mix=NMIX, seed=SEED)
+    # Simulate a prior fit that ratcheted both ceilings down at maxdecs.
+    m.rholrate = m.rholrate0 * 0.25
+    m.lrate_cap = m.lrate0 * 0.25
+    m._initialize_parameters()  # fit-start reset
+    assert m.rholrate == m.rholrate0
+    assert m.lrate_cap == m.lrate0