Skip to content

Commit 079a0d9

Browse files
committed
NF - Add 'truncate' and 'join' methods to colormaps.
This is based largely on https://gist.github.com/denis-bz/8052855 Which was based on http://stackoverflow.com/a/18926541/2121597
1 parent 07f6427 commit 079a0d9

8 files changed

+254
-1
lines changed

lib/matplotlib/colors.py

+187-1
Original file line numberDiff line numberDiff line change
@@ -742,6 +742,99 @@ def func_r(x):
742742

743743
return LinearSegmentedColormap(name, data_r, self.N, self._gamma)
744744

745+
def join(self, other, name=None, frac_self=None, N=None):
746+
"""
747+
Join colormap `self` to `other` and return the new colormap.
748+
749+
Parameters
750+
----------
751+
other : cmap
752+
The other colormap to be joined to this one.
753+
name : str, optional
754+
The name for the reversed colormap. If it's None the
755+
name will be ``self.name + '-' + other.name``.
756+
frac_self : float in the interval ``(0.0, 1.0)``, optional
757+
The fraction of the new colormap that should be occupied
758+
by self. By default, this is ``self.N / (self.N +
759+
other.N)``.
760+
761+
Returns
762+
-------
763+
LinearSegmentedColormap
764+
The joined colormap.
765+
766+
Examples
767+
--------
768+
import matplotlib.pyplat as plt
769+
cmap1 = plt.get_cmap('jet')
770+
cmap2 = plt.get_cmap('plasma_r')
771+
772+
joined_cmap = cmap1.join(cmap2)
773+
"""
774+
if N is None:
775+
N = self.N + other.N
776+
if frac_self is None:
777+
frac_self = self.N / (other.N + self.N)
778+
if name is None:
779+
name = '{}+{}'.format(self.name, other.name)
780+
assert 0 < frac_self and frac_self < 1, (
781+
"The parameter ``frac_self`` must be in the interval ``(0.0, 1.0)``."
782+
)
783+
map0 = self(np.linspace(0, 1, int(N * frac_self)))
784+
map1 = other(np.linspace(0, 1, int(N * (1 - frac_self))))
785+
# N is set by len of the vstack'd array:
786+
return LinearSegmentedColormap.from_list(name, np.vstack((map0, map1)))
787+
788+
def truncate(self, minval=0.0, maxval=1.0, N=None):
789+
"""
790+
Truncate a colormap.
791+
792+
Parameters
793+
----------
794+
minval : float in the interval ``(0.0, 1.0)``, optional
795+
The lower fraction of the colormap you want to truncate
796+
(default 0.0).
797+
798+
maxval : float in the interval ``(0.0, 1.0)``, optional
799+
The upper limit of the colormap you want to keep. i.e. truncate
800+
the section above this value (default 1.0).
801+
802+
N : int
803+
The number of entries in the map. The default is *None*,
804+
in which case the same color-step density is preserved,
805+
i.e.: N = ceil(N * (maxval - minval))
806+
807+
Returns
808+
-------
809+
LinearSegmentedColormap
810+
The truncated colormap.
811+
812+
Examples
813+
--------
814+
import matplotlib.pyplat as plt
815+
cmap = plt.get_cmap('jet')
816+
817+
# This will return the `jet` colormap with the bottom 20%,
818+
# and top 30% removed:
819+
cmap_trunc = cmap.truncate(0.2, 0.7)
820+
821+
"""
822+
assert minval < maxval, "``minval`` must be less than ``maxval``"
823+
assert 0 <= minval and minval < 1, (
824+
"The parameter ``minval`` must be in the interval ``(0.0, 1.0)``.")
825+
assert 0 < maxval and maxval <= 1, (
826+
"The parameter ``maxval`` must be in the interval ``(0.0, 1.0)``.")
827+
assert minval != 0 or maxval != 1, (
828+
"This is not a truncation.")
829+
# This was taken largely from
830+
# https://gist.github.com/denis-bz/8052855
831+
# Which, in turn was from @unutbu's SO answer:
832+
# http://stackoverflow.com/a/18926541/2121597
833+
if N is None:
834+
N = np.ceil(self.N * (maxval - minval))
835+
name = "trunc({},{:.2f},{:.2f})".format(self.name, minval, maxval)
836+
return LinearSegmentedColormap.from_list(name, self(np.linspace(minval, maxval, N)), N)
837+
745838

