Allows one to not clip to 0...1 in colors.MultivarColormap.__call__()

trygvrad · trygvrad · commit f5f5890b7a98 · 2024-07-03T13:20:28.000+02:00
Also adds a an improvement to colors.BIvarColormap.__getitem__() so that this returns a ListedColormap object instead of a Colormap object
diff --git a/lib/matplotlib/colors.py b/lib/matplotlib/colors.py
@@ -1295,7 +1295,7 @@ def __init__(self, name, colormaps, combination_mode):
         self.n_variates = len(colormaps)
         self._rgba_bad = (0.0, 0.0, 0.0, 0.0)  # If bad, don't paint anything.
 
-    def __call__(self, X, alpha=None, bytes=False):
+    def __call__(self, X, alpha=None, bytes=False, clip=True):
         r"""
         Parameters
         ----------
@@ -1315,6 +1315,8 @@ def __call__(self, X, alpha=None, bytes=False):
             If False (default), the returned RGBA values will be floats in the
             interval ``[0, 1]`` otherwise they will be `numpy.uint8`\s in the
             interval ``[0, 255]``.
+        clip : bool
+            If True, clip output to 0 to 1
 
         Returns
         -------
@@ -1340,17 +1342,24 @@ def __call__(self, X, alpha=None, bytes=False):
 
         rgba[mask_bad] = self.get_bad()
 
-        rgba = np.clip(rgba, 0, 1)
+        if clip:
+            rgba = np.clip(rgba, 0, 1)
 
         if alpha is not None:
-            alpha = np.clip(alpha, 0, 1)
+            if clip:
+                alpha = np.clip(alpha, 0, 1)
             if alpha.shape not in [(), np.array(X[0]).shape]:
                 raise ValueError(
                     f"alpha is array-like but its shape {alpha.shape} does "
                     f"not match that of X[0] {np.array(X[0]).shape}")
             rgba[..., -1] *= alpha
 
         if bytes:
+            if not clip:
+                raise ValueError(
+                    "clip cannot be false while bytes is true"
+                    " as uint8 does not support values below 0"
+                    " or above 255.")
             rgba = (rgba * 255).astype('uint8')
 
         if not np.iterable(X[0]):
@@ -1736,24 +1745,20 @@ def __getitem__(self, item):
         if not self._isinit:
             self._init()
         if item == 0:
-            cmap = Colormap(self.name+'0', self.N)
             one_d_lut = self._lut[:, self._origin[1]]
+            new_cmap = ListedColormap(one_d_lut, name=self.name+'_0', N=self.N)
+
         elif item == 1:
-            cmap = Colormap(self.name+'1', self.M)
             one_d_lut = self._lut[self._origin[0], :]
+            new_cmap = ListedColormap(one_d_lut, name=self.name+'_1', N=self.M)
         else:
             raise KeyError(f"only 0 or 1 are"
                            f" valid keys for BivarColormap, not {item!r}")
-        cmap._lut = np.zeros((self.N + 3, 4), float)
-        cmap._lut[:-3] = one_d_lut
-        cmap.set_bad(self._rgba_bad)
-        self._rgba_outside
+        new_cmap._rgba_bad = self._rgba_bad
         if self.shape in ['ignore', 'circleignore']:
-            cmap.set_under(self._rgba_outside)
-            cmap.set_over(self._rgba_outside)
-        cmap._set_extremes()
-        cmap._isinit = True
-        return cmap
+            new_cmap.set_over(self._rgba_outside)
+            new_cmap.set_under(self._rgba_outside)
+        return new_cmap
 
     def _repr_png_(self):
         """Generate a PNG representation of the BivarColormap."""
