From 331f9aa0db21a8255f9cece7ffc773810a4af5ff Mon Sep 17 00:00:00 2001 From: Wei Ji <23487320+weiji14@users.noreply.github.com> Date: Mon, 30 Sep 2024 10:38:57 +1300 Subject: [PATCH 1/3] Initial implementation of virtualfile_from_image Enable passing in 3-band images to GMT via a virtualfile mechanism instead of using a temporary GeoTIFF file which requires rioxarray to be installed. Implementation based on virtualfile_from_grid. Made some adjacent changes around put_matrix to handle ndim==3, though a segfault is happening now. --- pygmt/clib/conversion.py | 5 +-- pygmt/clib/session.py | 77 ++++++++++++++++++++++++++++++++++------ 2 files changed, 70 insertions(+), 12 deletions(-) diff --git a/pygmt/clib/conversion.py b/pygmt/clib/conversion.py index a296601e5a1..13d3d112505 100644 --- a/pygmt/clib/conversion.py +++ b/pygmt/clib/conversion.py @@ -84,9 +84,10 @@ def dataarray_to_matrix(grid): >>> print(inc) [2.0, 2.0] """ - if len(grid.dims) != 2: + if len(grid.dims) not in {2, 3}: raise GMTInvalidInput( - f"Invalid number of grid dimensions '{len(grid.dims)}'. Must be 2." + f"Invalid number of grid/image dimensions '{len(grid.dims)}'. " + "Must be 2 for grid, or 3 for image." ) # Extract region and inc from the grid region = [] diff --git a/pygmt/clib/session.py b/pygmt/clib/session.py index 95ed8c19d76..69dc442eabe 100644 --- a/pygmt/clib/session.py +++ b/pygmt/clib/session.py @@ -29,12 +29,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 @@ -868,6 +863,10 @@ def _check_dtype_and_dim(self, array, ndim): ... gmttype = ses._check_dtype_and_dim(data, ndim=2) ... gmttype == ses["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"] """ # Check that the array has the given number of dimensions if array.ndim != ndim: @@ -1006,9 +1005,9 @@ def put_strings(self, dataset, family, strings): f"Failed to put strings of type {strings.dtype} into dataset" ) - def put_matrix(self, dataset, matrix, pad=0): + def put_matrix(self, dataset, matrix, pad=0, ndim=2): """ - Attach a numpy 2-D 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``. @@ -1048,7 +1047,7 @@ def put_matrix(self, dataset, matrix, pad=0): 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 @@ -1610,6 +1609,64 @@ def virtualfile_from_grid(self, grid): with self.open_virtualfile(*args) as vfile: yield vfile + @contextlib.contextmanager + def virtualfile_from_image(self, image: xr.DataArray): + """ + Store a 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. + 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}", + ranges=region[0:4], # (xmin, xmax, ymin, ymax) only, leave out (zmin, zmax) + inc=inc[0:2], # (x-inc, y-inc) only, leave out z-inc + registration=_reg, + ) + 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): r""" @@ -1796,7 +1853,7 @@ def virtualfile_in( # noqa: PLR0912 "arg": 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, # Note: virtualfile_from_matrix is not used because a matrix can be # converted to vectors instead, and using vectors allows for better From 8c63ae6265768a2599c9a8b4b4632b0217dd836c Mon Sep 17 00:00:00 2001 From: Wei Ji <23487320+weiji14@users.noreply.github.com> Date: Mon, 30 Sep 2024 22:10:50 +1300 Subject: [PATCH 2/3] Update pygmt/clib/session.py Co-authored-by: Dongdong Tian --- pygmt/clib/session.py | 1 + 1 file changed, 1 insertion(+) diff --git a/pygmt/clib/session.py b/pygmt/clib/session.py index 69dc442eabe..87ecfeb2310 100644 --- a/pygmt/clib/session.py +++ b/pygmt/clib/session.py @@ -867,6 +867,7 @@ def _check_dtype_and_dim(self, array, ndim): >>> 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: From 78281533b43b50c0f8715ad6461ed8b7f8940b66 Mon Sep 17 00:00:00 2001 From: Wei Ji <23487320+weiji14@users.noreply.github.com> Date: Mon, 30 Sep 2024 22:18:42 +1300 Subject: [PATCH 3/3] Typo --- pygmt/clib/session.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pygmt/clib/session.py b/pygmt/clib/session.py index 87ecfeb2310..f0a6e1ff6fc 100644 --- a/pygmt/clib/session.py +++ b/pygmt/clib/session.py @@ -1613,7 +1613,7 @@ def virtualfile_from_grid(self, grid): @contextlib.contextmanager def virtualfile_from_image(self, image: xr.DataArray): """ - Store a image in a virtual file. + 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.