diff --git a/pygmt/datatypes/grid.py b/pygmt/datatypes/grid.py index 4be2599afd5..18ad679aeef 100644 --- a/pygmt/datatypes/grid.py +++ b/pygmt/datatypes/grid.py @@ -139,8 +139,8 @@ def to_dataarray(self) -> xr.DataArray: title: Produced by grdcut history: grdcut @earth_relief_01d_p -R-55/-47/-24/-10 -Gstatic_ea... description: Reduced by Gaussian Cartesian filtering (111.2 km fullwi... - long_name: elevation (m) actual_range: [190. 981.] + long_name: elevation (m) >>> da.coords["lon"] # doctest: +NORMALIZE_WHITESPACE, +ELLIPSIS ... array([-54.5, -53.5, -52.5, -51.5, -50.5, -49.5, -48.5, -47.5]) diff --git a/pygmt/datatypes/header.py b/pygmt/datatypes/header.py index d7806e50a24..197eaa26e45 100644 --- a/pygmt/datatypes/header.py +++ b/pygmt/datatypes/header.py @@ -210,17 +210,17 @@ def data_attrs(self) -> dict[str, Any]: GridFormat.NI, GridFormat.NF, GridFormat.ND, - }: # Only set the 'Conventions' attribute for netCDF. + }: # Set attributes specific to CF-1.7 conventions attrs["Conventions"] = "CF-1.7" - attrs["title"] = self.title.decode() - attrs["history"] = self.command.decode() - attrs["description"] = self.remark.decode() + attrs["title"] = self.title.decode() + attrs["history"] = self.command.decode() + attrs["description"] = self.remark.decode() + attrs["actual_range"] = np.array([self.z_min, self.z_max]) long_name, units = _parse_nameunits(self.z_units.decode()) if long_name: attrs["long_name"] = long_name if units: attrs["units"] = units - attrs["actual_range"] = np.array([self.z_min, self.z_max]) return attrs @property @@ -250,6 +250,16 @@ def gtype(self) -> int: "lon"/"lat" or have units "degrees_east"/"degrees_north", then the grid is assumed to be geographic. """ + gtype = 0 # Cartesian by default + dims = self.dims - gtype = 1 if dims[0] == "lat" and dims[1] == "lon" else 0 + if dims[0] == "lat" and dims[1] == "lon": + # Check dimensions for grids that following CF-conventions + gtype = 1 + elif self.ProjRefPROJ4 is not None: + # Check ProjRefPROJ4 for images imported via GDAL. + # The logic comes from GMT's `gmtlib_read_image_info` function. + projref = self.ProjRefPROJ4.decode() + if "longlat" in projref or "latlong" in projref: + gtype = 1 return gtype diff --git a/pygmt/datatypes/image.py b/pygmt/datatypes/image.py index aaccc75cdf5..743cc003681 100644 --- a/pygmt/datatypes/image.py +++ b/pygmt/datatypes/image.py @@ -5,7 +5,9 @@ import ctypes as ctp from typing import ClassVar -from pygmt.datatypes.grid import _GMT_GRID_HEADER +import numpy as np +import xarray as xr +from pygmt.datatypes.header import _GMT_GRID_HEADER class _GMT_IMAGE(ctp.Structure): # noqa: N801 @@ -63,6 +65,8 @@ class _GMT_IMAGE(ctp.Structure): # noqa: N801 array([-179.5, -178.5, ..., 178.5, 179.5]) >>> y # doctest: +NORMALIZE_WHITESPACE, +ELLIPSIS array([ 89.5, 88.5, ..., -88.5, -89.5]) + >>> data.dtype + dtype('uint8') >>> data.shape (180, 360, 3) >>> data.min(), data.max() @@ -91,3 +95,118 @@ class _GMT_IMAGE(ctp.Structure): # noqa: N801 # Book-keeping variables "hidden" from the API ("hidden", ctp.c_void_p), ] + + def to_dataarray(self) -> xr.DataArray: + """ + Convert a _GMT_IMAGE object to an :class:`xarray.DataArray` object. + + Returns + ------- + dataarray + A :class:`xarray.DataArray` object. + + Examples + -------- + >>> from pygmt.clib import Session + >>> with Session() as lib: + ... with lib.virtualfile_out(kind="image") as voutimg: + ... lib.call_module("read", ["@earth_day_01d_p", voutimg, "-Ti"]) + ... # Read the image from the virtual file + ... image = lib.read_virtualfile(voutimg, kind="image") + ... # Convert to xarray.DataArray and use it later + ... da = image.contents.to_dataarray() + >>> da # doctest: +NORMALIZE_WHITESPACE, +ELLIPSIS + ... + array([[[ 10, 10, 10, ..., 10, 10, 10], + [ 10, 10, 10, ..., 10, 10, 10], + [ 10, 10, 10, ..., 10, 10, 10], + ..., + [192, 193, 193, ..., 193, 192, 191], + [204, 206, 206, ..., 205, 206, 204], + [208, 210, 210, ..., 210, 210, 208]], + + [[ 10, 10, 10, ..., 10, 10, 10], + [ 10, 10, 10, ..., 10, 10, 10], + [ 10, 10, 10, ..., 10, 10, 10], + ..., + [186, 187, 188, ..., 187, 186, 185], + [196, 198, 198, ..., 197, 197, 196], + [199, 201, 201, ..., 201, 202, 199]], + + [[ 51, 51, 51, ..., 51, 51, 51], + [ 51, 51, 51, ..., 51, 51, 51], + [ 51, 51, 51, ..., 51, 51, 51], + ..., + [177, 179, 179, ..., 178, 177, 177], + [185, 187, 187, ..., 187, 186, 185], + [189, 191, 191, ..., 191, 191, 189]]], dtype=uint8) + Coordinates: + * y (y) float64... 89.5 88.5 87.5 86.5 ... -86.5 -87.5 -88.5 -89.5 + * x (x) float64... -179.5 -178.5 -177.5 -176.5 ... 177.5 178.5 179.5 + * band (band) uint8... 1 2 3 + Attributes: + long_name: z + + >>> da.coords["x"] # doctest: +NORMALIZE_WHITESPACE, +ELLIPSIS + ... + array([-179.5, -178.5, -177.5, ..., 177.5, 178.5, 179.5]) + Coordinates: + * x (x) float64... -179.5 -178.5 -177.5 -176.5 ... 177.5 178.5 179.5 + Attributes: + long_name: x + axis: X + actual_range: [-180. 180.] + >>> da.coords["y"] # doctest: +NORMALIZE_WHITESPACE, +ELLIPSIS + ... + array([ 89.5, 88.5, 87.5, 86.5, ..., -87.5, -88.5, -89.5]) + Coordinates: + * y (y) float64... 89.5 88.5 87.5 86.5 ... -86.5 -87.5 -88.5 -89.5 + Attributes: + long_name: y + axis: Y + actual_range: [-90. 90.] + >>> da.gmt.registration, da.gmt.gtype + (1, 1) + """ + # The image header + header = self.header.contents + + if header.n_bands != 3: + msg = ( + f"The raster image has {header.n_bands} band(s). " + "Currently only 3-band images are supported. " + "Please consider submitting a feature request to us. " + ) + raise NotImplementedError(msg) + + # Get dimensions and their attributes from the header. + dims, dim_attrs = header.dims, header.dim_attrs + # 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() + coords = [ + (dims[0], y, dim_attrs[0]), + (dims[1], x, dim_attrs[1]), + ("band", np.array([1, 2, 3], dtype=np.uint8), None), + ] + + # Get DataArray without padding + data = np.ctypeslib.as_array( + self.data, shape=(header.my, header.mx, header.n_bands) + ).copy() + pad = header.pad[:] + data = data[pad[2] : header.my - pad[3], pad[0] : header.mx - pad[1], :] + + # Create the xarray.DataArray object + image = xr.DataArray( + data=data, + coords=coords, + name=header.name, + attrs=header.data_attrs, + ).transpose("band", dims[0], dims[1]) + + # Set GMT accessors. + # Must put at the end, otherwise info gets lost after certain image operations. + image.gmt.registration = header.registration + image.gmt.gtype = header.gtype + return image diff --git a/pygmt/tests/test_clib_read_data.py b/pygmt/tests/test_clib_read_data.py index b2f92403df0..fd34fbced9f 100644 --- a/pygmt/tests/test_clib_read_data.py +++ b/pygmt/tests/test_clib_read_data.py @@ -4,6 +4,7 @@ from pathlib import Path +import numpy as np import pandas as pd import pytest import xarray as xr @@ -14,7 +15,7 @@ from pygmt.src import which try: - import rioxarray # noqa: F401 + import rioxarray _HAS_RIOXARRAY = True except ImportError: @@ -29,6 +30,18 @@ def fixture_expected_xrgrid(): return load_dataarray(which("@static_earth_relief.nc")) +@pytest.fixture(scope="module", name="expected_xrimage") +def fixture_expected_xrimage(): + """ + The expected xr.DataArray object for the @earth_day_01d_p file. + """ + if _HAS_RIOXARRAY: + with rioxarray.open_rasterio(which("@earth_day_01d_p")) as da: + dataarray = da.load().drop_vars("spatial_ref") + return dataarray + return None + + def test_clib_read_data_dataset(): """ Test the Session.read_data method for datasets. @@ -98,56 +111,63 @@ def test_clib_read_data_grid_two_steps(expected_xrgrid): # Read the data lib.read_data(infile, kind="grid", mode="GMT_DATA_ONLY", data=data_ptr) + + # Full check xrgrid = data_ptr.contents.to_dataarray() xr.testing.assert_equal(xrgrid, expected_xrgrid) -def test_clib_read_data_grid_actual_image(): +def test_clib_read_data_grid_actual_image(expected_xrimage): """ Test the Session.read_data method for grid, but actually the file is an image. """ with Session() as lib: - data_ptr = lib.read_data( - "@earth_day_01d_p", kind="grid", mode="GMT_CONTAINER_AND_DATA" - ) - 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] + image = lib.read_data("@earth_day_01d_p", kind="grid").contents # Explicitly check n_bands. Only one band is read for 3-band images. - assert header.n_bands == 1 + assert image.header.contents.n_bands == 1 + + xrimage = image.to_dataarray() + assert xrimage.shape == (180, 360) + assert xrimage.coords["x"].data.min() == -179.5 + assert xrimage.coords["x"].data.max() == 179.5 + assert xrimage.coords["y"].data.min() == -89.5 + assert xrimage.coords["y"].data.max() == 89.5 + assert xrimage.data.min() == 10.0 + assert xrimage.data.max() == 255.0 + # Data are stored as uint8 in images but are converted to float32 when reading + # into a GMT_GRID container. + assert xrimage.data.dtype == np.float32 if _HAS_RIOXARRAY: # Full check if rioxarray is installed. - xrimage = image.to_dataarray() - expected_xrimage = xr.open_dataarray( - which("@earth_day_01d_p"), engine="rasterio" - ) assert expected_xrimage.band.size == 3 # 3-band image. xr.testing.assert_equal( xrimage, - expected_xrimage.isel(band=0) - .drop_vars(["band", "spatial_ref"]) - .sortby("y"), + expected_xrimage.isel(band=0).drop_vars(["band"]).sortby("y"), ) -# Note: Simplify the tests for images after GMT_IMAGE.to_dataarray() is implemented. -def test_clib_read_data_image(): +def test_clib_read_data_image(expected_xrimage): """ 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 + + xrimage = image.to_dataarray() + assert xrimage.shape == (3, 180, 360) + assert xrimage.coords["x"].data.min() == -179.5 + assert xrimage.coords["x"].data.max() == 179.5 + assert xrimage.coords["y"].data.min() == -89.5 + assert xrimage.coords["y"].data.max() == 89.5 + assert xrimage.data.min() == 10 + assert xrimage.data.max() == 255 + assert xrimage.data.dtype == np.uint8 + + if _HAS_RIOXARRAY: # Full check if rioxarray is installed. + xr.testing.assert_equal(xrimage, expected_xrimage) -def test_clib_read_data_image_two_steps(): +def test_clib_read_data_image_two_steps(expected_xrimage): """ Test the Session.read_data method for images in two steps, first reading the header and then the data. @@ -166,7 +186,19 @@ def test_clib_read_data_image_two_steps(): # Read the data lib.read_data(infile, kind="image", mode="GMT_DATA_ONLY", data=data_ptr) - assert image.data + + xrimage = image.to_dataarray() + assert xrimage.shape == (3, 180, 360) + assert xrimage.coords["x"].data.min() == -179.5 + assert xrimage.coords["x"].data.max() == 179.5 + assert xrimage.coords["y"].data.min() == -89.5 + assert xrimage.coords["y"].data.max() == 89.5 + assert xrimage.data.min() == 10 + assert xrimage.data.max() == 255 + assert xrimage.data.dtype == np.uint8 + + if _HAS_RIOXARRAY: # Full check if rioxarray is installed. + xr.testing.assert_equal(xrimage, expected_xrimage) def test_clib_read_data_fails():