Skip to content

Refactor info and grdinfo to use virtualfile_from_data #961

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 10 commits into from
Mar 2, 2021
9 changes: 5 additions & 4 deletions doc/api/index.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
85 changes: 85 additions & 0 deletions pygmt/clib/session.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
GMTInvalidInput,
GMTVersionError,
)
from pygmt.helpers import data_kind, dummy_context

FAMILIES = [
"GMT_IS_DATASET",
Expand Down Expand Up @@ -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())
...
<vector memory>: 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.
Expand Down
11 changes: 1 addition & 10 deletions pygmt/src/grdinfo.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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]
Expand Down
25 changes: 2 additions & 23 deletions pygmt/src/info.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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(
Expand Down