@@ -0,0 +1,1991 @@
+from builtins import range as irange%0A%0Aimport numpy%0Aimport scipy.ndimage%0A%0Aimport dask%0Aimport dask.array%0A%0A%0Adef fourier_shift_wrap(array, shift):%0A    result = numpy.empty_like(array)%0A    for i in irange(len(array)):%0A        result%5Bi%5D = scipy.ndimage.fourier_shift(array%5Bi%5D, shift%5B0%5D%5Bi%5D)%0A    return result%0A%0A%0Adef find_best_match(matches):%0A    best_match = numpy.zeros(%0A        matches.shape%5B:1%5D,%0A        dtype=matches.dtype%0A    )%0A    if matches.size:%0A        i = numpy.argmin((matches ** 2).sum(axis=0))%0A        best_match = matches%5B:, i%5D%0A%0A    return best_match%0A%0A%0Adef compute_offset(match_mask):%0A    while type(match_mask) is list:%0A        match_mask = match_mask%5B0%5D%0A%0A    result = numpy.empty((len(match_mask), match_mask.ndim - 1), dtype=int)%0A    for i in irange(len(match_mask)):%0A        match_mask_i = match_mask%5Bi%5D%0A%0A        frame_shape = numpy.array(match_mask_i.shape)%0A        half_frame_shape = frame_shape // 2%0A%0A        matches = numpy.array(match_mask_i.nonzero())%0A        above = (matches %3E half_frame_shape%5B:, None%5D).astype(matches.dtype)%0A        matches -= above * frame_shape%5B:, None%5D%0A%0A        result%5Bi%5D = find_best_match(matches)%0A%0A    return result%0A%0A%0Adef roll_frames_chunk(frames, shifts):%0A    # Needed as Dask shares objects and we plan to write to it.%0A    # Also if there is only one refcount the old object is freed.%0A    frames = numpy.copy(frames)%0A%0A    for i in irange(len(frames)):%0A        frames%5Bi%5D = numpy.roll(%0A            frames%5Bi%5D,%0A            tuple(shifts%5Bi%5D),%0A            axis=tuple(irange(frames.ndim - 1))%0A        )%0A%0A    return frames%0A%0A%0Adef roll_frames(frames, shifts):%0A    frames = frames.rechunk(%7B%0A        k: v for k, v in enumerate(frames.shape%5B1:%5D, 1)%0A    %7D)%0A    shifts = shifts.rechunk(%7B1: shifts.shape%5B1%5D%7D)%0A%0A    rolled_frames = dask.array.atop(%0A        roll_frames_chunk, tuple(irange(frames.ndim)),%0A        frames, tuple(irange(frames.ndim)),%0A        shifts, (0, frames.ndim),%0A        dtype=frames.dtype,%0A        concatenate=True%0A    )%0A%0A    return rolled_frames%0A

Add registration functions to `reg`