Defer call to norm and cmap in imshow

patniharshit · patniharshit · commit 8a985f0400f3 · 2017-08-27T15:00:04.000+05:30
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -5137,11 +5137,6 @@ def imshow(self, X, cmap=None, norm=None, aspect=None,
         of pixel (0, 0).
 
         """
-
-        temp = np.asarray(X)
-        if temp.ndim == 3 and isinstance(norm, mcolors.BivariateNorm):
-                    X = norm(temp)
-
         if not self._hold:
             self.cla()
 
diff --git a/lib/matplotlib/axes/_subplots.py b/lib/matplotlib/axes/_subplots.py
@@ -177,7 +177,6 @@ def subplot_class_factory(axes_class=None):
                          (SubplotBase, axes_class),
                          {'_axes_class': axes_class})
         _subplot_classes[axes_class] = new_class
-
     return new_class
 
 # This is provided for backward compatibility
diff --git a/lib/matplotlib/image.py b/lib/matplotlib/image.py
@@ -253,7 +253,11 @@ def get_size(self):
         if self._A is None:
             raise RuntimeError('You must first set the image array')
 
-        return self._A.shape[:2]
+        if isinstance(self.norm, mcolors.BivariateNorm):
+            imshape = self._A.shape[1:]
+        else:
+            imshape = self._A.shape[:2]
+        return imshape
 
     def set_alpha(self, alpha):
         """
@@ -300,6 +304,12 @@ def _make_image(self, A, in_bbox, out_bbox, clip_bbox, magnification=1.0,
         `trans` is the affine transformation from the image to pixel
         space.
         """
+        if isinstance(self.norm, mcolors.BivariateNorm):
+            imwidth = A.shape[1]
+            imheight = A.shape[2]
+        else:
+            imwidth = A.shape[0]
+            imheight = A.shape[1]
         if A is None:
             raise RuntimeError('You must first set the image '
                                'array or the image attribute')
@@ -323,15 +333,15 @@ def _make_image(self, A, in_bbox, out_bbox, clip_bbox, magnification=1.0,
             # Flip the input image using a transform.  This avoids the
             # problem with flipping the array, which results in a copy
             # when it is converted to contiguous in the C wrapper
-            t0 = Affine2D().translate(0, -A.shape[0]).scale(1, -1)
+            t0 = Affine2D().translate(0, -imwidth).scale(1, -1)
         else:
             t0 = IdentityTransform()
 
         t0 += (
             Affine2D()
             .scale(
-                in_bbox.width / A.shape[1],
-                in_bbox.height / A.shape[0])
+                in_bbox.width / imheight,
+                in_bbox.height / imwidth)
             .translate(in_bbox.x0, in_bbox.y0)
             + self.get_transform())
 
@@ -362,17 +372,17 @@ def _make_image(self, A, in_bbox, out_bbox, clip_bbox, magnification=1.0,
             if A.ndim not in (2, 3):
                 raise ValueError("Invalid dimensions, got {}".format(A.shape))
 
-            if A.ndim == 2:
-                if not isinstance(self.norm, mcolors.BivariateNorm):
-                    A = self.norm(A)
+            if A.ndim == 2 or (A.ndim == 3 and
+                               isinstance(self.norm, mcolors.BivariateNorm)):
+                A = self.norm(A)
                 if A.dtype.kind == 'f':
                     # If the image is greyscale, convert to RGBA and
                     # use the extra channels for resizing the over,
                     # under, and bad pixels.  This is needed because
                     # Agg's resampler is very aggressive about
                     # clipping to [0, 1] and we use out-of-bounds
                     # values to carry the over/under/bad information
-                    rgba = np.empty((A.shape[0], A.shape[1], 4), dtype=A.dtype)
+                    rgba = np.empty((imwidth, imheight, 4), dtype=A.dtype)
                     rgba[..., 0] = A  # normalized data
                     # this is to work around spurious warnings coming
                     # out of masked arrays.
@@ -411,9 +421,10 @@ def _make_image(self, A, in_bbox, out_bbox, clip_bbox, magnification=1.0,
 
             if not created_rgba_mask:
                 # Always convert to RGBA, even if only RGB input
+                isBivari = (A.ndim == 2 and A.shape[0] == 2)
                 if A.shape[2] == 3:
                     A = _rgb_to_rgba(A)
-                elif A.shape[2] != 4:
+                elif A.shape[2] != 4 and not isBivari:
                     raise ValueError("Invalid dimensions, got %s" % (A.shape,))
 
                 output = np.zeros((out_height, out_width, 4), dtype=A.dtype)
@@ -596,8 +607,9 @@ def set_data(self, A):
                 not np.can_cast(self._A.dtype, float, "same_kind")):
             raise TypeError("Image data cannot be converted to float")
 
-        if not (self._A.ndim == 2
-                or self._A.ndim == 3 and self._A.shape[-1] in [3, 4]):
+        isRGB = (self._A.ndim == 3 and self._A.shape[-1] in [3, 4])
+        isBivari = (self._A.ndim == 3 and self._A.shape[0] == 2)
+        if not (self._A.ndim == 2 or isRGB or isBivari):
             raise TypeError("Invalid dimensions for image data")
 
         self._imcache = None
