Repository: pangeo-data/xESMF

Issue:
Very large Dask graph size with high resolution grids
Regridding from 0.1 to 0.25 degree global grids is creating a very large task graph with xesmf 0.6.1, causing problems when evaluating the result.

Sample code:
```python
import xarray
import xesmf
import dask.array

a = xesmf.util.grid_global(0.1,0.1)
b = xesmf.util.grid_global(0.25,0.25)
r = xesmf.Regridder(a,b,method='bilinear')

data = dask.array.random.random((720,a.dims['y'], a.dims['x']), chunks=(12,-1,-1))
da = xarray.DataArray(data, dims=['time','y','x'], coords={'y': a.y, 'x': a.x})

r(da).data
```

Evaluated versions at https://nbviewer.org/gist/ScottWales/4fe0e9a5725b5a3bf07ab94e80778846

With 0.6.0, the regridded data has 240 Dask tasks. With 0.6.1 the number has exploded to 409,841 tasks, and appears to scale with the grid resolution - low resolution regrid task graphs have a reasonable size.

Expected behaviour is that the Dask task graph created by xesmf is a reasonably small multiple of the number of chunks in the input data, and that it is independent of grid resolution. Having a too large task graph slows down Dask as it prepares the tasks to run and is causing excess memory use when the result is evaluated.


Discussion:
huard: @ScottWales Could you check this out and confirm that it works for you ?
dcherian: >  Sadly, that didn't go as well as planned : the performance was not better, sometimes it was equal, sometimes worst.

Well the other way of looking at it is that you've enabled a major new feature (chunked spatial dimensions) without much of a performance regression 😉 

Maybe it's possible to improve tensordot for the case where there is only one chunk along spatial dimensions?
aulemahal: AFAIK, the performance issues where not from `tensordot` itself. In my first test, I was doing everything a "duck type" manner, i.e. : all operations on `indata` in `_regrid_array` and `apply_weights` were chosen to work with numpy and dask arrays alike. This way, dask's tensordot was used, but it added many `reshape`, `transpose` and a `astype` task. I think it is this high number of very small tasks that cancelled the potential performance gain.
raphaeldussin: great job @aulemahal 

Patch:
diff --git a/CHANGES.rst b/CHANGES.rst
index b2989ee..2c74834 100644
--- a/CHANGES.rst
+++ b/CHANGES.rst
@@ -7,8 +7,12 @@ What's new
 Bug fixes
 ~~~~~~~~~
 - The introduction of `sparse`, with `numba` under the hood, restricted input data to little-endian dtypes. In those cases, xESMF switches back to using scipy (:pull:`125`). By `Pascal Bourgault <https://github.com/aulemahal>`_
-- Regridding datasets with dask-backed variables is fixed. Dtype of the outputs is changed for this specific case. By `Pascal Bourgault <https://github.com/aulemahal>`_
 - SpatialAverager did not compute the same weights as Regridder when source cell areas were not uniform (:pull:`128`). By `David Huard <https://github.com/huard>`_
+- Refactor of how the regridding is called internally, to fix a bug with dask and sparse (:pull:`135`). By `Pascal Bourgault <https://github.com/aulemahal>`_
+
+Internal changes
+~~~~~~~~~~~~~~~~
+- Deprecation of ``regrid_numpy`` and ``regrid_dask`` is scheduled for 0.7.0. All checks on shape, array layout and numba support are now done at call time, rather then at computation time (:pull:`135`).
 
 0.6.1 (23-09-2021)
 ------------------
diff --git a/xesmf/frontend.py b/xesmf/frontend.py
index 0414172..d485fbc 100644
--- a/xesmf/frontend.py
+++ b/xesmf/frontend.py
@@ -15,6 +15,7 @@ from .smm import (
     _parse_coords_and_values,
     add_nans_to_weights,
     apply_weights,
+    check_shapes,
     read_weights,
 )
 from .util import split_polygons_and_holes
