Repository: pangeo-data/xESMF

Issue:
Regridder fails when ds_in doesn't have (lat,lon) as last two coordinates
Following the example on the tutorial:

```python
import xarray as xr
import xesmf as xe
import numpy as np

ds = xr.tutorial.open_dataset('air_temperature')
ds_out = xr.Dataset({'lat': (['lat'], np.arange(16, 75, 1.0)),
                     'lon': (['lon'], np.arange(200, 330, 1.5)),})

```
This works (as it should):

```python
regridder = xe.Regridder(ds, ds_out, 'bilinear')
regridder

xESMF Regridder 
Regridding algorithm:       bilinear 
Weight filename:            bilinear_25x53_59x87.nc 
Reuse pre-computed weights? False 
Input grid shape:           (25, 53) 
Output grid shape:          (59, 87) 
Output grid dimension name: ('lat', 'lon') 
Periodic in longitude?      False

dr_out = regridder(ds)
```
Let's create another dataset:
```python
ds_t = ds.transpose('lat','lon','time')
```

This doesn't work:
```python
dr_out = regridder(ds_t)
```

it gives this error

```
using dimensions ('lon', 'time') from data variable air as the horizontal dimensions for this dataset.
/srv/conda/envs/notebook/lib/python3.8/site-packages/xesmf/smm.py:70: UserWarning: Input array is not C_CONTIGUOUS. Will affect performance.
  warnings.warn("Input array is not C_CONTIGUOUS. "
---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
<ipython-input-24-d76ce0875d6d> in <module>
----> 1 dr_out = regridder(ds_t)
      2 dr_out
      3 dr_out = regridder_t(ds.transpose('lat','lon','time'))
      4 dr_out

/srv/conda/envs/notebook/lib/python3.8/site-packages/xesmf/frontend.py in __call__(self, indata, keep_attrs)
    357             return self.regrid_dataarray(indata, keep_attrs=keep_attrs)
    358         elif isinstance(indata, xr.Dataset):
--> 359             return self.regrid_dataset(indata, keep_attrs=keep_attrs)
    360         else:
    361             raise TypeError(

/srv/conda/envs/notebook/lib/python3.8/site-packages/xesmf/frontend.py in regrid_dataset(self, ds_in, keep_attrs)
    475               )
    476 
--> 477         ds_out = xr.apply_ufunc(
    478             self.regrid_numpy, ds_in,
    479             input_core_dims=[input_horiz_dims],

/srv/conda/envs/notebook/lib/python3.8/site-packages/xarray/core/computation.py in apply_ufunc(func, input_core_dims, output_core_dims, exclude_dims, vectorize, join, dataset_join, dataset_fill_value, keep_attrs, kwargs, dask, output_dtypes, output_sizes, meta, dask_gufunc_kwargs, *args)
   1114     # feed datasets apply_variable_ufunc through apply_dataset_vfunc
   1115     elif any(is_dict_like(a) for a in args):
-> 1116         return apply_dataset_vfunc(
   1117             variables_vfunc,
   1118             *args,

/srv/conda/envs/notebook/lib/python3.8/site-packages/xarray/core/computation.py in apply_dataset_vfunc(func, signature, join, dataset_join, fill_value, exclude_dims, keep_attrs, *args)
    426     args = [getattr(arg, "data_vars", arg) for arg in args]
    427 
--> 428     result_vars = apply_dict_of_variables_vfunc(
    429         func, *args, signature=signature, join=dataset_join, fill_value=fill_value
    430     )

/srv/conda/envs/notebook/lib/python3.8/site-packages/xarray/core/computation.py in apply_dict_of_variables_vfunc(func, signature, join, fill_value, *args)
    371     result_vars = {}
    372     for name, variable_args in zip(names, grouped_by_name):
--> 373         result_vars[name] = func(*variable_args)
    374 
    375     if signature.num_outputs > 1:

/srv/conda/envs/notebook/lib/python3.8/site-packages/xarray/core/computation.py in apply_variable_ufunc(func, signature, exclude_dims, dask, output_dtypes, vectorize, keep_attrs, dask_gufunc_kwargs, *args)
    722             )
    723 
--> 724     result_data = func(*input_data)
    725 
    726     if signature.num_outputs == 1:

/srv/conda/envs/notebook/lib/python3.8/site-packages/xesmf/frontend.py in regrid_numpy(self, indata)
    369             indata = np.expand_dims(indata, axis=-2)
    370 
--> 371         outdata = apply_weights(self.weights, indata,
    372                                 self.shape_in, self.shape_out)
    373         return outdata

/srv/conda/envs/notebook/lib/python3.8/site-packages/xesmf/smm.py in apply_weights(weights, indata, shape_in, shape_out)
     75     extra_shape = indata.shape[0:-2]
     76 
---> 77     assert shape_horiz == shape_in, (
     78         'The horizontal shape of input data is {}, different from that of'
     79         'the regridder {}!'.format(shape_horiz, shape_in)

AssertionError: The horizontal shape of input data is (53, 2920), different from that ofthe regridder (25, 53)!
```

