Skip to content

Commit e28e846

Browse files
committed
Pass GMT_GRID virtual file to GMT C API
1 parent c783a79 commit e28e846

File tree

2 files changed

+56
-1
lines changed

2 files changed

+56
-1
lines changed

doc/api/index.rst

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -319,3 +319,4 @@ Low level access (these are mostly used by the :mod:`pygmt.clib` package):
319319
clib.Session.virtualfile_from_grid
320320
clib.Session.virtualfile_from_matrix
321321
clib.Session.virtualfile_from_vectors
322+
clib.Session.virtualfile_from_xrgrid

pygmt/clib/session.py

Lines changed: 55 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
)
2727
from pygmt.clib.loading import load_libgmt
2828
from pygmt.datatypes import _GMT_DATASET, _GMT_GRID
29+
from pygmt.datatypes.header import gmt_grdfloat
2930
from pygmt.exceptions import (
3031
GMTCLibError,
3132
GMTCLibNoSessionError,
@@ -325,6 +326,22 @@ def get_libgmt_func(self, name, argtypes=None, restype=None):
325326
function.restype = restype
326327
return function
327328

329+
def set_allocmode(self, family, obj):
330+
"""
331+
Set allocation mode of object to external.
332+
"""
333+
c_set_allocmode = self.get_libgmt_func(
334+
"GMT_Set_AllocMode",
335+
argtypes=[ctp.c_void_p, ctp.c_uint, ctp.c_void_p],
336+
restype=ctp.c_void_p,
337+
)
338+
family_int = self._parse_constant(family, valid=FAMILIES, valid_modifiers=VIAS)
339+
status = c_set_allocmode(self.session_pointer, family_int, obj)
340+
if status:
341+
raise GMTCLibError(
342+
f"Failed to set allocation mode of object to external:\n{self._error_message}"
343+
)
344+
328345
def create(self, name):
329346
"""
330347
Create a new GMT C API session.
@@ -1597,7 +1614,7 @@ def virtualfile_in( # noqa: PLR0912
15971614
"file": contextlib.nullcontext,
15981615
"arg": contextlib.nullcontext,
15991616
"geojson": tempfile_from_geojson,
1600-
"grid": self.virtualfile_from_grid,
1617+
"grid": self.virtualfile_from_xrgrid,
16011618
"image": tempfile_from_image,
16021619
# Note: virtualfile_from_matrix is not used because a matrix can be
16031620
# converted to vectors instead, and using vectors allows for better
@@ -1839,6 +1856,43 @@ def read_virtualfile(
18391856
dtype = {"dataset": _GMT_DATASET, "grid": _GMT_GRID}[kind]
18401857
return ctp.cast(pointer, ctp.POINTER(dtype))
18411858

1859+
@contextlib.contextmanager
1860+
def virtualfile_from_xrgrid(self, xrgrid):
1861+
"""
1862+
Create a virtual file from an xarray.DataArray object.
1863+
"""
1864+
family = "GMT_IS_GRID"
1865+
geometry = "GMT_IS_SURFACE"
1866+
matrix, region, inc = dataarray_to_matrix(xrgrid)
1867+
1868+
_gtype = {0: "GMT_GRID_IS_CARTESIAN", 1: "GMT_GRID_IS_GEO"}[xrgrid.gmt.gtype]
1869+
_reg = {0: "GMT_GRID_NODE_REG", 1: "GMT_GRID_PIXEL_REG"}[
1870+
xrgrid.gmt.registration
1871+
]
1872+
1873+
data = self.create_data(
1874+
family,
1875+
geometry,
1876+
mode=f"GMT_CONTAINER_ONLY|{_gtype}",
1877+
ranges=region,
1878+
inc=inc,
1879+
registration=_reg,
1880+
pad=self["GMT_PAD_DEFAULT"],
1881+
)
1882+
if Version(self._info["version"]) < Version("6.5.0"):
1883+
# Upstream bug fixed in GMT>=6.5.0
1884+
self.set_allocmode(family, data)
1885+
1886+
gmtgrid = ctp.cast(data, ctp.POINTER(_GMT_GRID))
1887+
header = gmtgrid.contents.header.contents
1888+
header.z_min, header.z_max = matrix.min(), matrix.max()
1889+
1890+
matrix = np.pad(matrix, self["GMT_PAD_DEFAULT"]).astype(np.float32)
1891+
gmtgrid.contents.data = matrix.ctypes.data_as(ctp.POINTER(gmt_grdfloat))
1892+
1893+
with self.open_virtualfile(family, geometry, "GMT_IN", gmtgrid) as vfile:
1894+
yield vfile
1895+
18421896
def virtualfile_to_dataset(
18431897
self,
18441898
vfname: str,

0 commit comments

Comments
 (0)