diff --git a/pygmt/clib/session.py b/pygmt/clib/session.py index 1a77659ac75..ec5bdcf3be2 100644 --- a/pygmt/clib/session.py +++ b/pygmt/clib/session.py @@ -26,7 +26,7 @@ vectors_to_arrays, ) from pygmt.clib.loading import load_libgmt -from pygmt.datatypes import _GMT_DATASET, _GMT_GRID +from pygmt.datatypes import _GMT_DATASET, _GMT_GRID, _GMT_IMAGE from pygmt.exceptions import ( GMTCLibError, GMTCLibNoSessionError, @@ -1071,7 +1071,7 @@ def put_matrix(self, dataset, matrix, pad=0): def read_data( self, infile: str, - kind: Literal["dataset", "grid"], + kind: Literal["dataset", "grid", "image"], family: str | None = None, geometry: str | None = None, mode: str = "GMT_READ_NORMAL", @@ -1089,8 +1089,8 @@ def read_data( infile The input file name. kind - The data kind of the input file. Valid values are ``"dataset"`` and - ``"grid"``. + The data kind of the input file. Valid values are ``"dataset"``, ``"grid"`` + and ``"image"``. 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 @@ -1141,6 +1141,7 @@ def read_data( _family, _geometry, dtype = { "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), }[kind] if family is None: family = _family @@ -1797,7 +1798,9 @@ def virtualfile_from_data( @contextlib.contextmanager def virtualfile_out( - self, kind: Literal["dataset", "grid"] = "dataset", fname: str | None = None + self, + kind: Literal["dataset", "grid", "image"] = "dataset", + fname: str | None = None, ) -> Generator[str, None, None]: r""" Create a virtual file or an actual file for storing output data. @@ -1810,8 +1813,8 @@ def virtualfile_out( Parameters ---------- kind - The data kind of the virtual file to create. Valid values are ``"dataset"`` - and ``"grid"``. Ignored if ``fname`` is specified. + The data kind of the virtual file to create. Valid values are ``"dataset"``, + ``"grid"``, and ``"image"``. Ignored if ``fname`` is specified. fname The name of the actual file to write the output data. No virtual file will be created. @@ -1854,8 +1857,10 @@ def virtualfile_out( family, geometry = { "dataset": ("GMT_IS_DATASET", "GMT_IS_PLP"), "grid": ("GMT_IS_GRID", "GMT_IS_SURFACE"), + "image": ("GMT_IS_IMAGE", "GMT_IS_SURFACE"), }[kind] - with self.open_virtualfile(family, geometry, "GMT_OUT", None) as vfile: + direction = "GMT_OUT|GMT_IS_REFERENCE" if kind == "image" else "GMT_OUT" + with self.open_virtualfile(family, geometry, direction, None) as vfile: yield vfile def inquire_virtualfile(self, vfname: str) -> int: @@ -1901,7 +1906,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"`` and ``None``. If ``None``, will return a ctypes void pointer. + ``"grid"``, ``"image"`` and ``None``. If ``None``, will return a ctypes void + pointer. Returns ------- @@ -1951,9 +1957,9 @@ def read_virtualfile( # _GMT_DATASET). if kind is None: # Return the ctypes void pointer return pointer - if kind in {"image", "cube"}: + if kind == "cube": raise NotImplementedError(f"kind={kind} is not supported yet.") - dtype = {"dataset": _GMT_DATASET, "grid": _GMT_GRID}[kind] + dtype = {"dataset": _GMT_DATASET, "grid": _GMT_GRID, "image": _GMT_IMAGE}[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..3489dd19d10 100644 --- a/pygmt/datatypes/__init__.py +++ b/pygmt/datatypes/__init__.py @@ -4,3 +4,4 @@ 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/image.py b/pygmt/datatypes/image.py new file mode 100644 index 00000000000..63c267b1ca0 --- /dev/null +++ b/pygmt/datatypes/image.py @@ -0,0 +1,94 @@ +""" +Wrapper for the GMT_IMAGE data type. +""" + +import ctypes as ctp +from typing import ClassVar + +from pygmt.datatypes.grid import _GMT_GRID_HEADER + + +class _GMT_IMAGE(ctp.Structure): # noqa: N801 + """ + GMT image data structure. + + Examples + -------- + >>> import numpy as np + >>> from pygmt.clib import Session + >>> with Session() as lib: + ... with lib.virtualfile_out(kind="image") as voutimg: + ... lib.call_module("read", ["@earth_day_01d", voutimg, "-Ti"]) + ... # Read the image from the virtual file + ... image = lib.read_virtualfile(vfname=voutimg, kind="image").contents + ... # The image header + ... header = image.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.title) + ... print(header.command) + ... print(header.remark) + ... 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.nan_value, header.xy_off) + ... # Image-specific attributes. + ... print(image.type, image.n_indexed_colors) + ... # The x and y coordinates + ... x = image.x[: header.n_columns] + ... y = image.y[: header.n_rows] + ... # The data array (with paddings) + ... data = np.reshape( + ... image.data[: header.n_bands * header.mx * header.my], + ... (header.my, header.mx, header.n_bands), + ... ) + ... # The data array (without paddings) + ... pad = header.pad[:] + ... data = data[pad[2] : header.my - pad[3], pad[0] : header.mx - pad[1], :] + 180 360 1 + [-180.0, 180.0, -90.0, 90.0] [1.0, 1.0] + 1.0 0.0 + b'x' b'y' b'z' + b'' + b'' + b'' + 64800 66976 0 + 0 3 364 184 + [2, 2, 2, 2] + b'BRPa' 0.0 0.5 + 1 0 + >>> x + [-179.5, -178.5, ..., 178.5, 179.5] + >>> y + [89.5, 88.5, ..., -88.5, -89.5] + >>> data.shape + (180, 360, 3) + >>> data.min(), data.max() + (10, 255) + """ + + _fields_: ClassVar = [ + # Data type, e.g. GMT_FLOAT + ("type", ctp.c_int), + # Array with color lookup values + ("colormap", ctp.POINTER(ctp.c_int)), + # Number of colors in a paletted image + ("n_indexed_colors", ctp.c_int), + # Pointer to full GMT header for the image + ("header", ctp.POINTER(_GMT_GRID_HEADER)), + # Pointer to actual image + ("data", ctp.POINTER(ctp.c_ubyte)), + # Pointer to an optional transparency layer stored in a separate variable + ("alpha", ctp.POINTER(ctp.c_ubyte)), + # Color interpolation + ("color_interp", ctp.c_char_p), + # Pointer to the x-coordinate vector + ("x", ctp.POINTER(ctp.c_double)), + # Pointer to the y-coordinate vector + ("y", ctp.POINTER(ctp.c_double)), + # Book-keeping variables "hidden" from the API + ("hidden", ctp.c_void_p), + ] diff --git a/pygmt/tests/test_clib_read_data.py b/pygmt/tests/test_clib_read_data.py index 43978b291c2..b2f92403df0 100644 --- a/pygmt/tests/test_clib_read_data.py +++ b/pygmt/tests/test_clib_read_data.py @@ -132,6 +132,43 @@ def test_clib_read_data_grid_actual_image(): ) +# Note: Simplify the tests for images after GMT_IMAGE.to_dataarray() is implemented. +def test_clib_read_data_image(): + """ + Test the Session.read_data method for images. + """ + with Session() as lib: + image = lib.read_data("@earth_day_01d_p", kind="image").contents + header = image.header.contents + assert header.n_rows == 180 + assert header.n_columns == 360 + assert header.n_bands == 3 + assert header.wesn[:] == [-180.0, 180.0, -90.0, 90.0] + assert image.data + + +def test_clib_read_data_image_two_steps(): + """ + Test the Session.read_data method for images in two steps, first reading the header + and then the data. + """ + infile = "@earth_day_01d_p" + with Session() as lib: + # Read the header first + data_ptr = lib.read_data(infile, kind="image", mode="GMT_CONTAINER_ONLY") + image = data_ptr.contents + header = image.header.contents + assert header.n_rows == 180 + assert header.n_columns == 360 + assert header.wesn[:] == [-180.0, 180.0, -90.0, 90.0] + assert header.n_bands == 3 # Explicitly check n_bands + assert not image.data # The data is not read yet + + # Read the data + lib.read_data(infile, kind="image", mode="GMT_DATA_ONLY", data=data_ptr) + assert image.data + + def test_clib_read_data_fails(): """ Test that the Session.read_data method raises an exception if there are errors.