746839
class ListedColormap(Colormap):
747840
"""Colormap object generated from a list of colors.
@@ -835,6 +928,99 @@ def reversed(self, name=None):
835928
colors_r = list(reversed(self.colors))
836929
return ListedColormap(colors_r, name=name, N=self.N)
837930

931+
def join(self, other, name=None, frac_self=None, N=None):
932+
"""
933+
Join colormap `self` to `other` and return the new colormap.
934+
935+
Parameters
936+
----------
937+
other : cmap
938+
The other colormap to be joined to this one.
939+
name : str, optional
940+
The name for the reversed colormap. If it's None the
941+
name will be ``self.name + '-' + other.name``.
942+
frac_self : float in the interval ``(0.0, 1.0)``, optional
943+
The fraction of the new colormap that should be occupied
944+
by self. By default, this is ``self.N / (self.N +
945+
other.N)``.
946+
947+
Returns
948+
-------
949+
ListedColormap
950+
The joined colormap.
951+
952+
Examples
953+
--------
954+
import matplotlib.pyplat as plt
955+
cmap1 = plt.get_cmap('viridis')
956+
cmap2 = plt.get_cmap('plasma_r')
957+
958+
joined_cmap = cmap1.join(cmap2)
959+
"""
960+
if N is None:
961+
N = self.N + other.N
962+
if frac_self is None:
963+
frac_self = self.N / (other.N + self.N)
964+
if name is None:
965+
name = '{}+{}'.format(self.name, other.name)
966+
assert 0 < frac_self and frac_self < 1, (
967+
"The parameter ``frac_self`` must be in the interval ``(0.0, 1.0)``."
968+
)
969+
map0 = self(np.linspace(0, 1, int(N * frac_self)))
970+
map1 = other(np.linspace(0, 1, int(N * (1 - frac_self))))
971+
# N is set by len of the vstack'd array:
972+
return ListedColormap(np.vstack((map0, map1)), name, )
973+
974+
def truncate(self, minval=0.0, maxval=1.0, N=None):
975+
"""
976+
Truncate a colormap.
977+
978+
Parameters
979+
----------
980+
minval : float in the interval ``(0.0, 1.0)``, optional
981+
The lower fraction of the colormap you want to truncate
982+
(default 0.0).
983+
984+
maxval : float in the interval ``(0.0, 1.0)``, optional
985+
The upper limit of the colormap you want to keep. i.e. truncate
986+
the section above this value (default 1.0).
987+
988+
N : int
989+
The number of entries in the map. The default is *None*,
990+
in which case the same color-step density is preserved,
991+
i.e.: N = ceil(N * (maxval - minval))
992+
993+
Returns
994+
-------
995+
ListedColormap
996+
The truncated colormap.
997+
998+
Examples
999+
--------
1000+
import matplotlib.pyplat as plt
1001+
cmap = plt.get_cmap('viridis')
1002+
1003+
# This will return the `viridis` colormap with the bottom 20%,
1004+
# and top 30% removed:
1005+
cmap_trunc = cmap.truncate(0.2, 0.7)
1006+
1007+
"""
1008+
assert minval < maxval, "``minval`` must be less than ``maxval``"
1009+
assert 0 <= minval and minval < 1, (
1010+
"The parameter ``minval`` must be in the interval ``(0.0, 1.0)``.")
1011+
assert 0 < maxval and maxval <= 1, (
1012+
"The parameter ``maxval`` must be in the interval ``(0.0, 1.0)``.")
1013+
assert minval != 0 or maxval != 1, (
1014+
"This is not a truncation.")
1015+
# This was taken largely from
1016+
# https://gist.github.com/denis-bz/8052855
1017+
# Which, in turn was from @unutbu's SO answer:
1018+
# http://stackoverflow.com/a/18926541/2121597
1019+
if N is None:
1020+
N = np.ceil(self.N * (maxval - minval))
1021+
name = "trunc({},{:.2f},{:.2f})".format(self.name, minval, maxval)
1022+
return ListedColormap(self(np.linspace(minval, maxval, N)), name)
1023+
8381024

8391025
class Normalize(object):
8401026
"""
@@ -1040,7 +1226,7 @@ class SymLogNorm(Normalize):
10401226
*linthresh* allows the user to specify the size of this range
10411227
(-*linthresh*, *linthresh*).
10421228
"""
1043-
def __init__(self, linthresh, linscale=1.0,
1229+
def __init__(self, linthresh, linscale=1.0,
10441230
vmin=None, vmax=None, clip=False):
10451231
"""
10461232
*linthresh*:

