__getitem__ for colors.BivarColormap

trygvrad · trygvrad · commit 991610f2188e · 2024-07-02T15:51:17.000+02:00
Adds support for __getitem__ on colors.BivarColormap, i.e.: BivarColormap[0] and BivarColormap[1], which returns (1D) Colormap objects along the selected axes
diff --git a/lib/matplotlib/_cm_bivar.py b/lib/matplotlib/_cm_bivar.py
@@ -1305,8 +1305,8 @@
 
 cmaps = {
     "BiPeak": SegmentedBivarColormap(
-        BiPeak, "BiPeak", 256, "square"),
+        BiPeak, "BiPeak", 256, "square", (128, 128)),
     "BiOrangeBlue": SegmentedBivarColormap(
-        BiOrangeBlue, "BiOrangeBlue", 256, "square"),
-    "BiCone": SegmentedBivarColormap(BiPeak, "BiCone", 256, "circle"),
+        BiOrangeBlue, "BiOrangeBlue", 256, "square", (0, 0)),
+    "BiCone": SegmentedBivarColormap(BiPeak, "BiCone", 256, "circle", (128, 128)),
 }
diff --git a/lib/matplotlib/colors.py b/lib/matplotlib/colors.py
@@ -1414,7 +1414,7 @@ class BivarColormap(ColormapBase):
     lookup table. To be used with `~matplotlib.cm.VectorMappable`.
     """
 
-    def __init__(self, name, N=256, M=256, shape='square'):
+    def __init__(self, name, N=256, M=256, shape='square', origin=(0, 0)):
         """
         Parameters
         ----------
@@ -1436,6 +1436,10 @@ def __init__(self, name, N=256, M=256, shape='square'):
             - If  'circleignore' a circular mask is applied, but the data is not
               clipped and instead assigned the 'outside' color
 
