Coverage for src/CSET/operators/_colormaps.py: 99%

244 statements  

« prev     ^ index     » next       coverage.py v7.15.2, created at 2026-07-20 13:56 +0000

1# © Crown copyright, Met Office (2022-2026) and CSET contributors. 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# http://www.apache.org/licenses/LICENSE-2.0 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14 

15"""Functions to support colormap settings for CSET plots.""" 

16 

17import functools 

18import importlib.resources 

19import itertools 

20import json 

21import logging 

22from typing import Literal 

23 

24import iris 

25import matplotlib as mpl 

26import matplotlib.colors as mcolors 

27import matplotlib.pyplot as plt 

28import numpy as np 

29 

30from CSET._common import ( 

31 combine_dicts, 

32 get_recipe_metadata, 

33 iter_maybe, 

34) 

35 

36DEFAULT_DISCRETE_COLORS = mpl.colormaps["tab10"].colors + mpl.colormaps["Accent"].colors 

37 

38 

39@functools.cache 

40def load_colorbar_map(user_colorbar_file: str = None) -> dict: 

41 """Load the colorbar definitions from a file. 

42 

43 This is a separate function to make it cacheable. 

44 """ 

45 colorbar_file = importlib.resources.files().joinpath("_colorbar_definition.json") 

46 with open(colorbar_file, "rt", encoding="UTF-8") as fp: 

47 colorbar = json.load(fp) 

48 

49 logging.debug("User colour bar file: %s", user_colorbar_file) 

50 override_colorbar = {} 

51 if user_colorbar_file: 

52 try: 

53 with open(user_colorbar_file, "rt", encoding="UTF-8") as fp: 

54 override_colorbar = json.load(fp) 

55 except FileNotFoundError: 

56 logging.warning("Colorbar file does not exist. Using default values.") 

57 

58 # Overwrite values with the user supplied colorbar definition. 

59 colorbar = combine_dicts(colorbar, override_colorbar) 

60 return colorbar 

61 

62 

63def get_model_colors_map(cubes: iris.cube.CubeList | iris.cube.Cube) -> dict: 

64 """Get an appropriate colors for model lines in line plots. 

65 

66 For each model in the list of cubes colors either from user provided 

67 color definition file (so-called style file) or from default colors are mapped 

68 to model_name attribute. 

69 

70 Parameters 

71 ---------- 

72 cubes: CubeList or Cube 

73 Cubes with model_name attribute 

74 

75 Returns 

76 ------- 

77 model_colors_map: 

78 Dictionary mapping model_name attribute to colors 

79 """ 

80 user_colorbar_file = get_recipe_metadata().get("style_file_path", None) 

81 colorbar = load_colorbar_map(user_colorbar_file) 

82 model_names = sorted( 

83 filter( 

84 lambda x: x is not None, 

85 (cube.attributes.get("model_name", None) for cube in iter_maybe(cubes)), 

86 ) 

87 ) 

88 if not model_names: 

89 return {} 

90 use_user_colors = all(mname in colorbar.keys() for mname in model_names) 

91 if use_user_colors: 91 ↛ 92line 91 didn't jump to line 92 because the condition on line 91 was never true

92 return {mname: colorbar[mname] for mname in model_names} 

93 

94 color_list = itertools.cycle(DEFAULT_DISCRETE_COLORS) 

95 return {mname: color for mname, color in zip(model_names, color_list, strict=False)} 

96 

97 

98def colorbar_map_levels(cube: iris.cube.Cube, axis: Literal["x", "y"] | None = None): 

