Alpha argument can be an array

efiring · efiring · commit 20b894bfaa90 · 2020-09-08T10:46:44.000-10:00
Previously this was supported only for images.  Now it works for
collections, and when directly calling a colormap or to_rgba_array.
diff --git a/doc/users/next_whats_new/alpha_array.rst b/doc/users/next_whats_new/alpha_array.rst
@@ -0,0 +1,31 @@
+Transparency (alpha) can be set as an array in collections
+----------------------------------------------------------
+Previously, the alpha value controlling tranparency in collections could be
+specified only as a scalar which was applied to all elements in the collection.
+For example, all the markers in a `~.Axes.scatter` plot, or all the
+quadrilaterals in a `~.Axes.pcolormesh` plot, would have the same alpha value.
+
+Now it is possible to supply alpha as an array with one value for each element
+(marker, quadrilateral, etc.) in a collection.
+
+.. plot::
+
+    x = np.arange(5, dtype=float)
+    y = np.arange(5, dtype=float)
+    # z and zalpha for demo pcolormesh
+    z = x[1:, np.newaxis] + y[np.newaxis, 1:]
+    zalpha = np.ones_like(z)
+    zalpha[::2, ::2] = 0.3  # alternate patches are partly transparent
+    # s and salpha for demo scatter
+    s = x
+    salpha = np.linspace(0.1, 0.9, len(x))  # just a ramp
+
+    fig, axs = plt.subplots(2, 2, constrained_layout=True)
+    axs[0, 0].pcolormesh(x, y, z, alpha=zalpha)
+    axs[0, 0].set_title("pcolormesh")
+    axs[0, 1].scatter(x, y, c=s, alpha=salpha)
+    axs[0, 1].set_title("color-mapped")
+    axs[1, 0].scatter(x, y, c='k', alpha=salpha)
+    axs[1, 0].set_title("c='k'")
+    axs[1, 1].scatter(x, y, c=['r', 'g', 'b', 'c', 'm'], alpha=salpha)
+    axs[1, 1].set_title("c=['r', 'g', 'b', 'c', 'm']")
diff --git a/lib/matplotlib/artist.py b/lib/matplotlib/artist.py
@@ -954,10 +954,37 @@ def set_alpha(self, alpha):
 
         Parameters
         ----------
-        alpha : float or None
+        alpha : scalar or None
+            *alpha* must be within the 0-1 range, inclusive.
         """
         if alpha is not None and not isinstance(alpha, Number):
-            raise TypeError('alpha must be a float or None')
+            raise TypeError(
+                f'alpha must be numeric or None, not {type(alpha)}')
+        if alpha is not None and not (0 <= alpha <= 1):
+            raise ValueError(f'alpha ({alpha}) is outside 0-1 range')
+        self._alpha = alpha
+        self.pchanged()
+        self.stale = True
+
+    def _set_alpha_for_array(self, alpha):
+        """
+        Set the alpha value used for blending - not supported on all backends.
+
+        Parameters
+        ----------
+        alpha : array-like or scalar or None
+            All values must be within the 0-1 range, inclusive.
+            Masked values and nans are not supported.
+        """
+        if isinstance(alpha, str):
+            raise TypeError("alpha must be numeric or None, not a string")
+        if not np.iterable(alpha):
+            Artist.set_alpha(self, alpha)
+            return
+        alpha = np.asarray(alpha)
+        if not (alpha.min() >= 0 and alpha.max() <= 1):
+            raise ValueError('alpha must be between 0 and 1, inclusive, '
+                             f'but min is {alpha.min()}, max is {alpha.max()}')
         self._alpha = alpha
         self.pchanged()
         self.stale = True
diff --git a/lib/matplotlib/collections.py b/lib/matplotlib/collections.py
@@ -832,12 +832,24 @@ def set_edgecolor(self, c):
         self._set_edgecolor(c)
 
     def set_alpha(self, alpha):
-        # docstring inherited
-        super().set_alpha(alpha)
+        """
+        Set the transparency of the collection.
+
+        Parameters
+        ----------
+        alpha: float or array of float
+            If not None, *alpha* values must be between 0 and 1, inclusive.
+            If an array is provided, it's length must match the number of
+            elements in the collection.  Masked values and nans are not
+            supported.
+        """
+        artist.Artist._set_alpha_for_array(self, alpha)
         self._update_dict['array'] = True
         self._set_facecolor(self._original_facecolor)
         self._set_edgecolor(self._original_edgecolor)
 
+    set_alpha.__doc__ = artist.Artist._set_alpha_for_array.__doc__
+
     def get_linewidth(self):
         return self._linewidths
 
@@ -848,11 +860,23 @@ def update_scalarmappable(self):
         """Update colors from the scalar mappable array, if it is not None."""
         if self._A is None:
             return
-        # QuadMesh can map 2d arrays
+        # QuadMesh can map 2d arrays (but pcolormesh supplies 1d array)
         if self._A.ndim > 1 and not isinstance(self, QuadMesh):
             raise ValueError('Collections can only map rank 1 arrays')
         if not self._check_update("array"):
             return
+        if np.iterable(self._alpha):
+            if self._alpha.size != self._A.size:
+                # This can occur with the deprecated behavior of 'flat'
+                # pcolormesh shading.  If we bring the current change in
+                # before that deprecated behavior is removed, we need to
+                # add the explanation to the message below.
+                raise ValueError(f'Data array shape, {self._A.shape} '
+                                 'is incompatible with alpha array shape, '
+                                 f'{self._alpha.shape}.')
+            # pcolormesh, scatter, maybe others flatten their _A
+            self._alpha = self._alpha.reshape(self._A.shape)
+
         if self._is_filled:
             self._facecolors = self.to_rgba(self._A, self._alpha)
         elif self._is_stroked:
diff --git a/lib/matplotlib/colors.py b/lib/matplotlib/colors.py
@@ -289,18 +289,40 @@ def to_rgba_array(c, alpha=None):
     """
     Convert *c* to a (n, 4) array of RGBA colors.
 
