diff --git a/src/CSET/operators/plot.py b/src/CSET/operators/plot.py index d95dcc9d7..23c6a5486 100644 --- a/src/CSET/operators/plot.py +++ b/src/CSET/operators/plot.py @@ -844,6 +844,60 @@ def _spatial_plot( _make_plot_html_page(complete_plot_index) +def _calculate_CFAD( + cube: iris.cube.Cube, vertical_coordinate: str, bin_edges: list[float] +) -> iris.cube.Cube: + """Calculate a Contour Frequency by Altitude Diagram (CFAD). + + Parameters + ---------- + cube: iris.cube.Cube + A cube of the data to be turned into a CFAD. It should be a minimum + of two dimensions with one being a user specified vertical coordinate. + vertical_coordinate: str + The vertical coordinate of the cube for the CFAD to be calculated over. + bin_edges: list[float] + The bin edges for the histogram. The bins need to be specified to + ensure consistency across the CFAD, otherwise it cannot be interpreted. + """ + # Setup empty array for containing the CFAD data. + CFAD_values = np.zeros( + (len(cube.coord(vertical_coordinate).points), len(bin_edges) - 1) + ) + + # Set iterator for CFAD values. + i = 0 + + # Calculate the CFAD as a histogram summing to one for each level. + for level_cube in cube.slices_over(vertical_coordinate): + # Note setting density to True does not produce the correct + # normalization for a CFAD, where each row must sum to one. + CFAD_values[i, :] = ( + np.histogram(level_cube.data.reshape(level_cube.data.size), bins=bin_edges)[ + 0 + ] + / level_cube.data.size + ) + i += 1 + # calculate central points for bins + bins = (np.array(bin_edges[:-1]) + np.array(bin_edges[1:])) / 2.0 + bin_bounds = np.array((bin_edges[:-1], bin_edges[1:])).T + # Now construct the coordinates for the cube. + vert_coord = cube.coord(vertical_coordinate) + bin_coord = iris.coords.DimCoord( + bins, bounds=bin_bounds, standard_name=cube.standard_name, units=cube.units + ) + # Now construct the cube that is to be output. + CFAD = iris.cube.Cube( + CFAD_values, + dim_coords_and_dims=[(vert_coord, 0), (bin_coord, 1)], + standard_name=cube.standard_name, + units="1", + ) + CFAD.attributes["type"] = "Contour Frequency by Altitude Diagram (CFAD)" + return CFAD + + #################### # Public functions # #################### diff --git a/tests/operators/test_plots.py b/tests/operators/test_plots.py index 908d6dfdd..c9a03d254 100644 --- a/tests/operators/test_plots.py +++ b/tests/operators/test_plots.py @@ -17,11 +17,18 @@ from pathlib import Path import iris.cube +import numpy as np import pytest from CSET.operators import collapse, plot, read +@pytest.fixture() +def xwind() -> iris.cube.Cube: + """Get regridded xwind to run tests on.""" + return iris.load_cube("tests/test_data/ageofair/aoa_in_rgd.nc", "x_wind") + + def test_check_single_cube(): """Conversion to a single cube, and rejection where not possible.""" cube = iris.cube.Cube([0.0]) @@ -318,3 +325,24 @@ def test_invalid_plotting_method_postage_stamp_spatial_plot(cube, tmp_working_di plot._plot_and_save_postage_stamp_spatial_plot( cube, "filename", "realization", "title", "invalid" ) + + +def test_calculate_CFAD(xwind): + """Test calculating a CFAD.""" + bins = np.array([-50, -40, -30, -20, -10, 0, 10, 20, 30, 40, 50]) + calculated_CFAD = np.zeros((len(xwind.coord("pressure").points), len(bins) - 1)) + j = 0 + for level_cube in xwind.slices_over("pressure"): + calculated_CFAD[j, :] = ( + np.histogram(level_cube.data.reshape(level_cube.data.size), bins=bins)[0] + / level_cube.data.size + ) + j += 1 + assert np.allclose( + plot._calculate_CFAD( + xwind, "pressure", [-50, -40, -30, -20, -10, 0, 10, 20, 30, 40, 50] + ).data, + calculated_CFAD, + rtol=1e-06, + atol=1e-02, + )