Interestingly, the order of the dimensions does not matter when I create the weights. In fact this works:

```python
regridder_t = xe.Regridder(ds_t, ds_out, 'bilinear')
regridder_t

xESMF Regridder 
Regridding algorithm:       bilinear 
Weight filename:            bilinear_25x53_59x87.nc 
Reuse pre-computed weights? False 
Input grid shape:           (25, 53) 
Output grid shape:          (59, 87) 
Output grid dimension name: ('lat', 'lon') 
Periodic in longitude?      False
```

And both `regridder` and `regridder_t` work on the dataset `ds`, both do not work on `ds_t`.

 I `pip` installed the latest xarray version ('0.16.2') and I have 0.3.0 for xesmf

I don't consider this a bug, but I guess an inconsistency, since the order of the dimensions does not matter in creating the weights but it does when I apply them.

Best - and thanks so much for this tool! It is great! 
Chiara

Discussion:

Patch:
diff --git a/xesmf/frontend.py b/xesmf/frontend.py
index 957687a..0b27a98 100644
--- a/xesmf/frontend.py
+++ b/xesmf/frontend.py
@@ -8,7 +8,7 @@ import cf_xarray as cfxr
 import numpy as np
 import scipy.sparse as sps
 import xarray as xr
-from xarray import DataArray
+from xarray import DataArray, Dataset
 
 from .backend import Grid, LocStream, Mesh, add_corner, esmf_regrid_build, esmf_regrid_finalize
 from .smm import _combine_weight_multipoly, add_nans_to_weights, apply_weights, read_weights
@@ -105,6 +105,13 @@ def ds_to_ESMFgrid(ds, need_bounds=False, periodic=None, append=None):
     # use np.asarray(dr) instead of dr.values, so it also works for dictionary
 
     lon, lat = _get_lon_lat(ds)
+    if hasattr(lon, 'dims'):
+        if lon.ndim == 1:
+            dim_names = lat.dims + lon.dims
+        else:
+            dim_names = lon.dims
+    else:
+        dim_names = None
     lon, lat = as_2d_mesh(np.asarray(lon), np.asarray(lat))
 
     if 'mask' in ds:
@@ -123,7 +130,7 @@ def ds_to_ESMFgrid(ds, need_bounds=False, periodic=None, append=None):
         lon_b, lat_b = as_2d_mesh(np.asarray(lon_b), np.asarray(lat_b))
         add_corner(grid, lon_b.T, lat_b.T)
 
-    return grid, lon.shape
+    return grid, lon.shape, dim_names
 
 
 def ds_to_ESMFlocstream(ds):
@@ -142,6 +149,10 @@ def ds_to_ESMFlocstream(ds):
     """
 
     lon, lat = _get_lon_lat(ds)
+    if hasattr(lon, 'dims'):
+        dim_names = lon.dims
+    else:
+        dim_names = None
     lon, lat = np.asarray(lon), np.asarray(lat)
 
     if len(lon.shape) > 1:
@@ -153,7 +164,7 @@ def ds_to_ESMFlocstream(ds):
 
     locstream = LocStream.from_xarray(lon, lat)
 
-    return locstream, (1,) + lon.shape
+    return locstream, (1,) + lon.shape, dim_names
 
 
 def polys_to_ESMFmesh(polys):
@@ -195,6 +206,7 @@ class BaseRegridder(object):
         extrap_num_src_pnts=None,
         weights=None,
         ignore_degenerate=None,
+        input_dims=None,
     ):
         """
         Base xESMF regridding class supporting ESMF objects: `Grid`, `Mesh` and `LocStream`.
