Hide keyboard shortcuts

Hot-keys on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

1""" 

2Builtin colormaps, colormap handling utilities, and the `ScalarMappable` mixin. 

3 

4.. seealso:: 

5 

6 :doc:`/gallery/color/colormap_reference` for a list of builtin colormaps. 

7 

8 :doc:`/tutorials/colors/colormap-manipulation` for examples of how to 

9 make colormaps. 

10 

11 :doc:`/tutorials/colors/colormaps` an in-depth discussion of 

12 choosing colormaps. 

13 

14 :doc:`/tutorials/colors/colormapnorms` for more details about data 

15 normalization. 

16""" 

17 

18import functools 

19 

20import numpy as np 

21from numpy import ma 

22 

23import matplotlib as mpl 

24import matplotlib.colors as colors 

25import matplotlib.cbook as cbook 

26from matplotlib._cm import datad 

27from matplotlib._cm_listed import cmaps as cmaps_listed 

28 

29 

30def _reverser(f, x): # Deprecated, remove this at the same time as revcmap. 

31 return f(1 - x) # Toplevel helper for revcmap ensuring cmap picklability. 

32 

33 

34@cbook.deprecated("3.2", alternative="Colormap.reversed()") 

35def revcmap(data): 

36 """Can only handle specification *data* in dictionary format.""" 

37 data_r = {} 

38 for key, val in data.items(): 

39 if callable(val): 

40 # Return a partial object so that the result is picklable. 

41 valnew = functools.partial(_reverser, val) 

42 else: 

43 # Flip x and exchange the y values facing x = 0 and x = 1. 

44 valnew = [(1.0 - x, y1, y0) for x, y0, y1 in reversed(val)] 

45 data_r[key] = valnew 

46 return data_r 

47 

48 

49LUTSIZE = mpl.rcParams['image.lut'] 

50 

51 

52def _gen_cmap_d(): 

53 """ 

54 Generate a dict mapping standard colormap names to standard colormaps, as 

55 well as the reversed colormaps. 

56 """ 

57 cmap_d = {**cmaps_listed} 

58 for name, spec in datad.items(): 

59 cmap_d[name] = ( # Precache the cmaps at a fixed lutsize.. 

60 colors.LinearSegmentedColormap(name, spec, LUTSIZE) 

61 if 'red' in spec else 

62 colors.ListedColormap(spec['listed'], name) 

63 if 'listed' in spec else 

64 colors.LinearSegmentedColormap.from_list(name, spec, LUTSIZE)) 

65 # Generate reversed cmaps. 

66 for cmap in list(cmap_d.values()): 

67 rmap = cmap.reversed() 

68 cmap_d[rmap.name] = rmap 

69 return cmap_d 

70 

71 

72cmap_d = _gen_cmap_d() 

73locals().update(cmap_d) 

74 

75 

76# Continue with definitions ... 

77 

78 

79def register_cmap(name=None, cmap=None, data=None, lut=None): 

80 """ 

81 Add a colormap to the set recognized by :func:`get_cmap`. 

82 

83 It can be used in two ways:: 

84 

85 register_cmap(name='swirly', cmap=swirly_cmap) 

86 

87 register_cmap(name='choppy', data=choppydata, lut=128) 

88 

89 In the first case, *cmap* must be a :class:`matplotlib.colors.Colormap` 

90 instance. The *name* is optional; if absent, the name will 

91 be the :attr:`~matplotlib.colors.Colormap.name` attribute of the *cmap*. 

92 

93 In the second case, the three arguments are passed to 

94 the :class:`~matplotlib.colors.LinearSegmentedColormap` initializer, 

95 and the resulting colormap is registered. 

96 """ 

97 cbook._check_isinstance((str, None), name=name) 

98 if name is None: 

99 try: 

100 name = cmap.name 

101 except AttributeError: 

102 raise ValueError("Arguments must include a name or a Colormap") 

103 if isinstance(cmap, colors.Colormap): 

104 cmap_d[name] = cmap 

105 return 

106 # For the remainder, let exceptions propagate. 

107 if lut is None: 

108 lut = mpl.rcParams['image.lut'] 

109 cmap = colors.LinearSegmentedColormap(name, data, lut) 

110 cmap_d[name] = cmap 

111 

112 

113def get_cmap(name=None, lut=None): 

