Fix bivariate handling with pcolormesh

patniharshit · patniharshit · commit b43e6c7abd1d · 2017-08-15T20:33:49.000+05:30
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -5182,8 +5182,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:
@@ -5201,8 +5202,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))
@@ -5599,25 +5601,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 not isinstance(self.norm, colors.BivariateNorm):
                 if x.shape[2] == 3:
                     if alpha is None:
                         alpha = 1
diff --git a/lib/matplotlib/collections.py b/lib/matplotlib/collections.py
@@ -730,7 +730,7 @@ 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
