diff options
Diffstat (limited to 'ep_run/holo_ep.py')
| -rw-r--r-- | ep_run/holo_ep.py | 44 |
1 files changed, 44 insertions, 0 deletions
diff --git a/ep_run/holo_ep.py b/ep_run/holo_ep.py index 31054e4..5485bc7 100644 --- a/ep_run/holo_ep.py +++ b/ep_run/holo_ep.py @@ -254,6 +254,50 @@ def holo_a_track(blk, zs, xin, y, r, T2max, eps, K=10, exit_mult=5.0): return a_best.detach(), t_best +def holo_a_track_fast(blk, zs, xin, y, r, T2max, eps, K=10, exit_mult=5.0): + """EXACT restructure of holo_a_track (same math, ~half the correction cost): the two phase-halves' + deviations from the common mode are exact negatives (v[B:] = -v[:B], since zbar is their mean) and + the Jacobian anchor zbar is shared, so the doubled-batch jvp/vjp computed [J v0; -J v0] redundantly. + Compute Jv/JTv once at batch B and mirror. Bit-equal up to fp nondeterminism.""" + 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 = a_best = None + inc_min, t_best = 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: # measurement brake: Tikhonov-regularized adjoint + f = f - kappa * (Z - zs2a) + v0 = (Z[:B] - zbar).contiguous() # v of the +r phase; the -r phase's v is exactly -v0 + _, 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() + if inc < inc_min: + inc_min, a_best, t_best = inc, a_t, t + elif inc > exit_mult * inc_min and t >= 3 * K: + break + a_prev = a_t + if a_best is None: + a_best = a_prev if a_prev is not None else (Z[B:] - Z[:B]) / (2 * r) + t_best = T2max + return a_best.detach(), t_best + + 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 |