diff --git a/doc/api/index.rst b/doc/api/index.rst index 0ef0b359b39..58a41dd8cec 100644 --- a/doc/api/index.rst +++ b/doc/api/index.rst @@ -186,14 +186,15 @@ the :meth:`~pygmt.clib.Session.call_module` method: clib.Session.call_module -Passing memory blocks between Python variables (:class:`numpy.ndarray`, -:class:`pandas.Series`, and :class:`xarray.DataArray`) and GMT happens through *virtual -files*. These methods are context managers that automate the conversion of Python -variables to GMT virtual files: +Passing memory blocks between Python data objects (e.g. :class:`numpy.ndarray`, +:class:`pandas.Series`, :class:`xarray.DataArray`, etc) and GMT happens through +*virtual files*. These methods are context managers that automate the +conversion of Python variables to GMT virtual files: .. autosummary:: :toctree: generated + clib.Session.virtualfile_from_data clib.Session.virtualfile_from_matrix clib.Session.virtualfile_from_vectors clib.Session.virtualfile_from_grid diff --git a/pygmt/clib/session.py b/pygmt/clib/session.py index 93734645e90..09919dc9394 100644 --- a/pygmt/clib/session.py +++ b/pygmt/clib/session.py @@ -25,6 +25,7 @@ GMTInvalidInput, GMTVersionError, ) +from pygmt.helpers import data_kind, dummy_context FAMILIES = [ "GMT_IS_DATASET", @@ -1359,6 +1360,90 @@ def virtualfile_from_grid(self, grid): with self.open_virtual_file(*args) as vfile: yield vfile + def virtualfile_from_data(self, check_kind=None, data=None, x=None, y=None, z=None): + """ + Store any data inside a virtual file. + + This convenience function automatically detects the kind of data passed + into it, and produces a virtualfile that can be passed into GMT later + on. + + Parameters + ---------- + check_kind : str + Used to validate the type of data that can be passed in. Choose + from 'raster', 'vector' or None. Default is None (no validation). + data : str, xarray.DataArray, 2d array, or None + Any raster or vector data format. This could be a file name, a + raster grid, a vector matrix/arrays, or other supported data input. + x/y/z : 1d arrays or None + x, y and z columns as numpy arrays. + + Returns + ------- + file_context : contextlib._GeneratorContextManager + The virtual file stored inside a context manager. Access the file + name of this virtualfile using ``with file_context as fname: ...``. + + Examples + -------- + >>> from pygmt.helpers import GMTTempFile + >>> import xarray as xr + >>> data = xr.Dataset( + ... coords={"index": [0, 1, 2]}, + ... data_vars={ + ... "x": ("index", [9, 8, 7]), + ... "y": ("index", [6, 5, 4]), + ... "z": ("index", [3, 2, 1]), + ... }, + ... ) + >>> with Session() as ses: + ... with ses.virtualfile_from_data( + ... check_kind="vector", data=data + ... ) as fin: + ... # Send the output to a file so that we can read it + ... with GMTTempFile() as fout: + ... ses.call_module("info", f"{fin} ->{fout.name}") + ... print(fout.read().strip()) + ... + : N = 3 <7/9> <4/6> <1/3> + """ + kind = data_kind(data, x, y, z) + + if check_kind == "raster" and kind not in ("file", "grid"): + raise GMTInvalidInput(f"Unrecognized data type: {type(data)}") + if check_kind == "vector" and kind not in ("file", "matrix", "vectors"): + raise GMTInvalidInput(f"Unrecognized data type: {type(data)}") + + # Decide which virtualfile_from_ function to use + _virtualfile_from = { + "file": dummy_context, + "grid": self.virtualfile_from_grid, + # Note: virtualfile_from_matrix is not used because a matrix can be + # converted to vectors instead, and using vectors allows for better + # handling of string type inputs (e.g. for datetime data types) + "matrix": self.virtualfile_from_vectors, + "vectors": self.virtualfile_from_vectors, + }[kind] + + # Ensure the data is an iterable (Python list or tuple) + if kind in ("file", "grid"): + _data = (data,) + elif kind == "vectors": + _data = (x, y, z) + elif kind == "matrix": # turn 2D arrays into list of vectors + try: + # pandas.DataFrame and xarray.Dataset types + _data = [array for _, array in data.items()] + except AttributeError: + # Python lists, tuples, and numpy ndarray types + _data = np.atleast_2d(np.asanyarray(data).T) + + # Finally create the virtualfile from the data, to be passed into GMT + file_context = _virtualfile_from(*_data) + + return file_context + def extract_region(self): """ Extract the WESN bounding box of the currently active figure. diff --git a/pygmt/src/grdinfo.py b/pygmt/src/grdinfo.py index 43be9292aab..be481758497 100644 --- a/pygmt/src/grdinfo.py +++ b/pygmt/src/grdinfo.py @@ -2,12 +2,9 @@ grdinfo - Retrieve info about grid file. """ from pygmt.clib import Session -from pygmt.exceptions import GMTInvalidInput from pygmt.helpers import ( GMTTempFile, build_arg_string, - data_kind, - dummy_context, fmt_docstring, kwargs_to_strings, use_alias, @@ -109,15 +106,9 @@ def grdinfo(grid, **kwargs): info : str A string with information about the grid. """ - kind = data_kind(grid, None, None) with GMTTempFile() as outfile: with Session() as lib: - if kind == "file": - file_context = dummy_context(grid) - elif kind == "grid": - file_context = lib.virtualfile_from_grid(grid) - else: - raise GMTInvalidInput("Unrecognized data type: {}".format(type(grid))) + file_context = lib.virtualfile_from_data(check_kind="raster", data=grid) with file_context as infile: arg_str = " ".join( [infile, build_arg_string(kwargs), "->" + outfile.name] diff --git a/pygmt/src/info.py b/pygmt/src/info.py index e8d3730aeb3..e6f2e9aac61 100644 --- a/pygmt/src/info.py +++ b/pygmt/src/info.py @@ -3,15 +3,7 @@ """ import numpy as np from pygmt.clib import Session -from pygmt.exceptions import GMTInvalidInput -from pygmt.helpers import ( - GMTTempFile, - build_arg_string, - data_kind, - dummy_context, - fmt_docstring, - use_alias, -) +from pygmt.helpers import GMTTempFile, build_arg_string, fmt_docstring, use_alias @fmt_docstring @@ -66,21 +58,8 @@ def info(table, **kwargs): - :class:`numpy.ndarray` if either of the above parameters are used. - str if none of the above parameters are used. """ - kind = data_kind(table) with Session() as lib: - if kind == "file": - file_context = dummy_context(table) - elif kind == "matrix": - try: - # pandas.DataFrame and xarray.Dataset types - arrays = [array for _, array in table.items()] - except AttributeError: - # Python lists, tuples, and numpy ndarray types - arrays = np.atleast_2d(np.asanyarray(table).T) - file_context = lib.virtualfile_from_vectors(*arrays) - else: - raise GMTInvalidInput(f"Unrecognized data type: {type(table)}") - + file_context = lib.virtualfile_from_data(data=table) with GMTTempFile() as tmpfile: with file_context as fname: arg_str = " ".join(