Correctly guess output chunks for SpatialAverager (#308)

* correctly guess output chunks for Savg

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: aulemahal

CHANGES.rst

@@ -1,14 +1,20 @@
 What's new
 ==========
 
+0.8.2 (unreleased)
+------------------
+
+Bug fixes
+~~~~~~~~~
+* Correct guess of output chunks for the :``SpatialAverager``.
+
 0.8.1 (2023-09-05)
 ------------------
 
 Bug fixes
 ~~~~~~~~~
 * Change import to support shapely 1 and 2.
 
-
 0.8.0 (2023-09-01)
 ------------------
 

xesmf/frontend.py

@@ -590,21 +590,24 @@ def regrid_array(self, indata, weights, skipna=False, na_thres=1.0, output_chunk
 
         kwargs.update(skipna=skipna, na_thres=na_thres)
 
+        weights = self.weights.data.reshape(self.shape_out + self.shape_in)
         if isinstance(indata, dask_array_type):  # dask
             if output_chunks is None:
-                output_chunks = indata.chunksize[-2:]
-            elif output_chunks is not None:
-                if len(output_chunks) != len(self.shape_out):
+                output_chunks = tuple(
+                    [min(shp, inchnk) for shp, inchnk in zip(self.shape_out, indata.chunksize[-2:])]
+                )
+            if len(output_chunks) != len(self.shape_out):
+                if len(output_chunks) == 1 and self.sequence_out:
+                    output_chunks = (1, output_chunks[0])
+                else:
                     raise ValueError(
                         f'output_chunks must have same dimension as ds_out,'
                         f' output_chunks dimension ({len(output_chunks)}) does not '
                         f'match ds_out dimension ({len(self.shape_out)})'
                     )
-            weights = da.from_array(self.w.data, chunks=(output_chunks + indata.chunksize[-2:]))
-
+            weights = da.from_array(weights, chunks=(output_chunks + indata.chunksize[-2:]))
             outdata = self._regrid(indata, weights, **kwargs)
         else:  # numpy
-            weights = self.w.data  # 4D weights
             outdata = self._regrid(indata, weights, **kwargs)
         return outdata
 

xesmf/tests/test_frontend.py

@@ -7,6 +7,7 @@
 import pytest
 import xarray as xr
 from numpy.testing import assert_allclose, assert_almost_equal, assert_equal
+from shapely import segmentize
 from shapely.geometry import MultiPolygon, Polygon
 
 import xesmf as xe
@@ -780,18 +781,22 @@ def test_ds_to_ESMFlocstream():
         locstream, shape, names = ds_to_ESMFlocstream(ds_bogus)
 
 
+@pytest.mark.parametrize('use_dask', [True, False])
 @pytest.mark.parametrize('poly,exp', list(zip(polys, exps_polys)))
-def test_spatial_averager(poly, exp):
+def test_spatial_averager(poly, exp, use_dask):
     if isinstance(poly, (Polygon, MultiPolygon)):
         poly = [poly]
-    savg = xe.SpatialAverager(ds_savg, poly, geom_dim_name='my_geom')
-    out = savg(ds_savg.abc)
+    if use_dask:
+        ds_in = ds_savg.chunk(lat=10)
+    else:
+        ds_in = ds_savg
+    savg = xe.SpatialAverager(ds_in, poly, geom_dim_name='my_geom')
+    out = savg(ds_in.abc)
     assert_allclose(out, exp, rtol=1e-3)
 
     assert 'my_geom' in out.dims
 
 
-@pytest.mark.xfail
 def test_spatial_averager_with_zonal_region():
     # We expect the spatial average for all regions to be one
     zonal_south = Polygon([(0, -90), (10, 0), (0, 0)])
@@ -800,6 +805,7 @@ def test_spatial_averager_with_zonal_region():
     zonal_full = Polygon([(0, -90), (10, 0), (0, 90), (0, 0)])  # This yields 0... why?
 
     polys = [zonal_south, zonal_north, zonal_short, zonal_full]
+    polys = segmentize(polys, 1)
 
     # Create field of ones on a global grid
     ds = xe.util.grid_global(20, 12, cf=True)