Wrap grd2xyz (#1284)

Co-authored-by: Meghan Jones <[email protected]>
Co-authored-by: Wei Ji <[email protected]>
Co-authored-by: Dongdong Tian <[email protected]>

Author: 4 people

doc/api/index.rst

@@ -89,6 +89,7 @@ Operations on grids:
 .. autosummary::
     :toctree: generated
 
+    grd2xyz
     grdclip
     grdcut
     grdfill

pygmt/__init__.py

@@ -35,6 +35,7 @@
     blockmode,
     config,
     grd2cpt,
+    grd2xyz,
     grdclip,
     grdcut,
     grdfill,

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

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

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")