From c7330da501ab1817f247451edd45937bc176754a Mon Sep 17 00:00:00 2001 From: Will Schlitzer <schlitzer90@gmail.com> Date: Tue, 18 May 2021 07:55:15 +0100 Subject: [PATCH] Wrap grdclip function (#1261) Wrap the gmt function grdclip and add tests for it. Co-authored-by: Wei Ji <23487320+weiji14@users.noreply.github.com> Co-authored-by: Dongdong Tian <seisman.info@gmail.com> Co-authored-by: Wei Ji <23487320+weiji14@users.noreply.github.com> --- doc/api/index.rst | 1 + pygmt/__init__.py | 1 + pygmt/src/__init__.py | 1 + pygmt/src/grdclip.py | 89 +++++++++++++++++++++++++++++++++++++ pygmt/tests/test_grdclip.py | 47 ++++++++++++++++++++ 5 files changed, 139 insertions(+) create mode 100644 pygmt/src/grdclip.py create mode 100644 pygmt/tests/test_grdclip.py diff --git a/doc/api/index.rst b/doc/api/index.rst index b6cb79a26d6..dec563a4b8d 100644 --- a/doc/api/index.rst +++ b/doc/api/index.rst @@ -88,6 +88,7 @@ Operations on grids: .. autosummary:: :toctree: generated + grdclip grdcut grdfilter grdtrack diff --git a/pygmt/__init__.py b/pygmt/__init__.py index 1dc2ed79ac4..f9a7d00e1b0 100644 --- a/pygmt/__init__.py +++ b/pygmt/__init__.py @@ -33,6 +33,7 @@ blockmedian, config, grd2cpt, + grdclip, grdcut, grdfilter, grdinfo, diff --git a/pygmt/src/__init__.py b/pygmt/src/__init__.py index 083a2970109..5189c096243 100644 --- a/pygmt/src/__init__.py +++ b/pygmt/src/__init__.py @@ -10,6 +10,7 @@ from pygmt.src.config import config from pygmt.src.contour import contour from pygmt.src.grd2cpt import grd2cpt +from pygmt.src.grdclip import grdclip from pygmt.src.grdcontour import grdcontour from pygmt.src.grdcut import grdcut from pygmt.src.grdfilter import grdfilter diff --git a/pygmt/src/grdclip.py b/pygmt/src/grdclip.py new file mode 100644 index 00000000000..f0b5dbd7004 --- /dev/null +++ b/pygmt/src/grdclip.py @@ -0,0 +1,89 @@ +""" +grdclip - Change the range and extremes of grid values. +""" + +import xarray as xr +from pygmt.clib import Session +from pygmt.helpers import ( + GMTTempFile, + build_arg_string, + fmt_docstring, + kwargs_to_strings, + use_alias, +) + + +@fmt_docstring +@use_alias( + G="outgrid", + R="region", + Sa="above", + Sb="below", + Si="between", + Sr="new", + V="verbose", +) +@kwargs_to_strings( + R="sequence", + Sa="sequence", + Sb="sequence", + Si="sequence", + Sr="sequence", +) +def grdclip(grid, **kwargs): + r""" + Sets values in a grid that meet certain criteria to a new value. + + Produce a clipped ``outgrid`` or :class:`xarray.DataArray` version of the + input ``grid`` file. + + The parameters ``above`` and ``below`` allow for a given value to be set + for values above or below a set amount, respectively. This allows for + extreme values in a grid, such as points below a certain depth when + plotting Earth relief, to all be set to the same value. + + Full option list at :gmt-docs:`grdclip.html` + + {aliases} + + Parameters + ---------- + 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. + {R} + above : str or list or tuple + [*high*, *above*]. + Set all data[i] > *high* to *above*. + below : str or list or tuple + [*low*, *below*]. + Set all data[i] < *low* to *below*. + between : str or list or tuple + [*low*, *high*, *between*]. + Set all data[i] >= *low* and <= *high* to *between*. + new : str or list or tuple + [*old*, *new*]. + Set all data[i] == *old* to *new*. This is mostly useful when + your data are known to be integer values. + {V} + """ + with GMTTempFile(suffix=".nc") as tmpfile: + with Session() as lib: + file_context = lib.virtualfile_from_data(check_kind="raster", data=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("grdclip", arg_str) + + if outgrid == tmpfile.name: # if user did not set outgrid, return DataArray + with xr.open_dataarray(outgrid) as dataarray: + result = dataarray.load() + _ = result.gmt # load GMTDataArray accessor information + else: + result = None # if user sets an outgrid, return None + + return result diff --git a/pygmt/tests/test_grdclip.py b/pygmt/tests/test_grdclip.py new file mode 100644 index 00000000000..af08406eee2 --- /dev/null +++ b/pygmt/tests/test_grdclip.py @@ -0,0 +1,47 @@ +""" +Tests for grdclip. +""" +import os + +import numpy.testing as npt +import pytest +from pygmt import grdclip, grdinfo +from pygmt.datasets import load_earth_relief +from pygmt.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(resolution="01d", region=[-5, 5, -5, 5]) + + +def test_grdclip_outgrid(grid): + """ + Test the below and above parameters for grdclip and creates a test outgrid. + """ + with GMTTempFile(suffix=".nc") as tmpfile: + result = grdclip( + grid=grid, outgrid=tmpfile.name, below=[-1500, -1800], above=[-200, 40] + ) + assert result is None # return value is None + assert os.path.exists(path=tmpfile.name) # check that outgrid exists + result = ( + grdinfo(grid=tmpfile.name, force_scan=0, per_column="n").strip().split() + ) + assert int(result[4]) == -1800 # minimum value + assert int(result[5]) == 40 # maximum value + + +def test_grdclip_no_outgrid(grid): + """ + Test the below and above parameters for grdclip with no set outgrid. + """ + temp_grid = grdclip(grid=grid, below=[-1500, -1800], above=[-200, 40]) + assert temp_grid.dims == ("lat", "lon") + assert temp_grid.gmt.gtype == 1 # Geographic grid + assert temp_grid.gmt.registration == 1 # Pixel registration + npt.assert_allclose(temp_grid.min(), -1800) + npt.assert_allclose(temp_grid.max(), 40)