Issue in AIM drift correction round 2

[HazenBabcock]

Claude Code identified an issue in the AIM drift correction algorithm as
implemented in the repository (and Picasso). The short explanation is that
the second round of drift correction is usually a NOP because each segment
is aligned against an image that includes itself, and due to the nature of the
min() function this means it will almost always align best to itself. The proposed
fix, leave one out correction, is included.

From Claude Code:

AIM round 2: reference includes the segment being aligned, so it can never revise round 1

This is a question about whether to diverge from picasso/aim.py, not a bug
report — webSMLM reproduces the reference implementation faithfully here.
Details and measurements below; a patch is included if you decide it's worth
diverging.

What I observed

In aimDrift2D() (MODULE: drift), round 2 builds its reference from every
segment and then aligns each segment against it:

const full = new Map();
for (let k = 0; k < nSeg; k++) addTo(full, seg[k], dx[k], dy[k]);
for (let k = 0; k < nSeg; k++) {
  const [sx, sy] = bestShift(seg[k], dx[k], dy[k], full);   // full contains seg[k]

Because bestShift scores Σ min(c_k, ref) and full ⊇ seg[k], at zero shift
every one of segment k's bins returns min(c_k, ref) = c_k. The score there is
Σ c_k — the largest value the score can take — so zero shift is always the
argmax, whatever dx[k] happens to be.

The code comment says round 2 "fixes segment 0 and limits error accumulation".
As written it can do neither: it only ever contributes the sub-pixel parabola
term.

Evidence

Driving the real aimDrift2D() in an instrumented copy, 20 segments × 300
localizations, molecules re-blinking across segments, known linear drift:

• Round 2 never moves a segment by a whole bin. Largest shift across all 20
  segments and both axes was 0.16 bins (2.4 nm), entirely sub-pixel.

• Injecting a 60 nm (4-bin) error into round 1's estimate for one segment:

  round 2	shift it applies	needed
  as shipped	−0.001 bins	−4
  leave-one-out	−4.00 bins	−4

  The shipped version leaves the error entirely in place.

Picasso does the same thing

picasso/aim.py's aim() passes the corrected set as both target and
reference for the second round:

# the second run is with the entire dataset as reference
x_pdc, y_pdc, drift_x2, drift_y2 = intersection_max(
    x_pdc, y_pdc,      # target
    x_pdc, y_pdc,      # reference
    ..., aim_round=2)

and _count_intersections scores with the same min()
(roi_cc[s] += c0 if c0 < cj else cj). With reference and target identical,
c0 == cj everywhere, so zero shift again attains the maximum.

So webSMLM is doing exactly what it documents itself as doing. This behaviour
originates upstream, and possibly in Ma et al. itself — I have not read the
paper's methods, so I can't say whether self-inclusion is intended there.

Measured impact — modest

End-to-end drift recovery against ground truth, RMS residual over all frames,
24 segments, non-monotonic (thermal-wander-like) drift, mean of 5 seeds:

localizations/segment	as shipped	leave-one-out
200	3.50 nm	3.54 nm
60	5.47 nm	4.47 nm
25	35.22 nm	30.46 nm
12	386 nm	386 nm

Leave-one-out helps by roughly 15% in a middling-sparsity window and not at all
elsewhere. With plenty of localizations round 1 is already right and the second
round has nothing to correct; at extreme sparsity everything has failed and it
cannot help either.

Worth being explicit: on well-conditioned data this costs essentially nothing.
It is a safety net that isn't attached rather than a live source of error. The
regime where it would matter — few localizations per segment, short segment
times — is the one AIM is specifically pitched at.

Also note the built-in fixture wouldn't surface this: generateSynthetic()'s
driftpx produces purely linear drift (driftAt = fi/(n-1)*driftPx), which
round 1's sequential chaining already handles well.

Proposed patch, if you want to diverge

Subtract the segment's own contribution before scoring and restore it after —
O(nSeg) overall, rather than O(nSeg²) to rebuild the map per segment.

@@ MODULE: drift — aimDrift2D()
   const KEY=(bx,by)=>(bx+16384)*32768+(by+16384);
+  // Inverse of addTo: decrement this list's own bin counts, dropping empties.
+  const subFrom=(map,list,ox,oy)=>{ for(const L of list){ const k=KEY(Math.floor((L.x+ox)*qd),Math.floor((L.y+oy)*qd));
+    const c=map.get(k); if(c===undefined) continue; if(c<=1) map.delete(k); else map.set(k,c-1); } return map; };
   const addTo=(map,list,ox,oy)=>{ ... };
@@ round 2
   const full=new Map(); for(let k=0;k<nSeg;k++) addTo(full, seg[k], dx[k], dy[k]);
   const dx2=new Float64Array(nSeg), dy2=new Float64Array(nSeg);
   for(let k=0;k<nSeg;k++){
-    const [sx,sy]=bestShift(seg[k], dx[k], dy[k], full);
+    // Leave-one-out: a segment matched against a reference that CONTAINS it
+    // scores min(c_k,ref)=c_k at every one of its own bins, i.e. the maximum
+    // the score can take, so zero shift always wins and round 2 can never
+    // revise round 1.
+    subFrom(full, seg[k], dx[k], dy[k]);
+    const [sx,sy]=bestShift(seg[k], dx[k], dy[k], full);
+    addTo(full, seg[k], dx[k], dy[k]);
     dx2[k]=dx[k]+sx; dy2[k]=dy[k]+sy;
     await yieldMaybe(0.5+0.5*k/nSeg);
   }

Verified: with this applied, the injected 60 nm error above is recovered to
within 0.2 bins (3 nm), and the RMS figures match the leave-one-out column.

Reproducing

Expose aimDrift2D on window, record dx2[k] - dx[k] per segment before the
mean-subtraction step, and call it with synthetic localizations carrying a known
drift. Round-2 shifts come back below one bin for every segment. Adding a hook
that perturbs one dx[k] between the rounds shows the error surviving round 2
untouched.

[jhohlbein]

Thanks, Hazen. I'll look into that and also inform the Picasso folks if you haven't done already. JH

[HazenBabcock]

I haven't informed the Picasso folks. If you agree that this is a real issue please feel free to let them know.