Skip to content

Commit

Permalink
Wrap grdcut (#492)
Browse files Browse the repository at this point in the history
Wrap the `grdcut` function.

Ideally, `grdcut` should take a grid file or a `xarray.DataArray` as input, and
store output as a grid file or return a `xarray.DataArray`.

Currently supported features:

- [X] input: grid file; output: grid file
- [X] input: grid file; output: `xarray.DataArray`
- [ ] input: `xarray.DataArray`; output: grid file
- [ ] input: `xarray.DataArray`; output: `xarray.DataArray`

Passing `xarray.DataArray` to `grdcut` is not implemented in GMT 6.0.0, 
only is only possible for GMT>=6.1.0.
See GenericMappingTools/gmt#3532 for
the feature request to the upstream GMT.

Co-authored-by: Wei Ji <23487320+weiji14@users.noreply.github.com>
  • Loading branch information
seisman and weiji14 authored Jun 28, 2020
1 parent 0afaa69 commit f7b6d6c
Show file tree
Hide file tree
Showing 4 changed files with 185 additions and 0 deletions.
1 change: 1 addition & 0 deletions doc/api/index.rst
Original file line number Diff line number Diff line change
Expand Up @@ -70,6 +70,7 @@ Operations on grids:
.. autosummary::
:toctree: generated

grdcut
grdinfo
grdtrack

Expand Down
1 change: 1 addition & 0 deletions pygmt/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
from .sampling import grdtrack
from .mathops import makecpt
from .modules import config, info, grdinfo, which
from .gridops import grdcut
from . import datasets


Expand Down
119 changes: 119 additions & 0 deletions pygmt/gridops.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,119 @@
"""
GMT modules for grid operations
"""

import xarray as xr


from .clib import Session
from .helpers import (
build_arg_string,
fmt_docstring,
kwargs_to_strings,
GMTTempFile,
use_alias,
data_kind,
dummy_context,
)
from .exceptions import GMTInvalidInput


@fmt_docstring
@use_alias(
G="outgrid",
R="region",
J="projection",
N="extend",
S="circ_subregion",
Z="z_subregion",
)
@kwargs_to_strings(R="sequence")
def grdcut(grid, **kwargs):
"""
Extract subregion from a grid.
Produce a new *outgrid* file which is a subregion of *grid*. The
subregion is specified with *region*; the specified range must not exceed
the range of *grid* (but see *extend*). If in doubt, run
:meth:`pygmt.grdinfo` to check range. Alternatively, define the subregion
indirectly via a range check on the node values or via distances from a
given point. Finally, you can give *projection* for oblique projections to
determine the corresponding rectangular *region* setting that will give a
grid that fully covers the oblique domain.
Full option list at :gmt-docs:`grdcut.html`
{aliases}
Parameters
----------
grid : str
The name of the input grid file.
outgrid : str or None
The name of the output netCDF file with extension .nc to store the grid
in.
{J}
{R}
extend : bool or int or float
Allow grid to be extended if new *region* exceeds existing boundaries.
Give a value to initialize nodes outside current region.
circ_subregion : str
``'lon/lat/radius[unit][+n]'``.
Specify an origin (*lon* and *lat*) and *radius*; append a distance
*unit* and we determine the corresponding rectangular region so that
all grid nodes on or inside the circle are contained in the subset.
If **+n** is appended we set all nodes outside the circle to NaN.
z_subregion : str
``'[min/max][+n|N|r]'``.
Determine a new rectangular region so that all nodes outside this
region are also outside the given z-range [-inf/+inf]. To indicate no
limit on *min* or *max* only, specify a hyphen (-). Normally, any NaNs
encountered are simply skipped and not considered in the
range-decision. Append **+n** to consider a NaN to be outside the given
z-range. This means the new subset will be NaN-free. Alternatively,
append **+r** to consider NaNs to be within the data range. In this
case we stop shrinking the boundaries once a NaN is found [Default
simply skips NaNs when making the range decision]. Finally, if your
core subset grid is surrounded by rows and/or columns that are all
NaNs, append **+N** to strip off such columns before (optionally)
considering the range of the core subset for further reduction of the
area.
Returns
-------
ret: xarray.DataArray or None
Return type depends on whether the *outgrid* parameter is set:
- xarray.DataArray if *outgrid* is not set
- None if *outgrid* is set (grid output will be stored in *outgrid*)
"""
kind = data_kind(grid)

with GMTTempFile(suffix=".nc") as tmpfile:
with Session() as lib:
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.
else:
raise GMTInvalidInput("Unrecognized data type: {}".format(type(grid)))

with file_context as infile:
if "G" not in kwargs.keys(): # if outgrid is unset, output to tempfile
kwargs.update({"G": tmpfile.name})
outgrid = kwargs["G"]
arg_str = " ".join([infile, build_arg_string(kwargs)])
lib.call_module("grdcut", arg_str)

if outgrid == tmpfile.name: # if user did not set outgrid, return DataArray
with xr.open_dataarray(outgrid) as dataarray:
result = dataarray.load()
else:
result = None # if user sets an outgrid, return None

return result
64 changes: 64 additions & 0 deletions pygmt/tests/test_grdcut.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,64 @@
"""
Tests for grdcut
"""
import os
import numpy as np
import pytest
import xarray as xr

from .. import grdcut, grdinfo
from ..datasets import load_earth_relief
from ..exceptions import GMTInvalidInput
from ..helpers import GMTTempFile


def test_grdcut_file_in_file_out():
"grduct an input grid file, and output to a grid file"
with GMTTempFile(suffix=".nc") as tmpfile:
result = grdcut("@earth_relief_01d", outgrid=tmpfile.name, region="0/180/0/90")
assert result is None # return value is None
assert os.path.exists(path=tmpfile.name) # check that outgrid exists
result = grdinfo(tmpfile.name, C=True)
assert result == "0 180 0 90 -8592.5 5559 1 1 181 91\n"


def test_grdcut_file_in_dataarray_out():
"grdcut an input grid file, and output as DataArray"
outgrid = grdcut("@earth_relief_01d", 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.0
assert outgrid.coords["lat"].data.max() == 90.0
assert outgrid.coords["lon"].data.min() == 0.0
assert outgrid.coords["lon"].data.max() == 180.0
assert outgrid.data.min() == -8592.5
assert outgrid.data.max() == 5559
assert outgrid.sizes["lat"] == 91
assert outgrid.sizes["lon"] == 181


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():
"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_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_fails():
"Check that grdcut fails correctly"
with pytest.raises(GMTInvalidInput):
grdcut(np.arange(10).reshape((5, 2)))

0 comments on commit f7b6d6c

Please sign in to comment.