99 """Get an appropriate colorbar for the given cube. 

100 

101 For the given variable the appropriate colorbar is looked up from a 

102 combination of the built-in CSET colorbar definitions, and any user supplied 

103 definitions. As well as varying on variables, these definitions may also 

104 exist for specific pressure levels to account for variables with 

105 significantly different ranges at different heights. The colorbars also exist 

106 for masks and mask differences for considering variable presence diagnostics. 

107 Specific variable ranges can be separately set in user-supplied definition 

108 for x- or y-axis limits, or indicate where automated range preferred. 

109 

110 Parameters 

111 ---------- 

112 cube: Cube 

113 Cube of variable for which the colorbar information is desired. 

114 axis: "x", "y", optional 

115 Select the levels for just this axis of a line plot. The min and max 

116 can be set by xmin/xmax or ymin/ymax respectively. For variables where 

117 setting a universal range is not desirable (e.g. temperature), users 

118 can set ymin/ymax values to "auto" in the colorbar definitions file. 

119 Where no additional xmin/xmax or ymin/ymax values are provided, the 

120 axis bounds default to use the vmin/vmax values provided. 

121 

122 Returns 

123 ------- 

124 cmap: 

125 Matplotlib colormap. 

126 levels: 

127 List of levels to use for plotting. For continuous plots the min and max 

128 should be taken as the range. 

129 norm: 

130 BoundaryNorm information. 

131 """ 

132 # Grab the colorbar file from the recipe global metadata. 

133 user_colorbar_file = get_recipe_metadata().get("style_file_path", None) 

134 colorbar = load_colorbar_map(user_colorbar_file) 

135 cmap = None 

136 

137 try: 

138 # We assume that pressure is a scalar coordinate here. 

139 pressure_level_raw = cube.coord("pressure").points[0] 

140 # Ensure pressure_level is a string, as it is used as a JSON key. 

141 pressure_level = str(int(pressure_level_raw)) 

142 except iris.exceptions.CoordinateNotFoundError: 

143 pressure_level = None 

144 

145 # First try long name, then standard name, then var name. This order is used 

146 # as long name is the one we correct between models, so it most likely to be 

147 # consistent. 

148 varnames = list(filter(None, [cube.long_name, cube.standard_name, cube.var_name])) 

149 # Treat observation-labelled var names consistently with model var names. 

150 varnames = [varname.replace("observed_", "") for varname in varnames] 

151 for varname in varnames: 

152 # Get the colormap for this variable. 

153 try: 

154 var_colorbar = colorbar[varname] 

155 cmap = plt.get_cmap(colorbar[varname]["cmap"], 51) 

156 varname_key = varname 

157 break 

158 except KeyError: 

159 logging.debug("Cube name %s has no colorbar definition.", varname) 

160 

161 # Get colormap if it is a mask. 

162 if any("mask_for_" in name for name in varnames): 

163 cmap, levels, norm = custom_colormap_mask(cube, axis=axis) 

164 return cmap, levels, norm 

165 # If winds on Beaufort Scale use custom colorbar and levels 

166 if any("Beaufort_Scale" in name for name in varnames): 

167 cmap, levels, norm = custom_beaufort_scale(cube, axis=axis) 

168 return cmap, levels, norm 

169 # If probability is plotted use custom colorbar and levels 

170 if any("probability_of_" in name for name in varnames): 

171 cmap, levels, norm = custom_colormap_probability(cube, axis=axis) 

172 return cmap, levels, norm 

173 # If aviation colour state use custom colorbar and levels 

174 if any("aviation_colour_state" in name for name in varnames): 

175 cmap, levels, norm = custom_colormap_aviation_colour_state(cube) 

176 return cmap, levels, norm 

177 # If verification scores use custom colorbar 

178 if any("RMSE_" in name for name in varnames): 

179 cmap, levels, norm = custom_colormap_scores(cube) 

180 return cmap, levels, norm 

181 

182 # If no valid colormap has been defined, use defaults and return. 

183 if not cmap: 

184 logging.warning("No colorbar definition exists for %s.", cube.name()) 

185 cmap, levels, norm = mpl.colormaps["viridis"], None, None 

186 return cmap, levels, norm 

187 

188 # Test if pressure-level specific settings are provided for cube. 

189 if pressure_level: 

190 try: 

191 var_colorbar = colorbar[varname_key]["pressure_levels"][pressure_level] 

192 except KeyError: 

193 logging.debug( 

194 "%s has no colorbar definition for pressure level %s.", 

195 varname, 

196 pressure_level, 

197 ) 

