Merge pull request #1484 from plutowing/enhance_wind_vector_plots

Add surface wind vector plots recipe and related function

Author: jfrost-mo

.github/workflows/pull-request-checks.yml

@@ -91,6 +91,7 @@ jobs:
           # This links to a hosted version of the HTML report.
           tar -czf coverage-report.tar.gz htmlcov/
           report_url="$(curl -sSf --data-binary @coverage-report.tar.gz https://tmpweb.net)"
+          echo "Report hosted at $report_url"
           badge_options="$(coverage json --fail-under=0 -qo - | jq -r .totals.percent_covered_display)%25-blue?style=for-the-badge"
           echo "[![Coverage](https://img.shields.io/badge/coverage-${badge_options})](${report_url})" >> ${{ runner.temp }}/cov-report.md
           # Edit last comment if it exists, else create new one.

src/CSET/operators/_colorbar_definition.json

@@ -153,6 +153,11 @@
     "max": 5,
     "min": -5
   },
+  "eastward_wind_northward_wind_magnitude": {
+    "cmap": "YlOrRd",
+    "max": 25.0,
+    "min": 0.0
+  },
   "exner_pressure_at_cell_interfaces": {
     "cmap": "coolwarm",
     "max": 1.1,

src/CSET/operators/plot.py

@@ -825,6 +825,134 @@ def _plot_and_save_scatter_plot(
     plt.close(fig)
 
 
+def _plot_and_save_vector_plot(
+    cube_u: iris.cube.Cube,
+    cube_v: iris.cube.Cube,
+    filename: str,
+    title: str,
+    method: Literal["contourf", "pcolormesh"],
+    **kwargs,
+):
+    """Plot and save a 2D vector plot.
+
+    Parameters
+    ----------
+    cube_u: Cube
+        2 dimensional Cube of u component of the data.
+    cube_v: Cube
+        2 dimensional Cube of v component of the data.
+    filename: str
+        Filename of the plot to write.
+    title: str
+        Plot title.
+    """
+    fig = plt.figure(figsize=(10, 10), facecolor="w", edgecolor="k")
+
+    # Create a cube containing the magnitude of the vector field.
+    cube_vec_mag = (cube_u**2 + cube_v**2) ** 0.5
+    cube_vec_mag.rename(f"{cube_u.name()}_{cube_v.name()}_magnitude")
+
+    # Specify the color bar
+    cmap, levels, norm = _colorbar_map_levels(cube_vec_mag)
+
+    if method == "contourf":
+        # Filled contour plot of the field.
+        plot = iplt.contourf(cube_vec_mag, cmap=cmap, levels=levels, norm=norm)
+    elif method == "pcolormesh":
+        try:
+            vmin = min(levels)
+            vmax = max(levels)
+        except TypeError:
+            vmin, vmax = None, None
+        # pcolormesh plot of the field and ensure to use norm and not vmin/vmax
+        # if levels are defined.
+        if norm is not None:
+            vmin = None
+            vmax = None
+        plot = iplt.pcolormesh(cube_vec_mag, cmap=cmap, norm=norm, vmin=vmin, vmax=vmax)
+    else:
+        raise ValueError(f"Unknown plotting method: {method}")
+
+    # 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:
+        lat_axis, lon_axis = get_cube_yxcoordname(cube_vec_mag)
+        axes.coastlines(resolution="10m")
+        x1 = np.min(cube_vec_mag.coord(lon_axis).points)
+        x2 = np.max(cube_vec_mag.coord(lon_axis).points)
+        y1 = np.min(cube_vec_mag.coord(lat_axis).points)
+        y2 = np.max(cube_vec_mag.coord(lat_axis).points)
+        # Adjust bounds within +/- 180.0 if x dimension extends beyond half-globe.
+        if (x2 - x1) > 180.0:
+            x1 = x1 - 180.0
+            x2 = x2 - 180.0
+        axes.set_extent([x1, x2, y1, y2])
+    except ValueError:
+        # Skip if no x and y map coordinates.
+        pass
+
+    # Check to see if transect, and if so, adjust y axis.
+    if is_transect(cube_vec_mag):
+        if "pressure" in [coord.name() for coord in cube_vec_mag.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_vec_mag.coords()}
+        ):
+            axes.set_yscale("log")
+
+        axes.set_title(
+            f"{title}\n"
+            f"Start Lat: {cube_vec_mag.attributes['transect_coords'].split('_')[0]}"
+            f" Start Lon: {cube_vec_mag.attributes['transect_coords'].split('_')[1]}"
+            f" End Lat: {cube_vec_mag.attributes['transect_coords'].split('_')[2]}"
+            f" End Lon: {cube_vec_mag.attributes['transect_coords'].split('_')[3]}",
+            fontsize=16,
+        )
+
+    else:
+        # Add title.
+        axes.set_title(title, fontsize=16)
+
+    # Add watermark with min/max/mean. Currently not user togglable.
+    # In the bbox dictionary, fc and ec are hex colour codes for grey shade.
+    axes.annotate(
+        f"Min: {np.min(cube_vec_mag.data):.3g} Max: {np.max(cube_vec_mag.data):.3g} Mean: {np.mean(cube_vec_mag.data):.3g}",
+        xy=(1, -0.05),
+        xycoords="axes fraction",
+        xytext=(-5, 5),
+        textcoords="offset points",
+        ha="right",
+        va="bottom",
+        size=11,
+        bbox=dict(boxstyle="round", fc="#cccccc", ec="#808080", alpha=0.9),
+    )
+
+    # Add colour bar.
+    cbar = fig.colorbar(plot, orientation="horizontal", pad=0.042, shrink=0.7)
+    cbar.set_label(label=f"{cube_vec_mag.name()} ({cube_vec_mag.units})", size=16)
+    # add ticks and tick_labels for every levels if less than 20 levels exist
+    if levels is not None and len(levels) < 20:
+        cbar.set_ticks(levels)
+        cbar.set_ticklabels([f"{level:.1f}" for level in levels])
+
+    # 30 barbs along the longest axis of the plot, or a barb per point for data
+    # with less than 30 points.
+    step = max(max(cube_u.shape) // 30, 1)
+    iplt.quiver(cube_u[::step, ::step], cube_v[::step, ::step], pivot="middle")
+
+    # Save plot.
+    fig.savefig(filename, bbox_inches="tight", dpi=_get_plot_resolution())
+    logging.info("Saved vector plot to %s", filename)
+    plt.close(fig)
+
+
 def _plot_and_save_histogram_series(
     cubes: iris.cube.Cube | iris.cube.CubeList,
     filename: str,
@@ -1670,6 +1798,62 @@ def scatter_plot(
     return iris.cube.CubeList([cube_x, cube_y])
 
 
+def vector_plot(
+    cube_u: iris.cube.Cube,
+    cube_v: iris.cube.Cube,
+    filename: str = None,
+    sequence_coordinate: str = "time",
+    **kwargs,
+) -> iris.cube.CubeList:
+    """Plot a vector plot based on the input u and v components."""
+    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)
+
+    # Cubes must have a matching sequence coordinate.
+    try:
+        # Check that the u and v cubes have the same sequence coordinate.
+        if cube_u.coord(sequence_coordinate) != cube_v.coord(sequence_coordinate):
+            raise ValueError("Coordinates do not match.")
+    except (iris.exceptions.CoordinateNotFoundError, ValueError) as err:
+        raise ValueError(
+            f"Cubes should have matching {sequence_coordinate} coordinate:\n{cube_u}\n{cube_v}"
+        ) from err
+
+    # Create a plot for each value of the sequence coordinate.
+    plot_index = []
+    for cube_u_slice, cube_v_slice in zip(
+        cube_u.slices_over(sequence_coordinate),
+        cube_v.slices_over(sequence_coordinate),
+        strict=True,
+    ):
+        # Use sequence value so multiple sequences can merge.
+        sequence_value = cube_u_slice.coord(sequence_coordinate).points[0]
+        plot_filename = f"{filename.rsplit('.', 1)[0]}_{sequence_value}.png"
+        coord = cube_u_slice.coord(sequence_coordinate)
+        # Format the coordinate value in a unit appropriate way.
+        title = f"{recipe_title}\n{coord.units.title(coord.points[0])}"
+        # Do the actual plotting.
+        _plot_and_save_vector_plot(
+            cube_u_slice,
+            cube_v_slice,
+            filename=plot_filename,
+            title=title,
+            method="contourf",
+        )
+        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 iris.cube.CubeList([cube_u, cube_v])
+
+
 def plot_histogram_series(
     cubes: iris.cube.Cube | iris.cube.CubeList,
     filename: str = None,

src/CSET/recipes/generic_surface_vector_plot_sequence.yaml

@@ -0,0 +1,32 @@
+category: Surface Vector Plot
+title: $MODEL_NAME Surface Vector Plot of $VARNAME_U and $VARNAME_V
+description: Extracts and plots the surface vector plot of $VARNAME_U and $VARNAME_V for all times in $MODEL_NAME.
+
+steps:
+  - operator: read.read_cubes
+    file_paths: $INPUT_PATHS
+
+  - operator: plot.vector_plot
+    cube_u:
+      operator: filters.filter_cubes
+      constraint:
+        operator: constraints.combine_constraints
+        variable_constraint:
+          operator: constraints.generate_var_constraint
+          varname: $VARNAME_U
+        level_constraint:
+          operator: constraints.generate_level_constraint
+          coordinate: $LEVELTYPE
+          levels: $LEVEL
+    cube_v:
+      operator: filters.filter_cubes
+      constraint:
+        operator: constraints.combine_constraints
+        variable_constraint:
+          operator: constraints.generate_var_constraint
+          varname: $VARNAME_V
+        level_constraint:
+          operator: constraints.generate_level_constraint
+          coordinate: $LEVELTYPE
+          levels: $LEVEL
+    sequence_coordinate: time

tests/conftest.py

@@ -17,6 +17,7 @@
 https://docs.pytest.org/en/latest/reference/fixtures.html#conftest-py-sharing-fixtures-across-multiple-files
 """
 
+import iris.cube
 import pytest
 
 from CSET.operators import constraints, filters, read
@@ -83,6 +84,35 @@ def vertical_profile_cube(vertical_profile_cube_readonly):
     return vertical_profile_cube_readonly.copy()
 
 
+@pytest.fixture(scope="session")
+def vector_cubes_readonly():
+    """Get vector plot cubes. It is NOT safe to modify."""
+    from CSET.operators import read
+
+    # Read the input cubes.
+    in_cubes = read.read_cubes("tests/test_data/u10_v10.nc")
+    # Generate constraints for the u and v components of the wind.
+    var_constraint_u = constraints.generate_var_constraint("eastward_wind")
+    var_constraint_v = constraints.generate_var_constraint("northward_wind")
+    lev_contraint = constraints.generate_level_constraint("height", 10)
+    combined_constraint_u = constraints.combine_constraints(
+        a=var_constraint_u, b=lev_contraint
+    )
+    combined_constraint_v = constraints.combine_constraints(
+        a=var_constraint_v, b=lev_contraint
+    )
+    # filter the cubes using the constraints.
+    cube_u = filters.filter_cubes(in_cubes, combined_constraint_u)
+    cube_v = filters.filter_cubes(in_cubes, combined_constraint_v)
+    return iris.cube.CubeList([cube_u, cube_v])
+
+
+@pytest.fixture()
+def vector_cubes(vector_cubes_readonly):
+    """Get vector plot cubes."""
+    return vector_cubes_readonly.copy()
+
+
 @pytest.fixture(scope="session")
 def histogram_cube_readonly():
     """Get a histogram Cube. It is NOT safe to modify."""

tests/operators/test_plots.py

@@ -259,6 +259,40 @@ def test_contour_plot_sequence(cube, tmp_working_dir):
     assert Path("untitled_462149.0.png").is_file()
 
 
+def test_vector_plot_with_filename(vector_cubes, tmp_working_dir):
+    """Plot a vector plot of u10 and v10 components."""
+    cube_u = vector_cubes[0].slices_over("time").next()
+    cube_v = vector_cubes[1].slices_over("time").next()
+    plot.vector_plot(cube_u, cube_v, filename="testvector")
+    assert Path("testvector_0.0.png").is_file()
+
+
+def test_vector_plot_sequence(vector_cubes, tmp_working_dir):
+    """Plot a sequence of vector plots."""
+    plot.vector_plot(
+        vector_cubes[0],
+        vector_cubes[1],
+        filename="testvectorseq",
+        sequence_coordinate="time",
+    )
+    assert Path("testvectorseq_0.0.png").is_file()
+    assert Path("testvectorseq_6.0.png").is_file()
+    assert Path("testvectorseq_12.0.png").is_file()
+
+
+def test_vector_plot_check(vector_cubes, tmp_working_dir):
+    """Check error when cubes has no time coordinate."""
+    vector_cubes[0].remove_coord("time")
+    vector_cubes[1].remove_coord("time")
+    with pytest.raises(ValueError):
+        plot.vector_plot(
+            vector_cubes[0],
+            vector_cubes[1],
+            filename="testvector",
+            sequence_coordinate="time",
+        )
+
+
 def test_postage_stamp_contour_plot(ensemble_cube, monkeypatch, tmp_path):
     """Plot postage stamp plots of ensemble data."""
     # Get a single time step.