1 files changed, 56 insertions, 0 deletions
diff --git a/ep_run/holo_ep.py b/ep_run/holo_ep.py
index 5485bc7..d8d2be5 100644
--- a/ep_run/holo_ep.py
+++ b/ep_run/holo_ep.py
@@ -298,6 +298,62 @@ def holo_a_track_fast(blk, zs, xin, y, r, T2max, eps, K=10, exit_mult=5.0):
     return a_best.detach(), t_best
 
 
+def holo_a_track_avg(blk, zs, xin, y, r, T2max, eps, K=10, exit_mult=5.0):
+    """track_fast + the semi-convergence fix (t2_probe 2026-07-05): the adjoint iteration on a
+    near-marginal operator SEMI-converges — error dips at a batch-dependent optimum then grows, and the
+    plain argmin-of-increment t_best gets fooled by rotating slow modes. Two changes:
+    (1) trend-aware stop: break after the increment rises on 2 consecutive checks past 2x inc_min
+        (instead of the blunt exit_mult=5 single-shot);
+    (2) plateau averaging: return the MEAN of the a_t whose increment <= 1.5x inc_min (the flat bottom
+        of the semi-convergence curve) — averages out the rotating error component around the optimum."""
+    import torch.func as tf
+    B = zs.size(0)
+    Z = torch.cat([zs, zs], 0)
+    X2 = torch.cat([xin, xin], 0)
+    y2 = torch.cat([y, y], 0)
+    sg = torch.cat([torch.full((B, 1, 1), r, device=zs.device), torch.full((B, 1, 1), -r, device=zs.device)], 0)
+    fnc = lambda zz: blk.nc_force(zz)
+    a_prev = None
+    hist = []                                              # (inc, a_t) at each K-checkpoint
+    inc_min, rise = float('inf'), 0
+    zs2a = torch.cat([zs, zs], 0)
+    kappa = getattr(blk, 'nbrake', 0.0)
+    for t in range(1, T2max + 1):
+        with torch.no_grad():
+            zbar = 0.5 * (Z[:B] + Z[B:])
+            f = rforce(blk, Z, X2) - sg * rgrad_ce(blk, Z, y2, denom=y.numel())
+            if kappa > 0:
+                f = f - kappa * (Z - zs2a)
+            v0 = (Z[:B] - zbar).contiguous()
+            _, Jv0 = tf.jvp(fnc, (zbar,), (v0,))
+            JTv0 = tf.vjp(fnc, zbar)[1](v0)[0]
+            corr0 = Jv0 - JTv0
+            Z = Z + eps * (f - torch.cat([corr0, -corr0], 0))
+        if t % K == 0 or t == T2max:
+            a_t = (Z[B:] - Z[:B]) / (2 * r)
+            if not torch.isfinite(a_t).all():
+                break
+            if a_prev is not None:
+                inc = (a_t - a_prev).norm().item()
+                hist.append((inc, a_t))
+                if inc < inc_min:
+                    inc_min, rise = inc, 0
+                elif inc > 2.0 * inc_min and t >= 3 * K:
+                    rise += 1                              # trend-aware: need 2 consecutive rising checks
+                    if rise >= 2:
+                        break
+                else:
+                    rise = 0
+            a_prev = a_t
+    if not hist:
+        return (a_prev if a_prev is not None else (Z[B:] - Z[:B]) / (2 * r)).detach(), T2max
+    flat = [a for inc, a in hist if inc <= 1.5 * inc_min]  # the semi-convergence plateau
+    if not flat:
+        flat = [min(hist, key=lambda p: p[0])[1]]
+    a_avg = torch.stack(flat).mean(0)
+    return a_avg.detach(), len(hist) * K
+
+
 def holo_a_lockin(blk, zs, xin, y, r, P, ncyc, eps):
     """True oscillatory EP / lock-in estimator (Laborieux–Zenke taken literally) — the
     noisy-physics form: ONE trajectory, sinusoidal nudge beta(t)=r·sin(2πt/P), in-phase