diff --git a/doc/api/index.rst b/doc/api/index.rst index 8e888419b57..61ed8498004 100644 --- a/doc/api/index.rst +++ b/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 ------------ diff --git a/pygmt/__init__.py b/pygmt/__init__.py index 96fe5cf908d..bd139af3d80 100644 --- a/pygmt/__init__.py +++ b/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 ( diff --git a/pygmt/datasets/earth_relief.py b/pygmt/datasets/earth_relief.py index dacd6333966..73c67273ade 100644 --- a/pygmt/datasets/earth_relief.py +++ b/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) diff --git a/pygmt/io.py b/pygmt/io.py new file mode 100644 index 00000000000..70bcfbfcdb9 --- /dev/null +++ b/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 diff --git a/pygmt/src/grdclip.py b/pygmt/src/grdclip.py index bead6dcae74..8c3a66d63c9 100644 --- a/pygmt/src/grdclip.py +++ b/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 diff --git a/pygmt/src/grdcut.py b/pygmt/src/grdcut.py index 353b82bc3a5..4842bfc0872 100644 --- a/pygmt/src/grdcut.py +++ b/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 diff --git a/pygmt/src/grdfill.py b/pygmt/src/grdfill.py index 6b553b964fe..92f06b7a1db 100644 --- a/pygmt/src/grdfill.py +++ b/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 diff --git a/pygmt/src/grdfilter.py b/pygmt/src/grdfilter.py index 46debc21534..33eb0afff72 100644 --- a/pygmt/src/grdfilter.py +++ b/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 diff --git a/pygmt/src/grdgradient.py b/pygmt/src/grdgradient.py index edbebcb402e..124adf882f5 100644 --- a/pygmt/src/grdgradient.py +++ b/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 diff --git a/pygmt/src/grdlandmask.py b/pygmt/src/grdlandmask.py index 8965d9d1a87..488607c4b69 100644 --- a/pygmt/src/grdlandmask.py +++ b/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 diff --git a/pygmt/src/grdproject.py b/pygmt/src/grdproject.py index c8a5d539421..d50c654438e 100644 --- a/pygmt/src/grdproject.py +++ b/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 diff --git a/pygmt/src/grdsample.py b/pygmt/src/grdsample.py index b1c807ccf94..d1416c34383 100644 --- a/pygmt/src/grdsample.py +++ b/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 diff --git a/pygmt/src/surface.py b/pygmt/src/surface.py index 857f8602c8e..200e68877ad 100644 --- a/pygmt/src/surface.py +++ b/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 diff --git a/pygmt/src/xyz2grd.py b/pygmt/src/xyz2grd.py index ef40abc889d..be08657f463 100644 --- a/pygmt/src/xyz2grd.py +++ b/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