lib/matplotlib/tests/test_colorbar.py

+67
Original file line numberDiff line numberDiff line change
@@ -197,6 +197,73 @@ def test_gridspec_make_colorbar():
197197
plt.subplots_adjust(top=0.95, right=0.95, bottom=0.2, hspace=0.25)
198198

199199

200+
@image_comparison(baseline_images=['colorbar_join_lsc',
201+
'colorbar_join_lsc_frac',
202+
'colorbar_join_listed',
203+
'colorbar_join_listed_frac', ],
204+
extensions=['png'], remove_text=True,
205+
savefig_kwarg={'dpi': 40})
206+
def test_join_colorbar():
207+
data = np.arange(1200).reshape(30, 40)
208+
levels = [0, 200, 400, 600, 800, 1000, 1200]
209+
210+
# Jet is a LinearSegmentedColormap
211+
cmap_lsc = plt.get_cmap('jet', 16)
212+
cmap_lst = plt.get_cmap('viridis', 16)
213+
214+
# join returns the same type of cmap as self
215+
# Thus, this is a lsc
216+
cmap = cmap_lsc.join(cmap_lst)
217+
218+
plt.figure()
219+
plt.contourf(data, levels=levels, cmap=cmap)
220+
plt.colorbar(orientation='vertical')
221+
222+
# Use the 'frac_self' kwarg for the lsc cmap
223+
cmap = cmap_lsc.join(cmap_lst, frac_self=0.7)
224+
225+
plt.figure()
226+
plt.contourf(data, levels=levels, cmap=cmap)
227+
plt.colorbar(orientation='vertical')
228+
229+
# This should be a listed colormap.
230+
cmap = cmap_lst.join(cmap_lsc)
231+
232+
plt.figure()
233+
plt.contourf(data, levels=levels, cmap=cmap)
234+
plt.colorbar(orientation='vertical')
235+
236+
# Use the 'frac_self' kwarg for the listed cmap
237+
cmap = cmap_lst.join(cmap_lsc, frac_self=0.7)
238+
239+
plt.figure()
240+
plt.contourf(data, levels=levels, cmap=cmap)
241+
plt.colorbar(orientation='vertical')
242+
243+
244+
@image_comparison(baseline_images=['colorbar_truncate_lsc',
245+
'colorbar_truncate_listed', ],
246+
extensions=['png'], remove_text=True,
247+
savefig_kwarg={'dpi': 40})
248+
def test_truncate_colorbar():
249+
data = np.arange(1200).reshape(30, 40)
250+
levels = [0, 200, 400, 600, 800, 1000, 1200]
251+
252+
# jet is a LinearSegmentedColormap
253+
cmap = plt.get_cmap('jet', 16).truncate(0.2, 0.7)
254+
255+
plt.figure()
256+
plt.contourf(data, levels=levels, cmap=cmap)
257+
plt.colorbar(orientation='vertical')
258+
259+
# viridis is a ListedColormap
260+
cmap = plt.get_cmap('viridis', 16).truncate(0.2, 0.7)
261+
262+
plt.figure()
263+
plt.contourf(data, levels=levels, cmap=cmap)
264+
plt.colorbar(orientation='vertical')
265+
266+
200267
@image_comparison(baseline_images=['colorbar_single_scatter'],
201268
extensions=['png'], remove_text=True,
202269
savefig_kwarg={'dpi': 40})

0 commit comments

Comments
 (0)