Repository: pangeo-data/xESMF

Issue:
Expose weights as a "normal" DataArray
The current weight matrix is flattened, with dimensions `(out_dim, in_dim)`. I think it would be useful, including for debugging, to reshape the weights so they're exposed as `lon_out, lat_out, lon_in, lat_in`. 

We could alternatively use the same dimension names as the input and output grids, but if they're both called lat and lon, this won't work. 

Discussion:
huard: @charlesgauthier-udm 

Useful ?
charlesgauthier-udm: Could be useful yes. However, with more testing I find that the performance related to using 4D weights seems to be highly dependent on the input data shape and chunk sizes (with Dask) with the 2D weights sometimes outperforming the 4D weights.
huard: Let's not merge this then, and wait for Charles' results. 
charlesgauthier-udm: I think we can merge it and I'll use the 4D w property to make another PR with the dask compatibility.
huard: Should I remove the non-related test ?
aulemahal: I think it's ok to keep it.

Patch:
diff --git a/CHANGES.rst b/CHANGES.rst
index 2f832ff..3c532af 100644
--- a/CHANGES.rst
+++ b/CHANGES.rst
@@ -4,6 +4,11 @@ What's new
 0.7.2 (unreleased)
 ------------------
 
+New features
+~~~~~~~~~~~~
+* Added a `w` property to the `Regridder` and `SpatialAverager` classes, returning the weights reshaped according to
+  the input and output grid dimensions. This is mostly intended for debugging and visualisation purposes. By `David Huard <https://github.com/huard>`_
+
 Internal changes
 ~~~~~~~~~~~~~~~~
 * Pin sphinx < 7.0 since it's not supported by `sphinx_rtd_theme`
diff --git a/xesmf/frontend.py b/xesmf/frontend.py
index 96e8145..e561bfd 100644
--- a/xesmf/frontend.py
+++ b/xesmf/frontend.py
@@ -106,7 +106,12 @@ def ds_to_ESMFgrid(ds, need_bounds=False, periodic=None, append=None):
 
     Returns
     -------
