Coverage for src/CSET/operators/plot.py: 87%
221 statements
« prev ^ index » next coverage.py v7.5.4, created at 2024-07-02 16:30 +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
23from typing import Union
25import iris
26import iris.coords
27import iris.cube
28import iris.exceptions
29import iris.plot as iplt
30import matplotlib as mpl
31import matplotlib.pyplot as plt
32import numpy as np
33from markdown_it import MarkdownIt
35from CSET._common import get_recipe_metadata, render_file, slugify
36from CSET.operators._utils import get_cube_yxcoordname
38############################
39# Private helper functions #
40############################
43def _append_to_plot_index(plot_index: list) -> list:
44 """Add plots into the plot index, returning the complete plot index."""
45 with open("meta.json", "r+t", encoding="UTF-8") as fp:
46 fcntl.flock(fp, fcntl.LOCK_EX)
47 fp.seek(0)
48 meta = json.load(fp)
49 complete_plot_index = meta.get("plots", [])
50 complete_plot_index = complete_plot_index + plot_index
51 meta["plots"] = complete_plot_index
52 fp.seek(0)
53 fp.truncate()
54 json.dump(meta, fp)
55 return complete_plot_index
58def _check_single_cube(
59 cube: Union[iris.cube.Cube, iris.cube.CubeList],
60) -> iris.cube.Cube:
61 """Ensure a single cube is given.
63 If a CubeList of length one is given that the contained cube is returned,
64 otherwise an error is raised.
66 Parameters
67 ----------
68 cube: Cube | CubeList
69 The cube to check.
71 Returns
72 -------
73 cube: Cube
74 The checked cube.
76 Raises
77 ------
78 TypeError
79 If the input cube is not a Cube or CubeList of a single Cube.
80 """
81 if isinstance(cube, iris.cube.Cube):
82 return cube
83 if isinstance(cube, iris.cube.CubeList):
84 if len(cube) == 1:
85 return cube[0]
86 raise TypeError("Must have a single cube", cube)
89def _make_plot_html_page(plots: list):
90 """Create a HTML page to display a plot image."""
91 # Debug check that plots actually contains some strings.
92 assert isinstance(plots[0], str)
94 # Load HTML template file.
95 # Importlib behaviour changed in 3.12 to avoid circular dependencies.
96 if sys.version_info.minor >= 12:
97 operator_files = importlib.resources.files()
98 else:
99 import CSET.operators
101 operator_files = importlib.resources.files(CSET.operators)
102 template_file = operator_files.joinpath("_plot_page_template.html")
104 # Get some metadata.
105 meta = get_recipe_metadata()
106 title = meta.get("title", "Untitled")
107 description = MarkdownIt().render(meta.get("description", "*No description.*"))
109 # Prepare template variables.
110 variables = {
111 "title": title,
112 "description": description,
113 "initial_plot": plots[0],
114 "plots": plots,
115 "title_slug": slugify(title),
116 }
118 # Render template.
119 html = render_file(template_file, **variables)
121 # Save completed HTML.
122 with open("index.html", "wt", encoding="UTF-8") as fp:
123 fp.write(html)
126def _colorbar_map_levels(varname: str, **kwargs):
127 """
128 Specify the color map and levels.
130 For the given variable name, from a colorbar dictionary file.
132 Parameters
133 ----------
134 colorbar_file: str
135 Filename of the colorbar dictionary to read.
136 varname: str
137 Variable name to extract from the dictionary
139 """
140 # Grab the colour bar file from the recipe global metadata. A non-existent
141 # placeholder path is used if not found.
142 colorbar_file = get_recipe_metadata().get(
143 "style_file_path", "/non-existent/NO_FILE_SPECIFIED"
144 )
145 try:
146 with open(colorbar_file, "rt", encoding="UTF-8") as fp: 146 ↛ 147, 146 ↛ 1502 missed branches: 1) line 146 didn't jump to line 147 because , 2) line 146 didn't jump to line 150 because
147 colorbar = json.load(fp)
149 # Specify the colormap for this variable
150 try:
151 cmap = colorbar[varname]["cmap"]
152 logging.debug("From color_bar dictionary: Using cmap")
153 except KeyError:
154 cmap = mpl.colormaps["viridis"]
156 # Specify the colorbar levels for this variable
157 try:
158 levels = colorbar[varname]["levels"]
160 actual_cmap = mpl.cm.get_cmap(cmap)
162 norm = mpl.colors.BoundaryNorm(levels, ncolors=actual_cmap.N)
163 logging.debug("From color_bar dictionary: Using levels")
164 except KeyError:
165 try:
166 vmin, vmax = colorbar[varname]["min"], colorbar[varname]["max"]
167 logging.debug("From color_bar dictionary: Using min and max")
168 levels = np.linspace(vmin, vmax, 10)
169 norm = None
170 except KeyError:
171 levels = None
172 norm = None
174 except FileNotFoundError:
175 logging.debug("Colour bar file: %s", colorbar_file)
176 logging.info("Colour bar file does not exist. Using default values.")
177 levels = None
178 norm = None
179 cmap = mpl.colormaps["viridis"]
181 return cmap, levels, norm
184def _plot_and_save_contour_plot(
185 cube: iris.cube.Cube,
186 filename: str,
187 title: str,
188 **kwargs,
189):
190 """Plot and save a contour plot.
192 Parameters
193 ----------
194 cube: Cube
195 2 dimensional (lat and lon) Cube of the data to plot.
196 filename: str
197 Filename of the plot to write.
198 title: str
199 Plot title.
201 """
202 # Setup plot details, size, resolution, etc.
203 fig = plt.figure(figsize=(15, 15), facecolor="w", edgecolor="k")
205 # Specify the color bar
206 cmap, levels, norm = _colorbar_map_levels(cube.name())
208 # Filled contour plot of the field.
209 contours = iplt.contourf(cube, cmap=cmap, levels=levels, norm=norm)
211 # Using pyplot interface here as we need iris to generate a cartopy GeoAxes.
212 axes = plt.gca()
214 # Add coastlines if cube contains x and y map coordinates.
215 try:
216 get_cube_yxcoordname(cube)
217 axes.coastlines(resolution="10m")
218 except ValueError:
219 pass
221 # Add title.
222 axes.set_title(title, fontsize=16)
224 # Add colour bar.
225 cbar = fig.colorbar(contours)
226 cbar.set_label(label=f"{cube.name()} ({cube.units})", size=20)
228 # Save plot.
229 fig.savefig(filename, bbox_inches="tight", dpi=150)
230 logging.info("Saved contour plot to %s", filename)
231 plt.close(fig)
234def _plot_and_save_postage_stamp_contour_plot(
235 cube: iris.cube.Cube,
236 filename: str,
237 stamp_coordinate: str,
238 title: str,
239 **kwargs,
240):
241 """Plot postage stamp contour plots from an ensemble.
243 Parameters
244 ----------
245 cube: Cube
246 Iris cube of data to be plotted. It must have the stamp coordinate.
247 filename: str
248 Filename of the plot to write.
249 stamp_coordinate: str
250 Coordinate that becomes different plots.
252 Raises
253 ------
254 ValueError
255 If the cube doesn't have the right dimensions.
256 """
257 # Use the smallest square grid that will fit the members.
258 grid_size = int(math.ceil(math.sqrt(len(cube.coord(stamp_coordinate).points))))
260 fig = plt.figure(figsize=(10, 10))
262 # Specify the color bar
263 cmap, levels, norm = _colorbar_map_levels(cube.name())
265 # Make a subplot for each member.
266 for member, subplot in zip(
267 cube.slices_over(stamp_coordinate), range(1, grid_size**2 + 1), strict=False
268 ):
269 # Implicit interface is much easier here, due to needing to have the
270 # cartopy GeoAxes generated.
271 plt.subplot(grid_size, grid_size, subplot)
272 plot = iplt.contourf(member, cmap=cmap, levels=levels, norm=norm)
273 ax = plt.gca()
274 ax.set_title(f"Member #{member.coord(stamp_coordinate).points[0]}")
275 ax.set_axis_off()
277 # Add coastlines if cube contains x and y map coordinates.
278 try:
279 get_cube_yxcoordname(cube)
280 ax.coastlines(resolution="10m")
281 except ValueError:
282 pass
284 # Put the shared colorbar in its own axes.
285 colorbar_axes = fig.add_axes([0.15, 0.07, 0.7, 0.03])
286 colorbar = fig.colorbar(plot, colorbar_axes, orientation="horizontal")
287 colorbar.set_label(f"{cube.name()} / {cube.units}")
289 # Overall figure title.
290 fig.suptitle(title)
292 fig.savefig(filename, bbox_inches="tight", dpi=150)
293 logging.info("Saved contour postage stamp plot to %s", filename)
294 plt.close(fig)
297def _plot_and_save_line_series(
298 cube: iris.cube.Cube, coord: iris.coords.Coord, filename: str, title: str, **kwargs
299):
300 """Plot and save a 1D line series.
302 Parameters
303 ----------
304 cube: Cube
305 1 dimensional Cube of the data to plot on y-axis.
306 coord: Coord
307 Coordinate to plot on x-axis.
308 filename: str
309 Filename of the plot to write.
310 title: str
311 Plot title.
312 """
313 fig = plt.figure(figsize=(8, 8), facecolor="w", edgecolor="k")
314 iplt.plot(coord, cube, "o-")
315 ax = plt.gca()
317 # Add some labels and tweak the style.
318 ax.set(
319 xlabel=f"{coord.name()} / {coord.units}",
320 ylabel=f"{cube.name()} / {cube.units}",
321 title=title,
322 )
323 ax.ticklabel_format(axis="y", useOffset=False)
324 ax.tick_params(axis="x", labelrotation=15)
325 ax.autoscale()
327 # Save plot.
328 fig.savefig(filename, bbox_inches="tight", dpi=150)
329 logging.info("Saved line plot to %s", filename)
330 plt.close(fig)
333def _plot_and_save_vertical_line_series(
334 cube: iris.cube.Cube,
335 coord: iris.coords.Coord,
336 filename: str,
337 series_coordinate: str,
338 title: str,
339 vmin: float,
340 vmax: float,
341 **kwargs,
342):
343 """Plot and save a 1D line series in vertical.
345 Parameters
346 ----------
347 cube: Cube
348 1 dimensional Cube of the data to plot on x-axis.
349 coord: Coord
350 Coordinate to plot on y-axis.
351 filename: str
352 Filename of the plot to write.
353 series_coordinate: str
354 Coordinate to use as vertical axis.
355 title: str
356 Plot title.
357 vmin: float
358 Minimum value for the x-axis.
359 vmax: float
360 Maximum value for the x-axis.
361 """
362 # plot the vertical pressure axis using log scale
363 fig = plt.figure(figsize=(8, 8), facecolor="w", edgecolor="k")
364 iplt.plot(cube, coord, "o-")
365 ax = plt.gca()
367 # Special handling for pressure level data.
368 if series_coordinate == "pressure": 368 ↛ 394line 368 didn't jump to line 394 because the condition on line 368 was always true
369 # Invert y-axis and set to log scale.
370 ax.invert_yaxis()
371 ax.set_yscale("log")
373 # Define y-ticks and labels for pressure log axis.
374 y_tick_labels = [
375 "1000",
376 "850",
377 "700",
378 "500",
379 "300",
380 "200",
381 "100",
382 "50",
383 "30",
384 "20",
385 "10",
386 ]
387 y_ticks = [1000, 850, 700, 500, 300, 200, 100, 50, 30, 20, 10]
389 # Set y-axis limits and ticks.
390 ax.set_ylim(1100, 100)
392 # test if series_coordinate is model level data. The um data uses model_level_number
393 # and lfric uses full_levels as coordinate.
394 elif series_coordinate in ("model_level_number", "full_levels", "half_levels"):
395 # Define y-ticks and labels for vertical axis.
396 y_ticks = cube.coord(series_coordinate).points
397 y_tick_labels = [str(int(i)) for i in y_ticks]
398 ax.set_ylim(min(y_ticks), max(y_ticks))
400 ax.set_yticks(y_ticks)
401 ax.set_yticklabels(y_tick_labels)
403 # set x-axis limits
404 ax.set_xlim(vmin, vmax)
406 # Add some labels and tweak the style.
407 ax.set(
408 ylabel=f"{coord.name()} / {coord.units}",
409 xlabel=f"{cube.name()} / {cube.units}",
410 title=title,
411 )
413 # Save plot.
414 fig.savefig(filename, bbox_inches="tight", dpi=150)
415 logging.info("Saved line plot to %s", filename)
416 plt.close(fig)
419####################
420# Public functions #
421####################
424def spatial_contour_plot(
425 cube: iris.cube.Cube,
426 filename: str = None,
427 sequence_coordinate: str = "time",
428 stamp_coordinate: str = "realization",
429 **kwargs,
430) -> iris.cube.Cube:
431 """Plot a spatial variable onto a map from a 2D, 3D, or 4D cube.
433 A 2D spatial field can be plotted, but if the sequence_coordinate is present
434 then a sequence of plots will be produced. Similarly if the stamp_coordinate
435 is present then postage stamp plots will be produced.
437 Parameters
438 ----------
439 cube: Cube
440 Iris cube of the data to plot. It should have two spatial dimensions,
441 such as lat and lon, and may also have a another two dimension to be
442 plotted sequentially and/or as postage stamp plots.
443 filename: str, optional
444 Name of the plot to write, used as a prefix for plot sequences. Defaults
445 to the recipe name.
446 sequence_coordinate: str, optional
447 Coordinate about which to make a plot sequence. Defaults to ``"time"``.
448 This coordinate must exist in the cube.
449 stamp_coordinate: str, optional
450 Coordinate about which to plot postage stamp plots. Defaults to
451 ``"realization"``.
453 Returns
454 -------
455 Cube
456 The original cube (so further operations can be applied).
458 Raises
459 ------
460 ValueError
461 If the cube doesn't have the right dimensions.
462 TypeError
463 If the cube isn't a single cube.
464 """
465 recipe_title = get_recipe_metadata().get("title", "Untitled")
467 # Ensure we have a name for the plot file.
468 if filename is None:
469 filename = slugify(recipe_title)
471 # Ensure we've got a single cube.
472 cube = _check_single_cube(cube)
474 # Make postage stamp plots if stamp_coordinate exists and has more than a
475 # single point.
476 plotting_func = _plot_and_save_contour_plot
477 try:
478 if cube.coord(stamp_coordinate).shape[0] > 1:
479 plotting_func = _plot_and_save_postage_stamp_contour_plot
480 except iris.exceptions.CoordinateNotFoundError:
481 pass
483 try:
484 cube.coord(sequence_coordinate)
485 except iris.exceptions.CoordinateNotFoundError as err:
486 raise ValueError(f"Cube must have a {sequence_coordinate} coordinate.") from err
488 # Create a plot for each value of the sequence coordinate.
489 plot_index = []
490 for cube_slice in cube.slices_over(sequence_coordinate):
491 # Use sequence value so multiple sequences can merge.
492 sequence_value = cube_slice.coord(sequence_coordinate).points[0]
493 plot_filename = f"{filename.rsplit('.', 1)[0]}_{sequence_value}.png"
494 coord = cube_slice.coord(sequence_coordinate)
495 # Format the coordinate value in a unit appropriate way.
496 title = f"{recipe_title} | {coord.units.title(coord.points[0])}"
497 # Do the actual plotting.
498 plotting_func(
499 cube_slice,
500 plot_filename,
501 stamp_coordinate=stamp_coordinate,
502 title=title,
503 )
504 plot_index.append(plot_filename)
506 # Add list of plots to plot metadata.
507 complete_plot_index = _append_to_plot_index(plot_index)
509 # Make a page to display the plots.
510 _make_plot_html_page(complete_plot_index)
512 return cube
515# TODO: Expand function to handle ensemble data.
516# line_coordinate: str, optional
517# Coordinate about which to plot multiple lines. Defaults to
518# ``"realization"``.
519def plot_line_series(
520 cube: iris.cube.Cube,
521 filename: str = None,
522 series_coordinate: str = "time",
523 # line_coordinate: str = "realization",
524 **kwargs,
525) -> iris.cube.Cube:
526 """Plot a line plot for the specified coordinate.
528 The cube must be 1D.
530 Parameters
531 ----------
532 cube: Cube
533 Iris cube of the data to plot. It should have a single dimension.
534 filename: str, optional
535 Name of the plot to write, used as a prefix for plot sequences. Defaults
536 to the recipe name.
537 series_coordinate: str, optional
538 Coordinate about which to make a series. Defaults to ``"time"``. This
539 coordinate must exist in the cube.
541 Returns
542 -------
543 Cube
544 The original cube (so further operations can be applied).
546 Raises
547 ------
548 ValueError
549 If the cube doesn't have the right dimensions.
550 TypeError
551 If the cube isn't a single cube.
552 """
553 # Check cube is right shape.
554 cube = _check_single_cube(cube)
555 try:
556 coord = cube.coord(series_coordinate)
557 except iris.exceptions.CoordinateNotFoundError as err:
558 raise ValueError(f"Cube must have a {series_coordinate} coordinate.") from err
559 if cube.ndim > 1:
560 raise ValueError("Cube must be 1D.")
562 # Ensure we have a name for the plot file.
563 title = get_recipe_metadata().get("title", "Untitled")
564 if filename is None:
565 filename = slugify(title)
567 # Add file extension.
568 plot_filename = f"{filename.rsplit('.', 1)[0]}.png"
570 # Do the actual plotting.
571 _plot_and_save_line_series(cube, coord, plot_filename, title)
573 # Add list of plots to plot metadata.
574 plot_index = _append_to_plot_index([plot_filename])
576 # Make a page to display the plots.
577 _make_plot_html_page(plot_index)
579 return cube
582def plot_vertical_line_series(
583 cube: iris.cube.Cube,
584 filename: str = None,
585 series_coordinate: str = "model_level_number",
586 sequence_coordinate: str = "time",
587 # line_coordinate: str = "realization",
588 **kwargs,
589) -> iris.cube.Cube:
590 """Plot a line plot against a type of vertical coordinate.
592 A 1D line plot with y-axis as pressure coordinate can be plotted, but if the sequence_coordinate is present
593 then a sequence of plots will be produced.
595 The cube must be 1D.
597 Parameters
598 ----------
599 cube: Cube
600 Iris cube of the data to plot. It should have a single dimension.
601 filename: str, optional
602 Name of the plot to write, used as a prefix for plot sequences. Defaults
603 to the recipe name.
604 series_coordinate: str, optional
605 Coordinate to plot on the y-axis. Can be ``pressure`` or
606 ``model_level_number`` for UM, or ``full_levels`` or ``half_levels``
607 for LFRic. Defaults to ``model_level_number``.
608 This coordinate must exist in the cube.
609 sequence_coordinate: str, optional
610 Coordinate about which to make a plot sequence. Defaults to ``"time"``.
611 This coordinate must exist in the cube.
613 Returns
614 -------
615 Cube
616 The original cube (so further operations can be applied).
618 Raises
619 ------
620 ValueError
621 If the cube doesn't have the right dimensions.
622 TypeError
623 If the cube isn't a single cube.
624 """
625 # Ensure we've got a single cube.
626 cube = _check_single_cube(cube)
628 # Test if series coordinate i.e. pressure level exist for any cube with cube.ndim >=1.
629 try:
630 coord = cube.coord(series_coordinate)
631 except iris.exceptions.CoordinateNotFoundError as err:
632 raise ValueError(f"Cube must have a {series_coordinate} coordinate.") from err
634 # If several individual vertical lines are plotted with time as sequence_coordinate
635 # for the time slider option.
636 try:
637 cube.coord(sequence_coordinate)
638 except iris.exceptions.CoordinateNotFoundError as err:
639 raise ValueError(f"Cube must have a {sequence_coordinate} coordinate.") from err
641 # Ensure we have a name for the plot file.
642 recipe_title = get_recipe_metadata().get("title", "Untitled")
643 if filename is None:
644 filename = slugify(recipe_title)
646 # Make vertical line plot
647 plotting_func = _plot_and_save_vertical_line_series
649 # set the lower and upper limit for the x-axis to ensure all plots
650 # have same range. This needs to read the whole cube over the range of
651 # the sequence and if applicable postage stamp coordinate.
652 # This only works if the plotting is done in the collate section of a
653 # recipe and not in the parallel section of a recipe.
654 vmin = np.floor((cube.data.min()))
655 vmax = np.ceil((cube.data.max()))
657 # Create a plot for each value of the sequence coordinate.
658 plot_index = []
659 for cube_slice in cube.slices_over(sequence_coordinate):
660 # Use sequence value so multiple sequences can merge.
661 sequence_value = cube_slice.coord(sequence_coordinate).points[0]
662 plot_filename = f"{filename.rsplit('.', 1)[0]}_{sequence_value}.png"
663 coord = cube_slice.coord(series_coordinate)
664 # Format the coordinate value in a unit appropriate way.
665 title = f"{recipe_title} | {coord.units.title(coord.points[0])}"
666 # Do the actual plotting.
667 plotting_func(
668 cube_slice,
669 coord,
670 plot_filename,
671 series_coordinate,
672 title=title,
673 vmin=vmin,
674 vmax=vmax,
675 )
676 plot_index.append(plot_filename)
678 # Add list of plots to plot metadata.
679 complete_plot_index = _append_to_plot_index(plot_index)
681 # Make a page to display the plots.
682 _make_plot_html_page(complete_plot_index)
684 return cube