-    If *alpha* is not ``None``, it forces the alpha value.  If *c* is
-    ``"none"`` (case-insensitive) or an empty list, an empty array is returned.
-    If *c* is a masked array, an ndarray is returned with a (0, 0, 0, 0)
-    row for each masked value or row in *c*.
+    Parameters
+    ----------
+    c : Matplotlib color or array of colors
+        If *c* is a masked array, an ndarray is returned with a (0, 0, 0, 0)
+        row for each masked value or row in *c*.
+
+    alpha : float or sequence of floats, optional
+        If *alpha* is not ``None``, it forces the alpha value, except if *c* is
+        ``"none"`` (case-insensitive), which always maps to ``(0, 0, 0, 0)``.
+        If *alpha* is a sequence and *c* is a single color, *c* will be
+        repeated to match the length of *alpha*.
+
+    Returns
+    -------
+    array
+        (n, 4) array of RGBA colors.
+
     """
     # Special-case inputs that are already arrays, for performance.  (If the
     # array has the wrong kind or shape, raise the error during one-at-a-time
     # conversion.)
+    if np.iterable(alpha):
+        alpha = np.asarray(alpha).ravel()
     if (isinstance(c, np.ndarray) and c.dtype.kind in "if"
             and c.ndim == 2 and c.shape[1] in [3, 4]):
         mask = c.mask.any(axis=1) if np.ma.is_masked(c) else None
         c = np.ma.getdata(c)
+        if np.iterable(alpha):
+            if c.shape[0] == 1 and alpha.shape[0] > 1:
+                c = np.tile(c, (alpha.shape[0], 1))
+            elif c.shape[0] != alpha.shape[0]:
+                raise ValueError("The number of colors must match the number"
+                                 " of alpha values if there are more than one"
+                                 " of each.")
         if c.shape[1] == 3:
             result = np.column_stack([c, np.zeros(len(c))])
             result[:, -1] = alpha if alpha is not None else 1.
@@ -320,7 +342,10 @@ def to_rgba_array(c, alpha=None):
     if cbook._str_lower_equal(c, "none"):
         return np.zeros((0, 4), float)
     try:
-        return np.array([to_rgba(c, alpha)], float)
+        if np.iterable(alpha):
+            return np.array([to_rgba(c, a) for a in alpha], float)
+        else:
+            return np.array([to_rgba(c, alpha)], float)
     except (ValueError, TypeError):
         pass
 
@@ -332,7 +357,10 @@ def to_rgba_array(c, alpha=None):
     if len(c) == 0:
         return np.zeros((0, 4), float)
     else:
-        return np.array([to_rgba(cc, alpha) for cc in c])
+        if np.iterable(alpha):
+            return np.array([to_rgba(cc, aa) for cc, aa in zip(c, alpha)])
+        else:
+            return np.array([to_rgba(cc, alpha) for cc in c])
 
 
 def to_rgb(c):
@@ -539,8 +567,9 @@ def __call__(self, X, alpha=None, bytes=False):
             return the RGBA values ``X*100`` percent along the Colormap line.
             For integers, X should be in the interval ``[0, Colormap.N)`` to
             return RGBA values *indexed* from the Colormap with index ``X``.
-        alpha : float, None
-            Alpha must be a scalar between 0 and 1, or None.
+        alpha : float, array-like, None
+            Alpha must be a scalar between 0 and 1, a sequence of such
+            floats with shape matching X, or None.
         bytes : bool
             If False (default), the returned RGBA values will be floats in the
             interval ``[0, 1]`` otherwise they will be uint8s in the interval
@@ -580,23 +609,29 @@ def __call__(self, X, alpha=None, bytes=False):
         else:
             lut = self._lut.copy()  # Don't let alpha modify original _lut.
 
+        rgba = np.empty(shape=xa.shape + (4,), dtype=lut.dtype)
+        lut.take(xa, axis=0, mode='clip', out=rgba)
+
         if alpha is not None:
+            if np.iterable(alpha):
+                alpha = np.asarray(alpha)
+                if not (alpha.shape == xa.shape):
+                    raise ValueError("alpha is array-like but it's shape"
+                                     " %s doesn't match that of X %s" %
+                                     (alpha.shape, xa.shape))
+
             alpha = np.clip(alpha, 0, 1)
             if bytes:
-                alpha = int(alpha * 255)
-            if (lut[-1] == 0).all():
-                lut[:-1, -1] = alpha
-                # All zeros is taken as a flag for the default bad
-                # color, which is no color--fully transparent.  We
-                # don't want to override this.
-            else:
-                lut[:, -1] = alpha
-                # If the bad value is set to have a color, then we
-                # override its alpha just as for any other value.
+                alpha = (alpha * 255).astype(np.uint8)
+            rgba[..., -1] = alpha
+
+            if (lut[-1] == 0).all() and mask_bad is not None:
+                if mask_bad.shape == xa.shape:
+                    rgba[mask_bad] = (0, 0, 0, 0)
+                elif mask_bad:
+                    rgba[..., :] = (0, 0, 0, 0)
 
-        rgba = lut[xa]
         if not np.iterable(X):
-            # Return a tuple if the input was a scalar
             rgba = tuple(rgba)
         return rgba
 
diff --git a/lib/matplotlib/image.py b/lib/matplotlib/image.py
@@ -274,16 +274,12 @@ def set_alpha(self, alpha):
 
         Parameters
         ----------
-        alpha : float
+        alpha : float or 2-d array or None
         """