198 

199 # Check for availability of x-axis or y-axis user-specific overrides 

200 # for setting level bounds for line plot types and return just levels. 

201 # Line plots do not need a colormap, and just use the data range. 

202 if axis: 

203 if axis == "x": 

204 try: 

205 vmin, vmax = var_colorbar["xmin"], var_colorbar["xmax"] 

206 except KeyError: 

207 vmin, vmax = var_colorbar["min"], var_colorbar["max"] 

208 if axis == "y": 

209 try: 

210 vmin, vmax = var_colorbar["ymin"], var_colorbar["ymax"] 

211 except KeyError: 

212 vmin, vmax = var_colorbar["min"], var_colorbar["max"] 

213 # Check if user-specified auto-scaling for this variable 

214 if vmin == "auto" or vmax == "auto": 

215 levels = None 

216 else: 

217 levels = [vmin, vmax] 

218 return None, levels, None 

219 # Get and use the colorbar levels for this variable if spatial or histogram. 

220 else: 

221 try: 

222 levels = var_colorbar["levels"] 

223 # Use discrete bins when levels are specified, rather 

224 # than a smooth range. 

225 norm = mpl.colors.BoundaryNorm(levels, ncolors=cmap.N) 

226 logging.debug("Using levels for %s colorbar.", varname) 

227 logging.info("Using levels: %s", levels) 

228 except KeyError: 

229 # Get the range for this variable. 

230 vmin, vmax = var_colorbar["min"], var_colorbar["max"] 

231 logging.debug("Using min and max for %s colorbar.", varname) 

232 # Calculate levels from range. 

233 if vmin == "auto" or vmax == "auto": 

234 levels = None 

235 else: 

236 levels = np.linspace(vmin, vmax, 101) 

237 norm = None 

238 

239 # Overwrite cmap, levels and norm for specific variables that 

240 # require custom colorbar_map as these can not be defined in the 

241 # JSON file. 

242 cmap, levels, norm = custom_colormap_precipitation(cube, cmap, levels, norm) 

243 cmap, levels, norm = custom_colourmap_nimrod_weights(cube, cmap, levels, norm) 

244 cmap, levels, norm = custom_colormap_visibility_in_air(cube, cmap, levels, norm) 

245 cmap, levels, norm = custom_colormap_celsius(cube, cmap, levels, norm) 

246 cmap, levels, norm = custom_colormap_feature_tracking(cube, cmap, levels, norm) 

247 return cmap, levels, norm 

248 

249 

250def custom_colormap_mask(cube: iris.cube.Cube, axis: Literal["x", "y"] | None = None): 

251 """Get colormap for mask. 

252 

253 If "mask_for_" appears anywhere in the name of a cube this function will be called 

254 regardless of the name of the variable to ensure a consistent plot. 

255 

256 Parameters 

257 ---------- 

258 cube: Cube 

259 Cube of variable for which the colorbar information is desired. 

260 axis: "x", "y", optional 

261 Select the levels for just this axis of a line plot. The min and max 

262 can be set by xmin/xmax or ymin/ymax respectively. For variables where 

263 setting a universal range is not desirable (e.g. temperature), users 

264 can set ymin/ymax values to "auto" in the colorbar definitions file. 

265 Where no additional xmin/xmax or ymin/ymax values are provided, the 

266 axis bounds default to use the vmin/vmax values provided. 

267 

268 Returns 

269 ------- 

270 cmap: 

271 Matplotlib colormap. 

272 levels: 

273 List of levels to use for plotting. For continuous plots the min and max 

274 should be taken as the range. 

275 norm: 

276 BoundaryNorm information. 

277 """ 

278 if "difference" not in cube.long_name: 

279 if axis: 

280 levels = [0, 1] 

281 # Complete settings based on levels. 

282 return None, levels, None 

283 else: 

284 # Define the levels and colors. 

285 levels = [0, 1, 2] 

286 colors = ["white", "dodgerblue"] 

287 # Create a custom color map. 

288 cmap = mcolors.ListedColormap(colors) 

