napari example demonstrating parametrically evaluating functions on arrays

napari_image_shift.py

import numpy as np
from scipy.ndimage.interpolation import shift
from dask import delayed
import dask.array as da
import napari
from itertools import product

# define a 2D periodic pattern
x = np.linspace(-np.pi, np.pi, 200) * 5
y = x
field = np.outer(np.sin(y) , np.cos(x))

# define a function that lazily shifts an array
def dashift(array, shifts):
    delshift = delayed(shift)
    return da.from_delayed(delshift(array, shifts), dtype='float64', shape=array.shape)

if __name__ == '__main__':
    # the range of x and y shifts we want to try
    x_shifts = np.arange(-10,10)
    y_shifts = x_shifts
    n_s = len(x_shifts)
    # build lazily shifted & unshifted arrays. Would be great to have a function for this based on da.tile,
    # but da.tile doesn't work for n > 1 dimensions
    shifted = (da.stack(dashift(field, arg)
                    for arg in product(y_shifts, x_shifts))
                ).reshape((n_s, n_s, *field.shape)).rechunk((1,1,-1,-1))
    
    with napari.gui_qt():
        napari.view(shifted - field)

Such a cool usage mode @d-v-b!