-        if alpha is not None and not isinstance(alpha, Number):
-            alpha = np.asarray(alpha)
-            if alpha.ndim != 2:
-                raise TypeError('alpha must be a float, two-dimensional '
-                                'array, or None')
-        self._alpha = alpha
-        self.pchanged()
-        self.stale = True
+        martist.Artist._set_alpha_for_array(self, alpha)
+        if np.ndim(alpha) not in (0, 2):
+            raise TypeError('alpha must be a float, two-dimensional '
+                            'array, or None')
         self._imcache = None
 
     def _get_scalar_alpha(self):
diff --git a/lib/matplotlib/tests/test_artist.py b/lib/matplotlib/tests/test_artist.py
@@ -277,3 +277,29 @@ def test_artist_inspector_get_aliases():
     ai = martist.ArtistInspector(mlines.Line2D)
     aliases = ai.get_aliases()
     assert aliases["linewidth"] == {"lw"}
+
+
+def test_set_alpha():
+    art = martist.Artist()
+    with pytest.raises(TypeError, match='^alpha must be numeric or None'):
+        art.set_alpha('string')
+    with pytest.raises(TypeError, match='^alpha must be numeric or None'):
+        art.set_alpha([1, 2, 3])
+    with pytest.raises(ValueError, match="outside 0-1 range"):
+        art.set_alpha(1.1)
+    with pytest.raises(ValueError, match="outside 0-1 range"):
+        art.set_alpha(np.nan)
+
+
+def test_set_alpha_for_array():
+    art = martist.Artist()
+    with pytest.raises(TypeError, match='^alpha must be numeric or None'):
+        art._set_alpha_for_array('string')
+    with pytest.raises(ValueError, match="outside 0-1 range"):
+        art._set_alpha_for_array(1.1)
+    with pytest.raises(ValueError, match="outside 0-1 range"):
+        art._set_alpha_for_array(np.nan)
+    with pytest.raises(ValueError, match="alpha must be between 0 and 1"):
+        art._set_alpha_for_array([0.5, 1.1])
+    with pytest.raises(ValueError, match="alpha must be between 0 and 1"):
+        art._set_alpha_for_array([0.5, np.nan])
diff --git a/lib/matplotlib/tests/test_collections.py b/lib/matplotlib/tests/test_collections.py
@@ -543,11 +543,38 @@ def test_cap_and_joinstyle_image():
 def test_scatter_post_alpha():
     fig, ax = plt.subplots()
     sc = ax.scatter(range(5), range(5), c=range(5))