289 # Normalize the levels. 

290 norm = mcolors.BoundaryNorm(levels, cmap.N) 

291 logging.debug("Colormap for %s.", cube.long_name) 

292 return cmap, levels, norm 

293 else: 

294 if axis: 

295 levels = [-1, 1] 

296 return None, levels, None 

297 else: 

298 # Search for if mask difference, set to +/- 0.5 as values plotted < 

299 # not <=. 

300 levels = [-2, -0.5, 0.5, 2] 

301 colors = ["goldenrod", "white", "teal"] 

302 cmap = mcolors.ListedColormap(colors) 

303 norm = mcolors.BoundaryNorm(levels, cmap.N) 

304 logging.debug("Colormap for %s.", cube.long_name) 

305 return cmap, levels, norm 

306 

307 

308def custom_beaufort_scale(cube: iris.cube.Cube, axis: Literal["x", "y"] | None = None): 

309 """Get a custom colorbar for a cube in the Beaufort Scale. 

310 

311 Specific variable ranges can be separately set in user-supplied definition 

312 for x- or y-axis limits, or indicate where automated range preferred. 

313 

314 Parameters 

315 ---------- 

316 cube: Cube 

317 Cube of variable with Beaufort Scale in name. 

318 axis: "x", "y", optional 

319 Select the levels for just this axis of a line plot. The min and max 

320 can be set by xmin/xmax or ymin/ymax respectively. For variables where 

321 setting a universal range is not desirable (e.g. temperature), users 

322 can set ymin/ymax values to "auto" in the colorbar definitions file. 

323 Where no additional xmin/xmax or ymin/ymax values are provided, the 

324 axis bounds default to use the vmin/vmax values provided. 

325 

326 Returns 

327 ------- 

328 cmap: 

329 Matplotlib colormap. 

330 levels: 

331 List of levels to use for plotting. For continuous plots the min and max 

332 should be taken as the range. 

333 norm: 

334 BoundaryNorm information. 

335 """ 

336 if "difference" not in cube.long_name: 

337 if axis: 

338 levels = [0, 12] 

339 return None, levels, None 

340 else: 

341 levels = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13] 

342 colors = [ 

343 "black", 

344 (0, 0, 0.6), 

345 "blue", 

346 "cyan", 

347 "green", 

348 "yellow", 

349 (1, 0.5, 0), 

350 "red", 

351 "pink", 

352 "magenta", 

353 "purple", 

354 "maroon", 

355 "white", 

356 ] 

357 cmap = mcolors.ListedColormap(colors) 

358 norm = mcolors.BoundaryNorm(levels, cmap.N) 

359 logging.info("change colormap for Beaufort Scale colorbar.") 

360 return cmap, levels, norm 

361 else: 

362 if axis: 

363 levels = [-4, 4] 

364 return None, levels, None 

365 else: 

366 levels = [ 

367 -3.5, 

368 -2.5, 

369 -1.5, 

370 -0.5, 

371 0.5, 

372 1.5, 

373 2.5, 

374 3.5, 

375 ] 

376 cmap = plt.get_cmap("bwr", 8) 

377 norm = mcolors.BoundaryNorm(levels, cmap.N) 

378 return cmap, levels, norm 

379 

380 

381def custom_colormap_celsius(cube: iris.cube.Cube, cmap, levels, norm): 

382 """Return altered colormap for temperature with change in units to Celsius. 

383 

384 If "Celsius" appears anywhere in the name of a cube this function will be called. 

385 

386 Parameters 

387 ---------- 

388 cube: Cube 

389 Cube of variable for which the colorbar information is desired. 

390 cmap: Matplotlib colormap. 

391 levels: List 

392 List of levels to use for plotting. For continuous plots the min and max 

393 should be taken as the range. 

394 norm: BoundaryNorm. 

395 

396 Returns 

397 ------- 

398 cmap: Matplotlib colormap. 

399 levels: List 

400 List of levels to use for plotting. For continuous plots the min and max 

401 should be taken as the range. 

402 norm: BoundaryNorm. 

403 """ 

