ENH: implement grid_mode paramter for zoom

cupyx/scipy/ndimage/_interp_kernels.py

@@ -34,15 +34,16 @@ def _get_coord_map(ndim):
     ops = []
     ops.append('ptrdiff_t ncoords = _ind.size();')
     for j in range(ndim):
-        ops.append(
-            '''
-    W c_{j} = coords[i + {j} * ncoords];'''.format(j=j))
+        ops.append(f'''
+    W c_{j} = coords[i + {j} * ncoords];''')
     return ops
 
 
 def _get_coord_zoom_and_shift(ndim):
     """Compute target coordinate based on a shift followed by a zoom.
 
+    This version zooms from the center of the edge pixels.
+
     Notes
     -----
     Assumes the following variables have been initialized on the device::
@@ -58,15 +59,41 @@ def _get_coord_zoom_and_shift(ndim):
     """
     ops = []
     for j in range(ndim):
-        ops.append(
-            '''
-    W c_{j} = zoom[{j}] * ((W)in_coord[{j}] - shift[{j}]);'''.format(j=j))
+        ops.append(f'''
+    W c_{j} = zoom[{j}] * ((W)in_coord[{j}] - shift[{j}]);''')
     return ops
 
 
-def _get_coord_zoom(ndim):
+def _get_coord_zoom_and_shift_grid(ndim):
+    """Compute target coordinate based on a shift followed by a zoom.
+
+    This version zooms from the outer edges of the grid pixels.
+
+    Notes
+    -----
+    Assumes the following variables have been initialized on the device::
+
+        in_coord[ndim]: array containing the source coordinate
+        zoom[ndim]: array containing the zoom for each axis
+        shift[ndim]: array containing the zoom for each axis
+
+    computes::
+
+        c_j = zoom[j] * (in_coord[j] - shift[j] + 0.5) - 0.5
+
+    """
+    ops = []
+    for j in range(ndim):
+        ops.append(f'''
+    W c_{j} = zoom[{j}] * ((W)in_coord[{j}] - shift[j] + 0.5) - 0.5;''')
+    return ops
+
+
+def _get_coord_zoom(ndim, nprepad=0):
     """Compute target coordinate based on a zoom.
 
+    This version zooms from the center of the edge pixels.
+
     Notes
     -----
     Assumes the following variables have been initialized on the device::
@@ -81,13 +108,36 @@ def _get_coord_zoom(ndim):
     """
     ops = []
     for j in range(ndim):
-        ops.append(
-            '''
-    W c_{j} = zoom[{j}] * (W)in_coord[{j}];'''.format(j=j))
+        ops.append(f'''
+    W c_{j} = zoom[{j}] * (W)in_coord[{j}];''')
     return ops
 
 
-def _get_coord_shift(ndim):
+def _get_coord_zoom_grid(ndim, nprepad=0):
+    """Compute target coordinate based on a zoom (grid_mode=True version).
+
+    This version zooms from the outer edges of the grid pixels.
+
+    Notes
+    -----
+    Assumes the following variables have been initialized on the device::
+
+        in_coord[ndim]: array containing the source coordinate
+        zoom[ndim]: array containing the zoom for each axis
+
+    computes::
+
+        c_j = zoom[j] * (in_coord[j] + 0.5) - 0.5
+
+    """
+    ops = []
+    for j in range(ndim):
+        ops.append(f'''
+    W c_{j} = zoom[{j}] * ((W)in_coord[{j}] + 0.5) - 0.5;''')
+    return ops
+
+
+def _get_coord_shift(ndim, nprepad=0):
     """Compute target coordinate based on a shift.
 
     Notes
@@ -104,9 +154,8 @@ def _get_coord_shift(ndim):
     """
     ops = []
     for j in range(ndim):
-        ops.append(
-            '''
-    W c_{j} = (W)in_coord[{j}] - shift[{j}];'''.format(j=j))
+        ops.append(f'''
+    W c_{j} = (W)in_coord[{j}] - shift[{j}];''')
     return ops
 
 
