Update is_bivar condition to expect either BivariateNorm or BivariateColormap

patniharshit · patniharshit · commit f98978faddee · 2017-08-27T15:00:04.000+05:30
diff --git a/examples/images_contours_and_fields/bivariate_demo.py b/examples/images_contours_and_fields/bivariate_demo.py
@@ -28,6 +28,8 @@
 
 bivariate = [air_temp, surf_pres]
 
+# to distinguish bivariate data either BivariateNorm or BivariateColormap must
+# be passed in as argument
 cax = ax.imshow(bivariate, norm=colors.BivariateNorm(),
                 cmap=colors.BivariateColormap())
 
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -5151,8 +5151,11 @@ def imshow(self, X, cmap=None, norm=None, aspect=None,
             raise ValueError(msg)
 
         temp = np.asarray(X)
-        is_bivari = (temp.ndim == 3 or temp.shape[0] == 2)
+        is_bivari = (isinstance(norm, mcolors.BivariateNorm) or
+                     isinstance(cmap, mcolors.BivariateColormap))
         if is_bivari:
+            if temp.ndim != 3 and temp.shape[0] != 2:
+                raise TypeError("Expected shape like (2, n, m)")
             if cmap is None:
                 cmap = mcolors.BivariateColormap()
             if norm is None:
@@ -5205,7 +5208,8 @@ def _pcolorargs(funcname, *args, **kw):
 
         if len(args) == 1:
             C = np.asanyarray(args[0])
-            is_bivari = (C.ndim == 3 or C.shape[0] == 2)
+            is_bivari = (isinstance(norm, mcolors.BivariateNorm) or
+                         isinstance(cmap, mcolors.BivariateColormap))
             if is_bivari:
                 numRows, numCols = C.shape[1:]
             else:
@@ -5220,12 +5224,9 @@ def _pcolorargs(funcname, *args, **kw):
 
         if len(args) == 3:
             X, Y, C = [np.asanyarray(a) for a in args]
-            is_bivari = (C.ndim == 3 or C.shape[0] == 2)
+            is_bivari = (isinstance(norm, mcolors.BivariateNorm) or
+                         isinstance(cmap, mcolors.BivariateColormap))
             if is_bivari:
-                if cmap is None:
-                    cmap = mcolors.BivariateColormap()
-                if norm is None:
-                    norm = mcolors.BivariateNorm()
                 numRows, numCols = C.shape[1:]
             else:
                 numRows, numCols = C.shape
@@ -5414,8 +5415,11 @@ def pcolor(self, *args, **kwargs):
         X, Y, C = self._pcolorargs('pcolor', *args, **kw)
         Ny, Nx = X.shape
 
-        is_bivari = (C.ndim == 3 or C.shape[0] == 2)
+        is_bivari = (isinstance(norm, mcolors.BivariateNorm) or
+                     isinstance(cmap, mcolors.BivariateColormap))
         if is_bivari:
+            if C.ndim != 3 and C.shape[0] != 2:
+                raise TypeError("Expected shape like (2, n, m)")
             if cmap is None:
                 cmap = mcolors.BivariateColormap()
             if norm is None:
@@ -5638,8 +5642,11 @@ def pcolormesh(self, *args, **kwargs):
         X, Y, C = self._pcolorargs('pcolormesh', *args, **kw)
         Ny, Nx = X.shape
 
-        is_bivari = (C.ndim == 3 or C.shape[0] == 2)
+        is_bivari = (isinstance(norm, mcolors.BivariateNorm) or
+                     isinstance(cmap, mcolors.BivariateColormap))
         if is_bivari:
+            if C.ndim != 3 and C.shape[0] != 2:
+                raise TypeError("Expected shape like (2, n, m)")
             if cmap is None:
                 cmap = mcolors.BivariateColormap()
             if norm is None:
@@ -5798,8 +5805,11 @@ def pcolorfast(self, *args, **kwargs):
 
         C = np.asarray(args[-1])
 
-        is_bivari = (C.ndim == 3 or C.shape[0] == 2)
+        is_bivari = (isinstance(norm, mcolors.BivariateNorm) or
+                     isinstance(cmap, mcolors.BivariateColormap))
         if is_bivari:
+            if C.ndim != 3 and C.shape[0] != 2:
+                raise TypeError("Expected shape like (2, n, m)")
             if cmap is None:
                 cmap = mcolors.BivariateColormap()
             if norm is None:
diff --git a/lib/matplotlib/tests/test_axes.py b/lib/matplotlib/tests/test_axes.py
@@ -5285,17 +5285,17 @@ def test_bivariates():
     bivariate = [air_temp, surf_pres]
 
     fig1, ax1 = plt.subplots()
-    cax1 = ax1.imshow(bivariate)
+    cax1 = ax1.imshow(bivariate, norm=mcolors.BivariateNorm())
     cbar = fig1.colorbar(cax1, xlabel='air_temp', ylabel='surf_pres')
 
     fig2, ax2 = plt.subplots()
-    cax2 = ax2.pcolor(bivariate)
+    cax2 = ax2.pcolor(bivariate, norm=mcolors.BivariateNorm())
     cbar = fig2.colorbar(cax2, xlabel='air_temp', ylabel='surf_pres')
 
     fig3, ax3 = plt.subplots()
-    cax3 = ax3.pcolormesh(bivariate)
+    cax3 = ax3.pcolormesh(bivariate, norm=mcolors.BivariateNorm())
     cbar = fig3.colorbar(cax3, xlabel='air_temp', ylabel='surf_pres')
 
     fig4, ax4 = plt.subplots()
-    cax4 = ax4.pcolorfast(bivariate)
+    cax4 = ax4.pcolorfast(bivariate, norm=mcolors.BivariateNorm())
     cbar = fig4.colorbar(cax4, xlabel='air_temp', ylabel='surf_pres')