404 varnames = filter(None, [cube.long_name, cube.standard_name, cube.var_name]) 

405 if any("temperature" in name for name in varnames) and "Celsius" == cube.units: 

406 levels = np.array(levels) 

407 levels -= 273 

408 levels = levels.tolist() 

409 else: 

410 # Do nothing keep the existing colourbar attributes 

411 levels = levels 

412 cmap = cmap 

413 norm = norm 

414 return cmap, levels, norm 

415 

416 

417def custom_colormap_probability( 

418 cube: iris.cube.Cube, axis: Literal["x", "y"] | None = None 

419): 

420 """Get a custom colorbar for a probability cube. 

421 

422 Specific variable ranges can be separately set in user-supplied definition 

423 for x- or y-axis limits, or indicate where automated range preferred. 

424 

425 Parameters 

426 ---------- 

427 cube: Cube 

428 Cube of variable with probability in name. 

429 axis: "x", "y", optional 

430 Select the levels for just this axis of a line plot. The min and max 

431 can be set by xmin/xmax or ymin/ymax respectively. For variables where 

432 setting a universal range is not desirable (e.g. temperature), users 

433 can set ymin/ymax values to "auto" in the colorbar definitions file. 

434 Where no additional xmin/xmax or ymin/ymax values are provided, the 

435 axis bounds default to use the vmin/vmax values provided. 

436 

437 Returns 

438 ------- 

439 cmap: 

440 Matplotlib colormap. 

441 levels: 

442 List of levels to use for plotting. For continuous plots the min and max 

443 should be taken as the range. 

444 norm: 

445 BoundaryNorm information. 

446 """ 

447 if axis: 

448 levels = [0, 1] 

449 return None, levels, None 

450 else: 

451 cmap = mcolors.ListedColormap( 

452 [ 

453 "#FFFFFF", 

454 "#636363", 

455 "#e1dada", 

456 "#B5CAFF", 

457 "#8FB3FF", 

458 "#7F97FF", 

459 "#ABCF63", 

460 "#E8F59E", 

461 "#FFFA14", 

462 "#FFD121", 

463 "#FFA30A", 

464 ] 

465 ) 

466 levels = [0.0, 0.01, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0] 

467 norm = mcolors.BoundaryNorm(levels, cmap.N) 

468 return cmap, levels, norm 

469 

470 

471def custom_colormap_precipitation(cube: iris.cube.Cube, cmap, levels, norm): 

472 """Return a custom colormap for the current recipe.""" 

473 varnames_lower = [ 

474 n.lower() for n in (cube.long_name, cube.standard_name, cube.var_name) if n 

475 ] 

476 

477 is_rainfall_var = any( 

478 key in name 

479 for name in varnames_lower 

480 for key in ( 

481 "surface_microphysical", 

482 "rainfall rate composite", 

483 "nimrod5min", 

484 "nimrod_5min", 

485 "rain_accumulation", 

486 "rain accumulation", 

487 ) 

488 ) 

489 

490 if is_rainfall_var: 

491 logging.debug( 

492 "Using custom precipitation colourmap due to varnames: %s", varnames_lower 

493 ) 

494 levels = [0, 0.125, 0.25, 0.5, 1, 2, 4, 8, 16, 32, 64, 128, 256] 

495 colors = [ 

496 "w", 

497 (0, 0, 0.6), 

498 "b", 

499 "c", 

500 "g", 

501 "y", 

502 (1, 0.5, 0), 

503 "r", 

504 "pink", 

505 "m", 

506 "purple", 

507 "maroon", 

508 "gray", 

509 ] 

510 # Create a custom colormap 

511 cmap = mcolors.ListedColormap(colors) 

512 # Normalize the levels 

513 norm = mcolors.BoundaryNorm(levels, cmap.N) 

514 logging.info("Using custom rainfall colourmap.") 

515 

516 # Set any Nan values to be plotted a light grey. 

517 cmap.set_bad("#dcdcdc") 

518 

519 return cmap, levels, norm 

