GMT_IMAGE: Implement the GMT_IMAGE.to_dataarray method for 3-band images

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,122 @@ 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
+        <xarray.DataArray 'z' (band: 3, y: 180, x: 360)>...
+        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]],
+        <BLANKLINE>
+               [[ 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]],
+        <BLANKLINE>
+               [[ 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:
+            title:
+            history:
+            description:
+            long_name:     z
+            actual_range:  [ 1.79769313e+308 -1.79769313e+308]
+
+        >>> da.coords["x"]  # doctest: +NORMALIZE_WHITESPACE, +ELLIPSIS
+        <xarray.DataArray 'x' (x: 360)>...
+        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
+        <xarray.DataArray 'y' (y: 180)>...
+        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, 0)
+        """
+        # 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

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():