Merge pull request #787 from MetOffice/add_pcolormesh_plot

Add pcolormesh plotting operator

src/CSET/operators/plot.py

@@ -20,7 +20,7 @@
 import logging
 import math
 import sys
-from typing import Union
+from typing import Literal
 
 import iris
 import iris.coords
@@ -55,9 +55,7 @@ def _append_to_plot_index(plot_index: list) -> list:
     return complete_plot_index
 
 
-def _check_single_cube(
-    cube: Union[iris.cube.Cube, iris.cube.CubeList],
-) -> iris.cube.Cube:
+def _check_single_cube(cube: iris.cube.Cube | iris.cube.CubeList) -> iris.cube.Cube:
     """Ensure a single cube is given.
 
     If a CubeList of length one is given that the contained cube is returned,
@@ -186,13 +184,14 @@ def _get_plot_resolution() -> int:
     return get_recipe_metadata().get("plot_resolution", 100)
 
 
-def _plot_and_save_contour_plot(
+def _plot_and_save_spatial_plot(
     cube: iris.cube.Cube,
     filename: str,
     title: str,
+    method: Literal["contourf", "pcolormesh"],
     **kwargs,
 ):
-    """Plot and save a contour plot.
+    """Plot and save a spatial plot.
 
     Parameters
     ----------
@@ -202,16 +201,23 @@ def _plot_and_save_contour_plot(
         Filename of the plot to write.
     title: str
         Plot title.
-
+    method: "contourf" | "pcolormesh"
+        The plotting method to use.
     """
     # Setup plot details, size, resolution, etc.
     fig = plt.figure(figsize=(8, 8), facecolor="w", edgecolor="k")
 
     # Specify the color bar
     cmap, levels, norm = _colorbar_map_levels(cube.name())
 
-    # Filled contour plot of the field.
-    contours = iplt.contourf(cube, cmap=cmap, levels=levels, norm=norm)
+    if method == "contourf":
+        # Filled contour plot of the field.
+        plot = iplt.contourf(cube, cmap=cmap, levels=levels, norm=norm)
+    elif method == "pcolormesh":
+        # pcolormesh plot of the field.
+        plot = iplt.pcolormesh(cube, cmap=cmap, norm=norm)
+    else:
+        raise ValueError(f"Unknown plotting method: {method}")
 
     # Using pyplot interface here as we need iris to generate a cartopy GeoAxes.
     axes = plt.gca()
@@ -230,6 +236,7 @@ def _plot_and_save_contour_plot(
             ]
         )
     except ValueError:
+        # Skip if no x and y map coordinates.
         pass
 
     # Check to see if transect, and if so, adjust y axis.
@@ -263,7 +270,7 @@ def _plot_and_save_contour_plot(
     )
 
     # Add colour bar.
-    cbar = fig.colorbar(contours)
+    cbar = fig.colorbar(plot)
     cbar.set_label(label=f"{cube.name()} ({cube.units})", size=20)
 
     # Save plot.
