diff --git a/pygmt/clib/session.py b/pygmt/clib/session.py index 9ab3e37574e..a9733be567e 100644 --- a/pygmt/clib/session.py +++ b/pygmt/clib/session.py @@ -23,7 +23,7 @@ vectors_to_arrays, ) from pygmt.clib.loading import get_gmt_version, load_libgmt -from pygmt.datatypes import _GMT_DATASET, _GMT_GRID, _GMT_IMAGE +from pygmt.datatypes import _GMT_CUBE, _GMT_DATASET, _GMT_GRID, _GMT_IMAGE from pygmt.exceptions import GMTCLibError, GMTCLibNoSessionError, GMTInvalidInput from pygmt.helpers import ( _validate_data_input, @@ -1112,7 +1112,7 @@ def put_matrix( def read_data( self, infile: str, - kind: Literal["dataset", "grid", "image"], + kind: Literal["dataset", "grid", "image", "cube"], family: str | None = None, geometry: str | None = None, mode: str = "GMT_READ_NORMAL", @@ -1130,8 +1130,8 @@ def read_data( infile The input file name. kind - The data kind of the input file. Valid values are ``"dataset"``, ``"grid"`` - and ``"image"``. + The data kind of the input file. Valid values are ``"dataset"``, ``"grid"``, + ``"image"`` and ``"cube"``. family A valid GMT data family name (e.g., ``"GMT_IS_DATASET"``). See the ``FAMILIES`` attribute for valid names. If ``None``, will determine the data @@ -1183,6 +1183,7 @@ def read_data( "dataset": ("GMT_IS_DATASET", "GMT_IS_PLP", _GMT_DATASET), "grid": ("GMT_IS_GRID", "GMT_IS_SURFACE", _GMT_GRID), "image": ("GMT_IS_IMAGE", "GMT_IS_SURFACE", _GMT_IMAGE), + "cube": ("GMT_IS_CUBE", "GMT_IS_VOLUME", _GMT_CUBE), }[kind] if family is None: family = _family @@ -1905,7 +1906,7 @@ def virtualfile_in( @contextlib.contextmanager def virtualfile_out( self, - kind: Literal["dataset", "grid", "image"] = "dataset", + kind: Literal["dataset", "grid", "image", "cube"] = "dataset", fname: str | None = None, ) -> Generator[str, None, None]: r""" @@ -1920,7 +1921,7 @@ def virtualfile_out( ---------- kind The data kind of the virtual file to create. Valid values are ``"dataset"``, - ``"grid"``, and ``"image"``. Ignored if ``fname`` is specified. + ``"grid"``, ``"image"`` and ``"cube"``. Ignored if ``fname`` is specified. fname The name of the actual file to write the output data. No virtual file will be created. @@ -1964,6 +1965,7 @@ def virtualfile_out( "dataset": ("GMT_IS_DATASET", "GMT_IS_PLP"), "grid": ("GMT_IS_GRID", "GMT_IS_SURFACE"), "image": ("GMT_IS_IMAGE", "GMT_IS_SURFACE"), + "cube": ("GMT_IS_CUBE", "GMT_IS_VOLUME"), }[kind] direction = "GMT_OUT|GMT_IS_REFERENCE" if kind == "image" else "GMT_OUT" with self.open_virtualfile(family, geometry, direction, None) as vfile: @@ -2012,8 +2014,8 @@ def read_virtualfile( Name of the virtual file to read. kind Cast the data into a GMT data container. Valid values are ``"dataset"``, - ``"grid"``, ``"image"`` and ``None``. If ``None``, will return a ctypes void - pointer. + ``"grid"``, ``"image"``, ``"cube"`` and ``None``. If ``None``, will return + a ctypes void pointer. Returns ------- @@ -2063,10 +2065,12 @@ def read_virtualfile( # _GMT_DATASET). if kind is None: # Return the ctypes void pointer return pointer - if kind == "cube": - msg = f"kind={kind} is not supported yet." - raise NotImplementedError(msg) - dtype = {"dataset": _GMT_DATASET, "grid": _GMT_GRID, "image": _GMT_IMAGE}[kind] + dtype = { + "dataset": _GMT_DATASET, + "grid": _GMT_GRID, + "image": _GMT_IMAGE, + "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 3489dd19d10..2c21418936e 100644 --- a/pygmt/datatypes/__init__.py +++ b/pygmt/datatypes/__init__.py @@ -2,6 +2,7 @@ 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 from pygmt.datatypes.image import _GMT_IMAGE diff --git a/pygmt/datatypes/cube.py b/pygmt/datatypes/cube.py new file mode 100644 index 00000000000..363febd1fc0 --- /dev/null +++ b/pygmt/datatypes/cube.py @@ -0,0 +1,105 @@ +""" +Wrapper for the GMT_CUBE data type. +""" + +import ctypes as ctp +from typing import ClassVar + +from pygmt.datatypes.header import ( + _GMT_GRID_HEADER, + GMT_GRID_UNIT_LEN80, + GMT_GRID_VARNAME_LEN80, + gmt_grdfloat, +) + + +class _GMT_CUBE(ctp.Structure): # noqa: N801 + """ + GMT cube data structure for 3-D data. + + The GMT_CUBE structure is a extension of the GMT_GRID structure to handle 3-D data + cubes. It requires a 2-D grid header and extended parameters for the 3rd dimension. + + header->n_bands is used for the number of layers in 3-D cubes. + + Examples + -------- + >>> import numpy as np + >>> from pygmt import which + >>> from pygmt.clib import Session + >>> cubefile = which("@cube.nc", download="c") + >>> with Session() as lib: + ... with lib.virtualfile_out(kind="cube") as voutcube: + ... lib.call_module("read", [cubefile, voutcube, "-Tu", "-Vd"]) + ... # Read the cube from the virtual file + ... cube = lib.read_virtualfile(vfname=voutcube, kind="cube").contents + ... # The cube header + ... header = cube.header.contents + ... # Access the header properties + ... print(header.n_rows, header.n_columns, header.registration) + ... print(header.wesn[:], header.inc[:]) + ... print(header.z_scale_factor, header.z_add_offset) + ... print(header.x_units, header.y_units, header.z_units) + ... print(header.nm, header.size, header.complex_mode) + ... print(header.type, header.n_bands, header.mx, header.my) + ... print(header.pad[:]) + ... print(header.mem_layout, header.xy_off) + ... # Cube-specific attributes. + ... print(cube.mode, cube.z_range[:], cube.z_inc, cube.name, cube.units) + ... # The x, y, and z coordinates + ... x = np.ctypeslib.as_array(cube.x, shape=(header.n_columns,)).copy() + ... y = np.ctypeslib.as_array(cube.y, shape=(header.n_rows,)).copy() + ... z = np.ctypeslib.as_array(cube.z, shape=(header.n_bands,)).copy() + ... # The data array (with paddings) + ... data = np.ctypeslib.as_array( + ... cube.data, shape=(header.my, header.mx, header.n_bands) + ... ).copy() + ... # The data array (without paddings) + ... pad = header.pad[:] + ... data = data[pad[2] : header.my - pad[3], pad[0] : header.mx - pad[1], :] + 11 11 0 + [0.0, 10.0, 0.0, 10.0] [1.0, 1.0] + 1.0 0.0 + b'x' b'y' b'cube' + 121 226 0 + 18 4 15 15 + [2, 2, 2, 2] + b'' 0.0 + 0 [1.0, 5.0] 0.0 b'' b'z' + >>> x + array([ 0., 1., 2., 3., 4., 5., 6., 7., 8., 9., 10.]) + >>> y + array([10., 9., 8., 7., 6., 5., 4., 3., 2., 1., 0.]) + >>> z + array([1., 2., 3., 5.]) + >>> data.shape + (11, 11, 4) + >>> # data.min(), data.max() # The min/max are wrong. Upstream bug? + >>> # (-29.4, 169.4) + """ + + _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), + # mode=GMT_CUBE_IS_STACK means the input dataset was a list of 2-D grids, rather + # than a single cube. + ("mode", ctp.c_uint), + # Minimum/maximum 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), + ] diff --git a/pygmt/helpers/caching.py b/pygmt/helpers/caching.py index ea6bed8d4cf..2a218e974f4 100644 --- a/pygmt/helpers/caching.py +++ b/pygmt/helpers/caching.py @@ -87,6 +87,7 @@ def cache_data() -> None: "@Table_5_11_mean.xyz", "@capitals.gmt", "@circuit.png", + "@cube.nc", "@earth_relief_20m_holes.grd", "@fractures_06.txt", "@hotspots.txt", diff --git a/pygmt/tests/test_clib_read_data.py b/pygmt/tests/test_clib_read_data.py index bfca1577ccc..2aad524858d 100644 --- a/pygmt/tests/test_clib_read_data.py +++ b/pygmt/tests/test_clib_read_data.py @@ -201,6 +201,24 @@ def test_clib_read_data_image_two_steps(expected_xrimage): xr.testing.assert_equal(xrimage, expected_xrimage) +def test_clib_read_data_cube_actual_grid(): + """ + Test the Session.read_data method for cube, but actually the file is a grid. + """ + with Session() as lib: + with pytest.raises(GMTCLibError): + lib.read_data("@earth_relief_01d_p", kind="cube", mode="GMT_CONTAINER_ONLY") + + +def test_clib_read_data_cube_actual_image(): + """ + Test the Session.read_data method for cube, but actually the file is an image. + """ + with Session() as lib: + with pytest.raises(GMTCLibError): + lib.read_data("@earth_day_01d", kind="cube", mode="GMT_CONTAINER_ONLY") + + def test_clib_read_data_fails(): """ Test that the Session.read_data method raises an exception if there are errors.