-    # this needs to be here to update internal state
-    fig.canvas.draw()
     sc.set_alpha(.1)
 
 
+def test_scatter_alpha_array():
+    x = np.arange(5)
+    alpha = x / 5
+    # With color mapping.
+    fig, (ax0, ax1) = plt.subplots(2)
+    sc0 = ax0.scatter(x, x, c=x, alpha=alpha)
+    sc1 = ax1.scatter(x, x, c=x)
+    sc1.set_alpha(alpha)
+    plt.draw()
+    assert_array_equal(sc0.get_facecolors()[:, -1], alpha)
+    assert_array_equal(sc1.get_facecolors()[:, -1], alpha)
+    # Without color mapping.
+    fig, (ax0, ax1) = plt.subplots(2)
+    sc0 = ax0.scatter(x, x, color=['r', 'g', 'b', 'c', 'm'], alpha=alpha)
+    sc1 = ax1.scatter(x, x, color='r', alpha=alpha)
+    plt.draw()
+    assert_array_equal(sc0.get_facecolors()[:, -1], alpha)
+    assert_array_equal(sc1.get_facecolors()[:, -1], alpha)
+    # Without color mapping, and set alpha afterward.
+    fig, (ax0, ax1) = plt.subplots(2)
+    sc0 = ax0.scatter(x, x, color=['r', 'g', 'b', 'c', 'm'])
+    sc0.set_alpha(alpha)
+    sc1 = ax1.scatter(x, x, color='r')
+    sc1.set_alpha(alpha)
+    plt.draw()
+    assert_array_equal(sc0.get_facecolors()[:, -1], alpha)
+    assert_array_equal(sc1.get_facecolors()[:, -1], alpha)
+
+
 def test_pathcollection_legend_elements():
     np.random.seed(19680801)
     x, y = np.random.rand(2, 10)
@@ -662,6 +689,39 @@ def test_quadmesh_set_array():
     assert np.array_equal(coll.get_array(), np.ones(9))
 
 
+def test_quadmesh_alpha_array():
+    x = np.arange(4)
+    y = np.arange(4)
+    z = np.arange(9).reshape((3, 3))
+    alpha = z / z.max()
+    alpha_flat = alpha.ravel()
+    # Provide 2-D alpha:
+    fig, (ax0, ax1) = plt.subplots(2)
+    coll1 = ax0.pcolormesh(x, y, z, alpha=alpha)
+    coll2 = ax1.pcolormesh(x, y, z)
+    coll2.set_alpha(alpha)
+    plt.draw()
+    assert_array_equal(coll1.get_facecolors()[:, -1], alpha_flat)
+    assert_array_equal(coll2.get_facecolors()[:, -1], alpha_flat)
+    # Or provide 1-D alpha:
+    fig, (ax0, ax1) = plt.subplots(2)
+    coll1 = ax0.pcolormesh(x, y, z, alpha=alpha_flat)
+    coll2 = ax1.pcolormesh(x, y, z)
+    coll2.set_alpha(alpha_flat)
+    plt.draw()
+    assert_array_equal(coll1.get_facecolors()[:, -1], alpha_flat)
+    assert_array_equal(coll2.get_facecolors()[:, -1], alpha_flat)
+
+
+def test_alpha_validation():
+    # Most of the relevant testing is in test_artist and test_colors.
+    fig, ax = plt.subplots()
+    pc = ax.pcolormesh(np.arange(12).reshape((3, 4)))
+    with pytest.raises(ValueError, match="^Data array shape"):
+        pc.set_alpha([0.5, 0.6])
+        pc.update_scalarmappable()
+
+
 def test_legend_inverse_size_label_relationship():
     """
     Ensure legend markers scale appropriately when label and size are
diff --git a/lib/matplotlib/tests/test_colors.py b/lib/matplotlib/tests/test_colors.py
diff --git a/lib/matplotlib/tests/test_image.py b/lib/matplotlib/tests/test_image.py