@@ -133,19 +182,13 @@ def _get_coord_affine(ndim):
     ops = []
     ncol = ndim + 1
     for j in range(ndim):
-        ops.append('''
-            W c_{j} = (W)0.0;
-            '''.format(j=j))
+        ops.append(f'''
+            W c_{j} = (W)0.0;''')
         for k in range(ndim):
-            m_index = ncol * j + k
-            ops.append(
-                '''
-            c_{j} += mat[{m_index}] * (W)in_coord[{k}];'''.format(
-                    j=j, k=k, m_index=m_index))
-        ops.append(
-            '''
-            c_{j} += mat[{m_index}];'''.format(
-                j=j, m_index=ncol * j + ndim))
+            ops.append(f'''
+            c_{j} += mat[{ncol * j + k}] * (W)in_coord[{k}];''')
+        ops.append(f'''
+            c_{j} += mat[{ncol * j + ndim}];''')
     return ops
 
 
@@ -155,16 +198,15 @@ def _unravel_loop_index(shape, uint_t='unsigned int'):
     This code assumes that the array is a C-ordered array.
     """
     ndim = len(shape)
-    code = [
-        '''
+    code = [f'''
         {uint_t} in_coord[{ndim}];
-        {uint_t} s, t, idx = i;'''.format(uint_t=uint_t, ndim=ndim)]
+        {uint_t} s, t, idx = i;''']
     for j in range(ndim - 1, 0, -1):
-        code.append('''
-        s = {size};
+        code.append(f'''
+        s = {shape[j]};
         t = idx / s;
         in_coord[{j}] = idx - t * s;
-        idx = t;'''.format(j=j, size=shape[j]))
+        idx = t;''')
     code.append('''
         in_coord[0] = idx;''')
     return '\n'.join(code)
@@ -413,18 +455,22 @@ def _get_shift_kernel(ndim, large_int, yshape, mode, cval=0.0, order=1,
 
 @cupy._util.memoize(for_each_device=True)
 def _get_zoom_shift_kernel(ndim, large_int, yshape, mode, cval=0.0, order=1,
-                           integer_output=False):
+                           integer_output=False, grid_mode=False):
     in_params = 'raw X x, raw W shift, raw W zoom'
     out_params = 'Y y'
+    if grid_mode:
+        zoom_shift_func = _get_coord_zoom_and_shift_grid
+    else:
+        zoom_shift_func = _get_coord_zoom_and_shift
     operation, name = _generate_interp_custom(
-        coord_func=_get_coord_zoom_and_shift,
+        coord_func=zoom_shift_func,
         ndim=ndim,
         large_int=large_int,
         yshape=yshape,
         mode=mode,
         cval=cval,
         order=order,
-        name='zoom_shift',
+        name="zoom_shift_grid" if grid_mode else "zoom_shift",
         integer_output=integer_output,
     )
     return cupy.ElementwiseKernel(in_params, out_params, operation, name,
@@ -433,18 +479,18 @@ def _get_zoom_shift_kernel(ndim, large_int, yshape, mode, cval=0.0, order=1,
 
 @cupy._util.memoize(for_each_device=True)
 def _get_zoom_kernel(ndim, large_int, yshape, mode, cval=0.0, order=1,
-                     integer_output=False):
+                     integer_output=False, grid_mode=False):
     in_params = 'raw X x, raw W zoom'
     out_params = 'Y y'
     operation, name = _generate_interp_custom(
-        coord_func=_get_coord_zoom,
+        coord_func=_get_coord_zoom_grid if grid_mode else _get_coord_zoom,
         ndim=ndim,
         large_int=large_int,
         yshape=yshape,
         mode=mode,
         cval=cval,
         order=order,
-        name='zoom',
+        name="zoom_grid" if grid_mode else "zoom",
         integer_output=integer_output,
     )
     return cupy.ElementwiseKernel(in_params, out_params, operation, name,

cupyx/scipy/ndimage/interpolation.py

@@ -576,7 +576,7 @@ def shift(input, shift, output=None, order=None, mode='constant', cval=0.0,
 
 
 def zoom(input, zoom, output=None, order=None, mode='constant', cval=0.0,
-         prefilter=True):
+         prefilter=True, *, grid_mode=False):
     """Zoom an array.
 
     The array is zoomed using spline interpolation of the requested order.