@@ -448,10 +449,8 @@ class BaseRegridder(object):
             for any other case, all outputs will be 'float64'.
 
         """
-        if isinstance(indata, np.ndarray):
-            return self.regrid_numpy(indata, skipna=skipna, na_thres=na_thres)
-        elif isinstance(indata, dask_array_type):
-            return self.regrid_dask(indata, skipna=skipna, na_thres=na_thres)
+        if isinstance(indata, dask_array_type + (np.ndarray,)):
+            return self.regrid_array(indata, self.weights.data, skipna=skipna, na_thres=na_thres)
         elif isinstance(indata, xr.DataArray):
             return self.regrid_dataarray(
                 indata, keep_attrs=keep_attrs, skipna=skipna, na_thres=na_thres
@@ -464,11 +463,7 @@ class BaseRegridder(object):
             raise TypeError('input must be numpy array, dask array, xarray DataArray or Dataset!')
 
     @staticmethod
-    def _regrid_array(indata, weights, *, shape_in, shape_out, sequence_in, skipna, na_thres):
-
-        if sequence_in:
-            indata = np.expand_dims(indata, axis=-2)
-
+    def _regrid(indata, weights, *, shape_in, shape_out, skipna, na_thres):
         # skipna: set missing values to zero
         if skipna:
             missing = np.isnan(indata)
@@ -487,62 +482,64 @@ class BaseRegridder(object):
 
         return outdata
 
-    @property
-    def _regrid_kwargs(self):
-        return {
-            'sequence_in': self.sequence_in,
+    def regrid_array(self, indata, weights, skipna=False, na_thres=1.0):
+        """See __call__()."""
+        if self.sequence_in:
+            indata = np.reshape(indata, (*indata.shape[:-1], 1, indata.shape[-1]))
+
+        kwargs = {
             'shape_in': self.shape_in,
             'shape_out': self.shape_out,
         }
 
-    def regrid_numpy(self, indata, skipna=False, na_thres=1.0):
-        """See __call__()."""
-        outdata = self._regrid_array(
-            indata, self.weights.data, skipna=skipna, na_thres=na_thres, **self._regrid_kwargs
-        )
-        return outdata
+        check_shapes(indata, weights, **kwargs)
 
-    def regrid_dask(self, indata, skipna=False, na_thres=1.0):
-        """See __call__()."""
-
-        extra_chunk_shape = indata.chunksize[0:-2]
+        kwargs.update(skipna=skipna, na_thres=na_thres)
 
-        output_chunk_shape = extra_chunk_shape + self.shape_out
+        if isinstance(indata, dask_array_type):  # dask
+            weights = da.from_array(weights, chunks=weights.shape)
+            output_chunks = indata.chunks[:-2] + ((self.shape_out[0],), (self.shape_out[1],))
+
+            outdata = da.map_blocks(
+                self._regrid,
+                indata,
+                weights,
+                dtype=indata.dtype,
+                chunks=output_chunks,
+                meta=np.array((), dtype=indata.dtype),
+                **kwargs,
+            )
+        else:  # numpy
+            outdata = self._regrid(indata, weights, **kwargs)
+        return outdata
 
-        outdata = da.map_blocks(
-            self._regrid_array,
-            indata,
-            self.weights.data,
-            dtype=indata.dtype,
-            chunks=output_chunk_shape,
-            skipna=skipna,
-            na_thres=na_thres,
-            **self._regrid_kwargs,
+    def regrid_numpy(self, indata, **kwargs):
+        warnings.warn(
+            '`regrid_numpy()` will be removed in xESMF 0.7, please use `regrid_array` instead.',
+            category=FutureWarning,
         )
+        return self.regrid_array(indata, self.weights.data, **kwargs)
 
-        return outdata
+    def regrid_dask(self, indata, **kwargs):
+        warnings.warn(
+            '`regrid_dask()` will be removed in xESMF 0.7, please use `regrid_array` instead.',
+            category=FutureWarning,
+        )
+        return self.regrid_array(indata, self.weights.data, **kwargs)
 
     def regrid_dataarray(self, dr_in, keep_attrs=False, skipna=False, na_thres=1.0):
         """See __call__()."""
 
         input_horiz_dims, temp_horiz_dims = self._parse_xrinput(dr_in)
-        kwargs = self._regrid_kwargs.copy()
-        kwargs.update(skipna=skipna, na_thres=na_thres)
+        kwargs = dict(skipna=skipna, na_thres=na_thres)
         dr_out = xr.apply_ufunc(
-            self._regrid_array,
+            self.regrid_array,
             dr_in,
             self.weights,
             kwargs=kwargs,
             input_core_dims=[input_horiz_dims, ('out_dim', 'in_dim')],
             output_core_dims=[temp_horiz_dims],
-            dask='parallelized',
-            output_dtypes=[dr_in.dtype],
-            dask_gufunc_kwargs={
-                'output_sizes': {
-                    temp_horiz_dims[0]: self.shape_out[0],
-                    temp_horiz_dims[1]: self.shape_out[1],
-                },
-            },
+            dask='allowed',
             keep_attrs=keep_attrs,
         )
 
@@ -554,8 +551,7 @@ class BaseRegridder(object):
         # get the first data variable to infer input_core_dims
         input_horiz_dims, temp_horiz_dims = self._parse_xrinput(ds_in)
 
-        kwargs = self._regrid_kwargs.copy()
-        kwargs.update(skipna=skipna, na_thres=na_thres)
+        kwargs = dict(skipna=skipna, na_thres=na_thres)
 
         non_regriddable = [
             name
@@ -564,41 +560,17 @@ class BaseRegridder(object):
         ]
         ds_in = ds_in.drop_vars(non_regriddable)
 
-        # Select dtype (only relevant for dask-backed arrays, this has no effect otherwise)
-        dtypes = set(da.dtype for da in ds_in.data_vars.values())
-        if len(dtypes) == 1:
-            out_dtype = dtypes.pop()
-        else:
-            size = 1  # init to 1 byte
-            for d in dtypes:
-                size = max(size, d.itemsize)  # use max number of bytes found in dtypes
-                out_dtype = np.dtype(f'f{size}')
-
         ds_out = xr.apply_ufunc(
-            self._regrid_array,
+            self.regrid_array,
             ds_in,
             self.weights,
             kwargs=kwargs,
             input_core_dims=[input_horiz_dims, ('out_dim', 'in_dim')],
             output_core_dims=[temp_horiz_dims],
-            dask='parallelized',
-            output_dtypes=[out_dtype],
-            dask_gufunc_kwargs={
-                'output_sizes': {
-                    temp_horiz_dims[0]: self.shape_out[0],
-                    temp_horiz_dims[1]: self.shape_out[1],
-                },
-            },
+            dask='allowed',
             keep_attrs=keep_attrs,
         )
 
-        # For dask-backed, we force the dtype to fit the input's
-        # (numpy-backed variables already have been handled in `apply_weights`)
-        for name, data in ds_out.data_vars.items():
-            indtype = ds_in[name].dtype
-            if xr.core.common.is_duck_dask_array(data.data) and data.dtype != indtype:
-                ds_out[name] = data.astype(indtype, keep_attrs=True)
-
         return self._format_xroutput(ds_out, temp_horiz_dims)
 
     def _parse_xrinput(self, dr_in):
diff --git a/xesmf/smm.py b/xesmf/smm.py
index 28a2b57..2586ae6 100644
--- a/xesmf/smm.py
+++ b/xesmf/smm.py
@@ -1,7 +1,6 @@
 """
 Sparse matrix multiplication (SMM) using scipy.sparse library.
 """
-
 import warnings
 from pathlib import Path
 
@@ -101,6 +100,61 @@ def _parse_coords_and_values(indata, n_in, n_out):
     return xr.DataArray(sps.COO(crds, s, (n_out, n_in)), dims=('out_dim', 'in_dim'), name='weights')
 
 
+def check_shapes(indata, weights, shape_in, shape_out):
+    """Compare the shapes of the input array, the weights and the regridder and raises
+    potential errors.
+
+    Parameters
+    ----------
+    indata : array
+      Input array with the two spatial dimensions at the end,
+      which should fit shape_in.
+    weights : array
+      Weights 2D array of shape (out_dim, in_dim).
+      First element should be the product of shape_out.
+      Second element should be the product of shape_in.
+    shape_in : 2-tuple of int
+      Shape of the input of the Regridder.
+    shape_out : 2-tuple of int
+      Shape of the output of the Regridder.
+
+    Raises
+    ------
+    ValueError
+      If any of the conditions is not respected.
+    """
+    # COO matrix is fast with F-ordered array but slow with C-array, so we
+    # take in a C-ordered and then transpose)
+    # (CSR or CRS matrix is fast with C-ordered array but slow with F-array)
+    if hasattr(indata, 'flags') and not indata.flags['C_CONTIGUOUS']:
+        warnings.warn('Input array is not C_CONTIGUOUS. ' 'Will affect performance.')
+
+    # Limitation from numba : some big-endian dtypes are not supported.
+    try:
+        nb.from_dtype(indata.dtype)
+        nb.from_dtype(weights.dtype)
+    except NotImplementedError:
+        warnings.warn(
+            'Input array has a dtype not supported by sparse and numba.'
+            'Computation will fall back to scipy.'
+        )
+
+    # get input shape information
+    shape_horiz = indata.shape[-2:]
+
+    if shape_horiz != shape_in:
+        raise ValueError(
+            f'The horizontal shape of input data is {shape_horiz}, different from that '
+            f'of the regridder {shape_in}!'
+        )
+
+    if shape_in[0] * shape_in[1] != weights.shape[1]:
+        raise ValueError('ny_in * nx_in should equal to weights.shape[1]')
+
+    if shape_out[0] * shape_out[1] != weights.shape[0]:
+        raise ValueError('ny_out * nx_out should equal to weights.shape[0]')
+
+
 def apply_weights(weights, indata, shape_in, shape_out):
     """
     Apply regridding weights to data.
