Bivariate works with imshow, pcolor, pcolormesh, pcolorfast

patniharshit · patniharshit · commit 8e93d2027c79 · 2017-07-07T18:20:28.000+05:30
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -5081,12 +5081,30 @@ def imshow(self, X, cmap=None, norm=None, aspect=None,
         if not self._hold:
             self.cla()
 
-        if norm is not None and not isinstance(norm, mcolors.Normalize):
-            msg = "'norm' must be an instance of 'mcolors.Normalize'"
+        isNorm = isinstance(norm, (mcolors.Normalize, mcolors.BivariateNorm))
+        if norm is not None and not isNorm:
+            msg = "'norm' must be an instance of 'mcolors.Normalize' " \
+                  "or 'mcolors.BivariateNorm'"
             raise ValueError(msg)
+
         if aspect is None:
             aspect = rcParams['image.aspect']
         self.set_aspect(aspect)
+
+        temp = np.asarray(X)
+        if (temp.ndim == 3 and isinstance(norm, mcolors.BivariateNorm) or
+                isinstance(cmap, mcolors.BivariateColormap)):
+            if cmap is None:
+                cmap = mcolors.BivariateColormap()
+            if norm is None:
+                norm = mcolors.BivariateNorm()
+            temp = norm(temp)
+            temp[0] = temp[0] * (cmap.N-1)
+            temp[1] = temp[1] * (cmap.N-1)
+            temp = temp.astype(int)
+            X = temp[0] + cmap.N * temp[1]
+            norm = mcolors.NoNorm()
+
         im = mimage.AxesImage(self, cmap, norm, interpolation, origin, extent,
                               filternorm=filternorm, filterrad=filterrad,
                               resample=resample, **kwargs)
@@ -5113,7 +5131,6 @@ def imshow(self, X, cmap=None, norm=None, aspect=None,
 
     @staticmethod
     def _pcolorargs(funcname, *args, **kw):
-        # This takes one kwarg, allmatch.
         # If allmatch is True, then the incoming X, Y, C must
         # have matching dimensions, taking into account that
         # X and Y can be 1-D rather than 2-D.  This perfect
@@ -5126,9 +5143,25 @@ def _pcolorargs(funcname, *args, **kw):
         # is False.
 
         allmatch = kw.pop("allmatch", False)
+        norm = kw.pop("norm", None)
+        cmap = kw.pop("cmap", None)
 
         if len(args) == 1:
             C = np.asanyarray(args[0])
+
+            if (C.ndim == 3 and isinstance(norm, mcolors.BivariateNorm) or
+                    isinstance(cmap, mcolors.BivariateColormap)):
+                if cmap is None:
+                    cmap = mcolors.BivariateColormap()
+                if norm is None:
+                    norm = mcolors.BivariateNorm()
+                C = norm(C)
+                C[0] = C[0] * (cmap.N-1)
+                C[1] = C[1] * (cmap.N-1)
+                C = C.astype(int)
+                C = C[0] + cmap.N * C[1]
+                C = np.array(C)
+
             numRows, numCols = C.shape
             if allmatch:
                 X, Y = np.meshgrid(np.arange(numCols), np.arange(numRows))
@@ -5140,6 +5173,18 @@ def _pcolorargs(funcname, *args, **kw):
 
         if len(args) == 3:
             X, Y, C = [np.asanyarray(a) for a in args]
+            if (C.ndim == 3 and isinstance(norm, mcolors.BivariateNorm) or
+                    isinstance(cmap, mcolors.BivariateColormap)):
+                if cmap is None:
+                    cmap = mcolors.BivariateColormap()
+                if norm is None:
+                    norm = mcolors.BivariateNorm()
+                C = norm(C)
+                C[0] = C[0] * (cmap.N-1)
+                C[1] = C[1] * (cmap.N-1)
+                C = C.astype(int)
+                C = C[0] + cmap.N * C[1]
+                C = np.array(C)
             numRows, numCols = C.shape
         else:
             raise TypeError(
@@ -5325,9 +5370,14 @@ def pcolor(self, *args, **kwargs):
         vmin = kwargs.pop('vmin', None)
         vmax = kwargs.pop('vmax', None)
 
-        X, Y, C = self._pcolorargs('pcolor', *args, allmatch=False)
+        kw = {'norm': norm, 'cmap': cmap, 'allmatch': False}
+        X, Y, C = self._pcolorargs('pcolor', *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)
@@ -5393,9 +5443,13 @@ def pcolor(self, *args, **kwargs):
 
         collection.set_alpha(alpha)
         collection.set_array(C)
-        if norm is not None and not isinstance(norm, mcolors.Normalize):
-            msg = "'norm' must be an instance of 'mcolors.Normalize'"
+
+        isNorm = isinstance(norm, (mcolors.Normalize, mcolors.BivariateNorm))
+        if norm is not None and not isNorm:
+            msg = "'norm' must be an instance of 'mcolors.Normalize' " \
+                  "or 'mcolors.BivariateNorm'"
             raise ValueError(msg)
+
         collection.set_cmap(cmap)
         collection.set_norm(norm)
         collection.set_clim(vmin, vmax)
@@ -5525,9 +5579,14 @@ def pcolormesh(self, *args, **kwargs):
 
         allmatch = (shading == 'gouraud')
 
-        X, Y, C = self._pcolorargs('pcolormesh', *args, allmatch=allmatch)
+        kw = {'norm': norm, 'cmap': cmap, 'allmatch': allmatch}
+        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)
@@ -5666,11 +5725,28 @@ def pcolorfast(self, *args, **kwargs):
         cmap = kwargs.pop('cmap', None)
         vmin = kwargs.pop('vmin', None)
         vmax = kwargs.pop('vmax', None)
-        if norm is not None and not isinstance(norm, mcolors.Normalize):
-            msg = "'norm' must be an instance of 'mcolors.Normalize'"
+        isNorm = isinstance(norm, (mcolors.Normalize, mcolors.BivariateNorm))
+        if norm is not None and not isNorm:
+            msg = "'norm' must be an instance of 'mcolors.Normalize' " \
+                  "or 'mcolors.BivariateNorm'"
             raise ValueError(msg)
 
-        C = args[-1]
+        C = np.asarray(args[-1])
+
+        if (C.ndim == 3 and isinstance(norm, mcolors.BivariateNorm) or
+                isinstance(cmap, mcolors.BivariateColormap)):
+            if cmap is None:
+                cmap = mcolors.BivariateColormap()
+            if norm is None:
+                norm = mcolors.BivariateNorm()
+            C = norm(C)
+            C[0] = C[0] * (cmap.N-1)
+            C[1] = C[1] * (cmap.N-1)
+            C = C.astype(int)
+            C = C[0] + cmap.N * C[1]
+            C = np.array(C)
+            norm = mcolors.NoNorm()
+
         nr, nc = C.shape
         if len(args) == 1:
             style = "image"
diff --git a/lib/matplotlib/colors.py b/lib/matplotlib/colors.py
@@ -855,6 +855,32 @@ def reversed(self, name=None):
         return ListedColormap(colors_r, name=name, N=self.N)
 
 
+class BivariateColormap(Colormap):
+    def __init__(self, name, N=256):
+        Colormap.__init__(self, name, N)
+
+    def _init(self):
+        red = np.linspace(0, 1, self.N)
+        green = np.linspace(0, 1, self.N)
+        red_mesh, green_mesh = np.meshgrid(red, green)
+        blue_mesh = np.zeros_like(red_mesh)
+        alpha_mesh = np.ones_like(red_mesh)
+        bivariate_cmap = np.dstack((red_mesh, green_mesh, blue_mesh, alpha_mesh))
+        self._lut = np.vstack(bivariate_cmap)
+        self._isinit = True
+        self.N = self.N * self.N
+        self._set_extremes()
+
+    def _resample(self, lutsize):
+        """
+        Return a new color map with *lutsize x lutsize* entries.
+        """
+        return BivariateColormap(self.name, lutsize)
+
+    def reversed(self, name=None):
+        raise NotImplementedError()
+
+
 class Normalize(object):
     """
     A class which, when called, can normalize data into
@@ -1387,8 +1413,8 @@ def __call__(self, values, clip=None):
         if clip is None:
             clip = [self.norm1.clip, self.norm2.clip]
 
-        return [self.norm1(values[0], clip=clip[0]),
-                self.norm2(values[1], clip=clip[1])]
+        return np.array([self.norm1(values[0], clip=clip[0]),
+                self.norm2(values[1], clip=clip[1])])
 
     def inverse(self, values):
         """
@@ -1403,6 +1429,7 @@ def inverse(self, values):
         """
         return [self.norm1.inverse(values[0]), self.norm.inverse(values[1])]
 
+
 def rgb_to_hsv(arr):
     """
     convert float rgb values (in the range [0, 1]), in a numpy array to hsv