@@ -254,6 +266,11 @@ class BaseRegridder(object):
             If False (default), raise error if grids contain degenerated cells
             (i.e. triangles or lines, instead of quadrilaterals)
 
+        input_dims : tuple of str, optional
+            A tuple of dimension names to look for when regridding DataArrays or Datasets.
+            If not given or if those are not found on the regridded object, regridding
+            uses the two last dimensions of the object (or the last one for input LocStreams and Meshes).
+
         Returns
         -------
         baseregridder : xESMF BaseRegridder object
@@ -271,6 +288,10 @@ class BaseRegridder(object):
         self.sequence_in = isinstance(self.grid_in, (LocStream, Mesh))
         self.sequence_out = isinstance(self.grid_out, (LocStream, Mesh))
 
+        if input_dims is not None and len(input_dims) != int(not self.sequence_in) + 1:
+            raise ValueError(f'Wrong number of dimension names in `input_dims` ({len(input_dims)}.')
+        self.in_horiz_dims = input_dims
+
         # record grid shape information
         # We need to invert Grid shapes to respect xESMF's convention (y, x).
         self.shape_in = self.grid_in.get_shape()[::-1]
@@ -355,7 +376,8 @@ class BaseRegridder(object):
         Parameters
         ----------
         indata : numpy array, dask array, xarray DataArray or Dataset.
-            The rightmost two dimensions must be the same as ``ds_in``.
+            If not an xarray object or if `input_dìms` was not given in the init,
+            the rightmost two dimensions must be the same as ``ds_in``.
             Can have arbitrary additional dimensions.
 
             Examples of valid shapes
@@ -363,8 +385,8 @@ class BaseRegridder(object):
             - (n_lat, n_lon), if ``ds_in`` has shape (n_lat, n_lon)
             - (n_time, n_lev, n_y, n_x), if ``ds_in`` has shape (Ny, n_x)
 
-            Transpose your input data if the horizontal dimensions are not
-            the rightmost two dimensions.
+            Either give `input_dims` or transpose your input data
+            if the horizontal dimensions are not the rightmost two dimensions
 
         keep_attrs : bool, optional
             Keep attributes for xarray DataArrays or Datasets.
@@ -459,19 +481,8 @@ class BaseRegridder(object):
     def regrid_dataset(self, ds_in, keep_attrs=False):
         """See __call__()."""
 
-        # most logic is the same as regrid_dataarray()
-        # the major caution is that some data variables might not contain
-        # the correct horizontal dimension names.
-
         # get the first data variable to infer input_core_dims
