8000 Some ideas for a colormap __getitem__ method. · matplotlib/matplotlib@fa9f56a · GitHub
[go: up one dir, main page]

Skip to content

Commit fa9f56a

Browse files
committed
Some ideas for a colormap __getitem__ method.
1 parent 71742d7 commit fa9f56a

File tree

1 file changed

+84
-0
lines changed

1 file changed

+84
-0
lines changed

lib/matplotlib/colors.py

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -770,6 +770,90 @@ def truncate(self, minval=0.0, maxval=1.0, name=None, N=None):
770770
name = "trunc({},{:.2f},{:.2f})".format(self.name, minval, maxval)
771771
return ListedColormap(self(np.linspace(minval, maxval, N)), name)
772772

773+
def __getitem__(self, item):
774+
"""Advanced indexing for colorbars.
775+
776+
777+
Examples
778+
--------
779+
import matplotlib.pyplat as plt
780+
cmap = plt.get_cmap('viridis', 128)
781+
782+
# ### float indexing
783+
# for float-style indexing, the values must be in [0.0, 1.0]
784+
# Truncate the colormap between 20 and 70%
785+
new_cm = cmap[0.2, 0.7]
786+
# equivalently:
787+
new_cm = cmap[0.2:0.7]
788+
789+
# Same as above, but specify the number of points
790+
new_cm = cmap[0.2, 0.7, 64]
791+
# equivalently, use `np.mgrid` complex-indexing:
792+
new_cm = cmap[0.2:0.7:1j * 64]
793+
794+
# ### Int-style indexing
795+
# for int-style indexing, the values must be in [0, self.N]
796+
new_cm = cmap[12:100]
797+
798+
# Same as above, but 4x fewer points
799+
new_cm = cmap[12:100:4]
800+
801+
# negative values are supported (same as above)
802+
new_cm = cmap[12:-28:4]
803+
804+
# And so is `np.mgrid` complex-indexing (same as above)
805+
new_cm = cmap[12:-28:1j * 22]
806+
"""
807+
if isinstance(item, tuple):
808+
if len(item) == 2:
809+
N = self.N
810+
elif len(item) == 3:
811+
N = item[2]
812+
else:
813+
raise IndexError("Invalid colorbar itemization")
814+
return self.truncate(item[0], item[1], N=N)
815+
elif isinstance(item, slice):
816+
name = self.name + '[{:s}]'.format(str(item))
817+
sss = [item.start, item.stop, item.step]
818+
if any([isinstance(s, int) for s in sss]):
819+
# This is an integer-style itemization
820+
if sss[0] is None:
821+
sss[0] = 0
822+
elif sss[0] < 0:
823+
sss[0] = sss[0] % self.N
824+
if sss[1] is None:
825+
sss[1] = self.N
826+
elif sss[1] < 0:
827+
sss[1] = sss[1] % self.N
828+
if sss[2] is None:
829+
sss[2] = 1
830+
sss[0] = sss[0] / self.N
831+
sss[1] = sss[1] / self.N
832+
if not isinstance(sss[2], complex):
833+
sss[2] = sss[2] / self.N
834+
else:
835+
if sss[0] is None:
836+
sss[0] = 0.0
837+
if sss[1] is None:
838+
sss[1] = 1.0
839+
if sss[2] is None:
840+
sss[2] = self.N * 1j
841+
if sss[0] < 0 or sss[0] >= 1 or sss[1] <= 0 or sss[1] > 1:
842+
raise IndexError("Invalid colorbar itemization - outside bounds")
843+
points = np.mgrid[slice(*sss)]
844+
elif isinstance(item, (list, np.ndarray)):
845+
name = self.name + '[...]'
846+
if isinstance(item, list):
847+
item = np.array(item)
848+
if item.dtype.kind in ('u', 'i'):
849+
item = item.astype('f') / self.N
850+
points = item
851+
else:
852+
raise IndexError("Invalid colorbar itemization")
853+
if len(points) <= 1:
854+
raise IndexError("Invalid colorbar itemization - too few points")
855+
return ListedColormap(self(points), name=name)
856+
773857

774858
class LinearSegmentedColormap(Colormap):
775859
"""Colormap objects based on lookup tables using linear segments.

0 commit comments

Comments
 (0)
0