diff --git a/pygmt/clib/session.py b/pygmt/clib/session.py index 1a699596eaa..d219af9c837 100644 --- a/pygmt/clib/session.py +++ b/pygmt/clib/session.py @@ -22,7 +22,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, @@ -1130,7 +1130,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", @@ -1148,8 +1148,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 @@ -1201,6 +1201,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 @@ -1904,7 +1905,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""" @@ -1919,7 +1920,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. @@ -1963,6 +1964,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: @@ -2027,8 +2029,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 ------- @@ -2078,10 +2080,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..27b65373f7a --- /dev/null +++ b/pygmt/datatypes/cube.py @@ -0,0 +1,257 @@ +""" +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): # ruff: ignore[invalid-class-name] + """ + 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. + + The ``data`` array is a stack of 2-D padded layers, i.e., layer ``k`` starts at + offset ``k * header.size``. Note that ``header.size`` is the allocated length of one + padded layer, which can exceed ``header.my * header.mx``. + + Examples + -------- + >>> import numpy as np + >>> from pygmt import which + >>> from pygmt.clib import Session + >>> from pygmt.datatypes import _GMT_CUBE + >>> 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"]) + ... # Read the cube from the virtual file + ... cube = lib.read_virtualfile(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 (one padded layer per row) + ... data = np.ctypeslib.as_array( + ... cube.data, shape=(header.n_bands, header.size) + ... ).copy() + ... # Reshape the layers to 2-D and strip the paddings + ... pad = header.pad[:] + ... # header.size can exceed header.my * header.mx because GMT rounds up to + ... # an even count. So we need to slice the data array to the actual size. + ... data = data[:, : header.my * header.mx].reshape( + ... header.n_bands, header.my, header.mx + ... ) + ... 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 + (4, 11, 11) + >>> # The first layer is a grid of X * Y. + >>> print(data[0, :, :]) + [[ 0. 10. 20. 30. 40. 50. 60. 70. 80. 90. 100.] + [ 0. 9. 18. 27. 36. 45. 54. 63. 72. 81. 90.] + [ 0. 8. 16. 24. 32. 40. 48. 56. 64. 72. 80.] + [ 0. 7. 14. 21. 28. 35. 42. 49. 56. 63. 70.] + [ 0. 6. 12. 18. 24. 30. 36. 42. 48. 54. 60.] + [ 0. 5. 10. 15. 20. 25. 30. 35. 40. 45. 50.] + [ 0. 4. 8. 12. 16. 20. 24. 28. 32. 36. 40.] + [ 0. 3. 6. 9. 12. 15. 18. 21. 24. 27. 30.] + [ 0. 2. 4. 6. 8. 10. 12. 14. 16. 18. 20.] + [ 0. 1. 2. 3. 4. 5. 6. 7. 8. 9. 10.] + [ 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]] + >>> # Other layers are the same, but scaled by 1.0, 1.1, 1.2, and 1.4, respectively. + >>> np.allclose(data[1, :, :], 1.1 * data[0, :, :]) + True + >>> np.allclose(data[2, :, :], 1.2 * data[0, :, :]) + True + >>> np.allclose(data[3, :, :], 1.4 * data[0, :, :]) + True + """ + + _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), + ] + + def _parse_z_dimension(self) -> tuple[str, dict]: + """ + Get the name and attributes of the 3rd dimension. + + Unlike the x/y dimensions, the 3rd dimension is described by the cube itself + rather than by the 2-D grid header. ``self.units`` holds the dimension's + "long_name [units]" string, and ``self.z_range`` its actual range. + """ + attrs: dict = {} + long_name, units = _parse_nameunits(self.units.decode()) + if long_name: + attrs["long_name"] = long_name + if units: + attrs["units"] = units + attrs["axis"] = "Z" + attrs["actual_range"] = np.array(self.z_range[:]) + return "z", attrs + + def to_xarray(self) -> xr.DataArray: + """ + Convert a _GMT_CUBE object to a :class:`xarray.DataArray` object. + + Returns + ------- + dataarray + A 3-D :class:`xr.DataArray` object with dimensions ordered as (z, y, x). + + 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"]) + ... # Read the cube from the virtual file + ... cube = lib.read_virtualfile(voutcube, kind="cube") + ... # Convert to xarray.DataArray and use it later + ... da = cube.contents.to_xarray() + >>> da.name, da.dims, da.shape + ('cube', ('z', 'y', 'x'), (4, 11, 11)) + >>> da.coords["x"] + Size: 88B + array([ 0., 1., 2., 3., 4., 5., 6., 7., 8., 9., 10.]) + Coordinates: + * x (x) float64 88B 0.0 1.0 2.0 3.0 4.0 5.0 6.0 7.0 8.0 9.0 10.0 + Attributes: + long_name: x + axis: X + actual_range: [ 0. 10.] + >>> da.coords["y"] + Size: 88B + array([ 0., 1., 2., 3., 4., 5., 6., 7., 8., 9., 10.]) + Coordinates: + * y (y) float64 88B 0.0 1.0 2.0 3.0 4.0 5.0 6.0 7.0 8.0 9.0 10.0 + Attributes: + long_name: y + axis: Y + actual_range: [ 0. 10.] + >>> da.coords["z"] + Size: 32B + array([1., 2., 3., 5.]) + Coordinates: + * z (z) float64 32B 1.0 2.0 3.0 5.0 + Attributes: + long_name: z + axis: Z + actual_range: [1. 5.] + >>> da.gmt.registration, da.gmt.gtype + (, ) + """ + # The cube header + header = self.header.contents + + # Get x/y dimensions and their attributes from the header. + dims, dim_attrs = header.dims, header.dim_attrs + # Get the 3rd dimension and its attributes from the cube itself. + zdim, zdim_attrs = self._parse_z_dimension() + + # The coordinates, given as a tuple of the form (dims, data, attrs) + x = np.ctypeslib.as_array(self.x, shape=(header.n_columns,)).copy() + y = np.ctypeslib.as_array(self.y, shape=(header.n_rows,)).copy() + z = np.ctypeslib.as_array(self.z, shape=(header.n_bands,)).copy() + coords = [ + (zdim, z, zdim_attrs), + (dims[0], y, dim_attrs[0]), + (dims[1], x, dim_attrs[1]), + ] + + # The cube is a stack of 2-D padded layers, i.e., layer k starts at offset + # k * header.size, which can exceed header.my * header.mx. + data = np.ctypeslib.as_array( + self.data, shape=(header.n_bands, header.size) + ).copy() + data = data[:, : header.my * header.mx].reshape( + header.n_bands, header.my, header.mx + ) + # The data array without paddings. + pad = header.pad[:] + data = data[:, pad[2] : header.my - pad[3], pad[0] : header.mx - pad[1]] + + # Create the xarray.DataArray object. + # The cube name is stored in the header's z_units attribute. + cube = xr.DataArray( + data, coords=coords, name=header.z_units.decode(), attrs=header.data_attrs + ) + + # Flip the coordinates and data if necessary so that coordinates are ascending. + # `cube.sortby(list(cube.dims))` sometimes causes crashes. + # The solution comes from https://github.com/pydata/xarray/discussions/6695. + for dim in cube.dims: + if cube[dim].size > 1 and cube[dim][0] > cube[dim][1]: + cube = cube.isel({dim: slice(None, None, -1)}) + + # Set GMT accessors. + # Must put at the end, otherwise info gets lost after certain grid operations. + cube.gmt.registration = header.registration + cube.gmt.gtype = header.gtype + return cube diff --git a/pygmt/datatypes/header.py b/pygmt/datatypes/header.py index afd60c08c28..10c12903296 100644 --- a/pygmt/datatypes/header.py +++ b/pygmt/datatypes/header.py @@ -108,8 +108,10 @@ class _GMT_GRID_HEADER(ctp.Structure): # ruff: ignore[invalid-class-name] # Below are items used internally by GMT # Number of data points (n_columns * n_rows) [paddings are excluded] ("nm", ctp.c_size_t), - # Actual number of items (not bytes) required to hold this grid (mx * my), - # per band (for images) + # Actual number of items (not bytes) required to hold one layer of this grid, + # paddings included. size >= mx * my, rounded up to an even number. + # For images and cubes this is also the stride between bands/layers, i.e., + # band/layer k starts at offset k * size. ("size", ctp.c_size_t), # Bits per data value (e.g., 32 for ints/floats; 8 for bytes). # Only used for ERSI ArcInfo ASCII Exchange grids. @@ -121,7 +123,7 @@ class _GMT_GRID_HEADER(ctp.Structure): # ruff: ignore[invalid-class-name] ("complex_mode", ctp.c_uint), # Grid format ("type", ctp.c_uint), - # Number of bands [1]. Used with GMT_IMAGE containers + # Number of bands for GMT_IMAGE or number of layers for GMT_CUBE [1]. ("n_bands", ctp.c_uint), # Actual x-dimension in memory. mx = n_columns + pad[0] + pad[1] ("mx", ctp.c_uint), diff --git a/pygmt/helpers/caching.py b/pygmt/helpers/caching.py index e315a0eddcc..0246d25b121 100644 --- a/pygmt/helpers/caching.py +++ b/pygmt/helpers/caching.py @@ -121,6 +121,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 66d9a7468dc..ea189b79cf9 100644 --- a/pygmt/tests/test_clib_read_data.py +++ b/pygmt/tests/test_clib_read_data.py @@ -43,6 +43,16 @@ def fixture_expected_xrimage(): return None +@pytest.fixture(scope="module", name="expected_xrcube") +def fixture_expected_xrcube(): + """ + The expected xr.DataArray object for the @cube.nc file. + """ + return xr.load_dataarray( + which("@cube.nc", download="c"), engine="gmt", raster_kind="cube" + ) + + def test_clib_read_data_dataset(): """ Test the Session.read_data method for datasets. @@ -202,6 +212,55 @@ def test_clib_read_data_image_two_steps(expected_xrimage): xr.testing.assert_equal(xrimage, expected_xrimage) +def test_clib_read_data_cube(expected_xrcube): + """ + Test the Session.read_data method for cubes. + """ + infile = which("@cube.nc", download="c") + with Session() as lib: + cube = lib.read_data(infile, kind="cube").contents + xrcube = cube.to_xarray() + xr.testing.assert_equal(xrcube, expected_xrcube) + + +def test_clib_read_data_cube_two_steps(expected_xrcube): + """ + Test the Session.read_data method for cubes in two steps, first reading the header + and then the data. + """ + infile = which("@cube.nc", download="c") + with Session() as lib: + # Read the header first + data_ptr = lib.read_data(infile, kind="cube", mode="GMT_CONTAINER_ONLY") + cube = data_ptr.contents + header = cube.header.contents + assert header.n_rows == 11 + assert header.n_columns == 11 + assert header.n_bands == 4 + assert header.wesn[:] == [0.0, 10.0, 0.0, 10.0] + assert header.z_min == 0.0 + assert header.z_max == 140.0 + assert not cube.data # The data is not read yet + + # Read the data + lib.read_data(infile, kind="cube", mode="GMT_DATA_ONLY", data=data_ptr) + + # Full check + xrcube = data_ptr.contents.to_xarray() + xr.testing.assert_equal(xrcube, expected_xrcube) + + +@pytest.mark.parametrize("infile", ["@earth_relief_01d_p", "@earth_day_01d"]) +def test_clib_read_data_cube_actual_grid_or_image(infile): + """ + Test the Session.read_data method for cube, but actually the file is a grid or an + image. + """ + with Session() as lib: + with pytest.raises(GMTCLibError): + lib.read_data(infile, 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. diff --git a/pygmt/xarray/backend.py b/pygmt/xarray/backend.py index cc5a0f6322c..a1ceb9582d7 100644 --- a/pygmt/xarray/backend.py +++ b/pygmt/xarray/backend.py @@ -115,7 +115,7 @@ def open_dataset( # type: ignore[override] filename_or_obj: PathLike, *, drop_variables=None, # ruff: ignore[unused-method-argument] - raster_kind: Literal["grid", "image"], + raster_kind: Literal["grid", "image", "cube"], region: Sequence[float] | str | None = None, # other backend specific keyword arguments # `chunks` and `cache` DO NOT go here, they are handled by xarray @@ -130,19 +130,25 @@ def open_dataset( # type: ignore[override] that can be read by GMT via the netCDF or GDAL C libraries. See also :gmt-docs:`reference/features.html#grid-file-format`. raster_kind - Whether to read the file as a "grid" (single-band) or "image" (multi-band). + Whether to read the file as a "grid" (single-band), "image" (multi-band), or + "cube" (stacks of 2-D grids). region The subregion of the grid or image to load, in the form of a sequence [*xmin*, *xmax*, *ymin*, *ymax*] or an ISO country code. """ - if raster_kind not in {"grid", "image"}: + if raster_kind not in {"grid", "image", "cube"}: raise GMTValueError( - raster_kind, description="raster kind", choices=["grid", "image"] + raster_kind, + description="raster kind", + choices=["grid", "image", "cube"], ) with Session() as lib: with lib.virtualfile_out(kind=raster_kind) as voutfile: - kwdict = {"R": region, "T": {"grid": "g", "image": "i"}[raster_kind]} + kwdict = { + "R": region, + "T": {"grid": "g", "image": "i", "cube": "u"}[raster_kind], + } lib.call_module( module="read", args=[filename_or_obj, voutfile, *build_arg_list(kwdict)],