diff --git a/pyproject.toml b/pyproject.toml index b1a6a83f22..b262819a89 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -34,6 +34,7 @@ dependencies = [ "scikit-image", "scores", "dask", + "simple-track", "xarray", "simple-track", ] diff --git a/src/CSET/operators/__init__.py b/src/CSET/operators/__init__.py index 01ccb8e575..741f1521c0 100644 --- a/src/CSET/operators/__init__.py +++ b/src/CSET/operators/__init__.py @@ -123,6 +123,7 @@ def get_operator(name: str): operator = CSET.operators for section in name_sections: operator = getattr(operator, section) + if callable(operator): return operator else: diff --git a/src/CSET/operators/_colorbar_definition.json b/src/CSET/operators/_colorbar_definition.json index 498d73f4b6..ae1b4ae7ad 100644 --- a/src/CSET/operators/_colorbar_definition.json +++ b/src/CSET/operators/_colorbar_definition.json @@ -199,6 +199,50 @@ "max": 1.1, "min": 0.5 }, + "feature_id": { + "cmap": "viridis", + "levels": [ + 1, + 50, + 100, + 150, + 200, + 250, + 300, + 350, + 400 + ], + "ymax": 1.0, + "ymin": 0.0 + }, + "feature_init": { + "cmap": "Blues", + "levels": [ + 0.5, + 1 + ], + "ymax": 1.0, + "ymin": 0.0 + }, + "feature_lifetime": { + "cmap": "YlGnBu", + "levels": [ + 1, + 2, + 3, + 4, + 5, + 6, + 7, + 8, + 9, + 10, + 11, + 12 + ], + "ymax": 1.0, + "ymin": 0.0 + }, "fog_fraction_at_screen_level": { "cmap": "viridis", "max": 1, diff --git a/src/CSET/operators/_colormaps.py b/src/CSET/operators/_colormaps.py index 9d37fc5db1..871c32c9cf 100644 --- a/src/CSET/operators/_colormaps.py +++ b/src/CSET/operators/_colormaps.py @@ -723,7 +723,11 @@ def custom_colormap_feature_tracking(cube: iris.cube.Cube): norm = mcolors.BoundaryNorm(levels, cmap.N) logger.info("change colormap for feature init variable colorbar.") - # Set all non-feature data to white. - cmap = cmap.with_extremes(under="white") + else: + cmap, levels, norm = None, None, None + + # Set all non-feature data to white + if cmap: + cmap = cmap.with_extremes(under="white") return cmap, levels, norm diff --git a/src/CSET/operators/collapse.py b/src/CSET/operators/collapse.py index 44e4b48af5..bea4e82dff 100644 --- a/src/CSET/operators/collapse.py +++ b/src/CSET/operators/collapse.py @@ -77,6 +77,7 @@ def collapse( raise ValueError("Must specify additional_percent") # Retain only common time points between different models if multiple model inputs. + # Do this only if "forecast_reference_time" and "forecast_period" are present in the cubes. if isinstance(cubes, iris.cube.CubeList) and len(cubes) > 1: logger.debug("Extracting common time points as multiple model inputs detected.") for cube in cubes: diff --git a/src/CSET/operators/feature.py b/src/CSET/operators/feature.py index cd34a0fcc2..394b173a64 100644 --- a/src/CSET/operators/feature.py +++ b/src/CSET/operators/feature.py @@ -21,8 +21,11 @@ import iris.cube import iris.util import numpy as np +from simpletrack.frame import Timeline from simpletrack.track import Tracker +from CSET._common import iter_maybe + logger = logging.getLogger(__name__) @@ -42,7 +45,6 @@ def track( ---------- cube: iris.cube.Cube The cube to identify features in. The cube must be 3D and contain a time coordinate - and horizontal coordinates of xy type (not latitude/longitude). threshold: float The threshold value for feature detection. under_threshold: bool, optional @@ -118,9 +120,6 @@ def track( >>> plt.show() """ - # Check that the input cube has horizontal coordinates of xy type, not latitude/longitude - _check_xy_coords(cube) - # Setup config tracker_config = { "FEATURE": { @@ -199,31 +198,430 @@ def track( return tracking_cubelist -def _check_xy_coords(cube: iris.cube.Cube) -> None: - """Check that the input cube has horizontal coordinates of xy type, not latitude/longitude. +def cell_stats( + cubes: iris.cube.Cube | iris.cube.CubeList, + threshold: float | list[float], + under_threshold: bool = False, + min_size: int = 4, + save_data: bool = False, +): + """Identify features in each timestep and output statistics. + + Parameters + ---------- + cubes: iris.cube.Cube | iris.cube.CubeList + An iris cube (single model) or cubelist (multiple models) containing 2D data to be + analysed. Cube must have horizontal coordinates on a regular grid. + The cube must also have a time coordinate, which is used to identify features in + each timestep. + threshold: float | list[float] + The threshold value(s) for feature detection. If a list is provided, each value + is used to identify features in the corresponding cube in the cubelist. Therefore, + the list should match the number of models. + under_threshold: bool, optional + If set to True, features are identified where the data is below the threshold. + If set to False, features are identified where the data is above the threshold. + Default is False. + min_size: int, optional + The minimum number of contiguous grid points required for a feature to be tracked. + Default is 4. + save_data: bool, optional + If set to True, all tracking data is saved to disk for further analysis (including csv + and txt files containing feature properties that are not returned in output cubes). + Default is False. + + Returns + ------- + cell_stats_cubelist: iris.cube.CubeList + An iris CubeList containing "feature_size", "feature_effective_diameter", "feature_mean", + and "feature_max" cubes. + + Notes + ----- + This operator uses the Simple-Track package with tracking disabled to identify features + in each timestep and compile cell statistics. Outputs cubes containing feature size (number + of grid points), effective diameter (in km), mean value within features, and maximum + value within features. + + Links + ---------- + .. https://github.com/ParaChute-UK/simple-track + + Examples + -------- + >>> cell_stats_cubes = feature.cell_stats(threshold=2) + >>> feature_size_cube = cell_stats_cubes.extract_cube("feature_size") + >>> plt.hist(feature_size_cube[-1]) + >>> plt.show() + + """ + # Check inputs + cubes = iter_maybe(cubes) + + # Require inputs to have a uniform grid + for cube in cubes: + _check_uniform_grid(cube) + + # Setup containing cube list + cell_stats_cubelist = iris.cube.CubeList() + + # If threshold is a list, check that it matches the number of cubes + if isinstance(threshold, list): + if len(threshold) != len(cubes): + raise ValueError( + f"Length of threshold list ({len(threshold)}) does not match " + f"number of cubes ({len(cubes)})." + ) + # else, make it iterable by repeating the same value for each cube + else: + threshold = [threshold] * len(cubes) + + # Run tracking on all input data + for cube, thresh in zip(cubes, threshold, strict=True): + model_name = cube.attributes.get("model_name", None) + # Setup config + tracker_config = { + "FEATURE": { + "threshold": thresh, + "under_threshold": under_threshold, + "min_size": min_size, + }, + "OUTPUT": { + "save_data": save_data, + "experiment_name": "feature_tracking", + "path": f"{os.getcwd()}/{model_name}/cell-stats_data", + "skip_tracking": True, + }, + } + logger.debug(f"Tracker config: {tracker_config}") + + # Get cube data into a dict to pass to Tracker + times = cube.coord("time").points + time_units = cube.coord("time").units + times_dt = [time_units.num2pydate(t) for t in times] + cube_dict = { + time: cube_slice.data + for time, cube_slice in zip(times_dt, cube.slices_over("time"), strict=True) + } + + # Run tracking, returning Timeline object + timeline = Tracker(tracker_config).run(cube_dict) + logger.debug(f"Tracking completed for {model_name}") + + # Get feature data from each frame of data + size_data, mean_data, max_data = _get_cell_stats_arrays_from_timeline( + timeline=timeline, expected_frame_times=times_dt + ) + + # Get effective diameter from feature size, using horizontal coordinate of input cube to estimate grid spacing + effective_diameter_data, grid_spacing = ( + _get_effective_diameter_from_feature_size( + size_data=size_data, cube_with_hzntl_coord=cube + ) + ) + + # Add grid_spacing as an attribute to the template_cube, so it is copied to + # output cubes in following function + cube.attributes["grid_spacing"] = grid_spacing + + # Set output cube properties + cube_properties = { + "feature_size": { + "data": size_data, + "long_name": "feature_size", + "units": 1, + }, + "feature_mean": { + "data": mean_data, + "long_name": "feature_mean", + "units": 1, + }, + "feature_max": {"data": max_data, "long_name": "feature_max", "units": 1}, + "feature_effective_diameter": { + "data": effective_diameter_data, + "long_name": "feature_effective_diameter", + "units": "km", + }, + } + + # Create cubes, add to existing cubelist + cell_stats_cubelist.extend( + _add_cell_stats_data_to_cubes( + data_and_metadata_dict=cube_properties, template_cube=cube + ) + ) + + return cell_stats_cubelist + + +def _check_uniform_grid(cube: iris.cube.Cube) -> bool: + """Check that the input cube has approximately uniform horizontal grid spacing. + + Prints warning if cube does not have uniform grid. Parameters ---------- cube: iris.cube.Cube An iris cube containing horizontal coordinates. + Returns + ------- + bool + True if the cube has a uniform grid, False otherwise. + Raises ------ ValueError - If the input cube has horizontal coordinates of latitude/longitude type. + If the input cube does not have a uniform grid. """ hzntl_coords = [ coord for coord in cube.coords() if iris.util.guess_coord_axis(coord) in ["X", "Y"] ] - invalid_coord_names = ["latitude", "longitude", "grid_latitude", "grid_longitude"] + for coord in hzntl_coords: - if coord.name() in invalid_coord_names and isinstance( - coord, iris.coords.DimCoord - ): - raise ValueError( - f"Input cube has horizontal coordinate {coord.name()} ({coord.units}), " - "which is a DimCoord not of xy type. Please provide a cube with horizontal " - "coordinates of xy type." + if not iris.util.is_regular(coord): + warning_msg = ( + f"Horizontal coordinate {coord} is not regular. " + "Feature statistics calculation may be inaccurate." ) + logger.warning(warning_msg) + print(warning_msg) + return False + return True + + +def _get_cell_stats_arrays_from_timeline( + timeline: Timeline, expected_frame_times: list +) -> list[np.ndarray]: + """Extract cell statistics data from a Simple-Track Timeline object. + + Parameters + ---------- + timeline: Timeline + A Simple-Track Timeline object containing tracked features. + + expected_frame_times: list + A list of expected frame times to extract data for. + + Returns + ------- + size_data: np.ndarray + A numpy array containing the size of each feature in grid points. + mean_data: np.ndarray + A numpy array containing the mean value of each feature. + max_data: np.ndarray + A numpy array containing the maximum value of each feature. + """ + size_data, mean_data, max_data = [], [], [] + number_of_features = [] + for time in expected_frame_times: + frame = timeline.get_frame(time) + features = frame.features + size_data.append([feature.get_size() for feature in features.values()]) + mean_data.append([feature.mean for feature in features.values()]) + max_data.append([feature.max for feature in features.values()]) + number_of_features.append(len(features)) + + # Pad data with NaNs to create arrays of consistent shape (max number of features across + # all timesteps) + arr_size = max(number_of_features) + + # Size data is integer, but we need to pad with NaNs (which is a float), so fill + # with invalid value first + size_data = np.array( + [ + np.pad(sizes, (0, arr_size - len(sizes)), constant_values=-100) + for sizes in size_data + ], + dtype=float, + ) + size_data[size_data == -100] = np.nan + + # Mean and max data are already float, so can be padded with NaNs directly. + mean_data = np.array( + [ + np.pad(means, (0, arr_size - len(means)), constant_values=np.nan) + for means in mean_data + ] + ) + max_data = np.array( + [ + np.pad(maxs, (0, arr_size - len(maxs)), constant_values=np.nan) + for maxs in max_data + ] + ) + + return size_data, mean_data, max_data + + +def _get_effective_diameter_from_feature_size( + size_data: np.ndarray, cube_with_hzntl_coord: iris.cube.Cube +) -> np.ndarray: + """Convert feature size in grid points to effective diameter in km. + + Parameters + ---------- + size_data: np.ndarray + An array containing "feature_size" data, in units of grid points. + cube_with_hzntl_coord: iris.cube.Cube + An iris cube containing a horizontal coordinate, which is used to + estimate the grid spacing for the effective radius calculation. + + Returns + ------- + effective_diameters_data: np.ndarray + An array containing "feature_effective_diameter" data, in units of km. + + grid_spacing: float + The estimated grid spacing in m, calculated from the horizontal coordinate of the input cube. + + Notes + ----- + This function assumes that the input cube has a horizontal coordinate system that is regular and + that the grid spacing can be estimated from the horizontal coordinates. The effective diameter is + calculated as the diameter of a circle with the same area as the feature size in grid points. + + """ + # Guess coord representing horizontal grid (choose first available) + hzntl_coord = next( + iter( + [ + coord + for coord in cube_with_hzntl_coord.coords() + if iris.util.guess_coord_axis(coord) in ["X", "Y"] + ] + ) + ) + + logger.debug(f"Attempting to convert to effective diameter using {hzntl_coord}") + + # Check coordinate is regular, but only warn if not, this is a naive estimate + # and will be inaccurate for irregular grids + if not iris.util.is_regular(hzntl_coord): + logger.warning( + f"Horizontal coordinate {hzntl_coord} is not regular. " + "Effective diameter calculation may be inaccurate." + ) + + # Get grid spacing in native coord units (degrees, m, km etc) + grid_spacing = iris.util.regular_step(hzntl_coord) + + if hzntl_coord.units == "m": + grid_spacing = grid_spacing / 1000 # Convert to km + + # If grid spacing is in degrees, convert to km using approximate conversion factor + if hzntl_coord.units == "degrees": + # Get latitude for better conversion to km + lat_coords = [ + coord + for coord in cube_with_hzntl_coord.coords() + if iris.util.guess_coord_axis(coord) in ["Y"] + ] + for coord in lat_coords: + if coord.units == "degrees": + lat_coord_for_conversion = coord + else: + lat_coord_for_conversion = None + + # Calculate conversion factor using latitude correction if available, + # otherwise use naive 111 km per degree conversion + if lat_coord_for_conversion is not None: + mean_latitude = np.mean(lat_coord_for_conversion.points) + else: + logger.warning( + "No latitude coordinate found for conversion to km. " + "Using naive conversion factor of 111 km per degree." + ) + mean_latitude = 0 + grid_spacing = grid_spacing * 111 * np.cos(np.radians(mean_latitude)) + + effective_diameters_data = np.sqrt(size_data * grid_spacing**2 / np.pi) * 2 + return effective_diameters_data, grid_spacing + + +def _add_cell_stats_data_to_cubes( + data_and_metadata_dict: dict, template_cube: iris.cube.Cube +) -> iris.cube.CubeList: + """Add data to cubes, using template cube for metadata. + + Parameters + ---------- + data_and_metadata_dict: dict + A dictionary containing data and metadata for each cube to be created. + The keys are the long names of the cubes, and the values are dictionaries + containing the data and units for each cube. + + template_cube: iris.cube.Cube + An iris cube to use as a template for the new cubes. The new cubes will + have the same attributes as the template cube. + + Returns + ------- + cubelist: iris.cube.CubeList + A list of iris cubes containing the added data. + + """ + cubelist = iris.cube.CubeList() + + # Construct coordinates for new cubes + time_coord = template_cube.coord("time").copy() + # To construct feature coordinate, look at the size of dimension 1 for each data + arr_size = max( + [data_and_metadata_dict[cb]["data"].shape[1] for cb in data_and_metadata_dict] + ) + feature_coord = iris.coords.DimCoord( + np.arange(arr_size), + long_name="feature_number", + var_name="feature_number", + units="1", + ) + coords = [time_coord, feature_coord] + coords_and_dims = [(coord, i) for i, coord in enumerate(coords)] + + # Get list of coords to copy from input cube to output cubes + copyable_coord_names = [ + "realization", + "hour", + "forecast_period", + "forecast_reference_time", + "model_name", + "cset_comparison_base", + ] + input_cube_coord_names = [] + for coord in template_cube.coords(): + input_cube_coord_names.append(coord.standard_name) + input_cube_coord_names.append(coord.long_name) + + coords_to_copy = [ + coord_name + for coord_name in copyable_coord_names + if coord_name in input_cube_coord_names + ] + + # Populate cubelist + for cb_props in data_and_metadata_dict.values(): + cell_stats_cube = iris.cube.Cube( + data=cb_props["data"], + long_name=cb_props["long_name"], + units=cb_props["units"], + dim_coords_and_dims=coords_and_dims, + ) + # Add other metadata from input cube + for coord_name in coords_to_copy: + coord = template_cube.coord(coord_name).copy() + # Check if this coord represents a dimension of data + dims = template_cube.coord_dims(coord) + if len(dims) > 0: + cell_stats_cube.add_aux_coord(coord, dims) + else: + cell_stats_cube.add_aux_coord(coord) + + # Copy over attributes + cell_stats_cube.attributes = template_cube.attributes + + # Add to cubelist + cubelist.append(cell_stats_cube) + + return cubelist diff --git a/src/CSET/operators/misc.py b/src/CSET/operators/misc.py index aae2fcd586..535f5311fe 100644 --- a/src/CSET/operators/misc.py +++ b/src/CSET/operators/misc.py @@ -703,3 +703,65 @@ def differentiate( return new_cubelist[0] else: return new_cubelist + + +def flatten( + cubes: iris.cube.Cube | iris.cube.CubeList, remove_nans: bool = False +) -> iris.cube.Cube | iris.cube.CubeList: + """Flatten a cube or cubelist along all dimensions. + + Flattened cube contains a single dimension coordinate named "flattened_index". + + Parameters + ---------- + cubes : iris.cube.Cube or iris.cube.CubeList + The input Cube or CubeList to flatten. + remove_nans : bool, optional + If True, remove NaN values from the flattened data. Default is True. + + Returns + ------- + iris.cube.Cube or iris.cube.CubeList + The flattened cube or cubelist. + """ + if isinstance(cubes, iris.cube.Cube): + cubes = iris.cube.CubeList([cubes]) + + if not isinstance(cubes, iris.cube.CubeList): + raise TypeError("Input must be an iris.cube.Cube or iris.cube.CubeList.") + + flattened_cubes = iris.cube.CubeList() + for cube in cubes: + # Remove NaN if required + cube_data = cube.data + if remove_nans: + cube_data = cube_data[~np.isnan(cube_data)] + + # Flatten the data + flattened_data = cube_data.flatten() + + # Create a new cube with the flattened data and the remaining coordinates + flat_coord = iris.coords.DimCoord( + np.arange(flattened_data.size), long_name="flattened_index", units="1" + ) + flattened_cube = iris.cube.Cube( + flattened_data, + standard_name=cube.standard_name, + long_name=cube.long_name, + var_name=cube.var_name, + units=cube.units, + attributes=cube.attributes, + dim_coords_and_dims=[(flat_coord, 0)], + ) + + # Add single time point as a scalar coord if it exists + if cube.coords("time"): + time_coord = cube.coord("time") + flattened_cube.add_aux_coord(time_coord[0]) + + flattened_cubes.append(flattened_cube) + + if len(flattened_cubes) == 1: + return flattened_cubes[0] + else: + return flattened_cubes diff --git a/src/CSET/operators/plot.py b/src/CSET/operators/plot.py index 372d2e7284..e8ba0ebd42 100644 --- a/src/CSET/operators/plot.py +++ b/src/CSET/operators/plot.py @@ -475,8 +475,8 @@ def _set_axis_range(cubes): break if levels is None: - vmin = min(cb.data.min() for cb in cubes) - vmax = max(cb.data.max() for cb in cubes) + vmin = min(np.nanmin(cb.data) for cb in cubes) + vmax = max(np.nanmax(cb.data) for cb in cubes) return vmin, vmax @@ -547,6 +547,9 @@ def _plot_and_save_spatial_plot( # Specify the color bar cmap, levels, norm = colorbar_map_levels(cube) + if "feature" in cube.long_name: + cmap.set_under("white") + # If overplotting, set required colorbars if overlay_cube: over_cmap, over_levels, over_norm = colorbar_map_levels(overlay_cube) @@ -1540,7 +1543,10 @@ def _plot_and_save_histogram_series( # Set default that histograms will produce probability density function # at each bin (integral over range sums to 1). - density = True + if "feature" in cubes[0].long_name: + density = False + else: + density = True for cube in iter_maybe(cubes): # Easier to check title (where var name originates) @@ -1568,6 +1574,22 @@ def _plot_and_save_histogram_series( ax.set_xscale("log") elif "lightning" in title: bins = [0, 1, 2, 3, 4, 5] + elif "feature_size" in cube.long_name: + bins = np.linspace(0, 500, 51) + elif "feature_effective_radius" in cube.long_name: + # TODO: use grid_spacing attribute in cubes to find min bin size + # for effective radius, rather than being hard coded + # Modified from RMED toolbox + bins = 10 ** (np.arange(0, 5.28, 0.12)) + bins = np.insert(bins, 0, 0) + vmin = bins[1] + vmax = bins[-1] + elif "feature_mean" in cube.long_name or "feature_max" in cube.long_name: + # From RMED toolbox + bins = 10 ** (np.arange(-1, 2.7, 0.12)) + bins = np.insert(bins, 0, 0) + vmin = bins[1] + vmax = bins[-1] else: bins = np.linspace(vmin, vmax, 51) logger.debug( @@ -1577,6 +1599,10 @@ def _plot_and_save_histogram_series( np.max(bins), ) + if "feature" in cube.long_name: + ax.set_yscale("log") + ax.set_xscale("log") + # Reshape cube data into a single array to allow for a single histogram. # Otherwise we plot xdim histograms stacked. cube_data_1d = (cube.data).flatten() @@ -1618,6 +1644,9 @@ def _plot_and_save_histogram_series( ax.set_ylabel( f"Contribution to mean ({iter_maybe(cubes)[0].units})", fontsize=14 ) + if "feature" in cubes[0].long_name: + ax.set_ylabel("Frequency", fontsize=14) + try: ax.set_xlim(vmin, vmax) except ValueError: diff --git a/src/CSET/recipes/example_recipes/example_feature_cell_stats.yaml b/src/CSET/recipes/example_recipes/example_feature_cell_stats.yaml new file mode 100755 index 0000000000..1d7908ba85 --- /dev/null +++ b/src/CSET/recipes/example_recipes/example_feature_cell_stats.yaml @@ -0,0 +1,34 @@ +category: Quick Look +title: Example running cell_stats and plotting histograms +description: | + Uses the feature.cell_stats operator to calculate cell size, mean cell values and max cell values for + identified features. Saves cell stats data to output and plots histograms + +steps: + - operator: read.read_cubes + file_paths: $INPUT_PATHS + model_names: $MODEL_NAME + + - operator: filters.filter_multiple_cubes + constraint: + operator: constraints.generate_var_constraint + varname: $VARNAME + + - operator: feature.cell_stats + threshold: 3 + save_data: True # save raw tracking data for further analysis + + # Filter tracking cubelist to one of "feature_mean", "feature_max", "feature_size", or "feature_effective_diameter" + - operator: filters.filter_multiple_cubes + constraint: + operator: constraints.generate_var_constraint + varname: feature_size + # varname: feature_mean + # varname: feature_max + # varname: feature_effective_diameter + + # Flatten data across case study period + - operator: misc.flatten + remove_nans: True + + - operator: plot.plot_histogram_series diff --git a/tests/operators/test_feature.py b/tests/operators/test_feature.py index 54e144f693..a125038fdf 100644 --- a/tests/operators/test_feature.py +++ b/tests/operators/test_feature.py @@ -14,6 +14,7 @@ """Tests for feature operators.""" import datetime as dt +import os import cf_units import iris @@ -21,17 +22,19 @@ import iris.cube import numpy as np import pytest +from simpletrack import Tracker +from simpletrack.frame import Timeline from CSET.operators import feature -@pytest.fixture +@pytest.fixture(scope="session") def feature_cube() -> iris.cube.Cube: """Set up three timesteps of data and place into cube.""" data_arr = np.zeros((3, 10, 10)) - data_arr[0, 2:6, 2:6] = 1 - data_arr[1, 3:7, 3:7] = 1 - data_arr[2, 4:8, 4:8] = 1 + data_arr[0, 2:6, 2:6] = 5 + data_arr[1, 3:7, 3:7] = 10 + data_arr[2, 4:8, 4:8] = 20 time_units = cf_units.Unit("days since 2000-01-01 00:00:00", calendar="gregorian") time_start = dt.datetime(2010, 1, 1, 0, 0, 0) @@ -40,6 +43,9 @@ def feature_cube() -> iris.cube.Cube: time_coord = iris.coords.DimCoord( points=time_points, standard_name="time", units=time_units ) + fcst_ref_coord = iris.coords.AuxCoord( + points=time_start, standard_name="forecast_reference_time", units=time_units + ) coord_system = iris.coord_systems.TransverseMercator( latitude_of_projection_origin=55, longitude_of_central_meridian=0 @@ -65,14 +71,51 @@ def feature_cube() -> iris.cube.Cube: coords = (time_coord, proj_y_coord, proj_x_coord) dim_coords_and_dims = [(coord, dim) for dim, coord in enumerate(coords)] + attributes = {"title": "Precipitation test"} cube = iris.cube.Cube( data=data_arr, dim_coords_and_dims=dim_coords_and_dims, long_name="Precipitation test", + attributes=attributes, + aux_coords_and_dims=[(fcst_ref_coord, None)], ) return cube +@pytest.fixture(scope="session") +def cell_stats_timeline(feature_cube) -> Timeline: + """Run feature tracking on feature_cube and return Timeline object.""" + # Setup config + tracker_config = { + "FEATURE": { + "threshold": 1, + "under_threshold": False, + "min_size": 4, + }, + "OUTPUT": { + "save_data": False, + "experiment_name": "feature_tracking", + "path": f"{os.getcwd()}/model/cell-stats_data", + "skip_tracking": True, + }, + } + + # Get cube data into a dict to pass to Tracker + times = feature_cube.coord("time").points + time_units = feature_cube.coord("time").units + times_dt = [time_units.num2pydate(t) for t in times] + cube_dict = { + time: cube_slice.data + for time, cube_slice in zip( + times_dt, feature_cube.slices_over("time"), strict=True + ) + } + + # Run tracking, returning Timeline object + timeline = Tracker(tracker_config).run(cube_dict) + return timeline + + def test_tracking_valid(feature_cube) -> None: """ Test feature tracking returns same cube shape as input cube. @@ -118,7 +161,53 @@ def test_save_data(feature_cube, tmp_working_dir) -> None: save_data=True, ) # Check expected lifetime field is created in output directory - output_directory = tmp_working_dir / "tracking_data" + output_directory = f"{tmp_working_dir}/tracking_data" + expected_file = f"{output_directory}/lifetime_20100101_0000.field" + assert os.path.isfile(expected_file) + + # Check expected csv file is created in output directory + expected_file = f"{output_directory}/frame_20100101_0000.csv" + assert os.path.isfile(expected_file) + + +def test_cell_stats_operator(feature_cube, tmp_working_dir): + """ + Test the cell_stats operator returns expected size, mean, and max values. + + The expected values are based on the feature_cube data and the threshold of 0.5. + """ + threshold = 0.5 + min_size = 1 + cubelist = feature.cell_stats( + cubes=feature_cube, threshold=threshold, min_size=min_size, save_data=True + ) + + # Extract data from cubelist, squeeze since there is only one feature per timestep + # in this test case + size_data = np.squeeze(cubelist.extract_cube("feature_size").data) + mean_data = np.squeeze(cubelist.extract_cube("feature_mean").data) + max_data = np.squeeze(cubelist.extract_cube("feature_max").data) + effective_diameter_data = np.squeeze( + cubelist.extract_cube("feature_effective_diameter").data + ) + + # Expected values based on the feature_cube data + expected_size_data = np.array([16, 16, 16]) # Each feature is a 4x4 square + expected_mean_data = np.array([5.0, 10.0, 20.0]) # Mean values of each feature + expected_max_data = np.array([5.0, 10.0, 20.0]) # Max values of each feature + + grid_spacing = 10 # Assuming grid spacing is 10 meters from test setup + expected_diameter_data = np.sqrt(expected_size_data * grid_spacing**2 / np.pi) * 2 + # Convert to km + expected_diameter_data = expected_diameter_data / 1000 + + np.testing.assert_array_equal(size_data, expected_size_data) + np.testing.assert_array_equal(mean_data, expected_mean_data) + np.testing.assert_array_equal(max_data, expected_max_data) + np.testing.assert_array_almost_equal( + effective_diameter_data, expected_diameter_data + ) + output_directory = tmp_working_dir / "None/cell-stats_data" expected_file = output_directory / "lifetime_20100101_0000.field" assert expected_file.is_file() @@ -127,21 +216,28 @@ def test_save_data(feature_cube, tmp_working_dir) -> None: assert expected_file.is_file() -def test_check_xy_coords_valid(feature_cube) -> None: - """Test that _check_xy_coords does not raise an error for valid xy coordinates.""" - try: - feature._check_xy_coords(feature_cube) - except ValueError: - pytest.fail("Unexpected ValueError raised for valid xy coordinates.") +def test_cell_stats_invalid_threshold_list(feature_cube): + """Test that cell_stats raises a ValueError if passed a list of wrong size.""" + invalid_thresholds = [0.5, 1.0] + with pytest.raises(ValueError): + feature.cell_stats(cubes=feature_cube, threshold=invalid_thresholds, min_size=1) -def test_check_xy_coords_invalid() -> None: - """Test that _check_xy_coords raises a ValueError for invalid latitude/longitude coordinates.""" - # Create a cube with latitude and longitude coordinates +def test_check_uniform_grid(feature_cube) -> None: + """Test that _check_uniform_grid does not raise an error for valid uniform grid.""" + result = feature._check_uniform_grid(feature_cube) + assert result is True + + +def test_check_uniform_grid_invalid() -> None: + """Test that _check_uniform_grid raises a ValueError for invalid uniform grid.""" + # Create a cube with non-uniform grid data_arr = np.zeros((10, 10)) + lat_points = np.linspace(-90, 90, 10) + lat_points[-1] = 95 lat_coord = iris.coords.DimCoord( - points=np.linspace(-90, 90, 10), + points=lat_points, standard_name="latitude", var_name="latitude", units="degrees", @@ -161,5 +257,179 @@ def test_check_xy_coords_invalid() -> None: long_name="Precipitation test", ) - with pytest.raises(ValueError): - feature._check_xy_coords(cube) + result = feature._check_uniform_grid(cube) + assert result is False + + +def test_get_cell_stats_arrays_from_timeline(cell_stats_timeline): + """Test that _get_cell_stats_arrays_from_timeline returns expected arrays.""" + expected_frame_times = [ + dt.datetime(2010, 1, 1, 0, 0, 0), + dt.datetime(2010, 1, 1, 0, 5, 0), + dt.datetime(2010, 1, 1, 0, 10, 0), + ] + + size_array, mean_array, max_array = feature._get_cell_stats_arrays_from_timeline( + cell_stats_timeline, expected_frame_times + ) + + # Expected values based on the feature_cube data and the threshold of 0.5 + expected_size_array = np.array([[16], [16], [16]]) # Each feature is a 4x4 square + expected_mean_array = np.array([[5], [10], [20]]) # Mean values of each feature + expected_max_array = np.array([[5], [10], [20]]) # Max values of each feature + + np.testing.assert_array_equal(size_array, expected_size_array) + np.testing.assert_array_equal(mean_array, expected_mean_array) + np.testing.assert_array_equal(max_array, expected_max_array) + + +def test_get_effective_diameter_from_feature_size(feature_cube): + """Test that _get_effective_diameter_from_feature_size returns expected values.""" + # Use the same size data from above test + size_data = np.array([[16], [16], [16]]) # Each feature is a 4x4 square + + effective_diameter_data, grid_spacing = ( + feature._get_effective_diameter_from_feature_size(size_data, feature_cube) + ) + + # Expected values based on the feature_cube data + grid_spacing = 10 # Assuming grid spacing is 10 meters from test setup + expected_diameter_data = np.sqrt(size_data * grid_spacing**2 / np.pi) * 2 + # Convert to km + expected_diameter_data = expected_diameter_data / 1000 + + np.testing.assert_array_almost_equal( + effective_diameter_data, expected_diameter_data + ) + assert grid_spacing == 10 + + +def test_get_effective_diameter_from_feature_size_km_input_cube(feature_cube): + """Test that _get_effective_diameter_from_feature_size returns expected values for km input cube.""" + # Use the same size data from above test + size_data = np.array([[16], [16], [16]]) # Each feature is a 4x4 square + + # Convert cube to km units + feature_cube_km = feature_cube.copy() + feature_cube_km.coord("projection_x_coordinate").convert_units("km") + feature_cube_km.coord("projection_y_coordinate").convert_units("km") + + effective_diameter_data, grid_spacing = ( + feature._get_effective_diameter_from_feature_size(size_data, feature_cube_km) + ) + + # Expected values based on the feature_cube data + grid_spacing = 0.01 # Assuming grid spacing is 10 meters (0.01 km) from test setup + expected_diameter_data = np.sqrt(size_data * grid_spacing**2 / np.pi) * 2 + + np.testing.assert_array_almost_equal( + effective_diameter_data, expected_diameter_data + ) + assert grid_spacing == 0.01 + + +def test_get_effective_diameter_from_feature_size_latlon_cube(): + """Test that _get_effective_diameter_from_feature_size returns expected values for lat/lon cube.""" + # Create a cube with lat/lon coordinates + data_arr = np.zeros((10, 10)) + lat_points = np.linspace(-1, 1, 10) + lon_points = np.linspace(-1, 1, 10) + + lat_coord = iris.coords.DimCoord( + points=lat_points, + standard_name="latitude", + var_name="latitude", + units="degrees", + ) + lon_coord = iris.coords.DimCoord( + points=lon_points, + standard_name="longitude", + var_name="longitude", + units="degrees", + ) + + coords = (lat_coord, lon_coord) + dim_coords_and_dims = [(coord, dim) for dim, coord in enumerate(coords)] + cube = iris.cube.Cube( + data=data_arr, + dim_coords_and_dims=dim_coords_and_dims, + long_name="Precipitation test", + ) + + # Use the same size data from above test + size_data = np.array([[2], [2], [2]]) # Each feature is a 4x4 square + + effective_diameter_data, grid_spacing = ( + feature._get_effective_diameter_from_feature_size(size_data, cube) + ) + + # Expected values based on the feature_cube data + # For lat/lon cube, grid spacing is calculated based on the distance between points + # Expected grid spacing is 2 degrees (222 km) / 9 intervals = 24.666 km + expected_grid_spacing = 24.666 + expected_diameter_data = np.sqrt(size_data * expected_grid_spacing**2 / np.pi) * 2 + + np.testing.assert_array_almost_equal( + effective_diameter_data, expected_diameter_data, decimal=3 + ) + np.testing.assert_almost_equal(grid_spacing, expected_grid_spacing, decimal=3) + + +def test_add_cell_stats_data_to_cubes(cell_stats_timeline, feature_cube): + """Test that _add_cell_stats_data_to_cubes adds expected data to cubelist.""" + expected_frame_times = [ + dt.datetime(2010, 1, 1, 0, 0, 0), + dt.datetime(2010, 1, 1, 0, 5, 0), + dt.datetime(2010, 1, 1, 0, 10, 0), + ] + + # Get feature data from each frame of data + size_data, mean_data, max_data = feature._get_cell_stats_arrays_from_timeline( + timeline=cell_stats_timeline, expected_frame_times=expected_frame_times + ) + + # Get effective diameter from feature size, using horizontal coordinate of input cube to estimate grid spacing + effective_diameter_data, __ = feature._get_effective_diameter_from_feature_size( + size_data=size_data, cube_with_hzntl_coord=feature_cube + ) + + cube_properties = { + "feature_size": { + "data": size_data, + "long_name": "feature_size", + "units": 1, + }, + "feature_mean": { + "data": mean_data, + "long_name": "feature_mean", + "units": 1, + }, + "feature_max": {"data": max_data, "long_name": "feature_max", "units": 1}, + "feature_effective_diameter": { + "data": effective_diameter_data, + "long_name": "feature_effective_diameter", + "units": "km", + }, + } + + cubelist = feature._add_cell_stats_data_to_cubes(cube_properties, feature_cube) + + # Test a cube of each name is produced + expected_cube_names = [ + "feature_size", + "feature_mean", + "feature_max", + "feature_effective_diameter", + ] + + for cube_name in expected_cube_names: + assert any(cube.long_name == cube_name for cube in cubelist) + + # Check each cube contains a forecast_reference_time coordinate, copied from + # feature_cube + for cube in cubelist: + assert cube.coords("forecast_reference_time") + + # test cube attributes copied to cell stats cubes + for cube in cubelist: + assert cube.attributes == feature_cube.attributes diff --git a/tests/operators/test_misc.py b/tests/operators/test_misc.py index 9b51e12ad5..5bbb1a3f6b 100644 --- a/tests/operators/test_misc.py +++ b/tests/operators/test_misc.py @@ -603,3 +603,55 @@ def test_not_remove_non_scalar_coord(): # Check it is still present cube_out = out[0] assert cube_out.coords("time") + + +def test_flatten_cube_no_nans(): + """Test misc.flatten without nans in single cube.""" + data_shape = (3, 4) + data = np.arange(12).reshape(data_shape) + coords = [ + iris.coords.DimCoord(np.arange(shape), long_name=f"test{shape}", units="1") + for shape in data_shape + ] + dim_coords_and_dims = [(coord, dim) for dim, coord in enumerate(coords)] + cube = iris.cube.Cube(data, dim_coords_and_dims=dim_coords_and_dims) + flattened_cube = misc.flatten(cube) + assert flattened_cube.shape == (12,) + assert np.allclose(flattened_cube.data, np.arange(12)) + + +def test_flatten_cubelist_no_nans(): + """Test misc.flatten without nans in CubeList.""" + data_shape = (3, 4) + data = np.arange(12).reshape(data_shape) + coords = [ + iris.coords.DimCoord(np.arange(shape), long_name=f"test{shape}", units="1") + for shape in data_shape + ] + dim_coords_and_dims = [(coord, dim) for dim, coord in enumerate(coords)] + cubes = iris.cube.CubeList( + [ + iris.cube.Cube(data, dim_coords_and_dims=dim_coords_and_dims) + for __ in range(3) + ] + ) + flattened_cubes = misc.flatten(cubes) + assert len(flattened_cubes) == 3 + for cube in flattened_cubes: + assert cube.shape == (12,) + assert np.allclose(cube.data, np.arange(12)) + + +def test_flatten_cube_nans_removed(): + """Test misc.flatten with nans removed in Cube.""" + data_shape = (3, 4) + data = np.arange(12, dtype=float).reshape(data_shape) + data[:, 0] = np.nan # 3 nans are inserted + coords = [ + iris.coords.DimCoord(np.arange(shape), long_name=f"test{shape}", units="1") + for shape in data_shape + ] + dim_coords_and_dims = [(coord, dim) for dim, coord in enumerate(coords)] + cube = iris.cube.Cube(data, dim_coords_and_dims=dim_coords_and_dims) + flattened_cube = misc.flatten(cube, remove_nans=True) + assert flattened_cube.shape == (9,)