diff --git a/pyproject.toml b/pyproject.toml index 4e169090c..4e4adf0ac 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 634506dd6..2a16bafa5 100644 --- a/src/CSET/operators/__init__.py +++ b/src/CSET/operators/__init__.py @@ -121,6 +121,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 2ae71543f..0236d87ba 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_lifetime": { + "cmap": "YlGnBu", + "levels": [ + 1, + 2, + 3, + 4, + 5, + 6, + 7, + 8, + 9, + 10, + 11, + 12 + ], + "ymax": 1.0, + "ymin": 0.0 + }, + "feature_init": { + "cmap": "Blues", + "levels": [ + 0.5, + 1 + ], + "ymax": 1.0, + "ymin": 0.0 + }, "fog_fraction_at_screen_level": { "cmap": "viridis", "max": 1, diff --git a/src/CSET/operators/collapse.py b/src/CSET/operators/collapse.py index dcd299484..cae6eacd5 100644 --- a/src/CSET/operators/collapse.py +++ b/src/CSET/operators/collapse.py @@ -75,18 +75,24 @@ 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: - logging.debug( - "Extracting common time points as multiple model inputs detected." - ) - for cube in cubes: - cube.coord("forecast_reference_time").bounds = None - cube.coord("forecast_period").bounds = None - cubes = cubes.extract_overlapping( - ["forecast_reference_time", "forecast_period"] + fcst_ref_time_check = all( + "forecast_reference_time" in cube.coords() for cube in cubes ) - if len(cubes) == 0: - raise ValueError("No overlapping times detected in input cubes.") + fcst_period_check = all("forecast_period" in cube.coords() for cube in cubes) + if fcst_ref_time_check and fcst_period_check: + logging.debug( + "Extracting common time points as multiple model inputs detected." + ) + for cube in cubes: + cube.coord("forecast_reference_time").bounds = None + cube.coord("forecast_period").bounds = None + cubes = cubes.extract_overlapping( + ["forecast_reference_time", "forecast_period"] + ) + if len(cubes) == 0: + raise ValueError("No overlapping times detected in input cubes.") collapsed_cubes = iris.cube.CubeList([]) with warnings.catch_warnings(): diff --git a/src/CSET/operators/feature.py b/src/CSET/operators/feature.py index 3eb594e4b..b7a0830a3 100755 --- a/src/CSET/operators/feature.py +++ b/src/CSET/operators/feature.py @@ -20,8 +20,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 + def track( cube: iris.cube.Cube, @@ -194,6 +197,160 @@ def track( return tracking_cubelist +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 of xy type, not latitude/longitude. + 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_cubes: iris.cube.CubeList + An iris CubeList containing "feature_size", "feature_effective_radius", "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 radius (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 horizontal coordinates of xy type, not latitude/longitude + for cube in cubes: + _check_xy_coords(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, + }, + } + logging.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) + logging.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 radius from feature size, using horizontal coordinate of input cube to estimate grid spacing + effective_radius_data, grid_spacing = _get_effective_radius_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_radius": { + "data": effective_radius_data, + "long_name": "feature_effective_radius", + "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_xy_coords(cube: iris.cube.Cube) -> None: """Check that the input cube has horizontal coordinates of xy type, not latitude/longitude. @@ -220,3 +377,200 @@ def _check_xy_coords(cube: iris.cube.Cube) -> None: "which is not of xy type. Please provide a cube with horizontal " "coordinates of xy type." ) + + +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. + + 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_radius_from_feature_size( + size_data: np.ndarray, cube_with_hzntl_coord: iris.cube.Cube +) -> np.ndarray: + """Convert feature size in grid points to effective radius 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_radii_data: np.ndarray + An array containing "feature_effective_radius" 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 radius is + calculated as the radius of a circle with the same area as the feature size in grid points. + + """ + # Guess coord representing horizontal grid (choose first available) + hzntl_coord = [ + coord + for coord in cube_with_hzntl_coord.coords() + if iris.util.guess_coord_axis(coord) in ["X", "Y"] + ][0] + logging.debug(f"Attempting to convert to effective radius 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): + logging.warning( + f"Horizontal coordinate {hzntl_coord} is not regular. " + "Effective radius calculation may be inaccurate." + ) + + grid_spacing = np.abs(np.mean(np.diff(hzntl_coord.points))) + effective_radii_data = np.sqrt(size_data * grid_spacing**2 / np.pi) + + return effective_radii_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 f667273eb..b8a5b0d5f 100644 --- a/src/CSET/operators/misc.py +++ b/src/CSET/operators/misc.py @@ -702,3 +702,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 7a71d498b..ec0278ff1 100644 --- a/src/CSET/operators/plot.py +++ b/src/CSET/operators/plot.py @@ -443,8 +443,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 @@ -515,6 +515,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) @@ -1506,7 +1509,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) @@ -1534,6 +1540,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) logging.debug( @@ -1543,6 +1565,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() @@ -1584,6 +1610,8 @@ 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) ax.set_xlim(vmin, vmax) ax.tick_params(axis="both", labelsize=12) 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 000000000..4f0f46153 --- /dev/null +++ b/src/CSET/recipes/example_recipes/example_feature_cell_stats.yaml @@ -0,0 +1,31 @@ +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", or "feature_size" + - operator: filters.filter_multiple_cubes + constraint: + operator: constraints.generate_var_constraint + varname: feature_size + + # 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 54e144f69..1fe1c8179 100644 --- a/tests/operators/test_feature.py +++ b/tests/operators/test_feature.py @@ -29,9 +29,9 @@ 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) @@ -118,6 +118,48 @@ def test_save_data(feature_cube, tmp_working_dir) -> None: save_data=True, ) # Check expected lifetime field is created in output directory + output_directory = f"{tmp_path}/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): + """ + 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 + ) + + # 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_radius_data = np.squeeze( + cubelist.extract_cube("feature_effective_radius").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_radius_data = np.sqrt(expected_size_data * grid_spacing**2 / np.pi) + + 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_equal(effective_radius_data, expected_radius_data) output_directory = tmp_working_dir / "tracking_data" expected_file = output_directory / "lifetime_20100101_0000.field" assert expected_file.is_file() diff --git a/tests/operators/test_misc.py b/tests/operators/test_misc.py index 72fc8d1d5..e5e230f4b 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,)