Allow passing an array as intensity for plot3d (#1109)

Co-authored-by: Dongdong Tian <[email protected]>

Author: core-man

pygmt/src/plot3d.py

@@ -111,11 +111,14 @@ def plot3d(
         Offset the plot symbol or line locations by the given amounts
         *dx*/*dy*\ [/*dz*] [Default is no offset].
     {G}
-    intensity : float or bool
-        Provide an *intens* value (nominally in the -1 to +1 range) to
-        modulate the fill color by simulating illumination [Default is None].
-        If using ``intensity=True``, we will instead read *intens* from the
-        first data column after the symbol parameters (if given).
+    intensity : float or bool or 1d array
+        Provide an *intensity* value (nominally in the -1 to +1 range) to
+        modulate the fill color by simulating illumination. If using
+        ``intensity=True``, we will instead read *intensity* from the first
+        data column after the symbol parameters (if given). *intensity* can
+        also be a 1d array to set varying intensity for symbols, but it is only
+        valid for ``x``/``y``/``z``.
+
     close : str
         [**+b**\|\ **d**\|\ **D**][**+xl**\|\ **r**\|\ *x0*]\
         [**+yl**\|\ **r**\|\ *y0*][**+p**\ *pen*].
@@ -183,9 +186,14 @@ def plot3d(
             )
         extra_arrays.append(sizes)
 
-    if "t" in kwargs and is_nonstr_iter(kwargs["t"]):
-        extra_arrays.append(kwargs["t"])
-        kwargs["t"] = ""
+    for flag in ["I", "t"]:
+        if flag in kwargs and is_nonstr_iter(kwargs[flag]):
+            if kind != "vectors":
+                raise GMTInvalidInput(
+                    f"Can't use arrays for {plot3d.aliases[flag]} if data is matrix or file."
+                )
+            extra_arrays.append(kwargs[flag])
+            kwargs[flag] = ""
 
     with Session() as lib:
         # Choose how data will be passed in to the module

pygmt/tests/baseline/test_plot3d_varying_intensity.png.dvc

@@ -0,0 +1,4 @@
+outs:
+- md5: 79f7d8062dbb6f29ffa0a3aaa7382f13
+  size: 24052
+  path: test_plot3d_varying_intensity.png

pygmt/tests/test_plot3d.py

@@ -130,30 +130,19 @@ def test_plot3d_fail_no_data(data, region):
         )
 
 
-def test_plot3d_fail_size_color(data, region):
+def test_plot3d_fail_color_size_intensity(data, region):
     """
-    Should raise an exception if array sizes and color are used with matrix.
+    Should raise an exception if array color, sizes and intensity are used with
+    matrix.
     """
     fig = Figure()
+    kwargs = dict(data=data, region=region, projection="X10c", frame="afg")
     with pytest.raises(GMTInvalidInput):
-        fig.plot3d(
-            data=data,
-            region=region,
-            projection="X4i",
-            style="c0.2c",
-            color=data[:, 2],
-            frame="afg",
-        )
+        fig.plot3d(style="c0.2c", color=data[:, 2], **kwargs)
     with pytest.raises(GMTInvalidInput):
-        fig.plot3d(
-            data=data,
-            region=region,
-            projection="X4i",
-            style="cc",
-            sizes=data[:, 2],
-            color="red",
-            frame="afg",
-        )
+        fig.plot3d(style="cc", sizes=data[:, 2], color="red", **kwargs)
+    with pytest.raises(GMTInvalidInput):
+        fig.plot3d(style="cc", intensity=data[:, 2], color="red", **kwargs)
 
 
 @check_figures_equal()
@@ -329,6 +318,33 @@ def test_plot3d_colors_sizes_proj(data, region):
     return fig_ref, fig_test
 
 
+@pytest.mark.mpl_image_compare
+def test_plot3d_varying_intensity():
+    """
+    Plot the data with array-like intensity.
+    """
+    x = np.arange(-1, 1.1, 0.1)
+    y = np.zeros(x.size)
+    z = y
+    intensity = x
+
+    fig = Figure()
+    fig.plot3d(
+        x=x,
+        y=y,
+        z=z,
+        region=[-1.1, 1.1, -0.5, 0.5, -0.5, 0.5],
+        projection="X15c/5c",
+        zsize="5c",
+        perspective=[135, 30],
+        frame=["Sltr", "xaf+lIntensity"],
+        style="c0.5c",
+        color="blue",
+        intensity=intensity,
+    )
+    return fig
+
+
 @check_figures_equal()
 def test_plot3d_transparency():
     """