diff --git a/pygmt/clib/conversion.py b/pygmt/clib/conversion.py index c7d713647ef..98bef46ba70 100644 --- a/pygmt/clib/conversion.py +++ b/pygmt/clib/conversion.py @@ -91,10 +91,12 @@ def dataarray_to_matrix( >>> print(inc) [2.0, 2.0] """ - if len(grid.dims) != 2: - msg = f"Invalid number of grid dimensions 'len({grid.dims})'. Must be 2." + if len(grid.dims) not in {2, 3}: + msg = ( + f"Invalid number of grid/image dimensions 'len({grid.dims})'. " + "Must be 2 for grid, or 3 for image." + ) raise GMTInvalidInput(msg) - # Extract region and inc from the grid region, inc = [], [] # Reverse the dims because it is rows, columns ordered. In geographic grids, this diff --git a/pygmt/clib/session.py b/pygmt/clib/session.py index 7680e81108e..90ba17bb850 100644 --- a/pygmt/clib/session.py +++ b/pygmt/clib/session.py @@ -25,12 +25,7 @@ from pygmt.clib.loading import get_gmt_version, load_libgmt from pygmt.datatypes import _GMT_DATASET, _GMT_GRID, _GMT_IMAGE from pygmt.exceptions import GMTCLibError, GMTCLibNoSessionError, GMTInvalidInput -from pygmt.helpers import ( - _validate_data_input, - data_kind, - tempfile_from_geojson, - tempfile_from_image, -) +from pygmt.helpers import _validate_data_input, data_kind, tempfile_from_geojson FAMILIES = [ "GMT_IS_DATASET", # Entity is a data table @@ -927,13 +922,18 @@ def _check_dtype_and_dim(self, array: np.ndarray, ndim: int) -> int: ... gmttype = lib._check_dtype_and_dim(data, ndim=2) ... gmttype == lib["GMT_FLOAT"] True + >>> data = np.ones((5, 3, 2), dtype="uint8") + >>> with Session() as ses: + ... gmttype = ses._check_dtype_and_dim(data, ndim=3) + ... gmttype == ses["GMT_UCHAR"] + True """ # Check that the array has the given number of dimensions. if array.ndim != ndim: msg = f"Expected a numpy {ndim}-D array, got {array.ndim}-D." raise GMTInvalidInput(msg) - # 1-D arrays can be numeric or text, 2-D arrays can only be numeric. + # 1-D arrays can be numeric or text, 2-D or 3-D arrays can only be numeric. valid_dtypes = DTYPES if ndim == 1 else DTYPES_NUMERIC if (dtype := array.dtype.type) not in valid_dtypes: msg = f"Unsupported numpy data type '{dtype}'." @@ -1058,10 +1058,10 @@ def put_strings( raise GMTCLibError(msg) def put_matrix( - self, dataset: ctp.c_void_p, matrix: np.ndarray, pad: int = 0 - ) -> None: + self, dataset: ctp.c_void_p, matrix: np.ndarray, pad: int = 0, ndim: int = 2 + ): """ - Attach a 2-D numpy array to a GMT dataset. + Attach a numpy n-D (2-D or 3-D) array to a GMT dataset. Use this function to attach numpy array data to a GMT dataset and pass it to GMT modules. Wraps ``GMT_Put_Matrix``. @@ -1088,6 +1088,9 @@ def put_matrix( pad The amount of padding that should be added to the matrix. Use when creating grids for modules that require padding. + ndim + Number of dimensions in the array, use 2 for grids, or 3 for images. Default + is 2. Raises ------ @@ -1100,7 +1103,7 @@ def put_matrix( restype=ctp.c_int, ) - gmt_type = self._check_dtype_and_dim(matrix, ndim=2) + gmt_type = self._check_dtype_and_dim(matrix, ndim=ndim) matrix_pointer = matrix.ctypes.data_as(ctp.c_void_p) status = c_put_matrix( self.session_pointer, dataset, gmt_type, pad, matrix_pointer @@ -1653,6 +1656,66 @@ def virtualfile_from_grid(self, grid: xr.DataArray) -> Generator[str, None, None ) as vfile: yield vfile + @contextlib.contextmanager + def virtualfile_from_image(self, image: xr.DataArray): + """ + Store an image in a virtual file. + + Use the virtual file name to pass in the data in your image to a GMT module. + Images must be :class:`xarray.DataArray` instances. + + Context manager (use in a ``with`` block). Yields the virtual file name that you + can pass as an argument to a GMT module call. Closes the virtual file upon exit + of the ``with`` block. + + The virtual file will contain the image as a ``GMT_MATRIX`` with extra metadata. + + Use this instead of creating a data container and virtual file by hand with + :meth:`pygmt.clib.Session.create_data`, :meth:`pygmt.clib.Session.put_matrix`, + and :meth:`pygmt.clib.Session.open_virtualfile`. + + The image data matrix must be C contiguous in memory. If it is not (e.g., it is + a slice of a larger array), the array will be copied to make sure it is. + + Parameters + ---------- + image : :class:`xarray.DataArray` + The image that will be included in the virtual file. + + Yields + ------ + fname : str + The name of virtual file. Pass this as a file name argument to a GMT module. + + """ + _gtype = {0: "GMT_GRID_IS_CARTESIAN", 1: "GMT_GRID_IS_GEO"}[image.gmt.gtype] + _reg = {0: "GMT_GRID_NODE_REG", 1: "GMT_GRID_PIXEL_REG"}[image.gmt.registration] + + # Conversion to a C-contiguous array needs to be done here and not in put_matrix + # because we need to maintain a reference to the copy while it is being used by + # the C API. Otherwise, the array would be garbage collected and the memory + # freed. Creating it in this context manager guarantees that the copy will be + # around until the virtual file is closed. The conversion is implicit in + # dataarray_to_matrix. + nbands, nrows, ncols = image.shape # CHW + matrix, _region, _inc = dataarray_to_matrix(image) + + family = "GMT_IS_IMAGE|GMT_VIA_MATRIX" + geometry = "GMT_IS_SURFACE" + gmt_image = self.create_data( + family, + geometry, + mode=f"GMT_CONTAINER_ONLY|{_gtype}", + dim=(ncols, nrows, nbands, self["GMT_UCHAR"]), + # ranges=_region[0:4], # (xmin, xmax, ymin, ymax) only, no (zmin, zmax) + # inc=_inc[0:2], # (x-inc, y-inc) only, leave out z-inc + registration=_reg, # type: ignore[arg-type] + ) + self.put_matrix(gmt_image, matrix, ndim=3) + args = (family, geometry, "GMT_IN|GMT_IS_REFERENCE", gmt_image) + with self.open_virtualfile(*args) as vfile: + yield vfile + @contextlib.contextmanager def virtualfile_from_stringio( self, stringio: io.StringIO @@ -1866,7 +1929,7 @@ def virtualfile_in( # noqa: PLR0912 "file": contextlib.nullcontext, "geojson": tempfile_from_geojson, "grid": self.virtualfile_from_grid, - "image": tempfile_from_image, + "image": self.virtualfile_from_image, "stringio": self.virtualfile_from_stringio, "matrix": self.virtualfile_from_matrix, "vectors": self.virtualfile_from_vectors,