diff --git a/doc/api/index.rst b/doc/api/index.rst index 5cf28bb3ebb..c5a260096ac 100644 --- a/doc/api/index.rst +++ b/doc/api/index.rst @@ -318,3 +318,4 @@ Low level access (these are mostly used by the :mod:`pygmt.clib` package): clib.Session.virtualfile_from_grid 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 b0aaff44ec3..94a276c8ac3 100644 --- a/pygmt/clib/session.py +++ b/pygmt/clib/session.py @@ -26,6 +26,7 @@ ) from pygmt.clib.loading import load_libgmt from pygmt.datatypes import _GMT_DATASET, _GMT_GRID +from pygmt.datatypes.header import gmt_grdfloat from pygmt.exceptions import ( GMTCLibError, GMTCLibNoSessionError, @@ -341,6 +342,22 @@ def get_libgmt_func(self, name, argtypes=None, restype=None): 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: + raise GMTCLibError( + f"Failed to set allocation mode of object to external:\n{self._error_message}" + ) + def create(self, name): """ Create a new GMT C API session. @@ -1612,7 +1629,7 @@ def virtualfile_in( # noqa: PLR0912 "file": contextlib.nullcontext, "arg": contextlib.nullcontext, "geojson": tempfile_from_geojson, - "grid": self.virtualfile_from_grid, + "grid": self.virtualfile_from_xrgrid, "image": tempfile_from_image, # Note: virtualfile_from_matrix is not used because a matrix can be # converted to vectors instead, and using vectors allows for better @@ -1854,6 +1871,43 @@ def read_virtualfile( dtype = {"dataset": _GMT_DATASET, "grid": _GMT_GRID}[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(self._info["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, "GMT_IN", gmtgrid) as vfile: + yield vfile + def virtualfile_to_dataset( self, vfname: str,