Fix bivariate handling with pcolormesh

patniharshit · patniharshit · commit 8f6928d3cb32 · 2017-08-27T15:00:04.000+05:30
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -5200,8 +5200,9 @@ def _pcolorargs(funcname, *args, **kw):
                     cmap = mcolors.BivariateColormap()
                 if norm is None:
                     norm = mcolors.BivariateNorm()
-                C = norm(C)
-            numRows, numCols = C.shape
+                numRows, numCols = C.shape[1:]
+            else:
+                numRows, numCols = C.shape
             if allmatch:
                 X, Y = np.meshgrid(np.arange(numCols), np.arange(numRows))
             else:
@@ -5219,8 +5220,9 @@ def _pcolorargs(funcname, *args, **kw):
                     cmap = mcolors.BivariateColormap()
                 if norm is None:
                     norm = mcolors.BivariateNorm()
-                C = norm(C)
-            numRows, numCols = C.shape
+                numRows, numCols = C.shape[1:]
+            else:
+                numRows, numCols = C.shape
         else:
             raise TypeError(
                 'Illegal arguments to %s; see help(%s)' % (funcname, funcname))
@@ -5614,25 +5616,25 @@ def pcolormesh(self, *args, **kwargs):
         X, Y, C = self._pcolorargs('pcolormesh', *args, **kw)
         Ny, Nx = X.shape
 
-        if (isinstance(norm, mcolors.BivariateNorm) or
-                isinstance(cmap, mcolors.BivariateColormap)):
-            norm = mcolors.NoNorm()
-
         # unit conversion allows e.g. datetime objects as axis values
         self._process_unit_info(xdata=X, ydata=Y, kwargs=kwargs)
         X = self.convert_xunits(X)
         Y = self.convert_yunits(Y)
 
-        # convert to one dimensional arrays
-        C = C.ravel()
+        # convert to one dimensional arrays if univariate
+        if isinstance(norm, mcolors.BivariateNorm):
+            C = np.asarray([C[0].ravel(), C[1].ravel()])
+        else:
+            C = C.ravel()
+
         coords = np.column_stack((X.flat, Y.flat)).astype(float, copy=False)
 
         collection = mcoll.QuadMesh(Nx - 1, Ny - 1, coords,
                                     antialiased=antialiased, shading=shading,
                                     **kwargs)
         collection.set_alpha(alpha)
         collection.set_array(C)
-        if norm is not None and not isinstance(norm, mcolors.Normalize):
+        if norm is not None and not isinstance(norm, mcolors.Norms):
             msg = "'norm' must be an instance of 'mcolors.Normalize'"
             raise ValueError(msg)
         collection.set_cmap(cmap)
diff --git a/lib/matplotlib/cm.py b/lib/matplotlib/cm.py
@@ -238,7 +238,7 @@ def to_rgba(self, x, alpha=None, bytes=False, norm=True):
         """
         # First check for special case, image input:
         try:
-            if x.ndim == 3:
+            if x.ndim == 3 and (x.shape[-1] == 3 or x.shape[-1] == 4):
                 if x.shape[2] == 3:
                     if alpha is None:
                         alpha = 1
diff --git a/lib/matplotlib/collections.py b/lib/matplotlib/collections.py
@@ -731,7 +731,8 @@ def update_scalarmappable(self):
         """
         if self._A is None:
             return
-        if self._A.ndim > 1:
+        if (self._A.ndim > 1 and
+                not isinstance(self.norm, mcolors.BivariateNorm)):
             raise ValueError('Collections can only map rank 1 arrays')
         if not self.check_update("array"):
             return