@@ -121,47 +175,26 @@ def apply_weights(weights, indata, shape_in, shape_out):
         Extra dimensions are the same as `indata`.
         If input data is C-ordered, output will also be C-ordered.
     """
+    extra_shape = indata.shape[0:-2]
+
+    # use flattened array for dot operation
+    indata_flat = indata.reshape(-1, shape_in[0] * shape_in[1])
+
     # Limitation from numba : some big-endian dtypes are not supported.
     try:
         nb.from_dtype(indata.dtype)
         nb.from_dtype(weights.dtype)
     except NotImplementedError:
-        warnings.warn(
-            'Input array has a dtype not supported by sparse and numba. Falling back to scipy.'
-        )
         weights = weights.to_scipy_sparse()
 
-    # COO matrix is fast with F-ordered array but slow with C-array, so we
-    # take in a C-ordered and then transpose)
-    # (CSR or CRS matrix is fast with C-ordered array but slow with F-array)
-    if not indata.flags['C_CONTIGUOUS']:
-        warnings.warn('Input array is not C_CONTIGUOUS. ' 'Will affect performance.')
-
-    # get input shape information
-    shape_horiz = indata.shape[-2:]
-    extra_shape = indata.shape[0:-2]
-
-    assert shape_horiz == shape_in, (
-        'The horizontal shape of input data is {}, different from that of'
-        'the regridder {}!'.format(shape_horiz, shape_in)
-    )
-
-    assert (
-        shape_in[0] * shape_in[1] == weights.shape[1]
-    ), 'ny_in * nx_in should equal to weights.shape[1]'
-
-    assert (
-        shape_out[0] * shape_out[1] == weights.shape[0]
-    ), 'ny_out * nx_out should equal to weights.shape[0]'
-
-    # use flattened array for dot operation
-    indata_flat = indata.reshape(-1, shape_in[0] * shape_in[1])
+    # Dot product
     outdata_flat = weights.dot(indata_flat.T).T
 
-    outdata_flat = outdata_flat.astype(indata_flat.dtype)
+    # Ensure same dtype as the input.
+    outdata_flat = outdata_flat.astype(indata.dtype)
 
     # unflattened output array
-    outdata = outdata_flat.reshape([*extra_shape, shape_out[0], shape_out[1]])
+    outdata = outdata_flat.reshape(*extra_shape, shape_out[0], shape_out[1])
     return outdata
 
 


Test patch:
diff --git a/xesmf/tests/test_frontend.py b/xesmf/tests/test_frontend.py
index bfcbaaf..ed668cc 100644
--- a/xesmf/tests/test_frontend.py
+++ b/xesmf/tests/test_frontend.py
@@ -411,13 +411,20 @@ def test_regrid_dataarray(use_cfxr):
         xr.testing.assert_identical(dr_out, dr_out_rn)
 
 
-def test_regrid_dataarray_endianess():
+@pytest.mark.parametrize('use_dask', [True, False])
+def test_regrid_dataarray_endianess(use_dask):
     # xarray.DataArray containing in-memory numpy array
     regridder = xe.Regridder(ds_in, ds_out, 'conservative')
 
     exp = regridder(ds_in['data'])  # Normal (little-endian)
-    with pytest.warns(UserWarning, match='Input array has a dtype not supported'):
-        out = regridder(ds_in['data'].astype('>f8'))  # big endian
+    # with pytest.warns(UserWarning, match='Input array has a dtype not supported'):
+
+    if use_dask:
+        indata = ds_in.data.astype('>f8').chunk()
+    else:
+        indata = ds_in.data.astype('>f8')
+
+    out = regridder(indata)  # big endian
 
     # Results should be the same
     assert_equal(exp.values, out.values)
@@ -470,14 +477,21 @@ def test_regrid_dask(request, scheduler):
     regridder = xe.Regridder(ds_in, ds_out, 'conservative')
 
     indata = ds_in_chunked['data4D'].data
-    outdata = regridder(indata)
+    # Use ridiculous small chunk size value to be sure it _isn't_ impacting computation.
+    with dask.config.set({'array.chunk-size': '1MiB'}):
+        outdata = regridder(indata)
 
     assert dask.is_dask_collection(outdata)
 
     # lazy dask arrays have incorrect shape attribute due to last chunk
-    assert outdata.compute().shape == indata.shape[:-2] + horiz_shape_out
+    assert outdata.shape == indata.shape[:-2] + horiz_shape_out
     assert outdata.chunksize == indata.chunksize[:-2] + horiz_shape_out
 
+    # Check that the number of tasks hasn't exploded.
+    n_task_in = len(indata.__dask_graph__().keys())
+    n_task_out = len(outdata.__dask_graph__().keys())
+    assert (n_task_out / n_task_in) < 3
+
     outdata_ref = ds_out['data4D_ref'].values
     rel_err = (outdata.compute() - outdata_ref) / outdata_ref
     assert np.max(np.abs(rel_err)) < 0.05