520 

521 

522def custom_colourmap_nimrod_weights(cube: iris.cube.Cube, cmap, levels, norm): 

523 """Return a custom colourmap for the current recipe.""" 

524 varnames = filter(None, [cube.long_name, cube.standard_name, cube.var_name]) 

525 if ( 

526 any("wts" in name for name in varnames) 

527 and "difference" not in cube.long_name 

528 and "mask" not in cube.long_name 

529 ): 

530 # Define the levels and colors. Remember the Nimrod weights vary over 

531 # the range [0,13] and should be integer values. Optimum value is 13. 

532 levels = [ 

533 -0.5, 

534 0.5, 

535 1.5, 

536 2.5, 

537 3.5, 

538 4.5, 

539 5.5, 

540 6.5, 

541 7.5, 

542 8.5, 

543 9.5, 

544 10.5, 

545 11.5, 

546 12.5, 

547 13.5, 

548 ] 

549 norm = mcolors.BoundaryNorm(levels, cmap.N) 

550 colours = [ 

551 "#dcdcdc", 

552 "#d10000", 

553 "purple", 

554 "#8f00d6", 

555 "#ff9700", 

556 "pink", 

557 "#ffff00", 

558 "#00007f", 

559 "#6c9ccd", 

560 "#aae8ff", 

561 "#37a648", 

562 "#8edc64", 

563 "#c5ffc5", 

564 "#ffffff", 

565 ] 

566 # Create a custom colormap. 

567 cmap = mcolors.ListedColormap(colours) 

568 # Normalize the levels. 

569 norm = mcolors.BoundaryNorm(levels, cmap.N) 

570 logging.info("Change colormap for Nimrod weights colorbar.") 

571 else: 

572 # Do nothing and keep existing colorbar attributes. 

573 cmap = cmap 

574 levels = levels 

575 norm = norm 

576 return cmap, levels, norm 

577 

578 

579def custom_colormap_aviation_colour_state(cube: iris.cube.Cube): 

580 """Return custom colormap for aviation colour state. 

581 

582 If "aviation_colour_state" appears anywhere in the name of a cube 

583 this function will be called. 

584 

585 Parameters 

586 ---------- 

587 cube: Cube 

588 Cube of variable for which the colorbar information is desired. 

589 

590 Returns 

591 ------- 

592 cmap: Matplotlib colormap. 

593 levels: List 

594 List of levels to use for plotting. For continuous plots the min and max 

595 should be taken as the range. 

596 norm: BoundaryNorm. 

597 """ 

598 levels = [-0.5, 0.5, 1.5, 2.5, 3.5, 4.5, 5.5, 6.5] 

599 colors = [ 

600 "#87ceeb", 

601 "#ffffff", 

602 "#8ced69", 

603 "#ffff00", 

604 "#ffd700", 

605 "#ffa500", 

606 "#fe3620", 

607 ] 

608 # Create a custom colormap 

609 cmap = mcolors.ListedColormap(colors) 

610 # Normalise the levels 

611 norm = mcolors.BoundaryNorm(levels, cmap.N) 

612 return cmap, levels, norm 

613 

614 

615def custom_colormap_visibility_in_air(cube: iris.cube.Cube, cmap, levels, norm): 

616 """Return a custom colormap for the current recipe.""" 

617 varnames = filter(None, [cube.long_name, cube.standard_name, cube.var_name]) 

618 if ( 

619 any("visibility_in_air" in name for name in varnames) 

620 and "difference" not in cube.long_name 

621 and "mask" not in cube.long_name 

622 ): 

623 # Define the levels and colors (in km) 

624 levels = [0, 0.05, 0.1, 0.2, 1.0, 2.0, 5.0, 10.0, 20.0, 30.0, 50.0, 70.0, 100.0] 

625 norm = mcolors.BoundaryNorm(levels, cmap.N) 

