Coverage for src/CSET/operators/plot.py: 88%
236 statements
« prev ^ index » next coverage.py v7.5.4, created at 2024-07-01 10:31 +0000
1# Copyright 2022 Met Office and 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.
15"""Operators to produce various kinds of plots."""
17import fcntl
18import importlib.resources
19import json
20import logging
21import math
22import sys
23import warnings
24from typing import Union
26import iris
27import iris.coords
28import iris.cube
29import iris.exceptions
30import iris.plot as iplt
31import matplotlib as mpl
32import matplotlib.pyplot as plt
33import numpy as np
34from markdown_it import MarkdownIt
36from CSET._common import get_recipe_metadata, render_file, slugify
37from CSET.operators._utils import get_cube_yxcoordname
39############################
40# Private helper functions #
41############################
44def _append_to_plot_index(plot_index: list) -> list:
45 """Add plots into the plot index, returning the complete plot index."""
46 with open("meta.json", "r+t", encoding="UTF-8") as fp:
47 fcntl.flock(fp, fcntl.LOCK_EX)
48 fp.seek(0)
49 meta = json.load(fp)
50 complete_plot_index = meta.get("plots", [])
51 complete_plot_index = complete_plot_index + plot_index
52 meta["plots"] = complete_plot_index
53 fp.seek(0)
54 fp.truncate()
55 json.dump(meta, fp)
56 return complete_plot_index
59def _check_single_cube(
60 cube: Union[iris.cube.Cube, iris.cube.CubeList],
61) -> iris.cube.Cube:
62 """Ensure a single cube is given.
64 If a CubeList of length one is given that the contained cube is returned,
65 otherwise an error is raised.
67 Parameters
68 ----------
69 cube: Cube | CubeList
70 The cube to check.
72 Returns
73 -------
74 cube: Cube
75 The checked cube.
77 Raises
78 ------
79 TypeError
80 If the input cube is not a Cube or CubeList of a single Cube.
81 """
82 if isinstance(cube, iris.cube.Cube):
83 return cube
84 if isinstance(cube, iris.cube.CubeList):
85 if len(cube) == 1:
86 return cube[0]
87 raise TypeError("Must have a single cube", cube)
90def _make_plot_html_page(plots: list):
91 """Create a HTML page to display a plot image."""
92 # Debug check that plots actually contains some strings.
93 assert isinstance(plots[0], str)
95 # Load HTML template file.
96 # Importlib behaviour changed in 3.12 to avoid circular dependencies.
97 if sys.version_info.minor >= 12:
98 operator_files = importlib.resources.files()
99 else:
100 import CSET.operators
102 operator_files = importlib.resources.files(CSET.operators)
103 template_file = operator_files.joinpath("_plot_page_template.html")
105 # Get some metadata.
106 meta = get_recipe_metadata()
107 title = meta.get("title", "Untitled")
108 description = MarkdownIt().render(meta.get("description", "*No description.*"))
110 # Prepare template variables.
111 variables = {
112 "title": title,
113 "description": description,
114 "initial_plot": plots[0],
115 "plots": plots,
116 "title_slug": slugify(title),
117 }
119 # Render template.
120 html = render_file(template_file, **variables)
122 # Save completed HTML.
123 with open("index.html", "wt", encoding="UTF-8") as fp:
124 fp.write(html)
127def _colorbar_map_levels(varname: str, **kwargs):
128 """
129 Specify the color map and levels.
131 For the given variable name, from a colorbar dictionary file.
133 Parameters
134 ----------
135 colorbar_file: str
136 Filename of the colorbar dictionary to read.
137 varname: str
138 Variable name to extract from the dictionary
140 """
141 # Grab the colour bar file from the recipe global metadata. A non-existent
142 # placeholder path is used if not found.
143 colorbar_file = get_recipe_metadata().get(
144 "style_file_path", "/non-existent/NO_FILE_SPECIFIED"
145 )
146 try:
147 with open(colorbar_file, "rt", encoding="UTF-8") as fp: 147 ↛ 148, 147 ↛ 1512 missed branches: 1) line 147 didn't jump to line 148 because , 2) line 147 didn't jump to line 151 because
148 colorbar = json.load(fp)
150 # Specify the colormap for this variable
151 try:
152 cmap = colorbar[varname]["cmap"]
153 logging.debug("From color_bar dictionary: Using cmap")
154 except KeyError:
155 cmap = mpl.colormaps["viridis"]
157 # Specify the colorbar levels for this variable
158 try:
159 levels = colorbar[varname]["levels"]
161 actual_cmap = mpl.cm.get_cmap(cmap)
163 norm = mpl.colors.BoundaryNorm(levels, ncolors=actual_cmap.N)
164 logging.debug("From color_bar dictionary: Using levels")
165 except KeyError:
166 try:
167 vmin, vmax = colorbar[varname]["min"], colorbar[varname]["max"]
168 logging.debug("From color_bar dictionary: Using min and max")
169 levels = np.linspace(vmin, vmax, 10)
170 norm = None
171 except KeyError:
172 levels = None
173 norm = None
175 except FileNotFoundError:
176 logging.debug("Colour bar file: %s", colorbar_file)
177 logging.info("Colour bar file does not exist. Using default values.")
178 levels = None
179 norm = None
180 cmap = mpl.colormaps["viridis"]
182 return cmap, levels, norm
185def _plot_and_save_contour_plot(
186 cube: iris.cube.Cube,
187 filename: str,
188 title: str,
189 **kwargs,
190):
191 """Plot and save a contour plot.
193 Parameters
194 ----------
195 cube: Cube
196 2 dimensional (lat and lon) Cube of the data to plot.
197 filename: str
198 Filename of the plot to write.
199 title: str
200 Plot title.
202 """
203 # Setup plot details, size, resolution, etc.
204 fig = plt.figure(figsize=(15, 15), facecolor="w", edgecolor="k")
206 # Specify the color bar
207 cmap, levels, norm = _colorbar_map_levels(cube.name())
209 # Filled contour plot of the field.
210 contours = iplt.contourf(cube, cmap=cmap, levels=levels, norm=norm)
212 # Using pyplot interface here as we need iris to generate a cartopy GeoAxes.
213 axes = plt.gca()
215 # Add coastlines if cube contains x and y map coordinates.
216 try:
217 get_cube_yxcoordname(cube)
218 axes.coastlines(resolution="10m")
219 except ValueError:
220 pass
222 # Add title.
223 axes.set_title(title, fontsize=16)
225 # Add colour bar.
226 cbar = fig.colorbar(contours)
227 cbar.set_label(label=f"{cube.name()} ({cube.units})", size=20)
229 # Save plot.
230 fig.savefig(filename, bbox_inches="tight", dpi=150)
231 logging.info("Saved contour plot to %s", filename)
232 plt.close(fig)
235def _plot_and_save_postage_stamp_contour_plot(
236 cube: iris.cube.Cube,
237 filename: str,
238 stamp_coordinate: str,
239 title: str,
240 **kwargs,
241):
242 """Plot postage stamp contour plots from an ensemble.
244 Parameters
245 ----------
246 cube: Cube
247 Iris cube of data to be plotted. It must have the stamp coordinate.
248 filename: str
249 Filename of the plot to write.
250 stamp_coordinate: str
251 Coordinate that becomes different plots.
253 Raises
254 ------
255 ValueError
256 If the cube doesn't have the right dimensions.
257 """
258 # Use the smallest square grid that will fit the members.
259 grid_size = int(math.ceil(math.sqrt(len(cube.coord(stamp_coordinate).points))))
261 fig = plt.figure(figsize=(10, 10))
263 # Specify the color bar
264 cmap, levels, norm = _colorbar_map_levels(cube.name())
266 # Make a subplot for each member.
267 for member, subplot in zip(
268 cube.slices_over(stamp_coordinate), range(1, grid_size**2 + 1), strict=False
269 ):
270 # Implicit interface is much easier here, due to needing to have the
271 # cartopy GeoAxes generated.
272 plt.subplot(grid_size, grid_size, subplot)
273 plot = iplt.contourf(member, cmap=cmap, levels=levels, norm=norm)
274 ax = plt.gca()
275 ax.set_title(f"Member #{member.coord(stamp_coordinate).points[0]}")
276 ax.set_axis_off()
278 # Add coastlines if cube contains x and y map coordinates.
279 try:
280 get_cube_yxcoordname(cube)
281 ax.coastlines(resolution="10m")
282 except ValueError:
283 pass
285 # Put the shared colorbar in its own axes.
286 colorbar_axes = fig.add_axes([0.15, 0.07, 0.7, 0.03])
287 colorbar = fig.colorbar(plot, colorbar_axes, orientation="horizontal")
288 colorbar.set_label(f"{cube.name()} / {cube.units}")
290 # Overall figure title.
291 fig.suptitle(title)
293 fig.savefig(filename, bbox_inches="tight", dpi=150)
294 logging.info("Saved contour postage stamp plot to %s", filename)
295 plt.close(fig)
298def _plot_and_save_line_series(
299 cube: iris.cube.Cube, coord: iris.coords.Coord, filename: str, title: str, **kwargs
300):
301 """Plot and save a 1D line series.
303 Parameters
304 ----------
305 cube: Cube
306 1 dimensional Cube of the data to plot on y-axis.
307 coord: Coord
308 Coordinate to plot on x-axis.
309 filename: str
310 Filename of the plot to write.
311 title: str
312 Plot title.
313 """
314 fig = plt.figure(figsize=(8, 8), facecolor="w", edgecolor="k")
315 iplt.plot(coord, cube, "o-")
316 ax = plt.gca()
318 # Add some labels and tweak the style.
319 ax.set(
320 xlabel=f"{coord.name()} / {coord.units}",
321 ylabel=f"{cube.name()} / {cube.units}",
322 title=title,
323 )
324 ax.ticklabel_format(axis="y", useOffset=False)
325 ax.tick_params(axis="x", labelrotation=15)
326 ax.autoscale()
328 # Save plot.
329 fig.savefig(filename, bbox_inches="tight", dpi=150)
330 logging.info("Saved line plot to %s", filename)
331 plt.close(fig)
334def _plot_and_save_vertical_line_series(
335 cube: iris.cube.Cube,
336 coord: iris.coords.Coord,
337 filename: str,
338 series_coordinate: str,
339 title: str,
340 vmin: float,
341 vmax: float,
342 **kwargs,
343):
344 """Plot and save a 1D line series in vertical.
346 Parameters
347 ----------
348 cube: Cube
349 1 dimensional Cube of the data to plot on x-axis.
350 coord: Coord
351 Coordinate to plot on y-axis.
352 filename: str
353 Filename of the plot to write.
354 series_coordinate: str
355 Coordinate to use as vertical axis.
356 title: str
357 Plot title.
358 vmin: float
359 Minimum value for the x-axis.
360 vmax: float
361 Maximum value for the x-axis.
362 """
363 # plot the vertical pressure axis using log scale
364 fig = plt.figure(figsize=(8, 8), facecolor="w", edgecolor="k")
365 iplt.plot(cube, coord, "o-")
366 ax = plt.gca()
368 # Special handling for pressure level data.
369 if series_coordinate == "pressure": 369 ↛ 395line 369 didn't jump to line 395 because the condition on line 369 was always true
370 # Invert y-axis and set to log scale.
371 ax.invert_yaxis()
372 ax.set_yscale("log")
374 # Define y-ticks and labels for pressure log axis.
375 y_tick_labels = [
376 "1000",
377 "850",
378 "700",
379 "500",
380 "300",
381 "200",
382 "100",
383 "50",
384 "30",
385 "20",
386 "10",
387 ]
388 y_ticks = [1000, 850, 700, 500, 300, 200, 100, 50, 30, 20, 10]
390 # Set y-axis limits and ticks.
391 ax.set_ylim(1100, 100)
393 # test if series_coordinate is model level data. The um data uses model_level_number
394 # and lfric uses full_levels as coordinate.
395 elif series_coordinate in ("model_level_number", "full_levels", "half_levels"):
396 # Define y-ticks and labels for vertical axis.
397 y_ticks = cube.coord(series_coordinate).points
398 y_tick_labels = [str(int(i)) for i in y_ticks]
399 ax.set_ylim(min(y_ticks), max(y_ticks))
401 ax.set_yticks(y_ticks)
402 ax.set_yticklabels(y_tick_labels)
404 # set x-axis limits
405 ax.set_xlim(vmin, vmax)
407 # Add some labels and tweak the style.
408 ax.set(
409 ylabel=f"{coord.name()} / {coord.units}",
410 xlabel=f"{cube.name()} / {cube.units}",
411 title=title,
412 )
414 # Save plot.
415 fig.savefig(filename, bbox_inches="tight", dpi=150)
416 logging.info("Saved line plot to %s", filename)
417 plt.close(fig)
420####################
421# Public functions #
422####################
425def spatial_contour_plot(
426 cube: iris.cube.Cube,
427 filename: str = None,
428 sequence_coordinate: str = "time",
429 stamp_coordinate: str = "realization",
430 **kwargs,
431) -> iris.cube.Cube:
432 """Plot a spatial variable onto a map from a 2D, 3D, or 4D cube.
434 A 2D spatial field can be plotted, but if the sequence_coordinate is present
435 then a sequence of plots will be produced. Similarly if the stamp_coordinate
436 is present then postage stamp plots will be produced.
438 Parameters
439 ----------
440 cube: Cube
441 Iris cube of the data to plot. It should have two spatial dimensions,
442 such as lat and lon, and may also have a another two dimension to be
443 plotted sequentially and/or as postage stamp plots.
444 filename: str, optional
445 Name of the plot to write, used as a prefix for plot sequences. Defaults
446 to the recipe name.
447 sequence_coordinate: str, optional
448 Coordinate about which to make a plot sequence. Defaults to ``"time"``.
449 This coordinate must exist in the cube.
450 stamp_coordinate: str, optional
451 Coordinate about which to plot postage stamp plots. Defaults to
452 ``"realization"``.
454 Returns
455 -------
456 Cube
457 The original cube (so further operations can be applied).
459 Raises
460 ------
461 ValueError
462 If the cube doesn't have the right dimensions.
463 TypeError
464 If the cube isn't a single cube.
465 """
466 recipe_title = get_recipe_metadata().get("title", "Untitled")
468 # Ensure we have a name for the plot file.
469 if filename is None:
470 filename = slugify(recipe_title)
472 # Ensure we've got a single cube.
473 cube = _check_single_cube(cube)
475 # Make postage stamp plots if stamp_coordinate exists and has more than a
476 # single point.
477 plotting_func = _plot_and_save_contour_plot
478 try:
479 if cube.coord(stamp_coordinate).shape[0] > 1:
480 plotting_func = _plot_and_save_postage_stamp_contour_plot
481 except iris.exceptions.CoordinateNotFoundError:
482 pass
484 try:
485 cube.coord(sequence_coordinate)
486 except iris.exceptions.CoordinateNotFoundError as err:
487 raise ValueError(f"Cube must have a {sequence_coordinate} coordinate.") from err
489 # Create a plot for each value of the sequence coordinate.
490 plot_index = []
491 for cube_slice in cube.slices_over(sequence_coordinate):
492 # Use sequence value so multiple sequences can merge.
493 sequence_value = cube_slice.coord(sequence_coordinate).points[0]
494 plot_filename = f"{filename.rsplit('.', 1)[0]}_{sequence_value}.png"
495 coord = cube_slice.coord(sequence_coordinate)
496 # Format the coordinate value in a unit appropriate way.
497 title = f"{recipe_title} | {coord.units.title(coord.points[0])}"
498 # Do the actual plotting.
499 plotting_func(
500 cube_slice,
501 plot_filename,
502 stamp_coordinate=stamp_coordinate,
503 title=title,
504 )
505 plot_index.append(plot_filename)
507 # Add list of plots to plot metadata.
508 complete_plot_index = _append_to_plot_index(plot_index)
510 # Make a page to display the plots.
511 _make_plot_html_page(complete_plot_index)
513 return cube
516# Deprecated
517def postage_stamp_contour_plot(
518 cube: iris.cube.Cube,
519 filename: str = None,
520 coordinate: str = "realization",
521 **kwargs,
522) -> iris.cube.Cube:
523 """Plot postage stamp contour plots from an ensemble.
525 Depreciated. Use spatial_contour_plot with a stamp_coordinate argument
526 instead.
528 Parameters
529 ----------
530 cube: Cube
531 Iris cube of data to be plotted. It must have a realization coordinate.
532 filename: pathlike, optional
533 The path of the plot to write. Defaults to the recipe name.
534 coordinate: str
535 The coordinate that becomes different plots. Defaults to "realization".
537 Returns
538 -------
539 Cube
540 The original cube (so further operations can be applied)
542 Raises
543 ------
544 ValueError
545 If the cube doesn't have the right dimensions.
546 TypeError
547 If cube isn't a Cube.
548 """
549 warnings.warn(
550 "postage_stamp_contour_plot is depreciated. Use spatial_contour_plot with a stamp_coordinate argument instead.",
551 DeprecationWarning,
552 stacklevel=2,
553 )
554 # Get suitable filename.
555 if filename is None:
556 filename = slugify(get_recipe_metadata().get("title", "Untitled"))
557 if not filename.endswith(".png"):
558 filename = filename + ".png"
560 # Check cube is suitable.
561 cube = _check_single_cube(cube)
562 try:
563 cube.coord(coordinate)
564 except iris.exceptions.CoordinateNotFoundError as err:
565 raise ValueError(f"Cube must have a {coordinate} coordinate.") from err
567 _plot_and_save_postage_stamp_contour_plot(cube, filename, coordinate, title="")
568 _make_plot_html_page([filename])
569 return cube
572# TODO: Expand function to handle ensemble data.
573# line_coordinate: str, optional
574# Coordinate about which to plot multiple lines. Defaults to
575# ``"realization"``.
576def plot_line_series(
577 cube: iris.cube.Cube,
578 filename: str = None,
579 series_coordinate: str = "time",
580 # line_coordinate: str = "realization",
581 **kwargs,
582) -> iris.cube.Cube:
583 """Plot a line plot for the specified coordinate.
585 The cube must be 1D.
587 Parameters
588 ----------
589 cube: Cube
590 Iris cube of the data to plot. It should have a single dimension.
591 filename: str, optional
592 Name of the plot to write, used as a prefix for plot sequences. Defaults
593 to the recipe name.
594 series_coordinate: str, optional
595 Coordinate about which to make a series. Defaults to ``"time"``. This
596 coordinate must exist in the cube.
598 Returns
599 -------
600 Cube
601 The original cube (so further operations can be applied).
603 Raises
604 ------
605 ValueError
606 If the cube doesn't have the right dimensions.
607 TypeError
608 If the cube isn't a single cube.
609 """
610 # Check cube is right shape.
611 cube = _check_single_cube(cube)
612 try:
613 coord = cube.coord(series_coordinate)
614 except iris.exceptions.CoordinateNotFoundError as err:
615 raise ValueError(f"Cube must have a {series_coordinate} coordinate.") from err
616 if cube.ndim > 1:
617 raise ValueError("Cube must be 1D.")
619 # Ensure we have a name for the plot file.
620 title = get_recipe_metadata().get("title", "Untitled")
621 if filename is None:
622 filename = slugify(title)
624 # Add file extension.
625 plot_filename = f"{filename.rsplit('.', 1)[0]}.png"
627 # Do the actual plotting.
628 _plot_and_save_line_series(cube, coord, plot_filename, title)
630 # Add list of plots to plot metadata.
631 plot_index = _append_to_plot_index([plot_filename])
633 # Make a page to display the plots.
634 _make_plot_html_page(plot_index)
636 return cube
639def plot_vertical_line_series(
640 cube: iris.cube.Cube,
641 filename: str = None,
642 series_coordinate: str = "model_level_number",
643 sequence_coordinate: str = "time",
644 # line_coordinate: str = "realization",
645 **kwargs,
646) -> iris.cube.Cube:
647 """Plot a line plot against a type of vertical coordinate.
649 A 1D line plot with y-axis as pressure coordinate can be plotted, but if the sequence_coordinate is present
650 then a sequence of plots will be produced.
652 The cube must be 1D.
654 Parameters
655 ----------
656 cube: Cube
657 Iris cube of the data to plot. It should have a single dimension.
658 filename: str, optional
659 Name of the plot to write, used as a prefix for plot sequences. Defaults
660 to the recipe name.
661 series_coordinate: str, optional
662 Coordinate to plot on the y-axis. Can be ``pressure`` or
663 ``model_level_number`` for UM, or ``full_levels`` or ``half_levels``
664 for LFRic. Defaults to ``model_level_number``.
665 This coordinate must exist in the cube.
666 sequence_coordinate: str, optional
667 Coordinate about which to make a plot sequence. Defaults to ``"time"``.
668 This coordinate must exist in the cube.
670 Returns
671 -------
672 Cube
673 The original cube (so further operations can be applied).
675 Raises
676 ------
677 ValueError
678 If the cube doesn't have the right dimensions.
679 TypeError
680 If the cube isn't a single cube.
681 """
682 # Ensure we've got a single cube.
683 cube = _check_single_cube(cube)
685 # Test if series coordinate i.e. pressure level exist for any cube with cube.ndim >=1.
686 try:
687 coord = cube.coord(series_coordinate)
688 except iris.exceptions.CoordinateNotFoundError as err:
689 raise ValueError(f"Cube must have a {series_coordinate} coordinate.") from err
691 # If several individual vertical lines are plotted with time as sequence_coordinate
692 # for the time slider option.
693 try:
694 cube.coord(sequence_coordinate)
695 except iris.exceptions.CoordinateNotFoundError as err:
696 raise ValueError(f"Cube must have a {sequence_coordinate} coordinate.") from err
698 # Ensure we have a name for the plot file.
699 recipe_title = get_recipe_metadata().get("title", "Untitled")
700 if filename is None: 700 ↛ 704line 700 didn't jump to line 704 because the condition on line 700 was always true
701 filename = slugify(recipe_title)
703 # Make vertical line plot
704 plotting_func = _plot_and_save_vertical_line_series
706 # set the lower and upper limit for the x-axis to ensure all plots
707 # have same range. This needs to read the whole cube over the range of
708 # the sequence and if applicable postage stamp coordinate.
709 # This only works if the plotting is done in the collate section of a
710 # recipe and not in the parallel section of a recipe.
711 vmin = np.floor((cube.data.min()))
712 vmax = np.ceil((cube.data.max()))
714 # Create a plot for each value of the sequence coordinate.
715 plot_index = []
716 for cube_slice in cube.slices_over(sequence_coordinate):
717 # Use sequence value so multiple sequences can merge.
718 sequence_value = cube_slice.coord(sequence_coordinate).points[0]
719 plot_filename = f"{filename.rsplit('.', 1)[0]}_{sequence_value}.png"
720 coord = cube_slice.coord(series_coordinate)
721 # Format the coordinate value in a unit appropriate way.
722 title = f"{recipe_title} | {coord.units.title(coord.points[0])}"
723 # Do the actual plotting.
724 plotting_func(
725 cube_slice,
726 coord,
727 plot_filename,
728 series_coordinate,
729 title=title,
730 vmin=vmin,
731 vmax=vmax,
732 )
733 plot_index.append(plot_filename)
735 # Add list of plots to plot metadata.
736 complete_plot_index = _append_to_plot_index(plot_index)
738 # Make a page to display the plots.
739 _make_plot_html_page(complete_plot_index)
741 return cube