matplotlib · trygvrad · Aug 7, 2024 · Aug 7, 2024 · Jan 6, 2025 · Jan 6, 2025
diff --git a/doc/api/colors_api.rst b/doc/api/colors_api.rst
@@ -53,6 +53,7 @@ Multivariate Colormaps
    BivarColormap
    SegmentedBivarColormap
    BivarColormapFromImage
+   MultivarColormap
 
 Other classes
 -------------

diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -5685,7 +5685,7 @@ def imshow(self, X, cmap=None, norm=None, *, aspect=None,
         """
         Display data as an image, i.e., on a 2D regular raster.
 
-        The input may either be actual RGB(A) data, or 2D scalar data, which
+        The input may either be actual RGB(A) data, or 2D data, which
         will be rendered as a pseudocolor image. For displaying a grayscale
         image, set up the colormapping using the parameters
         ``cmap='gray', vmin=0, vmax=255``.
@@ -5706,24 +5706,25 @@ def imshow(self, X, cmap=None, norm=None, *, aspect=None,
             - (M, N): an image with scalar data. The values are mapped to
               colors using normalization and a colormap. See parameters *norm*,
               *cmap*, *vmin*, *vmax*.
+            - (K, M, N): multiple images with scalar data. Must be used
+              with a multivariate or bivariate colormap that supports K channels.
             - (M, N, 3): an image with RGB values (0-1 float or 0-255 int).
             - (M, N, 4): an image with RGBA values (0-1 float or 0-255 int),
               i.e. including transparency.
 
-            The first two dimensions (M, N) define the rows and columns of
-            the image.
+            The dimensions M and N are the number of rows and columns of the image.
 
             Out-of-range RGB(A) values are clipped.
 
-        %(cmap_doc)s
+        %(multi_cmap_doc)s
 
             This parameter is ignored if *X* is RGB(A).
 
-        %(norm_doc)s
+        %(multi_norm_doc)s
 
             This parameter is ignored if *X* is RGB(A).
 
-        %(vmin_vmax_doc)s
+        %(multi_vmin_vmax_doc)s
 
             This parameter is ignored if *X* is RGB(A).
 
@@ -6062,9 +6063,10 @@ def pcolor(self, *args, shading=None, alpha=None, norm=None, cmap=None,
 
         Parameters
         ----------
-        C : 2D array-like
+        C : 2D array-like or list of 2D array-like objects
             The color-mapped values.  Color-mapping is controlled by *cmap*,
-            *norm*, *vmin*, and *vmax*.
+            *norm*, *vmin*, and *vmax*.  Multiple arrays are only supported
+            when a bivariate or mulivariate colormap is used, see *cmap*.
 
         X, Y : array-like, optional
             The coordinates of the corners of quadrilaterals of a pcolormesh::
@@ -6111,11 +6113,11 @@ def pcolor(self, *args, shading=None, alpha=None, norm=None, cmap=None,
             See :doc:`/gallery/images_contours_and_fields/pcolormesh_grids`
             for more description.
 
-        %(cmap_doc)s
+        %(multi_cmap_doc)s
 
-        %(norm_doc)s
+        %(multi_norm_doc)s
 
-        %(vmin_vmax_doc)s
+        %(multi_vmin_vmax_doc)s
 
         %(colorizer_doc)s
 
@@ -6190,6 +6192,26 @@ def pcolor(self, *args, shading=None, alpha=None, norm=None, cmap=None,
         if shading is None:
             shading = mpl.rcParams['pcolor.shading']
         shading = shading.lower()
+
+        mcolorizer.ColorizingArtist._check_exclusionary_keywords(colorizer,
+                                                                 vmin=vmin,
+                                                                 vmax=vmax)
+
+        # we need to get the colorizer object to know the number of
+        # n_variates that should exist in the array, we therefore get the
+        # colorizer here.
+        colorizer_obj = mcolorizer.ColorizingArtist._get_colorizer(norm=norm,
+                                                                   cmap=cmap,
+                                                                   colorizer=colorizer)
+        if colorizer_obj.norm.n_input > 1:
+            # Valid call signatures for pcolor are with args = (C) or args = (X, Y, C)
+            # If provided, _pcolorargs will check that X, Y and C have the same shape.
+            # Before this check, we need to convert C from shape (K, N, M), where K is
+            # the number of variates, to (N, M) with a data type with K fields.
+            data = mcolorizer._ensure_multivariate_data(args[-1],
+                                                        colorizer_obj.norm.n_input)
+            args = (*args[:-1], data)
+
         X, Y, C, shading = self._pcolorargs('pcolor', *args, shading=shading,
                                             kwargs=kwargs)
         linewidths = (0.25,)
@@ -6226,9 +6248,8 @@ def pcolor(self, *args, shading=None, alpha=None, norm=None, cmap=None,
         coords = stack([X, Y], axis=-1)
 
         collection = mcoll.PolyQuadMesh(
-            coords, array=C, cmap=cmap, norm=norm, colorizer=colorizer,
+            coords, array=C, colorizer=colorizer_obj,
             alpha=alpha, **kwargs)
-        collection._check_exclusionary_keywords(colorizer, vmin=vmin, vmax=vmax)
         collection._scale_norm(norm, vmin, vmax)
 
         coords = coords.reshape(-1, 2)  # flatten the grid structure; keep x, y
@@ -6266,12 +6287,13 @@ def pcolormesh(self, *args, alpha=None, norm=None, cmap=None, vmin=None,
             - (M, N) or M*N: a mesh with scalar data. The values are mapped to
               colors using normalization and a colormap. See parameters *norm*,
               *cmap*, *vmin*, *vmax*.
+            - (K, M, N): multiple layers with scalar data. Must be used with
+              a multivariate or bivariate colormap. See *cmap*.
             - (M, N, 3): an image with RGB values (0-1 float or 0-255 int).
             - (M, N, 4): an image with RGBA values (0-1 float or 0-255 int),
               i.e. including transparency.
 
-            The first two dimensions (M, N) define the rows and columns of
-            the mesh data.
+            Here M and N define the rows and columns of the mesh.
 
         X, Y : array-like, optional
             The coordinates of the corners of quadrilaterals of a pcolormesh::
@@ -6302,11 +6324,11 @@ def pcolormesh(self, *args, alpha=None, norm=None, cmap=None, vmin=None,
             expanded as needed into the appropriate 2D arrays, making a
             rectangular grid.
 
-        %(cmap_doc)s
+        %(multi_cmap_doc)s
 
-        %(norm_doc)s
+        %(multi_norm_doc)s
 
-        %(vmin_vmax_doc)s
+        %(multi_vmin_vmax_doc)s
 
         %(colorizer_doc)s
 
@@ -6429,6 +6451,24 @@ def pcolormesh(self, *args, alpha=None, norm=None, cmap=None, vmin=None,
         shading = mpl._val_or_rc(shading, 'pcolor.shading').lower()
         kwargs.setdefault('edgecolors', 'none')
 
+        mcolorizer.ColorizingArtist._check_exclusionary_keywords(colorizer,
+                                                                 vmin=vmin,
+                                                                 vmax=vmax)
+        # we need to get the colorizer object to know the number of
+        # n_variates that should exist in the array, we therefore get the
+        # colorizer here.
+        colorizer_obj = mcolorizer.ColorizingArtist._get_colorizer(norm=norm,
+                                                                   cmap=cmap,
+                                                                   colorizer=colorizer)
+        if colorizer_obj.norm.n_input > 1:
+            # Valid call signatures for pcolor are with args = (C) or args = (X, Y, C)
+            # If provided, _pcolorargs will check that X, Y and C have the same shape.
+            # Before this check, we need to convert C from shape (K, N, M), where K is
+            # the number of variates, to (N, M) with a data type with K fields.
+            data = mcolorizer._ensure_multivariate_data(args[-1],
+                                                        colorizer_obj.norm.n_input)
+            args = (*args[:-1], data)
+
         X, Y, C, shading = self._pcolorargs('pcolormesh', *args,
                                             shading=shading, kwargs=kwargs)
         coords = np.stack([X, Y], axis=-1)
@@ -6437,8 +6477,8 @@ def pcolormesh(self, *args, alpha=None, norm=None, cmap=None, vmin=None,
 
         collection = mcoll.QuadMesh(
             coords, antialiased=antialiased, shading=shading,
-            array=C, cmap=cmap, norm=norm, colorizer=colorizer, alpha=alpha, **kwargs)
-        collection._check_exclusionary_keywords(colorizer, vmin=vmin, vmax=vmax)
+            array=C, colorizer=colorizer_obj,
+            alpha=alpha, **kwargs)
         collection._scale_norm(norm, vmin, vmax)
 
         coords = coords.reshape(-1, 2)  # flatten the grid structure; keep x, y

diff --git a/lib/matplotlib/cbook.py b/lib/matplotlib/cbook.py
@@ -690,7 +690,17 @@ def safe_masked_invalid(x, copy=False):
     try:
         xm = np.ma.masked_where(~(np.isfinite(x)), x, copy=False)
     except TypeError:
-        return x
+        if len(x.dtype.descr) == 1:
+            return x
+        else:
+            # in case of a dtype with multiple fields:
+            try:
+                mask = np.empty(x.shape, dtype=np.dtype('bool, '*len(x.dtype.descr)))
+                for dd, dm in zip(x.dtype.descr, mask.dtype.descr):
+                    mask[dm[0]] = ~(np.isfinite(x[dd[0]]))
+                xm = np.ma.array(x, mask=mask, copy=False)
+            except TypeError:
+                return x
     return xm
 
 

diff --git a/lib/matplotlib/cm.py b/lib/matplotlib/cm.py
@@ -278,32 +278,3 @@ def get_cmap(name=None, lut=None):
         return _colormaps[name]
     else:
         return _colormaps[name].resampled(lut)
-
-
-def _ensure_cmap(cmap):
-    """
-    Ensure that we have a `.Colormap` object.
-
-    For internal use to preserve type stability of errors.
-
-    Parameters
-    ----------
-    cmap : None, str, Colormap
-
-        - if a `Colormap`, return it
-        - if a string, look it up in mpl.colormaps
-        - if None, look up the default color map in mpl.colormaps
-
-    Returns
-    -------
-    Colormap
-
-    """
-    if isinstance(cmap, colors.Colormap):
-        return cmap
-    cmap_name = mpl._val_or_rc(cmap, "image.cmap")
-    # use check_in_list to ensure type stability of the exception raised by
-    # the internal usage of this (ValueError vs KeyError)
-    if cmap_name not in _colormaps:
-        _api.check_in_list(sorted(_colormaps), cmap=cmap_name)
-    return mpl.colormaps[cmap_name]
diff --git a/lib/matplotlib/collections.py b/lib/matplotlib/collections.py
@@ -2611,10 +2611,16 @@ def _get_unmasked_polys(self):
         arr = self.get_array()
         if arr is not None:
             arr = np.ma.getmaskarray(arr)
-            if arr.ndim == 3:
+            if self.norm.n_input > 1:
+                # multivar case
+                for a in mcolors.MultiNorm._iterable_variates_in_data(
+                                                arr, self.norm.n_input):
+                    mask |= np.any(a, axis=0)
+            elif arr.ndim == 3:
                 # RGB(A) case
                 mask |= np.any(arr, axis=-1)
             elif arr.ndim == 2:
+                # scalar case
                 mask |= arr
             else:
                 mask |= arr.reshape(self._coordinates[:-1, :-1, :].shape[:2])