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)

How about

shifted = (da.stack(dashift(field, (y, x))
                    for y, x in product(y_shifts, x_shifts)
          ).reshape((ny, nx) + field.shape))

You don't need to stack unshifted at all — numpy's broadcasting machinery will do the right thing here.

Good idea! I updated the gist accordingly.

Very cool! This can also be a nice application of map_blocks. 🙂

How would this go exactly? I'm imagining 2D indexing into an array of shifts using block_id but that doesn't sound very nice to me 😟

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