Enable passing in a list of RGB xarray.DataArray grids into grdimage

Extends `grdimage` to accept a list of individual red, green and blue xarray.DataArrays for plotting. Had to revert using the 'sequence_space' decorator as it would not work nicely for multiple grids... Implementation uses the contextlib.ExitStack method to create multiple virtual grid files and enter multiple contexts (similar to 3448972). Also added one test case with dummy data, but could do with more to capture edge cases.

Author: weiji14

pygmt/base_plotting.py

@@ -2,6 +2,7 @@
 Base class with plot generating commands.
 Does not define any special non-GMT methods (savefig, show, etc).
 """
+import contextlib
 import csv
 import os
 import numpy as np
@@ -288,7 +289,7 @@ def grdcontour(self, grid, **kwargs):
 
     @fmt_docstring
     @use_alias(R="region", J="projection", W="pen", B="frame", I="shading", C="cmap")
-    @kwargs_to_strings(R="sequence", grid="sequence_space")
+    @kwargs_to_strings(R="sequence")
     def grdimage(self, grid, **kwargs):
         """
         Project grids or images and plot them on maps.
@@ -304,18 +305,41 @@ def grdimage(self, grid, **kwargs):
         ----------
         grid : str, xarray.DataArray or list
             The file name of the input grid or the grid loaded as a DataArray.
-            A list of red, green, blue grids can also be provided instead.
+            For plotting RGB grids, pass in a list made up of either file names or
+            DataArrays to the individual red, green and blue grids.
         """
         kwargs = self._preprocess(**kwargs)
-        kind = data_kind(grid, None, None)
+
+        if isinstance(grid, list):
+            if all([data_kind(g) == "file" for g in grid]):
+                kind = "file"
+                grid = " ".join(grid)
+            elif all([data_kind(g) == "grid" for g in grid]):
+                kind = "grid"
+        else:
+            kind = data_kind(grid)
+
         with Session() as lib:
             if kind == "file":
-                file_context = dummy_context(grid)
+                file_contexts = [dummy_context(grid)]
             elif kind == "grid":
-                file_context = lib.virtualfile_from_grid(grid)
+                if isinstance(grid, list):
+                    file_contexts = [
+                        lib.virtualfile_from_grid(grid[0]),
+                        lib.virtualfile_from_grid(grid[1]),
+                        lib.virtualfile_from_grid(grid[2]),
+                    ]
+                else:
+                    file_contexts = [lib.virtualfile_from_grid(grid)]
             else:
                 raise GMTInvalidInput("Unrecognized data type: {}".format(type(grid)))
-            with file_context as fname:
+            with contextlib.ExitStack() as stack:
+                fname = " ".join(
+                    [
+                        stack.enter_context(file_context)
+                        for file_context in file_contexts
+                    ]
+                )
                 arg_str = " ".join([fname, build_arg_string(kwargs)])
                 lib.call_module("grdimage", arg_str)
 

pygmt/tests/baseline/test_grdimage_rgb_grid.png

@@ -3,6 +3,7 @@
 """
 import numpy as np
 import pytest
+import xarray as xr
 
 from .. import Figure
 from ..exceptions import GMTInvalidInput
@@ -49,6 +50,30 @@ def test_grdimage_rgb_files():
     return fig
 
 
+@pytest.mark.mpl_image_compare
+def test_grdimage_rgb_grid():
+    "Plot an image using Red, Green, and Blue xarray.DataArray inputs"
+    red = xr.DataArray(
+        data=[[128, 0, 0], [128, 0, 0]],
+        dims=("lat", "lon"),
+        coords={"lat": [0, 1], "lon": [2, 3, 4]},
+    )
+    green = xr.DataArray(
+        data=[[0, 128, 0], [0, 128, 0]],
+        dims=("lat", "lon"),
+        coords={"lat": [0, 1], "lon": [2, 3, 4]},
+    )
+    blue = xr.DataArray(
+        data=[[0, 0, 128], [0, 0, 128]],
+        dims=("lat", "lon"),
+        coords={"lat": [0, 1], "lon": [2, 3, 4]},
+    )
+
+    fig = Figure()
+    fig.grdimage(grid=[red, green, blue], projection="x5c", frame=True)
+    return fig
+
+
 def test_grdimage_fails():
     "Should fail for unrecognized input"
     fig = Figure()