`DriftCorrection` shall use more than two images to compute loss in affine transformation (time-series)

[bobleesj]

Problem

Atm, https://github.com/electronmicroscopy/quantem/blob/dev/src/quantem/imaging/drift.py hardcodes two images to compute affine/translation loss:

im0, w0 = self.interpolator[0].warp_image(
    self.images[0].array,
    knot_0,
)
im1, w1 = self.interpolator[1].warp_image(
    self.images[1].array,
    knot_1,
)
# Cross correlation alignment
shifts, image_shift = cross_correlation_shift(
    im0,
    im1,
    ...
)
cost[a0] = np.mean(np.abs(im0 - image_shift))

Proposed solution

Extend affine() to use all images (not just 2)

Sth like below:

def _compute_cost(drift):
    # Warp all images with candidate drift
    warped_images = []
    for img_idx in range(self.num_images):
        knot = self.knots[img_idx].copy()
        scanline_offset = np.arange(knot.shape[1]) - (knot.shape[1] - 1) / 2
        knot[0] += drift[0] * scanline_offset[:, None]
        knot[1] += drift[1] * scanline_offset[:, None]
        ...
        warped_images.append(warped)

    # Compute cost: average of all images aligned to reference
    ref = warped_images[0]
    total_cost = 0
    for img_idx in range(1, self.num_images):
        _, aligned = cross_correlation_shift(
            ref, warped_images[img_idx],
            ...
        )
    ...

[bobleesj]

For each drift candidate:

1. Warp all images with that given drift
2. For all pairs (i, j):
   • Cross-correlate to align translation
   • Compute loss where BOTH images have valid data (no padding area)
3. Return average cost across all pairs

Seems to work:

Masked to valid regions only

[cophus]

Looks good! A few commentS:

1. Performing affine translation across all pairs seems a bit wasteful - it would be better to do it using a global alignment (not sure exactly how this will work in practice).
2. Limiting to only masked regions valid over all images won't work for all data, for example long time series where image 1 doesn't overlap with the final image. Plus if we run 0, 45, 90, 135, etc. all the corners are still valid for at least 2 images.