-    grid : ESMF.Grid object
+    grid
+        ESMF.Grid object
+    shape
+        Shape of the grid
+    dim_names
+        Dimension names of the grid
 
     """
     # use np.asarray(dr) instead of dr.values, so it also works for dictionary
@@ -354,6 +359,20 @@ class BaseRegridder(object):
         print(message)
         return self.weights
 
+    @property
+    def w(self) -> xr.DataArray:
+        """Return weights as a 4D DataArray with dimensions (y_out, x_out, y_in, x_in).
+
+        ESMF stores the weights in a 2D array with dimensions (out_dim, in_dim), the size of the output and input
+        grids respectively (ny x nx). This property returns the weights reshaped as a 4D array to simplify
+        comparisons with the original grids.
+        """
+        # TODO: Add coords ?
+        s = self.shape_out + self.shape_in
+        data = self.weights.data.reshape(s)
+        dims = 'y_out', 'x_out', 'y_in', 'x_in'
+        return xr.DataArray(data, dims=dims)
+
     def _get_default_filename(self):
         # e.g. bilinear_400x600_300x400.nc
         filename = '{0}_{1}x{2}_{3}x{4}'.format(
@@ -877,7 +896,7 @@ class SpatialAverager(BaseRegridder):
             Contain input and output grid coordinates. Look for variables
             ``lon``, ``lat``, ``lon_b`` and ``lat_b``.
 
-            Optionaly looks for ``mask``, in which case  the ESMF convention is used,
+            Optionally looks for ``mask``, in which case  the ESMF convention is used,
             where masked values are identified by 0, and non-masked values by 1.
 
             Shape can be 1D (n_lon,) and (n_lat,) for rectilinear grids,
@@ -1038,6 +1057,19 @@ class SpatialAverager(BaseRegridder):
         # Combine weights for all the subgeometries belonging to the same geometry
         return _combine_weight_multipoly(w, area, geom_indices).T
 
+    @property
+    def w(self) -> xr.DataArray:
+        """Return weights as a 3D DataArray with dimensions (geom, y_in, x_in).
+
+        ESMF stores the weights in a 2D array with dimensions (out_dim, in_dim), the size of the output and input
+        grids respectively (ny x nx). This property returns the weights reshaped as a 3D array to simplify
+        comparisons with the original grids.
+        """
+        s = self.shape_out[1:2] + self.shape_in
+        data = self.weights.data.reshape(s)
+        dims = self.geom_dim_name, 'y_in', 'x_in'
+        return xr.DataArray(data, dims=dims)
+
     def _get_default_filename(self):
         # e.g. bilinear_400x600_300x400.nc
         filename = 'spatialavg_{0}x{1}_{2}.nc'.format(


Test patch:
diff --git a/xesmf/tests/test_frontend.py b/xesmf/tests/test_frontend.py
index 4c23671..b5aa183 100644
--- a/xesmf/tests/test_frontend.py
+++ b/xesmf/tests/test_frontend.py
@@ -179,13 +179,35 @@ def test_existing_weights():
 
     # check fails on non-existent file
     with pytest.raises(OSError):
-        regridder_reuse = xe.Regridder(
-            ds_in, ds_out, method, reuse_weights=True, filename='fakewgts.nc'
-        )
+        xe.Regridder(ds_in, ds_out, method, reuse_weights=True, filename='fakewgts.nc')
 
     # check fails if no weights are provided
     with pytest.raises(ValueError):
-        regridder_reuse = xe.Regridder(ds_in, ds_out, method, reuse_weights=True)
+        xe.Regridder(ds_in, ds_out, method, reuse_weights=True)
+
+
+def test_regridder_w():
+    """Check that the `w` property for dimensioned weights works."""
+    regridder = xe.Regridder(ds_in, ds_out, method='bilinear')
+    w = regridder.w
+    assert w.shape == ds_out.lon.shape + ds_in.lon.shape
+
+    p = Polygon([(-10, -10), (10, -10), (10, 10), (-10, 10)])
+
+    averager = xe.SpatialAverager(ds_in, [p])
+    assert averager.w.shape == (1,) + ds_in.lon.shape
+
+    ds_in_cf = xe.util.grid_global(15, 15, cf=True)
+    ds_out_cf = xe.util.grid_global(30, 30, cf=True)
+
+    regridder_cf = xe.Regridder(ds_in_cf, ds_out_cf, method='bilinear')
+    w_cf = regridder_cf.w
+    assert w_cf.shape == (
+        ds_out_cf.lat.shape + ds_out_cf.lon.shape + ds_in_cf.lat.shape + ds_in_cf.lon.shape
+    )
+
+    averager = xe.SpatialAverager(ds_in_cf, [p, p])
+    assert averager.w.shape == (2,) + ds_in_cf.lat.shape + ds_in_cf.lon.shape
 
 
 def test_to_netcdf(tmp_path):
@@ -713,6 +735,27 @@ def test_spatial_averager(poly, exp):
     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)])
+    zonal_north = Polygon([(0, 90), (10, 0), (0, 0)])
+    zonal_short = Polygon([(0, -10), (10, -10), (10, 10), (0, 10)])
+    zonal_full = Polygon([(0, -90), (10, 0), (0, 90), (0, 0)])  # This yields 0... why?
+
+    polys = [zonal_south, zonal_north, zonal_short, zonal_full]
+
+    # Create field of ones on a global grid
+    ds = xe.util.grid_global(20, 12, cf=True)
+    ds['a'] = xr.DataArray(
+        np.ones((ds.lon.size, ds.lat.size)),
+        coords={'lat': ds.lat, 'lon': ds.lon},
+        dims=('lon', 'lat'),
+    )
+    out = xe.SpatialAverager(ds, polys)(ds.a)
+    assert_allclose(out, 1, rtol=1e-3)
+
+
 def test_compare_weights_from_poly_and_grid():
     """Confirm that the weights are identical when they are computed from a grid->grid and grid->poly."""