Add pcolormesh plotting operator

Also test transect plotting.

Fixes #779

src/CSET/operators/plot.py

@@ -316,6 +316,141 @@ def _plot_and_save_postage_stamp_contour_plot(
     plt.close(fig)
 
 
+def _plot_and_save_pcolormesh_plot(
+    cube: iris.cube.Cube,
+    filename: str,
+    title: str,
+    **kwargs,
+):
+    """Plot and save a pcolormesh plot.
+
+    Parameters
+    ----------
+    cube: Cube
+        2 dimensional (lat and lon) Cube of the data to plot.
+    filename: str
+        Filename of the plot to write.
+    title: str
+        Plot title.
+
+    """
+    # Setup plot details, size, resolution, etc.
+    fig = plt.figure(figsize=(15, 15), facecolor="w", edgecolor="k")
+
+    # Specify the color bar
+    cmap, levels, norm = _colorbar_map_levels(cube.name())
+
+    # Filled contour plot of the field.
+    colormesh = iplt.pcolormesh(cube, cmap=cmap, norm=norm)
+
+    # Using pyplot interface here as we need iris to generate a cartopy GeoAxes.
+    axes = plt.gca()
+
+    # Add coastlines if cube contains x and y map coordinates.
+    # If is spatial map, fix extent to keep plot tight.
+    try:
+        lataxis, lonaxis = get_cube_yxcoordname(cube)
+        axes.coastlines(resolution="10m")
+        axes.set_extent(
+            [
+                np.min(cube.coord(lonaxis).points),
+                np.max(cube.coord(lonaxis).points),
+                np.min(cube.coord(lataxis).points),
+                np.max(cube.coord(lataxis).points),
+            ]
+        )
+    except ValueError:
+        pass
+
+    # Check to see if transect, and if so, adjust y axis.
+    if is_transect(cube):
+        if "pressure" in [coord.name() for coord in cube.coords()]:
+            axes.invert_yaxis()
+            axes.set_yscale("log")
+            axes.set_ylim(1100, 100)
+        # If both model_level_number and level_height exists, iplt can construct
+        # plot as a function of height above orography (NOT sea level).
+        elif {"model_level_number", "level_height"}.issubset(
+            {coord.name() for coord in cube.coords()}
+        ):
+            axes.set_yscale("log")
+
+    # Add title.
+    axes.set_title(title, fontsize=16)
+
+    # Add colour bar.
+    cbar = fig.colorbar(colormesh)
+    cbar.set_label(label=f"{cube.name()} ({cube.units})", size=20)
+
+    # Save plot.
+    fig.savefig(filename, bbox_inches="tight", dpi=150)
+    logging.info("Saved contour plot to %s", filename)
+    plt.close(fig)
+
+
+def _plot_and_save_postage_stamp_pcolormesh_plot(
+    cube: iris.cube.Cube,
+    filename: str,
+    stamp_coordinate: str,
+    title: str,
+    **kwargs,
+):
+    """Plot postage stamp pcolormesh plots from an ensemble.
+
+    Parameters
+    ----------
+    cube: Cube
+        Iris cube of data to be plotted. It must have the stamp coordinate.
+    filename: str
+        Filename of the plot to write.
+    stamp_coordinate: str
+        Coordinate that becomes different plots.
+
+    Raises
+    ------
+    ValueError
+        If the cube doesn't have the right dimensions.
+    """
+    # Use the smallest square grid that will fit the members.
+    grid_size = int(math.ceil(math.sqrt(len(cube.coord(stamp_coordinate).points))))
+
+    fig = plt.figure(figsize=(10, 10))
+
+    # Specify the color bar
+    cmap, levels, norm = _colorbar_map_levels(cube.name())
+
+    # Make a subplot for each member.
+    for member, subplot in zip(
+        cube.slices_over(stamp_coordinate), range(1, grid_size**2 + 1), strict=False
+    ):
+        # Implicit interface is much easier here, due to needing to have the
+        # cartopy GeoAxes generated.
+        plt.subplot(grid_size, grid_size, subplot)
+        plot = iplt.pcolormesh(member, cmap=cmap, norm=norm)
+        ax = plt.gca()
+        ax.set_title(f"Member #{member.coord(stamp_coordinate).points[0]}")
+        ax.set_axis_off()
+
+        # Add coastlines if cube contains x and y map coordinates.
+        try:
+            get_cube_yxcoordname(cube)
+            ax.coastlines(resolution="10m")
+        except ValueError:
+            pass
+
+    # Put the shared colorbar in its own axes.
+    colorbar_axes = fig.add_axes([0.15, 0.07, 0.7, 0.03])
+    colorbar = fig.colorbar(plot, colorbar_axes, orientation="horizontal")
+    colorbar.set_label(f"{cube.name()} / {cube.units}")
+
+    # Overall figure title.
+    fig.suptitle(title)
+
+    fig.savefig(filename, bbox_inches="tight", dpi=150)
+    logging.info("Saved contour postage stamp plot to %s", filename)
+    plt.close(fig)
+
+
 def _plot_and_save_line_series(
     cube: iris.cube.Cube, coord: iris.coords.Coord, filename: str, title: str, **kwargs
 ):
@@ -761,6 +896,107 @@ def spatial_contour_plot(
     return cube
 
 
+def spatial_pcolormesh_plot(
+    cube: iris.cube.Cube,
+    filename: str = None,
+    sequence_coordinate: str = "time",
+    stamp_coordinate: str = "realization",
+    **kwargs,
+) -> iris.cube.Cube:
+    """Plot a spatial variable onto a map from a 2D, 3D, or 4D cube.
+
+    A 2D spatial field can be plotted, but if the sequence_coordinate is present
+    then a sequence of plots will be produced. Similarly if the stamp_coordinate
+    is present then postage stamp plots will be produced.
+
+    This function is significantly faster than ``spatial_contour_plot``,
+    especially at high resolutions, and should be preferred unless contiguous
+    contour areas are important.
+
+    Parameters
+    ----------
+    cube: Cube
+        Iris cube of the data to plot. It should have two spatial dimensions,
+        such as lat and lon, and may also have a another two dimension to be
+        plotted sequentially and/or as postage stamp plots.
+    filename: str, optional
+        Name of the plot to write, used as a prefix for plot sequences. Defaults
+        to the recipe name.
+    sequence_coordinate: str, optional
+        Coordinate about which to make a plot sequence. Defaults to ``"time"``.
+        This coordinate must exist in the cube.
+    stamp_coordinate: str, optional
+        Coordinate about which to plot postage stamp plots. Defaults to
+        ``"realization"``.
+
+    Returns
+    -------
+    Cube
+        The original cube (so further operations can be applied).
+
+    Raises
+    ------
+    ValueError
+        If the cube doesn't have the right dimensions.
+    TypeError
+        If the cube isn't a single cube.
+    """
+    recipe_title = get_recipe_metadata().get("title", "Untitled")
+
+    # Ensure we have a name for the plot file.
+    if filename is None:
+        filename = slugify(recipe_title)
+
+    # Ensure we've got a single cube.
+    cube = _check_single_cube(cube)
+
+    # Make postage stamp plots if stamp_coordinate exists and has more than a
+    # single point.
+    plotting_func = _plot_and_save_pcolormesh_plot
+    try:
+        if cube.coord(stamp_coordinate).shape[0] > 1:
+            plotting_func = _plot_and_save_postage_stamp_pcolormesh_plot
+    except iris.exceptions.CoordinateNotFoundError:
+        pass
+
+    # Must have a sequence coordinate.
+    try:
+        cube.coord(sequence_coordinate)
+    except iris.exceptions.CoordinateNotFoundError as err:
+        raise ValueError(f"Cube must have a {sequence_coordinate} coordinate.") from err
+
+    # Ensure dimension coordinates are bounded.
+    for dim_coord in cube.coords(dim_coords=True):
+        if not dim_coord.has_bounds():
+            dim_coord.guess_bounds()
+
+    # Create a plot for each value of the sequence coordinate.
+    plot_index = []
+    for cube_slice in cube.slices_over(sequence_coordinate):
+        # Use sequence value so multiple sequences can merge.
+        sequence_value = cube_slice.coord(sequence_coordinate).points[0]
+        plot_filename = f"{filename.rsplit('.', 1)[0]}_{sequence_value}.png"
+        coord = cube_slice.coord(sequence_coordinate)
+        # Format the coordinate value in a unit appropriate way.
+        title = f"{recipe_title} | {coord.units.title(coord.points[0])}"
+        # Do the actual plotting.
+        plotting_func(
+            cube_slice,
+            plot_filename,
+            stamp_coordinate=stamp_coordinate,
+            title=title,
+        )
+        plot_index.append(plot_filename)
+
+    # Add list of plots to plot metadata.
+    complete_plot_index = _append_to_plot_index(plot_index)
+
+    # Make a page to display the plots.
+    _make_plot_html_page(complete_plot_index)
+
+    return cube
+
+
 # TODO: Expand function to handle ensemble data.
 # line_coordinate: str, optional
 #     Coordinate about which to plot multiple lines. Defaults to

tests/operators/test_plots.py

@@ -71,6 +71,41 @@ def test_postage_stamp_contour_plot_sequence_coord_check(cube, tmp_working_dir):
         plot.spatial_contour_plot(cube)
 
 
+def test_spatial_pcolormesh_plot(cube, tmp_working_dir):
+    """Plot spatial pcolormesh plot of instant air temp."""
+    # Remove realization coord to increase coverage, and as its not needed.
+    cube.remove_coord("realization")
+    cube_2d = cube.slices_over("time").next()
+    plot.spatial_pcolormesh_plot(cube_2d, filename="plot")
+    assert Path("plot_462147.0.png").is_file()
+
+
+def test_pcolormesh_plot_sequence(cube, tmp_working_dir):
+    """Plot sequence of pcolormesh plots."""
+    plot.spatial_pcolormesh_plot(cube, sequence_coordinate="time")
+    assert Path("untitled_462147.0.png").is_file()
+    assert Path("untitled_462148.0.png").is_file()
+    assert Path("untitled_462149.0.png").is_file()
+
+
+def test_postage_stamp_pcolormesh_plot(monkeypatch, tmp_path):
+    """Plot postage stamp plots of ensemble data."""
+    ensemble_cube = read.read_cube("tests/test_data/exeter_em*.nc")
+    # Get a single time step.
+    ensemble_cube_3d = next(ensemble_cube.slices_over("time"))
+    monkeypatch.chdir(tmp_path)
+    plot.spatial_pcolormesh_plot(ensemble_cube_3d)
+    assert Path("untitled_463858.0.png").is_file()
+
+
+def test_postage_stamp_pcolormesh_plot_sequence_coord_check(cube, tmp_working_dir):
+    """Check error when cube has no time coordinate."""
+    # What does this even physically mean? No data?
+    cube.remove_coord("time")
+    with pytest.raises(ValueError):
+        plot.spatial_pcolormesh_plot(cube)
+
+
 def test_plot_line_series(cube, tmp_working_dir):
     """Save a line series plot."""
     cube = collapse.collapse(cube, ["grid_latitude", "grid_longitude"], "MEAN")

tests/operators/test_transect.py

@@ -14,13 +14,15 @@
 
 """Tests transect operator."""
 
+from pathlib import Path
+
 import iris
 import iris.coord_systems
 import iris.cube
 import numpy as np
 import pytest
 
-import CSET.operators.transect as transect
+from CSET.operators import plot, transect
 
 
 # Session scope fixtures, so the test data only has to be loaded once.
@@ -144,3 +146,33 @@ def test_transect_plotasfuncoflongitude(load_cube_pl):
         transect.calc_transect(
             load_cube_pl, startcoords=(-10.94, 19.06), endcoords=(-10.86, 19.18)
         ).coord("latitude")
+
+
+def test_transect_model_level_spatial_contour_plot(load_cube_ml_out, tmp_working_dir):
+    """Plot a contour plot of the transect model level data."""
+    plot.spatial_contour_plot(load_cube_ml_out, filename="plot")
+    assert Path("plot_449743.0.png").is_file()
+    assert Path("plot_449744.0.png").is_file()
+
+
+def test_transect_model_level_spatial_pcolormesh_plot(
+    load_cube_ml_out, tmp_working_dir
+):
+    """Plot a pcolormesh plot of the transect model level data."""
+    plot.spatial_pcolormesh_plot(load_cube_ml_out, filename="plot")
+    assert Path("plot_449743.0.png").is_file()
+    assert Path("plot_449744.0.png").is_file()
+
+
+def test_transect_pressure_spatial_contour_plot(load_cube_pl_out, tmp_working_dir):
+    """Plot a contour plot of the transect pressure data."""
+    plot.spatial_contour_plot(load_cube_pl_out, filename="plot")
+    assert Path("plot_449469.0.png").is_file()
+    assert Path("plot_449472.0.png").is_file()
+
+
+def test_transect_pressure_spatial_pcolormesh_plot(load_cube_pl_out, tmp_working_dir):
+    """Plot a pcolormesh plot of the transect pressure data."""
+    plot.spatial_pcolormesh_plot(load_cube_pl_out, filename="plot")
+    assert Path("plot_449469.0.png").is_file()
+    assert Path("plot_449472.0.png").is_file()