affine_transform: Remove inconsistencies with ndimage implementation. (…

…#205)

* affine_transform: Remove inconsistencies with ndimage implementation.
- change offset default parameter to float
- expect tuple for output_shape (and output_chunks)

* Simplify cupy imports by using pytest.importorskip

* These lines aren't necessary when I test this locally, maybe we can remove them entirely

Co-authored-by: GenevieveBuckley <30920819+GenevieveBuckley@users.noreply.github.com>

dask_image/ndinterp/__init__.py

@@ -25,7 +25,7 @@
 def affine_transform(
         image,
         matrix,
-        offset=None,
+        offset=0.0,
         output_shape=None,
         order=1,
         output_chunks=None,
@@ -54,16 +54,18 @@ def affine_transform(
         The image array.
     matrix : array (ndim,), (ndim, ndim), (ndim, ndim+1) or (ndim+1, ndim+1)
         Transformation matrix.
-    offset : array (ndim,)
-        Transformation offset.
-    output_shape : array (ndim,), optional
-        The size of the array to be returned.
+    offset : float or sequence, optional
+        The offset into the array where the transform is applied. If a float,
+        `offset` is the same for each axis. If a sequence, `offset` should
+        contain one value for each axis.
+    output_shape : tuple of ints, optional
+        The shape of the array to be returned.
     order : int, optional
         The order of the spline interpolation. Note that for order>1
         scipy's affine_transform applies prefiltering, which is not
         yet supported and skipped in this implementation.
-    output_chunks : array (ndim,), optional
-        The chunks of the output Dask Array.
+    output_chunks : tuple of ints, optional
+        The shape of the chunks of the output Dask Array.
 
     Returns
     -------
@@ -123,7 +125,8 @@ def affine_transform(
     image_shape = image.shape
 
     # calculate output array properties
-    normalized_chunks = da.core.normalize_chunks(output_chunks, output_shape)
+    normalized_chunks = da.core.normalize_chunks(output_chunks,
+                                                 tuple(output_shape))
     block_indices = product(*(range(len(bds)) for bds in normalized_chunks))
     block_offsets = [np.cumsum((0,) + bds[:-1]) for bds in normalized_chunks]
 
@@ -223,7 +226,7 @@ def affine_transform(
 
     transformed = da.Array(graph,
                            output_name,
-                           shape=output_shape,
+                           shape=tuple(output_shape),
                            # chunks=output_chunks,
                            chunks=normalized_chunks,
                            meta=meta)

tests/test_dask_image/test_ndinterp/test_affine_transformation.py

@@ -115,15 +115,7 @@ def test_affine_transform_cupy(n,
                                interp_order,
                                input_output_chunksize_per_dim,
                                random_seed):
-
-    pytest.importorskip("cupy", minversion="6.0.0")
-
-    # somehow, these lines are required for the first parametrized
-    # test to succeed
-    import cupy as cp
-    from dask_image.dispatch._dispatch_ndinterp import (
-        dispatch_affine_transform)
-    dispatch_affine_transform(cp.asarray([]))
+    cupy = pytest.importorskip("cupy", minversion="6.0.0")
 
     kwargs = dict()
     kwargs['n'] = n
@@ -180,6 +172,14 @@ def test_affine_transform_numpy_input():
     assert (image == image_t).min()
 
 
+def test_affine_transform_minimal_input():
+
+    image = np.ones((3, 3))
+    image_t = da_ndinterp.affine_transform(np.ones((3, 3)), np.eye(2))
+
+    assert image_t.shape == image.shape
+
+
 def test_affine_transform_type_consistency():
 
     image = da.ones((3, 3))
@@ -189,15 +189,15 @@ def test_affine_transform_type_consistency():
     assert isinstance(image[0, 0].compute(), type(image_t[0, 0].compute()))
 
 
+@pytest.mark.cupy
 def test_affine_transform_type_consistency_gpu():
 
-    pytest.importorskip("cupy", minversion="6.0.0")
+    cupy = pytest.importorskip("cupy", minversion="6.0.0")
 
     image = da.ones((3, 3))
     image_t = da_ndinterp.affine_transform(image, np.eye(2), [0, 0])
 
-    import cupy as cp
-    image.map_blocks(cp.asarray)
+    image.map_blocks(cupy.asarray)
 
     assert isinstance(image, type(image_t))
     assert isinstance(image[0, 0].compute(), type(image_t[0, 0].compute()))
@@ -236,17 +236,17 @@ def test_affine_transform_large_input_small_output_cpu():
     image_t[0, 0, 0].compute()
 
 
+@pytest.mark.cupy
 @pytest.mark.timeout(15)
 def test_affine_transform_large_input_small_output_gpu():
     """
     Make sure input array does not need to be computed entirely
     """
-    pytest.importorskip("cupy", minversion="6.0.0")
+    cupy = pytest.importorskip("cupy", minversion="6.0.0")
 
     # this array would occupy more than 24GB on a GPU
     image = da.random.random([2000] * 3, chunks=(50, 50, 50))
-    import cupy as cp
-    image.map_blocks(cp.asarray)
+    image.map_blocks(cupy.asarray)
 
     image_t = da_ndinterp.affine_transform(image, np.eye(3), [0, 0, 0],
                                            output_chunks=[1, 1, 1],