Implement support for ndimage.map_coordinates

dask_image/ndinterp/__init__.py

@@ -5,12 +5,15 @@
 from itertools import product
 import warnings
 
+from dask import compute, delayed
 import dask.array as da
 import numpy as np
 from dask.base import tokenize
 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 scipy.ndimage import labeled_comprehension as ndimage_labeled_comprehension
 
 from ..dispatch._dispatch_ndinterp import (
  dispatch_affine_transform,
@@ -381,3 +384,195 @@ def spline_filter1d(
  )
 
  return result
+
+
+def _map_single_coordinates_array_chunk(input, coordinates, order=3,
+ mode='constant', cval=0.0, prefilter=False):
+ """
+ Central helper function for implementing map_coordinates.
+
+ Receives 'input' as a dask array and computed coordinates.
+
+ Implementation details and steps:
+ 1) associate each coordinate in coordinates with the chunk
+ it maps to in the input
+ 2) for each input chunk that has been associated to at least one
+ coordinate, calculate the minimal slice required to map
+ all coordinates that are associated to it (note that resulting slice
+ coordinates can lie outside of the coordinate's chunk)
+ 3) for each previously obtained slice and its associated coordinates,
+ define a dask task and apply ndimage.map_coordinates
+ 4) outputs of ndimage.map_coordinates are rearranged to match input order
+ """
+
+ # STEP 1: Associate each coordinate in coordinates with
+ # the chunk it maps to in the input array
+
+ # get the input chunks each coordinates maps onto
+ coords_input_chunk_locations = coordinates.T // np.array(input.chunksize)
+
+ # all input chunk locations
+ input_chunk_locations = np.array([i for i in np.ndindex(input.numblocks)])
+
+ # linearize input chunk locations
+ coords_input_chunk_locations_linear = np.sum(coords_input_chunk_locations *\
+ np.array([np.product(input.numblocks[:dim])
+ for dim in range(input.ndim)])[::-1], axis=1, dtype=np.int64)
+
+ # determine the input chunks that have coords associated and
+ # count how many coords map onto each input chunk
+ chunk_indices_count = np.bincount(coords_input_chunk_locations_linear,
+ minlength=np.product(input.numblocks))
+ required_input_chunk_indices = np.where(chunk_indices_count > 0)[0]
+ required_input_chunks = input_chunk_locations[required_input_chunk_indices]
+ coord_rc_count = chunk_indices_count[required_input_chunk_indices]
+
+ # inverse mapping: input chunks to coordinates
+ required_input_chunk_coords_indices = \
+ [np.where(coords_input_chunk_locations_linear == irc)[0]
+ for irc in required_input_chunk_indices]
+
+ # STEP 2: for each input chunk that has been associated to at least
+ # one coordinate, calculate the minimal slice required to map all
+ # coordinates that are associated to it (note that resulting slice
+ # coordinates can lie outside of the coordinate's chunk)
+
+ # determine the slices of the input array that are required for
+ # mapping all coordinates associated to a given input chunk.
+ # Note that this slice can be larger than the given chunk when coords
+ # lie at chunk borders.
+ # (probably there's a more efficient way to do this)
+ input_slices_lower = np.array([np.clip(
+ ndimage_labeled_comprehension(
+ np.floor(coordinates[dim] - order // 2),
+ coords_input_chunk_locations_linear,
+ required_input_chunk_indices,
+ np.min, np.int64, 0),
+ 0, input.shape[dim] - 1)
+ for dim in range(input.ndim)])
+
+ input_slices_upper = np.array([np.clip(
+ ndimage_labeled_comprehension(
+ np.ceil(coordinates[dim] + order // 2) + 1,
+ coords_input_chunk_locations_linear,
+ required_input_chunk_indices,
+ np.max, np.int64, 0),
+ 0, input.shape[dim])
+ for dim in range(input.ndim)])
+
+ input_slices = np.array([input_slices_lower, input_slices_upper])\
+ .swapaxes(1, 2)
+
+ # STEP 3: For each previously obtained slice and its associated
+ # coordinates, define a dask task and apply ndimage.map_coordinates
+
+ # prepare building dask graph
+ # define one task per associated input chunk
+ name = "map_coordinates_chunk-%s" % tokenize(input,
+ coordinates,
+ order,
+ mode,
+ cval,
+ prefilter)
+
+ keys = [(name, i) for i in range(len(required_input_chunks))]
+
+ # pair map_coordinates calls with input slices and mapped coordinates
+ values = []
+ for irc in range(len(required_input_chunks)):
+
+ ric_slice = [slice(input_slices[0][irc][dim],
+ input_slices[1][irc][dim])
+ for dim in range(input.ndim)]
+ ric_offset = input_slices[0][irc]
+
+ values.append((ndimage_map_coordinates,
+ input[tuple(ric_slice)],
+ coordinates[:, required_input_chunk_coords_indices[irc]]
+ - ric_offset[:, None],
+ None,
+ order,
+ mode,
+ cval,
+ prefilter))
+
+ # build dask graph
+ dsk = dict(zip(keys, values))
+ ar = da.Array(dsk, name, tuple([list(coord_rc_count)]), input.dtype)
+
+ # STEP 4: rearrange outputs of ndimage.map_coordinates
+ # to match input order
+ orig_order = np.argsort(
+ [ic for ric_ci in required_input_chunk_coords_indices for ic in ric_ci])
+
+ # compute result and reorder
+ # (ordering first would probably unnecessarily inflate the task graph)
+ return ar.compute()[orig_order]
+
+
+def map_coordinates(input, coordinates, order=3,
+ mode='constant', cval=0.0, prefilter=False):
+ """
+ Wraps ndimage.map_coordinates.
+
+ Both the input and coordinate arrays can be dask arrays.
+
+ For each chunk in the coordinates array, the coordinates are computed
+ and mapped onto the required slices of the input array. One task is
+ is defined for each input array chunk that has been associated to at
+ least one coordinate. The outputs of the tasks are then rearranged to
+ match the input order. For more details see the docstring of
+ '_map_single_coordinates_array_chunk'.
+
+ input : array_like
+ The input array.
+ coordinates : array_like
+ The coordinates at which to sample the input.
+ order : int, optional
+ The order of the spline interpolation, default is 3. The order has to
+ be in the range 0-5.
+ mode : boundary behavior mode, optional
+ cval : float, optional
+ Value to fill past edges of input if mode is 'constant'. Default is 0.0
+ prefilter : bool, optional
+
+ Comments:
+ - in case of a small coordinate array, it might make sense to rechunk
+ it into a single chunk
+ """
+
+ # if coordinate array is not a dask array, convert it into one
+ if type(coordinates) is not da.Array:
+ coordinates = da.from_array(coordinates, chunks=coordinates.shape)
+ else:
+ # make sure indices are not split across chunks, i.e. that there's
+ # no chunking along the first dimension 
+ if len(coordinates.chunks[0]) > 1:
+ coordinates = da.rechunk(coordinates,
+ (-1,) + coordinates.chunks[1:])
+
+ # if the input array is not a dask array, convert it into one
+ if type(input) is not da.Array:
+ input = da.from_array(input, chunks=input.shape)
+
+ # Map each chunk of the coordinates array onto the entire input array.
+ # 'input' is passed to `_map_single_coordinates_array_chunk` using a bit of
+ # a dirty trick: it is split into its components and passed as a delayed
+ # object, which reconstructs the original array when the task is
+ # executed. Therefore two `compute` calls are required to obtain the
+ # final result, one of which is peformed by
+ # `_map_single_coordinates_array_chunk`
+ output = da.map_blocks(
+ _map_single_coordinates_array_chunk,
+ delayed(da.Array)(input.dask, input.name, input.chunks, input.dtype),
+ coordinates,
+ order=order,
+ mode=mode,
+ cval=cval,
+ prefilter=prefilter,
+ dtype=input.dtype,
+ chunks=coordinates.chunks[1:],
+ drop_axis=0,
+ )
+
+ return output

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,113 @@
+#!/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 input
+ np.random.seed(random_seed)
+ input = np.random.random([im_shape_per_dim] * n)
+ input_da = da.from_array(input, chunks=im_chunksize_per_dim)
+
+ # define test coordinates
+ coords = np.random.random((n, coord_len)) * im_shape_per_dim + coord_offset
+ coords_da = da.from_array(coords, chunks=(n, coord_chunksize))
+
+ # ndimage result
+ mapped_scipy = scipy.ndimage.map_coordinates(
+ input,
+ coords,
+ order=interp_order,
+ mode=interp_mode,
+ cval=0.0,
+ prefilter=prefilter)
+
+ # dask-image results
+ for input_array in [input, input_da]:
+ for coords_array in [coords, coords_da]:
+ mapped_dask = dask_image.ndinterp.map_coordinates(
+ input_array,
+ coords_array,
+ order=interp_order,
+ mode=interp_mode,
+ cval=0.0,
+ prefilter=prefilter)
+
+ mapped_dask_computed = mapped_dask.compute()
+
+ assert np.allclose(mapped_scipy, mapped_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():
+
+ """
+ This test assesses whether relatively large
+ inputs are processed before timeout.
+ """
+
+ # 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()