diff --git a/doc/api/index.rst b/doc/api/index.rst index 618217468c2..56d6ec72786 100644 --- a/doc/api/index.rst +++ b/doc/api/index.rst @@ -345,3 +345,4 @@ Low level access (these are mostly used by the :mod:`pygmt.clib` package): clib.Session.virtualfile_from_stringio clib.Session.virtualfile_from_matrix clib.Session.virtualfile_from_vectors + clib.Session.virtualfile_from_xrgrid diff --git a/pygmt/clib/session.py b/pygmt/clib/session.py index eba799660f3..17eda883522 100644 --- a/pygmt/clib/session.py +++ b/pygmt/clib/session.py @@ -16,6 +16,7 @@ import numpy as np import pandas as pd import xarray as xr +from packaging.version import Version from pygmt.clib.conversion import ( dataarray_to_matrix, sequence_to_ctypes_array, @@ -24,6 +25,7 @@ ) from pygmt.clib.loading import get_gmt_version, load_libgmt from pygmt.datatypes import _GMT_DATASET, _GMT_GRID, _GMT_IMAGE +from pygmt.datatypes.header import gmt_grdfloat from pygmt.exceptions import GMTCLibError, GMTCLibNoSessionError, GMTInvalidInput from pygmt.helpers import ( _validate_data_input, @@ -207,6 +209,24 @@ def info(self) -> dict[str, str]: } return self._info + def __init__(self, in_mode: str = "GMT_IN", grid_as_matrix: bool = False): + """ + Initialize to store session-level variables. + + Parameters + ---------- + in_mode + Mode when creating a virtualfile. Defaults to ``"GMT_IN"``. It's used in + :meth:`pygmt.clib.Session.virtualfile_from_xrgrid` only. For some modules + (e.g., ``grdgradient`` and ``grdsample``), we may need + ``"GMT_IN|GMT_IS_REFERENCE"`` due to potential upstream bugs in GMT. + grid_as_matrix + Whether to treat the grid as a matrix. Defaults to ``False``. If ``True``, + will use the :meth:`pygmt.clib.Session.virtualfile_from_grid` method instead + """ + self._in_mode = in_mode + self._grid_as_matrix = grid_as_matrix + def __enter__(self): """ Create a GMT API session. @@ -327,6 +347,21 @@ def get_libgmt_func( function.restype = restype return function + def set_allocmode(self, family, obj): + """ + Set allocation mode of object to external. + """ + c_set_allocmode = self.get_libgmt_func( + "GMT_Set_AllocMode", + argtypes=[ctp.c_void_p, ctp.c_uint, ctp.c_void_p], + restype=ctp.c_void_p, + ) + family_int = self._parse_constant(family, valid=FAMILIES, valid_modifiers=VIAS) + status = c_set_allocmode(self.session_pointer, family_int, obj) + if status: + msg = f"Failed to set allocation mode of object to external:\n{self._error_message}" + raise GMTCLibError(msg) + def create(self, name: str) -> None: """ Create a new GMT C API session. @@ -1874,13 +1909,18 @@ def virtualfile_in( # noqa: PLR0912 "empty": self.virtualfile_from_vectors, "file": contextlib.nullcontext, "geojson": tempfile_from_geojson, - "grid": self.virtualfile_from_grid, + "grid": self.virtualfile_from_xrgrid, "image": tempfile_from_image, "stringio": self.virtualfile_from_stringio, "matrix": self.virtualfile_from_matrix, "vectors": self.virtualfile_from_vectors, }[kind] + # Patch for upstream bugs where grdinfo -L doesn't work with + # virtualfile_from_xrgrid. + if kind == "grid" and self._grid_as_matrix is True: + _virtualfile_from = self.virtualfile_from_grid + # "_data" is the data that will be passed to the _virtualfile_from function. # "_data" defaults to "data" but should be adjusted for some cases. _data = data @@ -2093,6 +2133,43 @@ def read_virtualfile( dtype = {"dataset": _GMT_DATASET, "grid": _GMT_GRID, "image": _GMT_IMAGE}[kind] return ctp.cast(pointer, ctp.POINTER(dtype)) + @contextlib.contextmanager + def virtualfile_from_xrgrid(self, xrgrid): + """ + Create a virtual file from an xarray.DataArray object. + """ + family = "GMT_IS_GRID" + geometry = "GMT_IS_SURFACE" + matrix, region, inc = dataarray_to_matrix(xrgrid) + + _gtype = {0: "GMT_GRID_IS_CARTESIAN", 1: "GMT_GRID_IS_GEO"}[xrgrid.gmt.gtype] + _reg = {0: "GMT_GRID_NODE_REG", 1: "GMT_GRID_PIXEL_REG"}[ + xrgrid.gmt.registration + ] + + data = self.create_data( + family, + geometry, + mode=f"GMT_CONTAINER_ONLY|{_gtype}", + ranges=region, + inc=inc, + registration=_reg, + pad=self["GMT_PAD_DEFAULT"], + ) + if Version(__gmt_version__) < Version("6.5.0"): + # Upstream bug fixed in GMT>=6.5.0 + self.set_allocmode(family, data) + + gmtgrid = ctp.cast(data, ctp.POINTER(_GMT_GRID)) + header = gmtgrid.contents.header.contents + header.z_min, header.z_max = matrix.min(), matrix.max() + + matrix = np.pad(matrix, self["GMT_PAD_DEFAULT"]).astype(np.float32) + gmtgrid.contents.data = matrix.ctypes.data_as(ctp.POINTER(gmt_grdfloat)) + + with self.open_virtualfile(family, geometry, self._in_mode, gmtgrid) as vfile: + yield vfile + def virtualfile_to_dataset( self, vfname: str, diff --git a/pygmt/src/grdgradient.py b/pygmt/src/grdgradient.py index f2a87a4dc8a..aeb4c4246cb 100644 --- a/pygmt/src/grdgradient.py +++ b/pygmt/src/grdgradient.py @@ -170,7 +170,7 @@ def grdgradient( "azimuth, direction, or radiance." ) raise GMTInvalidInput(msg) - with Session() as lib: + with Session(in_mode="GMT_IN|GMT_IS_REFERENCE") as lib: with ( lib.virtualfile_in(check_kind="raster", data=grid) as vingrd, lib.virtualfile_out(kind="grid", fname=outgrid) as voutgrd, diff --git a/pygmt/src/grdinfo.py b/pygmt/src/grdinfo.py index 97eb070bb2d..40363876771 100644 --- a/pygmt/src/grdinfo.py +++ b/pygmt/src/grdinfo.py @@ -3,8 +3,9 @@ """ import xarray as xr +from packaging.version import Version from pygmt._typing import PathLike -from pygmt.clib import Session +from pygmt.clib import Session, __gmt_version__ from pygmt.helpers import ( GMTTempFile, build_arg_list, @@ -112,8 +113,13 @@ def grdinfo(grid: PathLike | xr.DataArray, **kwargs) -> str: info : str A string with information about the grid. """ + # Workaround for upstream bug https://github.com/GenericMappingTools/gmt/issues/8525 + grid_as_matrix = Version(__gmt_version__) <= Version("6.5.0") and bool( + kwargs.get("L") + ) + with GMTTempFile() as outfile: - with Session() as lib: + with Session(grid_as_matrix=grid_as_matrix) as lib: with lib.virtualfile_in(check_kind="raster", data=grid) as vingrd: lib.call_module( module="grdinfo", diff --git a/pygmt/src/grdsample.py b/pygmt/src/grdsample.py index c5d6ae548aa..9a2ff2e6d46 100644 --- a/pygmt/src/grdsample.py +++ b/pygmt/src/grdsample.py @@ -83,7 +83,7 @@ def grdsample( >>> # and set both x- and y-spacings to 0.5 arc-degrees >>> new_grid = pygmt.grdsample(grid=grid, translate=True, spacing=[0.5, 0.5]) """ - with Session() as lib: + with Session(in_mode="GMT_IN|GMT_IS_REFERENCE") as lib: with ( lib.virtualfile_in(check_kind="raster", data=grid) as vingrd, lib.virtualfile_out(kind="grid", fname=outgrid) as voutgrd, diff --git a/pygmt/tests/test_grdsample.py b/pygmt/tests/test_grdsample.py index 4b768ef22eb..023567d1b2e 100644 --- a/pygmt/tests/test_grdsample.py +++ b/pygmt/tests/test_grdsample.py @@ -43,11 +43,11 @@ def fixture_expected_grid(): """ return xr.DataArray( data=[ - [460.84375, 482.78125, 891.09375], - [680.46875, 519.09375, 764.9375], - [867.75, 579.03125, 852.53125], - [551.75, 666.6875, 958.21875], - [411.3125, 518.4375, 931.28125], + [460.84375, 482.78125, 848.6875], + [680.46875, 519.09375, 760.84375], + [867.75, 579.03125, 804.875], + [551.75, 666.6875, 934.09375], + [411.3125, 518.4375, 879.875], ], coords={ "lon": [-52, -50, -48],