Add support for custom coordinate names

xesmf/frontend.py

@@ -11,13 +11,20 @@
                        esmf_regrid_build, esmf_regrid_finalize)
 
 from . smm import read_weights, apply_weights
+from . util import get_lon_lat, get_lon_lat_bnds
 
 try:
     import dask.array as da
     dask_array_type = (da.Array,)  # for isinstance checks
 except ImportError:
     dask_array_type = ()
 
+default_var_names = {"lat": "lat",
+                     "lon": "lon",
+                     "lat_b": "lat_b",
+                     "lon_b": "lon_b"}
+
+
 def as_2d_mesh(lon, lat):
 
     if (lon.ndim, lat.ndim) == (2, 2):
@@ -30,7 +37,7 @@ def as_2d_mesh(lon, lat):
     return lon, lat
 
 
-def ds_to_ESMFgrid(ds, need_bounds=False, periodic=None, append=None):
+def ds_to_ESMFgrid(ds, need_bounds=False, periodic=None, append=None, var_names=default_var_names):
     '''
     Convert xarray DataSet or dictionary to ESMF.Grid object.
 
@@ -41,7 +48,7 @@ def ds_to_ESMFgrid(ds, need_bounds=False, periodic=None, append=None):
         and optionally ``lon_b``, ``lat_b`` if need_bounds=True.
 
         Shape should be ``(n_lat, n_lon)`` or ``(n_y, n_x)``,
-        as normal C or Python ordering. Will be then tranposed to F-ordered.
+        as normal C or Python ordering. Will be then transposed to F-ordered.
 
     need_bounds : bool, optional
         Need cell boundary values?
@@ -55,24 +62,22 @@ 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 = np.asarray(ds['lon'])
-    lat = np.asarray(ds['lat'])
+    # Need asarray ?
+    lon, lat = get_lon_lat(ds, var_names)
     lon, lat = as_2d_mesh(lon, lat)
 
     # tranpose the arrays so they become Fortran-ordered
     grid = esmf_grid(lon.T, lat.T, periodic=periodic)
 
     if need_bounds:
-        lon_b = np.asarray(ds['lon_b'])
-        lat_b = np.asarray(ds['lat_b'])
+        lon_b, lat_b = get_lon_lat_bnds(ds, var_names)
         lon_b, lat_b = as_2d_mesh(lon_b, lat_b)
         add_corner(grid, lon_b.T, lat_b.T)
 
     return grid, lon.shape
 
 
-def ds_to_ESMFlocstream(ds):
+def ds_to_ESMFlocstream(ds, var_names=default_var_names):
     '''
     Convert xarray DataSet or dictionary to ESMF.LocStream object.
 
@@ -86,9 +91,7 @@ def ds_to_ESMFlocstream(ds):
     locstream : ESMF.LocStream object
 
     '''
-
-    lon = np.asarray(ds['lon'])
-    lat = np.asarray(ds['lat'])
+    lon, lat = get_lon_lat(ds, var_names=var_names)
 
     if len(lon.shape) > 1:
         raise ValueError("lon can only be 1d")
@@ -105,7 +108,7 @@ def ds_to_ESMFlocstream(ds):
 class Regridder(object):
     def __init__(self, ds_in, ds_out, method, periodic=False,
                  filename=None, reuse_weights=False, ignore_degenerate=None,
-                 locstream_in=False, locstream_out=False):
+                 locstream_in=False, locstream_out=False, var_names=default_var_names):
         """
         Make xESMF regridder
 
@@ -185,17 +188,19 @@ def __init__(self, ds_in, ds_out, method, periodic=False,
 
         # construct ESMF grid, with some shape checking
         if locstream_in:
-            self._grid_in, shape_in = ds_to_ESMFlocstream(ds_in)
+            self._grid_in, shape_in = ds_to_ESMFlocstream(ds_in, var_names=var_names)
         else:
             self._grid_in, shape_in = ds_to_ESMFgrid(ds_in,
                                                      need_bounds=self.need_bounds,
-                                                     periodic=periodic
+                                                     periodic=periodic,
+                                                     var_names=var_names
                                                      )
         if locstream_out:
-            self._grid_out, shape_out = ds_to_ESMFlocstream(ds_out)
+            self._grid_out, shape_out = ds_to_ESMFlocstream(ds_out, var_names=var_names)
         else:
             self._grid_out, shape_out = ds_to_ESMFgrid(ds_out,
-                                                       need_bounds=self.need_bounds
+                                                       need_bounds=self.need_bounds,
+                                                       var_names=var_names
                                                        )
 
         # record output grid and metadata