diff --git a/src/CSET/operators/_utils.py b/src/CSET/operators/_utils.py index 7518f92e0..366d94e7b 100644 --- a/src/CSET/operators/_utils.py +++ b/src/CSET/operators/_utils.py @@ -247,6 +247,34 @@ def is_spatialdim(cube: iris.cube.Cube) -> bool: return False +def is_coorddim(cube: iris.cube.Cube, coord_name) -> bool: + """Determine whether a cube has specified dimension coordinates. + + Arguments + --------- + cube: iris.cube.Cube + An iris cube which will be checked to see if it contains coordinate + names that match a pre-defined list of acceptable coordinate names. + + coord_name: str + A cube dimension name + + Returns + ------- + bool + If true, then the cube has a spatial projection and thus can be plotted + as a map. + """ + # Get a list of dimension coordinate names for the cube + coord_names = [coord.name() for coord in cube.coords(dim_coords=True)] + + # Check if requested dimension is found in cube and get index + if coord_name in coord_names: + return True + else: + return False + + def is_transect(cube: iris.cube.Cube) -> bool: """Determine whether a cube is a transect. @@ -332,6 +360,40 @@ def fully_equalise_attributes(cubes: iris.cube.CubeList): return cubes +def slice_over_maybe(cube: iris.cube.Cube, coord_name, index): + """Test slicing over cube if exists. + + Return None if not existing. + + Arguments + --------- + cube: iris.cube.Cube + An iris cube which will be checked to see if it can be sliced over + given coordinate. + coord_name: coord + An iris coordinate over which to slice cube. + index: + Coordinate index value to extract + + Returns + ------- + cube_slice: iris.cube.Cube + A slice of iris cube, if available to slice. + """ + if cube: + if is_coorddim(cube, coord_name): + cube_slice = cube[index] + if ( + cube.ndim > 3 + ): ## More elegant way to subset?? Need to handle ensemble inputs + cube_slice = cube[:, index] + else: + cube_slice = cube + else: + cube_slice = None + return cube_slice + + def is_time_aggregatable(cube: iris.cube.Cube) -> bool: """Determine whether a cube can be aggregated in time. diff --git a/src/CSET/operators/plot.py b/src/CSET/operators/plot.py index 737e3c3d1..40ecc360a 100644 --- a/src/CSET/operators/plot.py +++ b/src/CSET/operators/plot.py @@ -50,6 +50,7 @@ fully_equalise_attributes, get_cube_yxcoordname, is_transect, + slice_over_maybe, ) from CSET.operators.collapse import collapse from CSET.operators.misc import _extract_common_time_points @@ -472,6 +473,8 @@ def _plot_and_save_spatial_plot( filename: str, title: str, method: Literal["contourf", "pcolormesh"], + overlay_cube: iris.cube.Cube | None = None, + contour_cube: iris.cube.Cube | None = None, **kwargs, ): """Plot and save a spatial plot. @@ -486,6 +489,10 @@ def _plot_and_save_spatial_plot( Plot title. method: "contourf" | "pcolormesh" The plotting method to use. + overlay_cube: Cube, optional + Optional 2 dimensional (lat and lon) Cube of data to overplot on top of base cube + contour_cube: Cube, optional + Optional 2 dimensional (lat and lon) Cube of data to overplot as contours over base cube """ # Setup plot details, size, resolution, etc. fig = plt.figure(figsize=(10, 10), facecolor="w", edgecolor="k") @@ -493,6 +500,12 @@ def _plot_and_save_spatial_plot( # Specify the color bar cmap, levels, norm = _colorbar_map_levels(cube) + # If overplotting, set required colorbars + if overlay_cube: + over_cmap, over_levels, over_norm = _colorbar_map_levels(overlay_cube) + if contour_cube: + cntr_cmap, cntr_levels, cntr_norm = _colorbar_map_levels(contour_cube) + # Setup plot map projection, extent and coastlines. axes = _setup_spatial_map(cube, fig, cmap) @@ -516,6 +529,36 @@ def _plot_and_save_spatial_plot( else: raise ValueError(f"Unknown plotting method: {method}") + # Overplot overlay field, if required + if overlay_cube: + try: + over_vmin = min(over_levels) + over_vmax = max(over_levels) + except TypeError: + over_vmin, over_vmax = None, None + if over_norm is not None: + over_vmin = None + over_vmax = None + iplt.pcolormesh( + overlay_cube, + cmap=over_cmap, + norm=over_norm, + alpha=0.8, + vmin=over_vmin, + vmax=over_vmax, + ) + # Overplot contour field, if required + if contour_cube: + iplt.contour( + contour_cube, + colors="darkgray", + levels=cntr_levels, + norm=cntr_norm, + alpha=0.5, + linestyles="--", + linewidths=1, + ) + # 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()]: @@ -579,6 +622,8 @@ def _plot_and_save_postage_stamp_spatial_plot( stamp_coordinate: str, title: str, method: Literal["contourf", "pcolormesh"], + overlay_cube: iris.cube.Cube | None = None, + contour_cube: iris.cube.Cube | None = None, **kwargs, ): """Plot postage stamp spatial plots from an ensemble. @@ -593,6 +638,10 @@ def _plot_and_save_postage_stamp_spatial_plot( Coordinate that becomes different plots. method: "contourf" | "pcolormesh" The plotting method to use. + overlay_cube: Cube, optional + Optional 2 dimensional (lat and lon) Cube of data to overplot on top of base cube + contour_cube: Cube, optional + Optional 2 dimensional (lat and lon) Cube of data to overplot as contours over base cube Raises ------ @@ -606,6 +655,11 @@ def _plot_and_save_postage_stamp_spatial_plot( # Specify the color bar cmap, levels, norm = _colorbar_map_levels(cube) + # If overplotting, set required colorbars + if overlay_cube: + over_cmap, over_levels, over_norm = _colorbar_map_levels(overlay_cube) + if contour_cube: + cntr_cmap, cntr_levels, cntr_norm = _colorbar_map_levels(contour_cube) # Make a subplot for each member. for member, subplot in zip( @@ -634,6 +688,36 @@ def _plot_and_save_postage_stamp_spatial_plot( plot = iplt.pcolormesh(member, cmap=cmap, norm=norm, vmin=vmin, vmax=vmax) else: raise ValueError(f"Unknown plotting method: {method}") + + # Overplot overlay field, if required + if overlay_cube: + try: + over_vmin = min(over_levels) + over_vmax = max(over_levels) + except TypeError: + over_vmin, over_vmax = None, None + if over_norm is not None: + over_vmin = None + over_vmax = None + iplt.pcolormesh( + overlay_cube[member.coord(stamp_coordinate).points[0]], + cmap=over_cmap, + norm=over_norm, + alpha=0.6, + vmin=over_vmin, + vmax=over_vmax, + ) + # Overplot contour field, if required + if contour_cube: + iplt.contour( + contour_cube[member.coord(stamp_coordinate).points[0]], + colors="darkgray", + levels=cntr_levels, + norm=cntr_norm, + alpha=0.6, + linestyles="--", + linewidths=1, + ) axes.set_title(f"Member #{member.coord(stamp_coordinate).points[0]}") axes.set_axis_off() @@ -1480,6 +1564,8 @@ def _spatial_plot( filename: str | None, sequence_coordinate: str, stamp_coordinate: str, + overlay_cube: iris.cube.Cube | None = None, + contour_cube: iris.cube.Cube | None = None, ): """Plot a spatial variable onto a map from a 2D, 3D, or 4D cube. @@ -1487,6 +1573,9 @@ def _spatial_plot( then a sequence of plots will be produced. Similarly if the stamp_coordinate is present then postage stamp plots will be produced. + If an overlay_cube and/or contour_cube are specified, multiple variables can + be overplotted on the same figure. + Parameters ---------- method: "contourf" | "pcolormesh" @@ -1504,6 +1593,10 @@ def _spatial_plot( stamp_coordinate: str Coordinate about which to plot postage stamp plots. Defaults to ``"realization"``. + overlay_cube: Cube | None, optional + Optional 2 dimensional (lat and lon) Cube of data to overplot on top of base cube + contour_cube: Cube | None, optional + Optional 2 dimensional (lat and lon) Cube of data to overplot as contours over base cube Raises ------ @@ -1539,7 +1632,8 @@ def _spatial_plot( # Create a plot for each value of the sequence coordinate. plot_index = [] nplot = np.size(cube.coord(sequence_coordinate).points) - for cube_slice in cube.slices_over(sequence_coordinate): + + for iseq, cube_slice in enumerate(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" @@ -1550,6 +1644,11 @@ def _spatial_plot( if nplot == 1 and coord.has_bounds: if np.size(coord.bounds) > 1: title = f"{recipe_title}\n [{coord.units.title(coord.bounds[0][0])} to {coord.units.title(coord.bounds[0][1])}]" + + # Extract sequence slice for overlay_cube and contour_cube if required. + overlay_slice = slice_over_maybe(overlay_cube, sequence_coordinate, iseq) + contour_slice = slice_over_maybe(contour_cube, sequence_coordinate, iseq) + # Do the actual plotting. plotting_func( cube_slice, @@ -1557,6 +1656,8 @@ def _spatial_plot( stamp_coordinate=stamp_coordinate, title=title, method=method, + overlay_cube=overlay_slice, + contour_cube=contour_slice, ) plot_index.append(plot_filename) @@ -2038,6 +2139,79 @@ def spatial_pcolormesh_plot( return cube +def spatial_multi_pcolormesh_plot( + cube: iris.cube.Cube, + overlay_cube: iris.cube.Cube, + contour_cube: iris.cube.Cube, + filename: str = None, + sequence_coordinate: str = "time", + stamp_coordinate: str = "realization", + **kwargs, +) -> iris.cube.Cube: + """Plot a set of spatial variables onto a map from a 2D, 3D, or 4D cube. + + A 2D basis cube 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. + + If specified, a masked overlay_cube can be overplotted on top of the base cube. + + If specified, contours of a contour_cube can be overplotted on top of those. + + For single-variable equivalent of this routine, use spatial_pcolormesh_plot. + + 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. + overlay_cube: Cube + Iris cube of the data to plot as an overlay on top of basis cube. 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. This is likely to be a masked cube in order not to hide the underlying basis cube. + contour_cube: Cube + Iris cube of the data to plot as a contour overlay on top of basis cube and overlay_cube. 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, + overlay_cube=overlay_cube, + contour_cube=contour_cube, + ) + return cube, overlay_cube, contour_cube + + # TODO: Expand function to handle ensemble data. # line_coordinate: str, optional # Coordinate about which to plot multiple lines. Defaults to diff --git a/src/CSET/recipes/surface_fields/multi_surface_spatial_plot_sequence.yaml b/src/CSET/recipes/surface_fields/multi_surface_spatial_plot_sequence.yaml new file mode 100644 index 000000000..374660737 --- /dev/null +++ b/src/CSET/recipes/surface_fields/multi_surface_spatial_plot_sequence.yaml @@ -0,0 +1,64 @@ +category: Multivar Spatial Plot +title: $MODEL_NAME $VARNAME_BASE m$METHOD +description: Extracts and overplots the $METHOD of $VARNAME_BASE, $VARNAME_OVER and $VARNAME_CONTOUR from all times in $MODEL_NAME. + +steps: + + - operator: read.read_cubes + file_paths: $INPUT_PATHS + subarea_type: $SUBAREA_TYPE + subarea_extent: $SUBAREA_EXTENT + + - operator: collapse.collapse + coordinate: [time] + method: $METHOD + + - operator: plot.spatial_multi_pcolormesh_plot + cube: + operator: filters.filter_cubes + constraint: + operator: constraints.combine_constraints + variable_constraint: + operator: constraints.generate_var_constraint + varname: $VARNAME_BASE + cell_methods_constraint: + operator: constraints.generate_cell_methods_constraint + cell_methods: [] + varname: $VARNAME_BASE + + overlay_cube: + operator: filters.apply_mask + original_field: + operator: filters.filter_cubes + constraint: + operator: constraints.combine_constraints + variable_constraint: + operator: constraints.generate_var_constraint + varname: $VARNAME_OVER + cell_methods_constraint: + operator: constraints.generate_cell_methods_constraint + cell_methods: [] + varname: $VARNAME_OVER + mask: + operator: filters.generate_mask + mask_field: + operator: filters.filter_cubes + constraint: + operator: constraints.combine_constraints + variable_constraint: + operator: constraints.generate_var_constraint + varname: $VARNAME_OVER + cell_methods_constraint: + operator: constraints.generate_cell_methods_constraint + cell_methods: [] + varname: $VARNAME_OVER + condition: '>=' + value: $OVERLAY_MASK_VALUE + + contour_cube: + operator: filters.filter_cubes + constraint: + operator: constraints.combine_constraints + variable_constraint: + operator: constraints.generate_var_constraint + varname: $VARNAME_CONTOUR diff --git a/tests/operators/test_utils.py b/tests/operators/test_utils.py index d229d6281..086b4374d 100644 --- a/tests/operators/test_utils.py +++ b/tests/operators/test_utils.py @@ -238,6 +238,16 @@ def test_is_spatialdim_true(transect_source_cube): assert operator_utils.is_spatialdim(transect_source_cube) +def test_is_coordim_false(transect_source_cube): + """Check that is coorddim test returns false if cube does not contain specified spatial coordinate.""" + assert not operator_utils.is_coorddim(transect_source_cube, "time") + + +def test_is_coordim_true(transect_source_cube): + """Check that is coorddim test returns true if cube contains specified spatial coordinate.""" + assert operator_utils.is_coorddim(transect_source_cube, "latitude") + + def test_fully_equalise_attributes_remove_unique_attributes(): """Check unique attributes are removed.""" original = iris.cube.Cube( @@ -301,6 +311,38 @@ def test_fully_equalise_attributes_equalise_coords(): assert "shared_attribute" not in cube.coord("foo").attributes +def test_slice_over_maybe_False(): + """Check that a missing cube returns None.""" + assert operator_utils.slice_over_maybe(None, "time", 0) is None + + +def test_slice_over_deterministic(long_forecast): + """Check that a sliced cube is returned.""" + assert operator_utils.slice_over_maybe(long_forecast, "time", 0) == long_forecast[0] + assert ( + operator_utils.slice_over_maybe(long_forecast, "time", 10) == long_forecast[10] + ) + + +def test_slice_over_ensemble(long_forecast): + """Check that a slice cube is returned from multi-time ensemble inputs.""" + ensemble_cube_em01 = long_forecast.copy() + ensemble_cube_em02 = long_forecast.copy() + ensemble_cube_em02.coord("realization").points = 1 + ensemble_cube_4d = iris.cube.CubeList( + [ensemble_cube_em01, ensemble_cube_em02] + ).merge_cube() + #### NEED TO UPDATE/IMPROVE THIS TESTING ASSUMPTION - DEPENDS ON ORDERING OF CUBE DIMS ONLY #### + assert ( + operator_utils.slice_over_maybe(ensemble_cube_4d, "time", 0) + == ensemble_cube_4d[:, 0, :, :] + ) + assert ( + operator_utils.slice_over_maybe(ensemble_cube_4d, "realization", 1) + == ensemble_cube_4d[:, 1, :, :] + ) + + def test_is_time_aggregatable_False(long_forecast): """Check that a cube that is not time aggregatable returns False.""" assert not operator_utils.is_time_aggregatable(long_forecast)