diff --git a/pygmt/clib/session.py b/pygmt/clib/session.py index 8a8b52df8e5..2ddca449490 100644 --- a/pygmt/clib/session.py +++ b/pygmt/clib/session.py @@ -25,7 +25,7 @@ vectors_to_arrays, ) from pygmt.clib.loading import load_libgmt -from pygmt.datatypes import _GMT_DATASET, _GMT_GRID +from pygmt.datatypes import _GMT_CUBE, _GMT_DATASET, _GMT_GRID from pygmt.exceptions import ( GMTCLibError, GMTCLibNoSessionError, @@ -1649,7 +1649,9 @@ def virtualfile_in( # noqa: PLR0912 @contextlib.contextmanager def virtualfile_out( - self, kind: Literal["dataset", "grid"] = "dataset", fname: str | None = None + self, + kind: Literal["dataset", "grid", "cube"] = "dataset", + fname: str | None = None, ): r""" Create a virtual file or an actual file for storing output data. @@ -1706,6 +1708,7 @@ def virtualfile_out( family, geometry = { "dataset": ("GMT_IS_DATASET", "GMT_IS_PLP"), "grid": ("GMT_IS_GRID", "GMT_IS_SURFACE"), + "cube": ("GMT_IS_CUBE", "GMT_IS_VOLUME"), }[kind] with self.open_virtualfile(family, geometry, "GMT_OUT", None) as vfile: yield vfile @@ -1801,9 +1804,9 @@ def read_virtualfile( # _GMT_DATASET). if kind is None: # Return the ctypes void pointer return pointer - if kind in ["image", "cube"]: + if kind == "image": raise NotImplementedError(f"kind={kind} is not supported yet.") - dtype = {"dataset": _GMT_DATASET, "grid": _GMT_GRID}[kind] + dtype = {"dataset": _GMT_DATASET, "grid": _GMT_GRID, "cube": _GMT_CUBE}[kind] return ctp.cast(pointer, ctp.POINTER(dtype)) def virtualfile_to_dataset( diff --git a/pygmt/datatypes/__init__.py b/pygmt/datatypes/__init__.py index 237a050a9f7..16627ed5798 100644 --- a/pygmt/datatypes/__init__.py +++ b/pygmt/datatypes/__init__.py @@ -2,5 +2,6 @@ Wrappers for GMT data types. """ +from pygmt.datatypes.cube import _GMT_CUBE from pygmt.datatypes.dataset import _GMT_DATASET from pygmt.datatypes.grid import _GMT_GRID diff --git a/pygmt/datatypes/cube.py b/pygmt/datatypes/cube.py new file mode 100644 index 00000000000..b51341c36ab --- /dev/null +++ b/pygmt/datatypes/cube.py @@ -0,0 +1,93 @@ +""" +Wrapper for the GMT_CUBE data type. +""" + +import ctypes as ctp +from typing import ClassVar + +import numpy as np +import xarray as xr +from pygmt.datatypes.header import ( + _GMT_GRID_HEADER, + GMT_GRID_UNIT_LEN80, + GMT_GRID_VARNAME_LEN80, + _parse_nameunits, + gmt_grdfloat, +) + + +class _GMT_CUBE(ctp.Structure): # noqa: N801 + """ + GMT cube data structure for 3D data. + """ + + _fields_: ClassVar = [ + # Pointer to full GMT 2-D header for a layer (common to all layers) + ("header", ctp.POINTER(_GMT_GRID_HEADER)), + # Pointer to the gmt_grdfloat 3-D cube - a stack of 2-D padded grids + ("data", ctp.POINTER(gmt_grdfloat)), + # Vector of x coordinates common to all layers + ("x", ctp.POINTER(ctp.c_double)), + # Vector of y coordinates common to all layers + ("y", ctp.POINTER(ctp.c_double)), + # Low-level information for GMT use only + ("hidden", ctp.c_void_p), + # GMT_CUBE_IS_STACK if input dataset was a list of 2-D grids rather than a + # single cube + ("mode", ctp.c_uint), + # Minimum/max z values (complements header->wesn[4]) + ("z_range", ctp.c_double * 2), + # z increment (complements inc[2]) (0 if variable z spacing) + ("z_inc", ctp.c_double), + # Array of z values (complements x, y) + ("z", ctp.POINTER(ctp.c_double)), + # Name of the 3-D variable, if read from file (or empty if just one) + ("name", ctp.c_char * GMT_GRID_VARNAME_LEN80), + # Units in 3rd direction (complements x_units, y_units, z_units) + ("units", ctp.c_char * GMT_GRID_UNIT_LEN80), + ] + + def to_dataarray(self): + """ + Convert the GMT_CUBE to an xarray.DataArray. + + Returns + ------- + xarray.DataArray: The data array representation of the GMT_CUBE. + """ + # The grid header + header = self.header.contents + + name = "cube" + # Dimensions and attributes + dims = header.dims + dim_attrs = header.dim_attrs + + # Patch for the 3rd dimension + dims.append("z") + z_attrs = {"actual_range": np.array(self.z_range[:]), "axis": "Z"} + long_name, units = _parse_nameunits(self.units.decode()) + if long_name: + z_attrs["long_name"] = long_name + if units: + z_attrs["units"] = units + dim_attrs.append(z_attrs) + + # The coordinates, given as a tuple of the form (dims, data, attrs) + coords = [ + (dims[0], self.y[: header.n_rows], dim_attrs[0]), + (dims[1], self.x[: header.n_columns], dim_attrs[1]), + # header->n_bands is used for the number of layers for 3-D cubes + (dims[2], self.z[: header.n_bands], dim_attrs[1]), + ] + + # The data array without paddings + pad = header.pad[:] + data = np.reshape( + self.data[: header.mx * header.my * header.n_bands], + (header.my, header.mx, header.n_bands), + )[pad[2] : header.my - pad[3], pad[0] : header.mx - pad[1], :] + + # Create the xarray.DataArray object + grid = xr.DataArray(data, coords=coords, name=name, attrs=header.dataA_attrs) + return grid