-        name, dr_in = next(iter(ds_in.items()))
-        input_horiz_dims, temp_horiz_dims = self._parse_xrinput(dr_in)
-
-        # help user debugging invalid horizontal dimensions
-        print(
-            'using dimensions {} from data variable {} '
-            'as the horizontal dimensions for this dataset.'.format(input_horiz_dims, name)
-        )
+        input_horiz_dims, temp_horiz_dims = self._parse_xrinput(ds_in)
 
         ds_out = xr.apply_ufunc(
             self._regrid_array,
@@ -491,11 +502,30 @@ class BaseRegridder(object):
         return self._format_xroutput(ds_out, temp_horiz_dims)
 
     def _parse_xrinput(self, dr_in):
+        # dr could be a DataArray or a Dataset
         # Get input horiz dim names and set output horiz dim names
-        if self.sequence_in:
-            input_horiz_dims = dr_in.dims[-1:]
+        if self.in_horiz_dims is not None and all(dim in dr_in.dims for dim in self.in_horiz_dims):
+            input_horiz_dims = self.in_horiz_dims
         else:
-            input_horiz_dims = dr_in.dims[-2:]
+            if isinstance(dr_in, Dataset):
+                name, dr_in = next(iter(dr_in.items()))
+            else:
+                # For warning purposes
+                name = dr_in.name
+
+            if self.sequence_in:
+                input_horiz_dims = dr_in.dims[-1:]
+            else:
+                input_horiz_dims = dr_in.dims[-2:]
+
+            # help user debugging invalid horizontal dimensions
+            warnings.warn(
+                (
+                    f'Using dimensions {input_horiz_dims} from data variable {name} '
+                    'as the horizontal dimensions for the regridding.'
+                ),
+                UserWarning,
+            )
 
         if self.sequence_out:
             temp_horiz_dims = ['dummy', 'locations']
@@ -654,16 +684,18 @@ class Regridder(BaseRegridder):
 
         # construct ESMF grid, with some shape checking
         if locstream_in:
-            grid_in, shape_in = ds_to_ESMFlocstream(ds_in)
+            grid_in, shape_in, input_dims = ds_to_ESMFlocstream(ds_in)
         else:
-            grid_in, shape_in = ds_to_ESMFgrid(ds_in, need_bounds=need_bounds, periodic=periodic)
+            grid_in, shape_in, input_dims = ds_to_ESMFgrid(
+                ds_in, need_bounds=need_bounds, periodic=periodic
+            )
         if locstream_out:
-            grid_out, shape_out = ds_to_ESMFlocstream(ds_out)
+            grid_out, shape_out, _ = ds_to_ESMFlocstream(ds_out)
         else:
-            grid_out, shape_out = ds_to_ESMFgrid(ds_out, need_bounds=need_bounds)
+            grid_out, shape_out, _ = ds_to_ESMFgrid(ds_out, need_bounds=need_bounds)
 
         # Create the BaseRegridder
-        super().__init__(grid_in, grid_out, method, **kwargs)
+        super().__init__(grid_in, grid_out, method, input_dims=input_dims, **kwargs)
 
         # record output grid and metadata
         lon_out, lat_out = _get_lon_lat(ds_out)
@@ -809,11 +841,9 @@ class SpatialAverager(BaseRegridder):
         """
         self.ignore_holes = ignore_holes
         self.polys = polys
-        self.ignore_degenerate = ignore_degenerate
         self.geom_dim_name = geom_dim_name
 
-        grid_in, shape_in = ds_to_ESMFgrid(ds_in, need_bounds=True, periodic=periodic)
-        self.grid_in = grid_in
+        grid_in, shape_in, input_dims = ds_to_ESMFgrid(ds_in, need_bounds=True, periodic=periodic)
 
         # Create an output locstream so that the regridder knows the output shape and coords.
         # Latitude and longitude coordinates are the polygon centroid.
@@ -832,13 +862,14 @@ class SpatialAverager(BaseRegridder):
         # We put names 'lon' and 'lat' so ds_to_ESMFlocstream finds them easily.
         # _lon_out_name and _lat_out_name are used on the output anyway.
         ds_out = {'lon': self._lon_out, 'lat': self._lat_out}
-        locstream_out, shape_out = ds_to_ESMFlocstream(ds_out)
+        locstream_out, shape_out, _ = ds_to_ESMFlocstream(ds_out)
 
         # BaseRegridder with custom-computed weights and dummy out grid
         super().__init__(
             grid_in,
             locstream_out,
             'conservative',
+            input_dims=input_dims,
             weights=weights,
             filename=filename,
             reuse_weights=reuse_weights,


Test patch:
diff --git a/xesmf/tests/test_frontend.py b/xesmf/tests/test_frontend.py
index 44b8f98..8b6d5ab 100644
--- a/xesmf/tests/test_frontend.py
+++ b/xesmf/tests/test_frontend.py
@@ -43,11 +43,11 @@ ds_locs['lon'] = xr.DataArray(data=[0, 5, 10, 15], dims=('locations',))
 ds_savg = xr.Dataset(
     coords={
         'lat': (('lat',), [0.5, 1.5]),
-        'lon': (('lon',), [0.5, 1.5]),
+        'lon': (('lon',), [0.5, 1.5, 2.5]),
         'lat_b': (('lat_b',), [0, 1, 2]),
-        'lon_b': (('lon_b',), [0, 1, 2]),
+        'lon_b': (('lon_b',), [0, 1, 2, 3]),
     },
-    data_vars={'abc': (('lon', 'lat'), [[1, 2], [3, 4]])},
+    data_vars={'abc': (('lon', 'lat'), [[1, 2], [3, 4], [2, 4]])},
 )
 polys = [
     Polygon([[0.5, 0.5], [0.5, 1.5], [1.5, 0.5]]),  # Simple triangle polygon
@@ -78,7 +78,7 @@ polys = [
         ]
     ),  # Long multifaceted polygon
 ]
-exps_polys = [1.75, 3.5, 1.5716, 4, 0, 2.5]
+exps_polys = [1.75, 3, 2.1429, 4, 0, 2.5]
 
 
 def test_as_2d_mesh():
@@ -305,6 +305,15 @@ def test_regrid_dataarray(use_cfxr):
     xr.testing.assert_identical(dr_out_4D['time'], ds_in2['time'])
     xr.testing.assert_identical(dr_out_4D['lev'], ds_in2['lev'])
 
+    # test transposed
+    dr_out_4D_t = regridder(ds_in2['data4D'].transpose(..., 'time', 'lev'))
+    xr.testing.assert_identical(dr_out_4D, dr_out_4D_t)
+
+    # test renamed dim
+    if not use_cfxr:
+        dr_out_rn = regridder(ds_in2.rename(y='why')['data'])
+        xr.testing.assert_identical(dr_out, dr_out_rn)
+
 
 def test_regrid_dataarray_to_locstream():
     # xarray.DataArray containing in-memory numpy array
@@ -313,9 +322,11 @@ def test_regrid_dataarray_to_locstream():
 
     outdata = regridder(ds_in['data'].values)  # pure numpy array
     dr_out = regridder(ds_in['data'])  # xarray DataArray
+    dr_out_t = regridder(ds_in['data'].transpose())  # Transpose to test name matching
 
     # DataArray and numpy array should lead to the same result
     assert_equal(outdata.squeeze(), dr_out.values)
+    assert_equal(outdata.squeeze(), dr_out_t.values)
 
     with pytest.raises(ValueError):
         regridder = xe.Regridder(ds_in, ds_locs, 'conservative', locstream_out=True)
@@ -328,9 +339,12 @@ def test_regrid_dataarray_from_locstream():
 
     outdata = regridder(ds_locs['lat'].values)  # pure numpy array
     dr_out = regridder(ds_locs['lat'])  # xarray DataArray
+    # New dim and transpose to test name-matching
+    dr_out_2D = regridder(ds_locs['lat'].expand_dims(other=[1]).transpose('locations', 'other'))
 
     # DataArray and numpy array should lead to the same result
     assert_equal(outdata, dr_out.values)
+    assert_equal(outdata, dr_out_2D.sel(other=1).values)
 
     with pytest.raises(ValueError):
         regridder = xe.Regridder(ds_locs, ds_in, 'bilinear', locstream_in=True)
@@ -496,18 +510,19 @@ def test_ds_to_ESMFlocstream():
 
     from xesmf.frontend import ds_to_ESMFlocstream
 
-    locstream, shape = ds_to_ESMFlocstream(ds_locs)
+    locstream, shape, names = ds_to_ESMFlocstream(ds_locs)
     assert isinstance(locstream, ESMF.LocStream)
     assert shape == (
         1,
         4,
     )
+    assert names == ('locations',)
     with pytest.raises(ValueError):
-        locstream, shape = ds_to_ESMFlocstream(ds_in)
+        locstream, shape, names = ds_to_ESMFlocstream(ds_in)
     ds_bogus = ds_in.copy()
     ds_bogus['lon'] = ds_locs['lon']
     with pytest.raises(ValueError):
-        locstream, shape = ds_to_ESMFlocstream(ds_bogus)
+        locstream, shape, names = ds_to_ESMFlocstream(ds_bogus)
 
 
 @pytest.mark.parametrize('poly,exp', list(zip(polys, exps_polys)))