Coverage for src/CSET/operators/plot.py: 90%
231 statements
« prev ^ index » next coverage.py v7.5.4, created at 2024-07-01 09:02 +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 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 title: str
354 Plot title.
355 vmin: float
356 Minimum value for the x-axis.
357 vmax: float
358 Maximum value for the x-axis.
359 """
360 # plot the vertical pressure axis using log scale
361 fig = plt.figure(figsize=(8, 8), facecolor="w", edgecolor="k")
362 iplt.plot(cube, coord, "o-")
363 ax = plt.gca()
364 ax.invert_yaxis()
365 ax.set_yscale("log")
367 # Define y-ticks and labels for pressure log axis
368 y_tick_labels = [
369 "1000",
370 "850",
371 "700",
372 "500",
373 "300",
374 "200",
375 "100",
376 "50",
377 "30",
378 "20",
379 "10",
380 ]
381 y_ticks = [1000, 850, 700, 500, 300, 200, 100, 50, 30, 20, 10]
383 # Set y-axis limits and ticks
384 ax.set_ylim(1100, 100)
385 ax.set_yticks(y_ticks)
386 ax.set_yticklabels(y_tick_labels)
388 # set x-axis limits
389 ax.set_xlim(vmin, vmax)
391 # Add some labels and tweak the style.
392 ax.set(
393 ylabel=f"{coord.name()} / {coord.units}",
394 xlabel=f"{cube.name()} / {cube.units}",
395 title=title,
396 )
398 # Save plot.
399 fig.savefig(filename, bbox_inches="tight", dpi=150)
400 logging.info("Saved line plot to %s", filename)
401 plt.close(fig)
404####################
405# Public functions #
406####################
409def spatial_contour_plot(
410 cube: iris.cube.Cube,
411 filename: str = None,
412 sequence_coordinate: str = "time",
413 stamp_coordinate: str = "realization",
414 **kwargs,
415) -> iris.cube.Cube:
416 """Plot a spatial variable onto a map from a 2D, 3D, or 4D cube.
418 A 2D spatial field can be plotted, but if the sequence_coordinate is present
419 then a sequence of plots will be produced. Similarly if the stamp_coordinate
420 is present then postage stamp plots will be produced.
422 Parameters
423 ----------
424 cube: Cube
425 Iris cube of the data to plot. It should have two spatial dimensions,
426 such as lat and lon, and may also have a another two dimension to be
427 plotted sequentially and/or as postage stamp plots.
428 filename: str, optional
429 Name of the plot to write, used as a prefix for plot sequences. Defaults
430 to the recipe name.
431 sequence_coordinate: str, optional
432 Coordinate about which to make a plot sequence. Defaults to ``"time"``.
433 This coordinate must exist in the cube.
434 stamp_coordinate: str, optional
435 Coordinate about which to plot postage stamp plots. Defaults to
436 ``"realization"``.
438 Returns
439 -------
440 Cube
441 The original cube (so further operations can be applied).
443 Raises
444 ------
445 ValueError
446 If the cube doesn't have the right dimensions.
447 TypeError
448 If the cube isn't a single cube.
449 """
450 recipe_title = get_recipe_metadata().get("title", "Untitled")
452 # Ensure we have a name for the plot file.
453 if filename is None:
454 filename = slugify(recipe_title)
456 # Ensure we've got a single cube.
457 cube = _check_single_cube(cube)
459 # Make postage stamp plots if stamp_coordinate exists and has more than a
460 # single point.
461 plotting_func = _plot_and_save_contour_plot
462 try:
463 if cube.coord(stamp_coordinate).shape[0] > 1:
464 plotting_func = _plot_and_save_postage_stamp_contour_plot
465 except iris.exceptions.CoordinateNotFoundError:
466 pass
468 try:
469 cube.coord(sequence_coordinate)
470 except iris.exceptions.CoordinateNotFoundError as err:
471 raise ValueError(f"Cube must have a {sequence_coordinate} coordinate.") from err
473 # Create a plot for each value of the sequence coordinate.
474 plot_index = []
475 for cube_slice in cube.slices_over(sequence_coordinate):
476 # Use sequence value so multiple sequences can merge.
477 sequence_value = cube_slice.coord(sequence_coordinate).points[0]
478 plot_filename = f"{filename.rsplit('.', 1)[0]}_{sequence_value}.png"
479 coord = cube_slice.coord(sequence_coordinate)
480 # Format the coordinate value in a unit appropriate way.
481 title = f"{recipe_title} | {coord.units.title(coord.points[0])}"
482 # Do the actual plotting.
483 plotting_func(
484 cube_slice,
485 plot_filename,
486 stamp_coordinate=stamp_coordinate,
487 title=title,
488 )
489 plot_index.append(plot_filename)
491 # Add list of plots to plot metadata.
492 complete_plot_index = _append_to_plot_index(plot_index)
494 # Make a page to display the plots.
495 _make_plot_html_page(complete_plot_index)
497 return cube
500# Deprecated
501def postage_stamp_contour_plot(
502 cube: iris.cube.Cube,
503 filename: str = None,
504 coordinate: str = "realization",
505 **kwargs,
506) -> iris.cube.Cube:
507 """Plot postage stamp contour plots from an ensemble.
509 Depreciated. Use spatial_contour_plot with a stamp_coordinate argument
510 instead.
512 Parameters
513 ----------
514 cube: Cube
515 Iris cube of data to be plotted. It must have a realization coordinate.
516 filename: pathlike, optional
517 The path of the plot to write. Defaults to the recipe name.
518 coordinate: str
519 The coordinate that becomes different plots. Defaults to "realization".
521 Returns
522 -------
523 Cube
524 The original cube (so further operations can be applied)
526 Raises
527 ------
528 ValueError
529 If the cube doesn't have the right dimensions.
530 TypeError
531 If cube isn't a Cube.
532 """
533 warnings.warn(
534 "postage_stamp_contour_plot is depreciated. Use spatial_contour_plot with a stamp_coordinate argument instead.",
535 DeprecationWarning,
536 stacklevel=2,
537 )
538 # Get suitable filename.
539 if filename is None:
540 filename = slugify(get_recipe_metadata().get("title", "Untitled"))
541 if not filename.endswith(".png"):
542 filename = filename + ".png"
544 # Check cube is suitable.
545 cube = _check_single_cube(cube)
546 try:
547 cube.coord(coordinate)
548 except iris.exceptions.CoordinateNotFoundError as err:
549 raise ValueError(f"Cube must have a {coordinate} coordinate.") from err
551 _plot_and_save_postage_stamp_contour_plot(cube, filename, coordinate, title="")
552 _make_plot_html_page([filename])
553 return cube
556# TODO: Expand function to handle ensemble data.
557# line_coordinate: str, optional
558# Coordinate about which to plot multiple lines. Defaults to
559# ``"realization"``.
560def plot_line_series(
561 cube: iris.cube.Cube,
562 filename: str = None,
563 series_coordinate: str = "time",
564 # line_coordinate: str = "realization",
565 **kwargs,
566) -> iris.cube.Cube:
567 """Plot a line plot for the specified coordinate.
569 The cube must be 1D.
571 Parameters
572 ----------
573 cube: Cube
574 Iris cube of the data to plot. It should have a single dimension.
575 filename: str, optional
576 Name of the plot to write, used as a prefix for plot sequences. Defaults
577 to the recipe name.
578 series_coordinate: str, optional
579 Coordinate about which to make a series. Defaults to ``"time"``. This
580 coordinate must exist in the cube.
582 Returns
583 -------
584 Cube
585 The original cube (so further operations can be applied).
587 Raises
588 ------
589 ValueError
590 If the cube doesn't have the right dimensions.
591 TypeError
592 If the cube isn't a single cube.
593 """
594 # Check cube is right shape.
595 cube = _check_single_cube(cube)
596 try:
597 coord = cube.coord(series_coordinate)
598 except iris.exceptions.CoordinateNotFoundError as err:
599 raise ValueError(f"Cube must have a {series_coordinate} coordinate.") from err
600 if cube.ndim > 1:
601 raise ValueError("Cube must be 1D.")
603 # Ensure we have a name for the plot file.
604 title = get_recipe_metadata().get("title", "Untitled")
605 if filename is None:
606 filename = slugify(title)
608 # Add file extension.
609 plot_filename = f"{filename.rsplit('.', 1)[0]}.png"
611 # Do the actual plotting.
612 _plot_and_save_line_series(cube, coord, plot_filename, title)
614 # Add list of plots to plot metadata.
615 plot_index = _append_to_plot_index([plot_filename])
617 # Make a page to display the plots.
618 _make_plot_html_page(plot_index)
620 return cube
623def plot_vertical_line_series(
624 cube: iris.cube.Cube,
625 filename: str = None,
626 series_coordinate: str = "pressure",
627 sequence_coordinate: str = "time",
628 # line_coordinate: str = "realization",
629 **kwargs,
630) -> iris.cube.Cube:
631 """Plot a line plot against a type of vertical coordinate.
633 A 1D line plot with y-axis as pressure coordinate can be plotted, but if the sequence_coordinate is present
634 then a sequence of plots will be produced.
636 The cube must be 1D.
638 Parameters
639 ----------
640 cube: Cube
641 Iris cube of the data to plot. It should have a single dimension.
642 filename: str, optional
643 Name of the plot to write, used as a prefix for plot sequences. Defaults
644 to the recipe name.
645 series_coordinate: str, optional
646 Coordinate to plot on the y-axis. Defaults to ``pressure``.
647 This coordinate must exist in the cube.
648 sequence_coordinate: str, optional
649 Coordinate about which to make a plot sequence. Defaults to ``"time"``.
650 This coordinate must exist in the cube.
652 Returns
653 -------
654 Cube
655 The original cube (so further operations can be applied).
657 Raises
658 ------
659 ValueError
660 If the cube doesn't have the right dimensions.
661 TypeError
662 If the cube isn't a single cube.
663 """
664 # Ensure we've got a single cube.
665 cube = _check_single_cube(cube)
667 # Test if series coordinate i.e. pressure level exist for any cube with cube.ndim >=1.
668 try:
669 coord = cube.coord(series_coordinate)
670 except iris.exceptions.CoordinateNotFoundError as err:
671 raise ValueError(f"Cube must have a {series_coordinate} coordinate.") from err
673 # If several individual vertical lines are plotted with time as sequence_coordinate
674 # for the time slider option.
675 try:
676 cube.coord(sequence_coordinate)
677 except iris.exceptions.CoordinateNotFoundError as err:
678 raise ValueError(f"Cube must have a {sequence_coordinate} coordinate.") from err
680 # Ensure we have a name for the plot file.
681 recipe_title = get_recipe_metadata().get("title", "Untitled")
682 if filename is None: 682 ↛ 686line 682 didn't jump to line 686 because the condition on line 682 was always true
683 filename = slugify(recipe_title)
685 # Make vertical line plot
686 plotting_func = _plot_and_save_vertical_line_series
688 # set the lower and upper limit for the x-axis to ensure all plots
689 # have same range. This needs to read the whole cube over the range of
690 # the sequence and if applicable postage stamp coordinate.
691 # This only works if the plotting is done in the collate section of a
692 # recipe and not in the parallel section of a recipe.
693 vmin = np.floor((cube.data.min()))
694 vmax = np.ceil((cube.data.max()))
696 # Create a plot for each value of the sequence coordinate.
697 plot_index = []
698 for cube_slice in cube.slices_over(sequence_coordinate):
699 # Use sequence value so multiple sequences can merge.
700 sequence_value = cube_slice.coord(sequence_coordinate).points[0]
701 plot_filename = f"{filename.rsplit('.', 1)[0]}_{sequence_value}.png"
702 coord = cube_slice.coord(series_coordinate)
703 # Format the coordinate value in a unit appropriate way.
704 title = f"{recipe_title} | {coord.units.title(coord.points[0])}"
705 # Do the actual plotting.
706 plotting_func(
707 cube_slice,
708 coord,
709 plot_filename,
710 title=title,
711 vmin=vmin,
712 vmax=vmax,
713 )
714 plot_index.append(plot_filename)
716 # Add list of plots to plot metadata.
717 complete_plot_index = _append_to_plot_index(plot_index)
719 # Make a page to display the plots.
720 _make_plot_html_page(complete_plot_index)
722 return cube