bilinear interpolation in dask

[miguelcarcamov]

anyone have ever programmed a bilinear interpolation using dask? I have this function where im is 2D dask array (an image) and uv is a tuple of 2 1D dask arrays, say (u,v) which are in an irregular space:

import dask.array as da
import numpy as np


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

The output of this function must have the same shape of one of the 1D arrays.

The problem is that this part of the code has a problem since these lines

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

are throwing the dask does not support fancy indexing. Has anyone ever been able to make a bilinear interpolation with dask bypassing the fancy indexing error? Cheers

[phobson]

@miguelcarcamov thanks for the issue and clear example.

If I understand correctly, this function would work properly with in-memory numpy arrays, and so this question is more about how to work around dask's limitation in fancy indexing and less about the implementation details of array-based bilinear interpolation.

Is that right?

[miguelcarcamov]

@phobson Yes, exactly - it's about to work around dask's fancy indexing. I have been told to use map_overlap but I'm not quite sure how to use it to do the same....

[phobson]

Thanks for clarifying. I see the conversation on slack that alludes to that. I agree with Martin that the first thing I'd try is to use .map_overlap.

I also wonder if xarray supports this better, though it's hard to imagine you won't run into the same issue if the DataArray is backed by dask.

[miguelcarcamov]

@phobson I have tried with map_blocks and it seems to work. However, I have changed my bilinear_interpolation function a little bit.

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


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

Then I have used map_blocks such that:
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 you can give me a hand @phobson, making map_overlap work, I would appreciate it :)

[phobson]

@miguelcarcamov Can you provide code to generate sample input from scratch?

[miguelcarcamov]

@phobson

import numpy as np
import dask.array as da

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)

[phobson]

Hey @miguelcarcamov -- Thanks for the additional details. With such a good reproducible example, at this point I think it's best to migrate this question to the dask discourse. There, more community members can engage and benefit from this.

[miguelcarcamov]

@phobson I have already posted this question on the dask discourse, but nobody seems to answer :)