Merge pull request JiaweiZhuang#68 from raphaeldussin/test_ocean

start adding testing for commonly used ocean models

ci/environment-upstream-dev.yml

@@ -2,6 +2,7 @@ name: xesmf
 channels:
   - conda-forge
 dependencies:
+  - cftime
   - codecov
   - dask
   - esmpy

ci/environment.yml

@@ -3,6 +3,7 @@ channels:
   - conda-forge
 dependencies:
   - cf_xarray>=0.3.1
+  - cftime
   - codecov
   - dask
   - esmpy

setup.cfg

@@ -8,7 +8,7 @@ extend-ignore = E203,E501,E402,W605
 
 [isort]
 known_first_party=xesmf
-known_third_party=ESMF,cf_xarray,dask,numpy,pytest,scipy,setuptools,shapely,xarray
+known_third_party=ESMF,cf_xarray,cftime,dask,numpy,pytest,scipy,setuptools,shapely,xarray
 multi_line_output=3
 include_trailing_comma=True
 force_grid_wrap=0

xesmf/tests/test_oceanmodels.py

@@ -1,17 +1,90 @@
-import cf_xarray  # noqa
+import cftime
 import numpy as np
 import pytest
 import xarray as xr
 
 import xesmf
 
+mom6like = xr.Dataset(
+    data_vars=dict(
+        tos=(['time', 'yh', 'xh'], np.random.rand(2, 180, 360)),
+    ),
+    coords=dict(
+        xq=xr.DataArray(
+            np.arange(-300, 60 + 1),
+            dims=['xq'],
+            attrs={
+                'long_name': 'q point nominal longitude',
+                'units': 'degrees_east',
+                'cartesian_axis': 'X',
+            },
+        ),
+        yq=xr.DataArray(
+            np.arange(-90, 90 + 1),
+            dims=['yq'],
+            attrs={
+                'long_name': 'q point nominal latitude',
+                'units': 'degrees_north',
+                'cartesian_axis': 'Y',
+            },
+        ),
+        xh=xr.DataArray(
+            0.5 + np.arange(-300, 60),
+            dims=['xh'],
+            attrs={
+                'long_name': 'h point nominal longitude',
+                'units': 'degrees_east',
+                'cartesian_axis': 'X',
+            },
+        ),
+        yh=xr.DataArray(
+            0.5 + np.arange(-90, 90),
+            dims=['yh'],
+            attrs={
+                'long_name': 'h point nominal latitude',
+                'units': 'degrees_north',
+                'cartesian_axis': 'Y',
+            },
+        ),
+        time=xr.DataArray(
+            [
+                cftime.DatetimeNoLeap(2007, 1, 16, 12, 0, 0, 0),
+                cftime.DatetimeNoLeap(2007, 2, 15, 0, 0, 0, 0),
+            ],
+            dims=['time'],
+        ),
+        reference_time=cftime.DatetimeNoLeap(1901, 1, 1, 0, 0, 0, 0),
+    ),
+    attrs=dict(description='Synthetic MOM6 data'),
+)
+
+
+def test_mom6like_to_5x5():
+    """regression test for MOM6 grid"""
+
+    grid_5x5 = xr.Dataset()
+    grid_5x5['lon'] = xr.DataArray(data=0.5 + np.arange(0, 360, 5), dims=('x'))
+    grid_5x5['lat'] = xr.DataArray(data=0.5 - 90 + np.arange(0, 180, 5), dims=('y'))
+
+    # multiple definition for lon/lat results in failure to determine
+    # which coordinate set to use.
+    with pytest.raises(ValueError):
+        regrid_to_5x5 = xesmf.Regridder(mom6like, grid_5x5, 'bilinear', periodic=True)
+
+    regrid_to_5x5 = xesmf.Regridder(
+        mom6like.rename({'xh': 'lon', 'yh': 'lat'}), grid_5x5, 'bilinear', periodic=True
+    )
+
+    tos_regridded = regrid_to_5x5(mom6like['tos'])
+    assert tos_regridded.shape == ((2, 36, 72))
+
 
 @pytest.mark.testcases
-def test_MOM6_to_1x1():
+def test_mom6_to_1x1():
     """ check regridding of MOM6 to regular lon/lat """
 
     dataurl = 'http://35.188.34.63:8080/thredds/dodsC/OM4p5/'
-    MOM6 = xr.open_dataset(
+    mom6 = xr.open_dataset(
         f'{dataurl}/ocean_monthly_z.200301-200712.nc4',
         chunks={'time': 1, 'z_l': 1},
         drop_variables=['average_DT', 'average_T1', 'average_T2'],
@@ -24,8 +97,8 @@ def test_MOM6_to_1x1():
     grid_1x1['lat'] = xr.DataArray(data=0.5 - 90 + np.arange(180), dims=('y'))
 
     regrid_to_1x1 = xesmf.Regridder(
-        MOM6.rename({'geolon': 'lon', 'geolat': 'lat'}), grid_1x1, 'bilinear', periodic=True
+        mom6.rename({'geolon': 'lon', 'geolat': 'lat'}), grid_1x1, 'bilinear', periodic=True
     )
 
-    thetao_1x1 = regrid_to_1x1(MOM6['thetao'])
+    thetao_1x1 = regrid_to_1x1(mom6['thetao'])
     assert np.allclose(thetao_1x1.isel(time=0, z_l=0, x=200, y=100).values, 27.15691922)