Remove invalid dimension checking in axes_rgb.

anntzer · anntzer · commit ab9610ef1e08 · 2017-01-13T22:33:16.000-08:00
Just let `dstack` error out if dimensions don't match.
diff --git a/lib/mpl_toolkits/axes_grid1/axes_rgb.py b/lib/mpl_toolkits/axes_grid1/axes_rgb.py
@@ -202,23 +202,14 @@ def imshow_rgb(self, r, g, b, **kwargs):
         g : matplotlib.image.AxesImage
         b : matplotlib.image.AxesImage
         """
-        ny, nx = r.shape
-        if not ((nx, ny) == g.shape == b.shape):
-            raise ValueError('Input shapes do not match.'
-                             '\nr.shape = {0}'
-                             '\ng.shape = {1}'
-                             '\nb.shape = {2}'
-                             ''.format(r.shape, g.shape, b.shape))
-
-        R = np.zeros([ny, nx, 3], dtype="d")
+        RGB = np.dstack([r, g, b])
+        R = np.zeros_like(RGB)
         R[:,:,0] = r
-        G = np.zeros_like(R)
+        G = np.zeros_like(RGB)
         G[:,:,1] = g
-        B = np.zeros_like(R)
+        B = np.zeros_like(RGB)
         B[:,:,2] = b
 
-        RGB = R + G + B
-
         im_rgb = self.RGB.imshow(RGB, **kwargs)
         im_r = self.R.imshow(R, **kwargs)
         im_g = self.G.imshow(G, **kwargs)