114 """ 

115 Get a colormap instance, defaulting to rc values if *name* is None. 

116 

117 Colormaps added with :func:`register_cmap` take precedence over 

118 built-in colormaps. 

119 

120 Parameters 

121 ---------- 

122 name : `matplotlib.colors.Colormap` or str or None, default: None 

123 If a `Colormap` instance, it will be returned. Otherwise, the name of 

124 a colormap known to Matplotlib, which will be resampled by *lut*. The 

125 default, None, means :rc:`image.cmap`. 

126 lut : int or None, default: None 

127 If *name* is not already a Colormap instance and *lut* is not None, the 

128 colormap will be resampled to have *lut* entries in the lookup table. 

129 """ 

130 if name is None: 

131 name = mpl.rcParams['image.cmap'] 

132 if isinstance(name, colors.Colormap): 

133 return name 

134 cbook._check_in_list(sorted(cmap_d), name=name) 

135 if lut is None: 

136 return cmap_d[name] 

137 else: 

138 return cmap_d[name]._resample(lut) 

139 

140 

141class ScalarMappable: 

142 """ 

143 This is a mixin class to support scalar data to RGBA mapping. 

144 The ScalarMappable makes use of data normalization before returning 

145 RGBA colors from the given colormap. 

146 

147 """ 

148 def __init__(self, norm=None, cmap=None): 

149 """ 

150 

151 Parameters 

152 ---------- 

153 norm : :class:`matplotlib.colors.Normalize` instance 

154 The normalizing object which scales data, typically into the 

155 interval ``[0, 1]``. 

156 If *None*, *norm* defaults to a *colors.Normalize* object which 

157 initializes its scaling based on the first data processed. 

158 cmap : str or :class:`~matplotlib.colors.Colormap` instance 

159 The colormap used to map normalized data values to RGBA colors. 

160 """ 

161 

162 self.callbacksSM = cbook.CallbackRegistry() 

163 

164 if cmap is None: 

165 cmap = get_cmap() 

166 if norm is None: 

167 norm = colors.Normalize() 

168 

169 self._A = None 

170 #: The Normalization instance of this ScalarMappable. 

171 self.norm = norm 

172 #: The Colormap instance of this ScalarMappable. 

173 self.cmap = get_cmap(cmap) 

174 #: The last colorbar associated with this ScalarMappable. May be None. 

175 self.colorbar = None 

176 self.update_dict = {'array': False} 

177 

178 def to_rgba(self, x, alpha=None, bytes=False, norm=True): 

179 """ 

180 Return a normalized rgba array corresponding to *x*. 

181 

182 In the normal case, *x* is a 1-D or 2-D sequence of scalars, and 

183 the corresponding ndarray of rgba values will be returned, 

184 based on the norm and colormap set for this ScalarMappable. 

185 

186 There is one special case, for handling images that are already 

187 rgb or rgba, such as might have been read from an image file. 

188 If *x* is an ndarray with 3 dimensions, 

189 and the last dimension is either 3 or 4, then it will be 

190 treated as an rgb or rgba array, and no mapping will be done. 

191 The array can be uint8, or it can be floating point with 

192 values in the 0-1 range; otherwise a ValueError will be raised. 

193 If it is a masked array, the mask will be ignored. 

194 If the last dimension is 3, the *alpha* kwarg (defaulting to 1) 

195 will be used to fill in the transparency. If the last dimension 

196 is 4, the *alpha* kwarg is ignored; it does not 

197 replace the pre-existing alpha. A ValueError will be raised 

198 if the third dimension is other than 3 or 4. 

199 

200 In either case, if *bytes* is *False* (default), the rgba 

201 array will be floats in the 0-1 range; if it is *True*, 

202 the returned rgba array will be uint8 in the 0 to 255 range. 

203 

204 If norm is False, no normalization of the input data is 

205 performed, and it is assumed to be in the range (0-1). 

206 

207 """ 

208 # First check for special case, image input: 

209 try: 

210 if x.ndim == 3: 

211 if x.shape[2] == 3: 

212 if alpha is None: 

213 alpha = 1 

214 if x.dtype == np.uint8: 

215 alpha = np.uint8(alpha * 255) 

216 m, n = x.shape[:2] 

217 xx = np.empty(shape=(m, n, 4), dtype=x.dtype) 

218 xx[:, :, :3] = x 

219 xx[:, :, 3] = alpha 

220 elif x.shape[2] == 4: 

221 xx = x 

222 else: 

223 raise ValueError("third dimension must be 3 or 4") 

224 if xx.dtype.kind == 'f': 

225 if norm and (xx.max() > 1 or xx.min() < 0): 

226 raise ValueError("Floating point image RGB values " 

227 "must be in the 0..1 range.") 

228 if bytes: 

