Bilinear interpolation using dask

Dask Array
1

Hi there,

I’m trying to develop a bilinear interpolation using dask.

I have this piece of code taken from stack overflow and I changed the numpy calls to dask.array.

def bilinear_interpolate(im: da.Array, uv: tuple) -> da.Array:

    u0 = da.floor(uv[0]).astype(np.int32)
    u1 = u0 + 1
    v0 = da.floor(uv[1]).astype(np.int32)
    v1 = v0 + 1

    u0 = da.clip(u0, 0, im.shape[1] - 1)
    u1 = da.clip(u1, 0, im.shape[1] - 1)
    v0 = da.clip(v0, 0, im.shape[0] - 1)
    v1 = da.clip(v1, 0, im.shape[0] - 1)

    Ia = im[v0, u0]
    Ib = im[v1, u0]
    Ic = im[v0, u1]
    Id = im[v1, u1]

    wa = (u1 - uv[0]) * (v1 - uv[1])
    wb = (u1 - uv[0]) * (uv[1] - v0)
    wc = (uv[0] - u0) * (v1 - uv[1])
    wd = (uv[0] - u0) * (uv[1] - v0)

    return wa * Ia + wb * Ib + wc * Ic + wd * Id

where im is a dask.array image, and uv is a tuple of dask array coordinates, say (x, y).

The main problem here are these lines:

Ia = im[v0, u0]
Ib = im[v1, u0]
Ic = im[v0, u1]
Id = im[v1, u1]

since dask does not support fancy indexing.

Does anyone know how can I make this work? - One possibility that I don’t like is to compute im and the coordinates uv. Is there another way to do this?

2

I have been able to work this out using map_blocks.

import dask.array as da
import numpy as np
from typing import Union

grid_fft = da.random.random((3000,3000)) + 1j * da.random.random((3000,3000))
u = da.random.randint(0, 3000, 48959) + da.random.random(48959)
v = da.random.randint(0, 3000, 48959) + da.random.random(48959)
uv = (u,v)


def bilinear_interpolation(u: Union[float, np.ndarray, da.Array], v: Union[float, np.ndarray, da.Array], im: Union[np.ndarray, da.Array]) -> Union[float, np.ndarray, da.Array]:

    u0 = np.floor(u).astype(np.int32)
    u1 = u0 + 1
    v0 = np.floor(v).astype(np.int32)
    v1 = v0 + 1

    u0 = np.clip(u0, 0, im.shape[1] - 1)
    u1 = np.clip(u1, 0, im.shape[1] - 1)
    v0 = np.clip(v0, 0, im.shape[0] - 1)
    v1 = np.clip(v1, 0, im.shape[0] - 1)

    Ia = im[v0, u0]
    Ib = im[v1, u0]
    Ic = im[v0, u1]
    Id = im[v1, u1]

    wa = (u1 - u) * (v1 - v)
    wb = (u1 - u) * (v - v0)
    wc = (u - u0) * (v1 - v)
    wd = (u - u0) * (v - v0)

    estimated_data = wa * Ia + wb * Ib + wc * Ic + wd * Id

    return estimated_data

data = da.map_blocks(bilinear_interpolation, uv[0], uv[1], grid_fft, meta=np.array(uv[0].shape, dtype=grid_fft.dtype))

This seems to work. However, I have tried to use map_overlap since I know there might be the case that when a coordinate needs to look to a neighbour pixel the value might be in another block. And I’m getting this error:

data = da.map_overlap(bilinear_interpolation, uv[0], uv[1], grid_fft,
                             depth=1, boundary=0, align_arrays=True, allow_rechunk=True, meta=np.array(uv[0].shape, dtype=grid_fft.dtype))


ValueError: Chunks do not add up to shape. Got chunks=((3000,), (48959,)), shape=(3000, 3000)

If anyone can give me a hand making map_overlap work, I would appreciate it :slight_smile: