diff --git a/doc/api/index.rst b/doc/api/index.rst index 4d9a2616d96..3fcabeedd25 100644 --- a/doc/api/index.rst +++ b/doc/api/index.rst @@ -89,6 +89,7 @@ Operations on grids: .. autosummary:: :toctree: generated + grd2xyz grdclip grdcut grdfill diff --git a/pygmt/__init__.py b/pygmt/__init__.py index ec4934baddf..656c730aefc 100644 --- a/pygmt/__init__.py +++ b/pygmt/__init__.py @@ -35,6 +35,7 @@ blockmode, config, grd2cpt, + grd2xyz, grdclip, grdcut, grdfill, diff --git a/pygmt/src/__init__.py b/pygmt/src/__init__.py index a7f1e1c74ef..09ce03fe886 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.grd2xyz import grd2xyz from pygmt.src.grdclip import grdclip from pygmt.src.grdcontour import grdcontour from pygmt.src.grdcut import grdcut diff --git a/pygmt/src/grd2xyz.py b/pygmt/src/grd2xyz.py new file mode 100644 index 00000000000..c6404febf63 --- /dev/null +++ b/pygmt/src/grd2xyz.py @@ -0,0 +1,117 @@ +""" +grd2xyz - Convert grid to data table +""" +import warnings + +import pandas as pd +import xarray as xr +from pygmt.clib import Session +from pygmt.exceptions import GMTInvalidInput +from pygmt.helpers import ( + GMTTempFile, + build_arg_string, + fmt_docstring, + kwargs_to_strings, + use_alias, +) + + +@fmt_docstring +@use_alias( + R="region", + V="verbose", + o="outcols", +) +@kwargs_to_strings(R="sequence", o="sequence_comma") +def grd2xyz(grid, output_type="pandas", outfile=None, **kwargs): + r""" + Convert grid to data table. + + Read a grid and output xyz-triplets as a :class:`numpy.ndarray`, + :class:`pandas.DataFrame`, or ASCII file. + + Full option list at :gmt-docs:`grd2xyz.html` + + {aliases} + + Parameters + ---------- + grid : str or xarray.DataArray + The file name of the input grid or the grid loaded as a + :class:`xarray.DataArray`. This is the only required parameter. + output_type : str + Determine the format the xyz data will be returned in [Default is + ``pandas``]: + + - ``numpy`` - :class:`numpy.ndarray` + - ``pandas``- :class:`pandas.DataFrame` + - ``file`` - ASCII file (requires ``outfile``) + outfile : str + The file name for the output ASCII file. + {R} + Adding ``region`` will select a subsection of the grid. If this + subsection exceeds the boundaries of the grid, only the common region + will be output. + {V} + {o} + + Returns + ------- + ret : pandas.DataFrame or numpy.ndarray or None + Return type depends on ``outfile`` and ``output_type``: + + - None if ``outfile`` is set (output will be stored in file set by + ``outfile``) + - :class:`pandas.DataFrame` or :class:`numpy.ndarray` if ``outfile`` is + not set (depends on ``output_type``) + + """ + if output_type not in ["numpy", "pandas", "file"]: + raise GMTInvalidInput( + "Must specify 'output_type' either as 'numpy', 'pandas' or 'file'." + ) + + if outfile is not None and output_type != "file": + msg = ( + f"Changing 'output_type' of grd2xyz from '{output_type}' to 'file' " + "since 'outfile' parameter is set. Please use output_type='file' " + "to silence this warning." + ) + warnings.warn(message=msg, category=RuntimeWarning, stacklevel=2) + output_type = "file" + elif outfile is None and output_type == "file": + raise GMTInvalidInput("Must specify 'outfile' for ASCII output.") + + if "o" in kwargs and output_type == "pandas": + raise GMTInvalidInput( + "If 'outcols' is specified, 'output_type' must be either 'numpy'" + "or 'file'." + ) + + # Set the default column names for the pandas dataframe header + dataframe_header = ["x", "y", "z"] + # Let output pandas column names match input DataArray dimension names + if isinstance(grid, xr.DataArray) and output_type == "pandas": + # Reverse the dims because it is rows, columns ordered. + dataframe_header = [grid.dims[1], grid.dims[0], grid.name] + + with GMTTempFile() as tmpfile: + with Session() as lib: + file_context = lib.virtualfile_from_data(check_kind="raster", data=grid) + with file_context as infile: + if outfile is None: + outfile = tmpfile.name + arg_str = " ".join([infile, build_arg_string(kwargs), "->" + outfile]) + lib.call_module("grd2xyz", arg_str) + + # Read temporary csv output to a pandas table + if outfile == tmpfile.name: # if user did not set outfile, return pd.DataFrame + result = pd.read_csv( + tmpfile.name, sep="\t", names=dataframe_header, comment=">" + ) + elif outfile != tmpfile.name: # return None if outfile set, output in outfile + result = None + + if output_type == "numpy": + result = result.to_numpy() + return result diff --git a/pygmt/tests/test_grd2xyz.py b/pygmt/tests/test_grd2xyz.py new file mode 100644 index 00000000000..230132dc548 --- /dev/null +++ b/pygmt/tests/test_grd2xyz.py @@ -0,0 +1,97 @@ +""" +Tests for grd2xyz. +""" +import os + +import numpy as np +import pandas as pd +import pytest +from pygmt import grd2xyz +from pygmt.datasets import load_earth_relief +from pygmt.exceptions import GMTInvalidInput +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=[-1, 1, 3, 5]) + + +def test_grd2xyz(grid): + """ + Make sure grd2xyz works as expected. + """ + xyz_data = grd2xyz(grid=grid, output_type="numpy") + assert xyz_data.shape == (4, 3) + + +def test_grd2xyz_format(grid): + """ + Test that correct formats are returned. + """ + lon = -0.5 + lat = 3.5 + orig_val = grid.sel(lon=lon, lat=lat).to_numpy() + xyz_default = grd2xyz(grid=grid) + xyz_val = xyz_default[(xyz_default["lon"] == lon) & (xyz_default["lat"] == lat)][ + "elevation" + ].to_numpy() + assert isinstance(xyz_default, pd.DataFrame) + assert orig_val.size == 1 + assert xyz_val.size == 1 + np.testing.assert_allclose(orig_val, xyz_val) + xyz_array = grd2xyz(grid=grid, output_type="numpy") + assert isinstance(xyz_array, np.ndarray) + xyz_df = grd2xyz(grid=grid, output_type="pandas") + assert isinstance(xyz_df, pd.DataFrame) + assert list(xyz_df.columns) == ["lon", "lat", "elevation"] + + +def test_grd2xyz_file_output(grid): + """ + Test that grd2xyz returns a file output when it is specified. + """ + with GMTTempFile(suffix=".xyz") as tmpfile: + result = grd2xyz(grid=grid, outfile=tmpfile.name, output_type="file") + assert result is None # return value is None + assert os.path.exists(path=tmpfile.name) # check that outfile exists + + +def test_grd2xyz_invalid_format(grid): + """ + Test that grd2xyz fails with incorrect format. + """ + with pytest.raises(GMTInvalidInput): + grd2xyz(grid=grid, output_type=1) + + +def test_grd2xyz_no_outfile(grid): + """ + Test that grd2xyz fails when a string output is set with no outfile. + """ + with pytest.raises(GMTInvalidInput): + grd2xyz(grid=grid, output_type="file") + + +def test_grd2xyz_outfile_incorrect_output_type(grid): + """ + Test that grd2xyz raises a warning when an outfile filename is set but the + output_type is not set to 'file'. + """ + with pytest.warns(RuntimeWarning): + with GMTTempFile(suffix=".xyz") as tmpfile: + result = grd2xyz(grid=grid, outfile=tmpfile.name, output_type="numpy") + assert result is None # return value is None + assert os.path.exists(path=tmpfile.name) # check that outfile exists + + +def test_grd2xyz_pandas_output_with_o(grid): + """ + Test that grd2xyz fails when outcols is set and output_type is set to + 'pandas'. + """ + with pytest.raises(GMTInvalidInput): + grd2xyz(grid=grid, output_type="pandas", outcols="2")