Implement support for ndimage.map_coordinates

dask_image/ndinterp/__init__.py

@@ -11,6 +11,7 @@
 from dask.highlevelgraph import HighLevelGraph
 import scipy
 from scipy.ndimage import affine_transform as ndimage_affine_transform
+from scipy.ndimage import map_coordinates as ndimage_map_coordinates
 
 from ..dispatch._dispatch_ndinterp import (
     dispatch_affine_transform,
@@ -381,3 +382,68 @@ def spline_filter1d(
     )
 
     return result
+
+
+def map_coordinates(image, coordinates, order=3,
+                    mode='constant', cval=0.0, prefilter=False):
+    """
+    Pre-map coordinates onto the chunks of 'image' and execute
+    `ndimage.map_coordinates` for every chunk.
+    """
+
+    coordinates = np.asarray(coordinates)
+
+    coord_chunk_locations = [np.searchsorted(np.cumsum(image.chunks[dim]),
+                                             coordinates[dim])
+                             for dim in range(image.ndim)]
+
+    coord_chunk_locations = np.array(coord_chunk_locations).T
+
+    required_chunks, coord_rc_inds = np.unique(coord_chunk_locations,
+                                               axis=0, return_inverse=True)
+
+    rc_coord_inds = [[] for _ in range(len(required_chunks))]
+
+    rois = np.ones((2, len(required_chunks), image.ndim)) * np.nan
+    for ic, c in enumerate(coordinates.T):
+
+        rois[0][coord_rc_inds[ic]] = np.nanmin([rois[0][coord_rc_inds[ic]],
+                                                np.floor(c - order // 2)], 0)
+        rois[1][coord_rc_inds[ic]] = np.nanmax([rois[1][coord_rc_inds[ic]],
+                                                np.ceil(c + order // 2)], 0)
+
+        rc_coord_inds[coord_rc_inds[ic]] += [ic]
+
+    name = "map_coordinates-%s" % tokenize(image,
+                                           coordinates,
+                                           order,
+                                           mode,
+                                           cval,
+                                           prefilter)
+
+    keys = [(name, i) for i in range(len(required_chunks))]
+
+    values = []
+    for irc, rc in enumerate(required_chunks):
+        offset = np.min([np.max([[0] * image.ndim,
+                                 rois[0][irc].astype(np.int64)], 0),
+                         np.array(image.shape) - 1], 0)
+        sl = [slice(offset[dim],
+                    min(image.shape[dim], int(rois[1][irc][dim]) + 1))
+              for dim in range(image.ndim)]
+
+        values.append((ndimage_map_coordinates,
+                       image[tuple(sl)],
+                       coordinates[:, rc_coord_inds[irc]] - offset[:, None],
+                       None,
+                       order,
+                       mode,
+                       cval,
+                       prefilter))
+
+    dsk = dict(zip(keys, values))
+    ar = da.Array(dsk, name, tuple([[len(rc_ci) for rc_ci in rc_coord_inds]]), image.dtype)
+
+    orig_order = np.argsort([ic for rc_ci in rc_coord_inds for ic in rc_ci])
+
+    return ar[orig_order]

docs/coverage.rst

@@ -145,7 +145,7 @@ This table shows which SciPy ndimage functions are supported by dask-image.
      - ✓
    * - ``map_coordinates``
      - ✓
-     -
+     - ✓
    * - ``maximum``
      - ✓
      - ✓

tests/test_dask_image/test_ndinterp/test_map_coordinates.py

@@ -0,0 +1,106 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+from packaging import version
+
+import dask.array as da
+import numpy as np
+import pytest
+import scipy
+import scipy.ndimage
+
+import dask_image.ndinterp
+
+# mode lists for the case with prefilter = False
+_supported_modes = ['constant', 'nearest']
+_unsupported_modes = ['wrap', 'reflect', 'mirror']
+
+# mode lists for the case with prefilter = True
+_supported_prefilter_modes = ['constant']
+_unsupported_prefilter_modes = _unsupported_modes + ['nearest']
+
+have_scipy16 = version.parse(scipy.__version__) >= version.parse('1.6.0')
+
+# additional modes are present in SciPy >= 1.6.0
+if have_scipy16:
+    _supported_modes += ['grid-constant']
+    _unsupported_modes += ['grid-mirror', 'grid-wrap']
+    _unsupported_prefilter_modes += ['grid-constant', 'grid-mirror',
+                                     'grid-wrap']
+
+
+def validate_map_coordinates_general(n=2,
+                                     interp_order=1,
+                                     interp_mode='constant',
+                                     coord_len=12,
+                                     coord_chunksize=6,
+                                     coord_offset=0.,
+                                     im_shape_per_dim=12,
+                                     im_chunksize_per_dim=6,
+                                     random_seed=0,
+                                     prefilter=False,
+                                     ):
+
+    if interp_order > 1 and interp_mode == 'nearest' and not have_scipy16:
+        # not clear on the underlying cause, but this fails on older SciPy
+        pytest.skip("requires SciPy >= 1.6.0")
+
+    # define test image
+    np.random.seed(random_seed)
+    image = np.random.random([im_shape_per_dim] * n)
+    image_da = da.from_array(image, chunks=im_chunksize_per_dim)
+
+    # define test coordinates
+    coords = np.random.random((n, coord_len)) * im_shape_per_dim + coord_offset
+
+    # ndimage result
+    ints_scipy = scipy.ndimage.map_coordinates(
+        image,
+        coords,
+        output=None,
+        order=interp_order,
+        mode=interp_mode,
+        cval=0.0,
+        prefilter=prefilter)
+
+    # dask-image result
+    ints_dask = dask_image.ndinterp.map_coordinates(
+        image_da,
+        coords,
+        order=interp_order,
+        mode=interp_mode,
+        cval=0.0,
+        prefilter=prefilter)
+
+    ints_dask_computed = ints_dask.compute()
+
+    assert np.allclose(ints_scipy, ints_dask_computed)
+
+
+@pytest.mark.parametrize("n",
+                         [1, 2, 3, 4])
+@pytest.mark.parametrize("random_seed",
+                         range(2))
+def test_map_coordinates_basic(n,
+                               random_seed,
+                               ):
+
+    kwargs = dict()
+    kwargs['n'] = n
+    kwargs['random_seed'] = random_seed
+
+    validate_map_coordinates_general(**kwargs)
+
+
+@pytest.mark.timeout(3)
+def test_map_coordinates_large_input():
+
+    # define large test image
+    image_da = da.random.random([1000] * 3, chunks=200)
+
+    # define sparse test coordinates
+    coords = np.random.random((3, 2)) * 1000
+
+    # dask-image result
+    dask_image.ndinterp.map_coordinates(
+        image_da,
+        coords).compute()