@@ -600,6 +600,21 @@ def zoom(input, zoom, output=None, order=None, mode='constant', cval=0.0,
             0.0
         prefilter (bool): It is not used yet. It just exists for compatibility
             with :mod:`scipy.ndimage`.
+        grid_mode (bool, optional): If False, the distance from the pixel
+            centers is zoomed. Otherwise, the distance including the full pixel
+            extent is used. For example, a 1d signal of length 5 is considered
+            to have length 4 when ``grid_mode`` is False, but length 5 when
+            ``grid_mode`` is True. See the following visual illustration:
+
+            .. code-block:: text
+
+                    | pixel 1 | pixel 2 | pixel 3 | pixel 4 | pixel 5 |
+                         |<-------------------------------------->|
+                                            vs.
+                    |<----------------------------------------------->|
+
+            The starting point of the arrow in the diagram above corresponds to
+            coordinate location 0 in each mode.
 
     Returns:
         cupy.ndarray or None:
@@ -644,10 +659,25 @@ def zoom(input, zoom, output=None, order=None, mode='constant', cval=0.0,
         if order is None:
             order = 1
 
+        if grid_mode:
+            # warn about modes that may have surprising behavior
+            suggest_mode = None
+            if mode == 'constant':
+                suggest_mode = 'grid-constant'
+            elif mode == 'wrap':
+                suggest_mode = 'grid-wrap'
+            if suggest_mode is not None:
+                warnings.warn(
+                    f'It is recommended to use mode = {suggest_mode} instead '
+                    f'of {mode} when grid_mode is True.')
+
         zoom = []
         for in_size, out_size in zip(input.shape, output_shape):
             if out_size > 1:
-                zoom.append(float(in_size - 1) / (out_size - 1))
+                if grid_mode:
+                    zoom.append(in_size / out_size)
+                else:
+                    zoom.append((in_size - 1) / (out_size - 1))
             else:
                 zoom.append(0)
 
@@ -659,7 +689,7 @@ def zoom(input, zoom, output=None, order=None, mode='constant', cval=0.0,
         large_int = max(_prod(input.shape), _prod(output_shape)) > 1 << 31
         kern = _interp_kernels._get_zoom_kernel(
             input.ndim, large_int, output_shape, mode, order=order,
-            integer_output=integer_output)
+            integer_output=integer_output, grid_mode=grid_mode)
         zoom = cupy.asarray(zoom, dtype=cupy.float64)
         kern(input, zoom, output)
     return output

tests/cupyx_tests/scipy_tests/ndimage_tests/test_interpolation.py

@@ -579,6 +579,28 @@ def test_zoom_int(self, xp, scp, dtype):
         return out
 
 
+@testing.parameterize(*testing.product({
+    'shape': [(2, 3), (4, 4)],
+    'zoom': [(1, 1), (3, 5), (8, 2), (8, 8)],
+    'mode': ['nearest', 'reflect', 'mirror', 'grid-wrap', 'grid-constant'],
+}))
+@testing.gpu
+class TestZoomIntegerGrid():
+
+    def test_zoom_grid_by_int_order0(self):
+        # When grid_mode is True,  order 0 zoom should be the same as
+        # replication via a Kronecker product. The only exceptions to this are
+        # the non-grid modes 'constant' and 'wrap'.
+        size = numpy.prod(self.shape)
+        x = cupy.arange(size, dtype=float).reshape(self.shape)
+        testing.assert_array_almost_equal(
+            cupyx.scipy.ndimage.zoom(
+                x, self.zoom, order=0, mode=self.mode, grid_mode=True
+            ),
+            cupy.kron(x, cupy.ones(self.zoom)),
+        )
+
+
 @testing.parameterize(
     {'zoom': 3},
     {'zoom': 0.3},