Let grdcut() accept xarray.DataArray as input (#541)

Author: seisman

pygmt/gridops.py

@@ -47,8 +47,8 @@ def grdcut(grid, **kwargs):
 
     Parameters
     ----------
-    grid : str
-        The name of the input grid file.
+    grid : str or xarray.DataArray
+        The file name of the input grid or the grid loaded as a DataArray.
     outgrid : str or None
         The name of the output netCDF file with extension .nc to store the grid
         in.
@@ -94,12 +94,7 @@ def grdcut(grid, **kwargs):
             if kind == "file":
                 file_context = dummy_context(grid)
             elif kind == "grid":
-                raise NotImplementedError(
-                    "xarray.DataArray is not supported as the input grid yet!"
-                )
-                # file_context = lib.virtualfile_from_grid(grid)
-                # See https://github.com/GenericMappingTools/gmt/pull/3532
-                # for a feature request.
+                file_context = lib.virtualfile_from_grid(grid)
             else:
                 raise GMTInvalidInput("Unrecognized data type: {}".format(type(grid)))
 

pygmt/tests/test_grdcut.py

@@ -12,6 +12,12 @@
 from ..helpers import GMTTempFile
 
 
+@pytest.fixture(scope="module", name="grid")
+def fixture_grid():
+    "Load the grid data from the sample earth_relief file"
+    return load_earth_relief(registration="pixel")
+
+
 def test_grdcut_file_in_file_out():
     "grduct an input grid file, and output to a grid file"
     with GMTTempFile(suffix=".nc") as tmpfile:
@@ -41,23 +47,30 @@ def test_grdcut_file_in_dataarray_out():
     assert outgrid.sizes["lon"] == 180
 
 
-def test_grdcut_dataarray_in_file_out():
-    "grdcut an input DataArray, and output to a grid file"
-    # Not supported yet.
-    # See https://github.com/GenericMappingTools/gmt/pull/3532
-
-
-def test_grdcut_dataarray_in_dataarray_out():
+def test_grdcut_dataarray_in_file_out(grid):
     "grdcut an input DataArray, and output to a grid file"
-    # Not supported yet.
-    # See https://github.com/GenericMappingTools/gmt/pull/3532
+    with GMTTempFile(suffix=".nc") as tmpfile:
+        result = grdcut(grid, outgrid=tmpfile.name, region="0/180/0/90")
+        assert result is None  # grdcut returns None if output to a file
+        result = grdinfo(tmpfile.name, C=True)
+        assert result == "0 180 0 90 -8182 5651.5 1 1 180 90 1 1\n"
 
 
-def test_grdcut_dataarray_in_fail():
-    "Make sure that grdcut fails correctly if DataArray is the input grid"
-    with pytest.raises(NotImplementedError):
-        grid = load_earth_relief()
-        grdcut(grid, region="0/180/0/90")
+def test_grdcut_dataarray_in_dataarray_out(grid):
+    "grdcut an input DataArray, and output as DataArray"
+    outgrid = grdcut(grid, region="0/180/0/90")
+    assert isinstance(outgrid, xr.DataArray)
+    # check information of the output grid
+    # the '@earth_relief_01d' is in pixel registration, so the grid range is
+    # not exactly 0/180/0/90
+    assert outgrid.coords["lat"].data.min() == 0.5
+    assert outgrid.coords["lat"].data.max() == 89.5
+    assert outgrid.coords["lon"].data.min() == 0.5
+    assert outgrid.coords["lon"].data.max() == 179.5
+    assert outgrid.data.min() == -8182.0
+    assert outgrid.data.max() == 5651.5
+    assert outgrid.sizes["lat"] == 90
+    assert outgrid.sizes["lon"] == 180
 
 
 def test_grdcut_fails():