@@ -272,14 +279,15 @@ def _plot_and_save_contour_plot(
     plt.close(fig)
 
 
-def _plot_and_save_postage_stamp_contour_plot(
+def _plot_and_save_postage_stamp_spatial_plot(
     cube: iris.cube.Cube,
     filename: str,
     stamp_coordinate: str,
     title: str,
+    method: Literal["contourf", "pcolormesh"],
     **kwargs,
 ):
-    """Plot postage stamp contour plots from an ensemble.
+    """Plot postage stamp spatial plots from an ensemble.
 
     Parameters
     ----------
@@ -289,6 +297,8 @@ def _plot_and_save_postage_stamp_contour_plot(
         Filename of the plot to write.
     stamp_coordinate: str
         Coordinate that becomes different plots.
+    method: "contourf" | "pcolormesh"
+        The plotting method to use.
 
     Raises
     ------
@@ -310,16 +320,33 @@ def _plot_and_save_postage_stamp_contour_plot(
         # Implicit interface is much easier here, due to needing to have the
         # cartopy GeoAxes generated.
         plt.subplot(grid_size, grid_size, subplot)
-        plot = iplt.contourf(member, cmap=cmap, levels=levels, norm=norm)
+        if method == "contourf":
+            # Filled contour plot of the field.
+            plot = iplt.contourf(member, cmap=cmap, levels=levels, norm=norm)
+        elif method == "pcolormesh":
+            # pcolormesh plot of the field.
+            plot = iplt.pcolormesh(member, cmap=cmap, norm=norm)
+        else:
+            raise ValueError(f"Unknown plotting method: {method}")
         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.
+        # If is spatial map, fix extent to keep plot tight.
         try:
-            get_cube_yxcoordname(cube)
+            lataxis, lonaxis = get_cube_yxcoordname(cube)
             ax.coastlines(resolution="10m")
+            ax.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:
+            # Skip if no x and y map coordinates.
             pass
 
     # Put the shared colorbar in its own axes.
@@ -681,18 +708,13 @@ def _plot_and_save_postage_stamps_in_single_plot_histogram_series(
     plt.close(fig)
 
 
-####################
-# Public functions #
-####################
-
-
-def spatial_contour_plot(
+def _spatial_plot(
+    method: Literal["contourf", "pcolormesh"],
     cube: iris.cube.Cube,
-    filename: str = None,
-    sequence_coordinate: str = "time",
-    stamp_coordinate: str = "realization",
-    **kwargs,
-) -> iris.cube.Cube:
+    filename: str | None,
+    sequence_coordinate: str,
+    stamp_coordinate: str,
+):
     """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
@@ -701,25 +723,22 @@ def spatial_contour_plot(
 
     Parameters
     ----------
+    method: "contourf" | "pcolormesh"
+        The plotting method to use.
     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
+    filename: str | None
+        Name of the plot to write, used as a prefix for plot sequences. If None
+        uses the recipe name.
+    sequence_coordinate: str
         Coordinate about which to make a plot sequence. Defaults to ``"time"``.
         This coordinate must exist in the cube.
-    stamp_coordinate: str, optional
+    stamp_coordinate: str
         Coordinate about which to plot postage stamp plots. Defaults to
         ``"realization"``.
 
-    Returns
-    -------
-    Cube
-        The original cube (so further operations can be applied).
-
     Raises
     ------
     ValueError
@@ -738,10 +757,10 @@ def spatial_contour_plot(
 
     # Make postage stamp plots if stamp_coordinate exists and has more than a
     # single point.
-    plotting_func = _plot_and_save_contour_plot
+    plotting_func = _plot_and_save_spatial_plot
     try:
         if cube.coord(stamp_coordinate).shape[0] > 1:
-            plotting_func = _plot_and_save_postage_stamp_contour_plot
+            plotting_func = _plot_and_save_postage_stamp_spatial_plot
     except iris.exceptions.CoordinateNotFoundError:
         pass
 
@@ -763,9 +782,10 @@ def spatial_contour_plot(
         # Do the actual plotting.
         plotting_func(
             cube_slice,
-            plot_filename,
+            filename=plot_filename,
             stamp_coordinate=stamp_coordinate,
             title=title,
+            method=method,
         )
         plot_index.append(plot_filename)
 
@@ -775,6 +795,103 @@ def spatial_contour_plot(
     # Make a page to display the plots.
     _make_plot_html_page(complete_plot_index)
 
+
+####################
+# Public functions #
+####################
+
+
+def spatial_contour_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.
+
+    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.
+    """
+    _spatial_plot("contourf", cube, filename, sequence_coordinate, stamp_coordinate)
+    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.
+    """
+    _spatial_plot("pcolormesh", cube, filename, sequence_coordinate, stamp_coordinate)
     return cube
 
 
@@ -954,7 +1071,7 @@ def scatter_plot(
     filename: str = None,
     one_to_one: bool = True,
     **kwargs,
-) -> (iris.cube.Cube, iris.cube.Cube):
+) -> tuple[iris.cube.Cube, iris.cube.Cube]:
     """Plot a scatter plot between two variables.
 
     Both cubes must be 1D.

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")
@@ -269,3 +304,17 @@ def test_get_plot_resolution_unset(tmp_working_dir):
     """Test getting the default plot resolution when unset."""
     resolution = plot._get_plot_resolution()
     assert resolution == 100
+
+
+def test_invalid_plotting_method_spatial_plot(cube, tmp_working_dir):
+    """Test plotting a spatial plot with an invalid method."""
+    with pytest.raises(ValueError, match="Unknown plotting method"):
+        plot._plot_and_save_spatial_plot(cube, "filename", "title", "invalid")
+
+
+def test_invalid_plotting_method_postage_stamp_spatial_plot(cube, tmp_working_dir):
+    """Test plotting a postage stamp spatial plot with an invalid method."""
+    with pytest.raises(ValueError, match="Unknown plotting method"):
+        plot._plot_and_save_postage_stamp_spatial_plot(
+            cube, "filename", "realization", "title", "invalid"
+        )