626 colours = [ 

627 "#8f00d6", 

628 "#d10000", 

629 "#ff9700", 

630 "#ffff00", 

631 "#00007f", 

632 "#6c9ccd", 

633 "#aae8ff", 

634 "#37a648", 

635 "#8edc64", 

636 "#c5ffc5", 

637 "#dcdcdc", 

638 "#ffffff", 

639 ] 

640 # Create a custom colormap 

641 cmap = mcolors.ListedColormap(colours) 

642 # Normalize the levels 

643 norm = mcolors.BoundaryNorm(levels, cmap.N) 

644 logging.info("change colormap for visibility_in_air variable colorbar.") 

645 else: 

646 # do nothing and keep existing colorbar attributes 

647 cmap = cmap 

648 levels = levels 

649 norm = norm 

650 return cmap, levels, norm 

651 

652 

653def custom_colormap_scores(cube: iris.cube.Cube): 

654 """Return altered colormap for statistical metrics. 

655 

656 Parameters 

657 ---------- 

658 cube: Cube 

659 Cube of variable for which the colorbar information is desired. 

660 

661 Returns 

662 ------- 

663 cmap: Matplotlib colormap. 

664 levels: List 

665 List of levels to use for plotting. For continuous plots the min and max 

666 should be taken as the range. 

667 norm: BoundaryNorm. 

668 """ 

669 varnames = filter(None, [cube.long_name, cube.standard_name, cube.var_name]) 

670 cmap, levels, norm = None, None, None 

671 if any("RMSE_" in name for name in varnames): 671 ↛ 673line 671 didn't jump to line 673 because the condition on line 671 was always true

672 cmap = plt.get_cmap("PuRd", 51) 

673 return cmap, levels, norm 

674 

675 

676def custom_colormap_feature_tracking(cube: iris.cube.Cube, cmap, levels, norm): 

677 """Return altered colormap for feature tracking. 

678 

679 Parameters 

680 ---------- 

681 cube: Cube 

682 Cube of variable for which the colorbar information is desired. 

683 

684 Returns 

685 ------- 

686 cmap: Matplotlib colormap. 

687 levels: List 

688 List of levels to use for plotting. For continuous plots the min and max 

689 should be taken as the range. 

690 norm: BoundaryNorm. 

691 """ 

692 varnames = list(filter(None, [cube.long_name, cube.standard_name, cube.var_name])) 

693 if ( 

694 any("feature_id" in name for name in varnames) 

695 and "difference" not in cube.long_name 

696 and "mask" not in cube.long_name 

697 ): 

698 # Define the levels and colors 

699 levels = np.linspace(1, np.ma.max(cube.data), 10) 

700 cmap = plt.get_cmap("viridis") 

701 # Normalize the levels 

702 norm = mcolors.BoundaryNorm(levels, cmap.N) 

703 logging.info("change colormap for feature id variable colorbar.") 

704 elif ( 

705 any("feature_lifetime" in name for name in varnames) 

706 and "difference" not in cube.long_name 

707 and "mask" not in cube.long_name 

708 ): 

709 # Define the levels and colors 

710 levels = np.linspace(1, np.ma.max(cube.data), 10) 

711 cmap = plt.get_cmap("YlGnBu") 

712 # Normalize the levels 

713 norm = mcolors.BoundaryNorm(levels, cmap.N) 

714 logging.info("change colormap for feature lifetime variable colorbar.") 

715 elif ( 

716 any("feature_init" in name for name in varnames) 

717 and "difference" not in cube.long_name 

718 and "mask" not in cube.long_name 

719 ): 

720 # Define the levels and colors 

721 levels = np.array([0.5, 1]) 

722 cmap = plt.get_cmap("Blues") 

723 # Normalize the levels 

724 norm = mcolors.BoundaryNorm(levels, cmap.N) 

725 logging.info("change colormap for feature init variable colorbar.") 

726 

727 else: 

728 # do nothing and keep existing colorbar attributes 

729 cmap = cmap 

730 levels = levels 

731 norm = norm 

732 

733 # Set all non-feature data to white 

734 if any("feature" in name for name in varnames): 

735 cmap.with_extremes(under="white") 

736 

737 return cmap, levels, norm