Add pygmt.load_dataarray function and refactor modules that return None or xarray.dataarray (#1439)

Add new function pygmt.load_datarray
Use pygmt.load_dataarrary for modules that return grids
Add documentation

Co-authored-by: Wei Ji <[email protected]>

Author: Meghan Jones

doc/api/index.rst

@@ -107,6 +107,14 @@ Crossover analysis with x2sys:
     x2sys_init
     x2sys_cross
 
+Input/output
+------------
+
+.. autosummary::
+    :toctree: generated
+
+    load_dataarray
+
 GMT Defaults
 ------------
 

pygmt/__init__.py

@@ -26,6 +26,7 @@
 from pygmt import datasets
 from pygmt.accessors import GMTDataArrayAccessor
 from pygmt.figure import Figure, set_display
+from pygmt.io import load_dataarray
 from pygmt.session_management import begin as _begin
 from pygmt.session_management import end as _end
 from pygmt.src import (

pygmt/datasets/earth_relief.py

@@ -4,9 +4,9 @@
 
 The grids are available in various resolutions.
 """
-import xarray as xr
 from pygmt.exceptions import GMTInvalidInput
 from pygmt.helpers import kwargs_to_strings
+from pygmt.io import load_dataarray
 from pygmt.src import grdcut, which
 
 
@@ -133,9 +133,7 @@ def load_earth_relief(resolution="01d", region=None, registration=None, use_srtm
                 f"'region' is required for Earth relief resolution '{resolution}'."
             )
         fname = which(f"@earth_relief_{resolution}{reg}", download="a")
-        with xr.open_dataarray(fname, engine="netcdf4") as dataarray:
-            grid = dataarray.load()
-            _ = grid.gmt  # load GMTDataArray accessor information
+        grid = load_dataarray(fname, engine="netcdf4")
     else:
         grid = grdcut(f"@{earth_relief_prefix}{resolution}{reg}", region=region)
 

pygmt/io.py

@@ -0,0 +1,46 @@
+"""
+PyGMT input/output (I/O) utilities.
+"""
+import xarray as xr
+
+
+def load_dataarray(filename_or_obj, **kwargs):
+    """
+    Open, load into memory, and close a DataArray from a file or file-like
+    object containing a single data variable.
+
+    This is a thin wrapper around :py:func:`xarray.open_dataarray`. It differs
+    from :py:func:`xarray.open_dataarray` in that it loads the DataArray into
+    memory, gets GMT specific metadata about the grid via
+    :py:meth:`GMTDataArrayAccessor`, closes the file, and returns the
+    DataArray. In contrast, :py:func:`xarray.open_dataarray` keeps the file
+    handle open and lazy loads its contents. All parameters are passed directly
+    to :py:func:`xarray.open_dataarray`. See that documentation for further
+    details.
+
+    Parameters
+    ----------
+    filename_or_obj : str or pathlib.Path or file-like or DataStore
+        Strings and Path objects are interpreted as a path to a netCDF file
+        or an OpenDAP URL and opened with python-netCDF4, unless the filename
+        ends with .gz, in which case the file is gunzipped and opened with
+        scipy.io.netcdf (only netCDF3 supported). Byte-strings or file-like
+        objects are opened by scipy.io.netcdf (netCDF3) or h5py (netCDF4/HDF).
+
+    Returns
+    -------
+    datarray : xarray.DataArray
+        The newly created DataArray.
+
+    See Also
+    --------
+    xarray.open_dataarray
+    """
+    if "cache" in kwargs:
+        raise TypeError("cache has no effect in this context")
+
+    with xr.open_dataarray(filename_or_obj, **kwargs) as dataarray:
+        result = dataarray.load()
+        _ = result.gmt  # load GMTDataArray accessor information
+
+    return result

pygmt/src/grdclip.py

@@ -2,7 +2,6 @@
 grdclip - Change the range and extremes of grid values.
 """
 
-import xarray as xr
 from pygmt.clib import Session
 from pygmt.helpers import (
     GMTTempFile,
@@ -11,6 +10,7 @@
     kwargs_to_strings,
     use_alias,
 )
+from pygmt.io import load_dataarray
 
 
 @fmt_docstring
@@ -88,11 +88,4 @@ def grdclip(grid, **kwargs):
                 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
+        return load_dataarray(outgrid) if outgrid == tmpfile.name else None

pygmt/src/grdcut.py

@@ -2,7 +2,6 @@
 grdcut - Extract subregion from a grid.
 """
 
-import xarray as xr
 from pygmt.clib import Session
 from pygmt.helpers import (
     GMTTempFile,
@@ -11,6 +10,7 @@
     kwargs_to_strings,
     use_alias,
 )
+from pygmt.io import load_dataarray
 
 
 @fmt_docstring
@@ -98,11 +98,4 @@ def grdcut(grid, **kwargs):
                 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()
-                _ = result.gmt  # load GMTDataArray accessor information
-        else:
-            result = None  # if user sets an outgrid, return None
-
-        return result
+        return load_dataarray(outgrid) if outgrid == tmpfile.name else None

pygmt/src/grdfill.py

@@ -2,7 +2,6 @@
 grdfill - Fill blank areas from a grid.
 """
 
-import xarray as xr
 from pygmt.clib import Session
 from pygmt.exceptions import GMTInvalidInput
 from pygmt.helpers import (
@@ -12,6 +11,7 @@
     kwargs_to_strings,
     use_alias,
 )
+from pygmt.io import load_dataarray
 
 
 @fmt_docstring
@@ -75,11 +75,4 @@ def grdfill(grid, **kwargs):
                 arg_str = " ".join([infile, build_arg_string(kwargs)])
                 lib.call_module("grdfill", 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
+        return load_dataarray(outgrid) if outgrid == tmpfile.name else None

pygmt/src/grdfilter.py

@@ -2,7 +2,6 @@
 grdfilter - Filter a grid in the space (or time) domain.
 """
 
-import xarray as xr
 from pygmt.clib import Session
 from pygmt.helpers import (
     GMTTempFile,
@@ -11,6 +10,7 @@
     kwargs_to_strings,
     use_alias,
 )
+from pygmt.io import load_dataarray
 
 
 @fmt_docstring
@@ -151,11 +151,4 @@ def grdfilter(grid, **kwargs):
                 arg_str = " ".join([infile, build_arg_string(kwargs)])
                 lib.call_module("grdfilter", 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
+        return load_dataarray(outgrid) if outgrid == tmpfile.name else None

pygmt/src/grdgradient.py

@@ -2,7 +2,6 @@
 grdgradient - Compute directional gradients from a grid.
 """
 
-import xarray as xr
 from pygmt.clib import Session
 from pygmt.exceptions import GMTInvalidInput
 from pygmt.helpers import (
@@ -13,6 +12,7 @@
     kwargs_to_strings,
     use_alias,
 )
+from pygmt.io import load_dataarray
 
 
 @fmt_docstring
@@ -117,11 +117,4 @@ def grdgradient(grid, **kwargs):
                 arg_str = " ".join([infile, build_arg_string(kwargs)])
                 lib.call_module("grdgradient", 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
+        return load_dataarray(outgrid) if outgrid == tmpfile.name else None

pygmt/src/grdlandmask.py

@@ -2,7 +2,6 @@
 grdlandmask - Create a "wet-dry" mask grid from shoreline data base
 """
 
-import xarray as xr
 from pygmt.clib import Session
 from pygmt.exceptions import GMTInvalidInput
 from pygmt.helpers import (
@@ -12,6 +11,7 @@
     kwargs_to_strings,
     use_alias,
 )
+from pygmt.io import load_dataarray
 
 
 @fmt_docstring
@@ -104,11 +104,4 @@ def grdlandmask(**kwargs):
             arg_str = build_arg_string(kwargs)
             lib.call_module("grdlandmask", 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
+        return load_dataarray(outgrid) if outgrid == tmpfile.name else None

pygmt/src/grdproject.py

@@ -2,7 +2,6 @@
 grdproject - Forward and inverse map transformation of grids.
 """
 
-import xarray as xr
 from pygmt.clib import Session
 from pygmt.exceptions import GMTInvalidInput
 from pygmt.helpers import (
@@ -12,6 +11,7 @@
     kwargs_to_strings,
     use_alias,
 )
+from pygmt.io import load_dataarray
 
 
 @fmt_docstring
@@ -83,11 +83,4 @@ def grdproject(grid, **kwargs):
                 arg_str = " ".join([infile, build_arg_string(kwargs)])
                 lib.call_module("grdproject", 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
+        return load_dataarray(outgrid) if outgrid == tmpfile.name else None

pygmt/src/grdsample.py

@@ -2,7 +2,6 @@
 grdsample - Resample a grid onto a new lattice
 """
 
-import xarray as xr
 from pygmt.clib import Session
 from pygmt.helpers import (
     GMTTempFile,
@@ -11,6 +10,7 @@
     kwargs_to_strings,
     use_alias,
 )
+from pygmt.io import load_dataarray
 
 
 @fmt_docstring
@@ -85,11 +85,4 @@ def grdsample(grid, **kwargs):
                 arg_str = " ".join([infile, build_arg_string(kwargs)])
                 lib.call_module("grdsample", 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
+        return load_dataarray(outgrid) if outgrid == tmpfile.name else None

pygmt/src/surface.py

@@ -2,7 +2,6 @@
 surface - Grids table data using adjustable tension continuous curvature
 splines.
 """
-import xarray as xr
 from pygmt.clib import Session
 from pygmt.exceptions import GMTInvalidInput
 from pygmt.helpers import (
@@ -15,6 +14,7 @@
     kwargs_to_strings,
     use_alias,
 )
+from pygmt.io import load_dataarray
 
 
 @fmt_docstring
@@ -101,11 +101,4 @@ def surface(x=None, y=None, z=None, data=None, **kwargs):
                 arg_str = " ".join([infile, build_arg_string(kwargs)])
                 lib.call_module(module="surface", args=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
-        elif outgrid != tmpfile.name:  # if user sets an outgrid, return None
-            result = None
-
-    return result
+        return load_dataarray(outgrid) if outgrid == tmpfile.name else None

pygmt/src/xyz2grd.py

@@ -1,7 +1,6 @@
 """
 xyz2grd - Convert data table to a grid.
 """
-import xarray as xr
 from pygmt.clib import Session
 from pygmt.helpers import (
     GMTTempFile,
@@ -10,6 +9,7 @@
     kwargs_to_strings,
     use_alias,
 )
+from pygmt.io import load_dataarray
 
 
 @fmt_docstring
@@ -64,11 +64,4 @@ def xyz2grd(table, **kwargs):
                 arg_str = " ".join([infile, arg_str])
                 lib.call_module("xyz2grd", 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
+        return load_dataarray(outgrid) if outgrid == tmpfile.name else None