+        origin: (int, int)
+            The relative origin of the colormap. Typically (0, 0), for colormaps
+            that are linear on both axis, and (int(N*0.5), int(.5*M) for
+            circular colormaps.
         """
 
         self.name = name
@@ -1446,6 +1450,7 @@ def __init__(self, name, N=256, M=256, shape='square'):
         self._rgba_outside = (1.0, 0.0, 1.0, 1.0)
         self._isinit = False
         self.n_variates = 2
+        self._origin = origin
         '''#: When this colormap exists on a scalar mappable and colorbar_extend
         #: is not False, colorbar creation will pick up ``colorbar_extend`` as
         #: the default value for the ``extend`` keyword in the
@@ -1692,6 +1697,30 @@ def _clip(self, X):
                 X[0][mask_outside] = -1
                 X[1][mask_outside] = -1
 
+    def __getitem__(self, item):
+        """Creates and returns a colorbar along the selected axis"""
+        if not self._isinit:
+            self._init()
+        if item == 0:
+            cmap = Colormap(self.name+'0', self.N)
+            one_d_lut = self._lut[:, self._origin[1]]
+        elif item == 1:
+            cmap = Colormap(self.name+'1', self.M)
+            one_d_lut = self._lut[self._origin[0], :]
+        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
+        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
+
     def _repr_png_(self):
         """Generate a PNG representation of the BivarColormap."""
         if not self._isinit:
@@ -1769,11 +1798,15 @@ class SegmentedBivarColormap(BivarColormap):
           'outside' color
         - If 'circleignore' a circular mask is applied, but the data is not clipped
 
+    origin: (int, int)
+        The relative origin of the colormap. Typically (0, 0), for colormaps
+        that are linear on both axis, and (int(N*0.5), int(.5*M) for
+        circular colormaps.
     """
 
-    def __init__(self, patch, name, N=256, shape='square'):
+    def __init__(self, patch, name, N=256, shape='square', origin=(0, 0)):
         self.patch = patch
-        super().__init__(name, N, N, shape)
+        super().__init__(name, N, N, shape, origin)
 
     def _init(self):
         s = self.patch.shape
@@ -1809,9 +1842,13 @@ class BivarColormapFromImage(BivarColormap):
           'outside' color
         - If 'circleignore' a circular mask is applied, but the data is not clipped
 
+    origin: (int, int)
+        The relative origin of the colormap. Typically (0, 0), for colormaps
+        that are linear on both axis, and (int(N*0.5), int(.5*M) for
+        circular colormaps.
     """
 
-    def __init__(self, lut, name='', shape='square'):
+    def __init__(self, lut, name='', shape='square', origin=(0, 0)):
         # We can allow for a PIL.Image as unput in the following way, but importing
         # matplotlib.image.pil_to_array() results in a circular import
         # For now, this function only accepts numpy arrays.
@@ -1822,7 +1859,7 @@ def __init__(self, lut, name='', shape='square'):
             raise ValueError("The lut must be an array of shape (n, m, 3) or (n, m, 4)",
                              " or a PIL.image encoded as RGB or RGBA")
         self._lut = lut
-        super().__init__(name, lut.shape[0], lut.shape[1], shape)
+        super().__init__(name, lut.shape[0], lut.shape[1], shape, origin)
 
     def _init(self):
         self._isinit = True
diff --git a/lib/matplotlib/colors.pyi b/lib/matplotlib/colors.pyi
@@ -173,7 +173,8 @@ class BivarColormap(ColormapBase):
     M: int
     shape: str
     n_variates: int
-    def __init__(self, name: str, N: int = ..., M: int | None = ..., shape: str = ...) -> None: ...
+    def __init__(self, name: str, N: int = ..., M: int | None = ..., shape: str = ..., origin: tuple[int, int] = ...
+    ) -> None: ...
     @overload
     def __call__(
         self, X: Sequence[Sequence[float]] | np.ndarray, alpha: ArrayLike | None = ..., bytes: bool = ...
@@ -198,11 +199,12 @@ class BivarColormap(ColormapBase):
 
 class SegmentedBivarColormap(BivarColormap):
     def __init__(
-        self, patch: np.ndarray, name: str, N: int = ..., shape: str = ...,
+        self, patch: np.ndarray, name: str, N: int = ..., shape: str = ..., origin: tuple[int, int] = ...
     ) -> None: ...
 
 class BivarColormapFromImage(BivarColormap):
-    def __init__(self, lut: np.ndarray, name: str = ..., shape: str = ...) -> None: ...
+    def __init__(self, lut: np.ndarray, name: str = ..., shape: str = ..., origin: tuple[int, int] = ...
+    ) -> None: ...
 
 class Normalize:
     callbacks: cbook.CallbackRegistry
diff --git a/lib/matplotlib/tests/test_multivariate_colormaps.py b/lib/matplotlib/tests/test_multivariate_colormaps.py
@@ -281,7 +281,30 @@ def test_bivar_cmap_call():
                        match="only implemented for use with with floats"):
         cs = cmap([(0, 5, 9, 0, 0, 9), (0, 0, 0, 5, 11, 11)])
 
-
+def test_bivar_getitem():
+    '''Test __getitem__  on BivarColormap'''
+    xA = ([.0, .25, .5, .75, 1., -1, 2], [.5]*7)
+    xB = ([.5]*7, [.0, .25, .5, .75, 1., -1, 2])
+
+    cmaps = mpl.bivar_colormaps['BiPeak']
+    assert_array_equal(cmaps(xA), cmaps[0](xA[0]))
+    assert_array_equal(cmaps(xB), cmaps[1](xB[1]))
+
+    cmaps.shape = 'ignore'
+    assert_array_equal(cmaps(xA), cmaps[0](xA[0]))
+    assert_array_equal(cmaps(xB), cmaps[1](xB[1]))
+
+    xA = ([.0, .25, .5, .75, 1., -1, 2], [.0]*7)
+    xB = ([.0]*7, [.0, .25, .5, .75, 1., -1, 2])
+    cmaps = mpl.bivar_colormaps['BiOrangeBlue']
+    assert_array_equal(cmaps(xA), cmaps[0](xA[0]))
+    assert_array_equal(cmaps(xB), cmaps[1](xB[1]))
+
+    cmaps.shape = 'ignore'
+    assert_array_equal(cmaps(xA), cmaps[0](xA[0]))
+    assert_array_equal(cmaps(xB), cmaps[1](xB[1]))
+
+    
 def test_bivar_cmap_bad_shape():
     """
     Tests calling a bivariate colormap with integer values