229 xx = (xx * 255).astype(np.uint8) 

230 elif xx.dtype == np.uint8: 

231 if not bytes: 

232 xx = xx.astype(np.float32) / 255 

233 else: 

234 raise ValueError("Image RGB array must be uint8 or " 

235 "floating point; found %s" % xx.dtype) 

236 return xx 

237 except AttributeError: 

238 # e.g., x is not an ndarray; so try mapping it 

239 pass 

240 

241 # This is the normal case, mapping a scalar array: 

242 x = ma.asarray(x) 

243 if norm: 

244 x = self.norm(x) 

245 rgba = self.cmap(x, alpha=alpha, bytes=bytes) 

246 return rgba 

247 

248 def set_array(self, A): 

249 """Set the image array from numpy array *A*. 

250 

251 Parameters 

252 ---------- 

253 A : ndarray 

254 """ 

255 self._A = A 

256 self.update_dict['array'] = True 

257 

258 def get_array(self): 

259 'Return the array' 

260 return self._A 

261 

262 def get_cmap(self): 

263 'return the colormap' 

264 return self.cmap 

265 

266 def get_clim(self): 

267 'return the min, max of the color limits for image scaling' 

268 return self.norm.vmin, self.norm.vmax 

269 

270 def set_clim(self, vmin=None, vmax=None): 

271 """ 

272 Set the norm limits for image scaling. 

273 

274 Parameters 

275 ---------- 

276 vmin, vmax : float 

277 The limits. 

278 

279 The limits may also be passed as a tuple (*vmin*, *vmax*) as a 

280 single positional argument. 

281 

282 .. ACCEPTS: (vmin: float, vmax: float) 

283 """ 

284 if vmax is None: 

285 try: 

286 vmin, vmax = vmin 

287 except (TypeError, ValueError): 

288 pass 

289 if vmin is not None: 

290 self.norm.vmin = colors._sanitize_extrema(vmin) 

291 if vmax is not None: 

292 self.norm.vmax = colors._sanitize_extrema(vmax) 

293 self.changed() 

294 

295 def get_alpha(self): 

296 """ 

297 Returns 

298 ------- 

299 alpha : float 

300 Always returns 1. 

301 """ 

302 # This method is intended to be overridden by Artist sub-classes 

303 return 1. 

304 

305 def set_cmap(self, cmap): 

306 """ 

307 set the colormap for luminance data 

308 

309 Parameters 

310 ---------- 

311 cmap : colormap or registered colormap name 

312 """ 

313 cmap = get_cmap(cmap) 

314 self.cmap = cmap 

315 self.changed() 

316 

317 def set_norm(self, norm): 

318 """Set the normalization instance. 

319 

320 Parameters 

321 ---------- 

322 norm : `.Normalize` 

323 

324 Notes 

325 ----- 

326 If there are any colorbars using the mappable for this norm, setting 

327 the norm of the mappable will reset the norm, locator, and formatters 

328 on the colorbar to default. 

329 

330 """ 

331 cbook._check_isinstance((colors.Normalize, None), norm=norm) 

332 if norm is None: 

333 norm = colors.Normalize() 

334 self.norm = norm 

335 self.changed() 

336 

337 def autoscale(self): 

338 """ 

339 Autoscale the scalar limits on the norm instance using the 

340 current array 

341 """ 

342 if self._A is None: 

343 raise TypeError('You must first set_array for mappable') 

344 self.norm.autoscale(self._A) 

345 self.changed() 

346 

347 def autoscale_None(self): 

348 """ 

349 Autoscale the scalar limits on the norm instance using the 

350 current array, changing only limits that are None 

351 """ 

352 if self._A is None: 

353 raise TypeError('You must first set_array for mappable') 

354 self.norm.autoscale_None(self._A) 

355 self.changed() 

356 

357 def add_checker(self, checker): 

358 """ 

359 Add an entry to a dictionary of boolean flags 

360 that are set to True when the mappable is changed. 

361 """ 

362 self.update_dict[checker] = False 

363 

364 def check_update(self, checker): 

365 """ 

366 If mappable has changed since the last check, 

367 return True; else return False 

368 """ 

369 if self.update_dict[checker]: 

370 self.update_dict[checker] = False 

371 return True 

372 return False 

373 

374 def changed(self): 

375 """ 

376 Call this whenever the mappable is changed to notify all the 

377 callbackSM listeners to the 'changed' signal 

378 """ 

379 self.callbacksSM.process('changed', self) 

380 

381 for key in self.update_dict: 

382 self.update_dict[key] = True 

383 self.stale = True