Skip to content

Commit

Permalink
Pass GMT_GRID virtual file to GMT C API
Browse files Browse the repository at this point in the history
  • Loading branch information
seisman committed May 6, 2024
1 parent c783a79 commit dc0d109
Show file tree
Hide file tree
Showing 2 changed files with 58 additions and 1 deletion.
1 change: 1 addition & 0 deletions doc/api/index.rst
Original file line number Diff line number Diff line change
Expand Up @@ -319,3 +319,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
58 changes: 57 additions & 1 deletion pygmt/clib/session.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -325,6 +326,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.
Expand Down Expand Up @@ -1597,7 +1614,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
Expand Down Expand Up @@ -1839,6 +1856,45 @@ 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|GMT_IS_REFERENCE", gmtgrid
) as vfile:
yield vfile

def virtualfile_to_dataset(
self,
vfname: str,
Expand Down

0 comments on commit dc0d109

Please sign in to comment.