Coverage for src/CSET/operators/feature.py: 94%
134 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-27 13:04 +0000
1# © Crown copyright, Met Office (2022-2026) and CSET contributors.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14"""Operators for identifying and tracking features."""
16import logging
17import os
19import iris
20import iris.coords
21import iris.cube
22import iris.util
23import numpy as np
24from simpletrack.frame import Timeline
25from simpletrack.track import Tracker
27from CSET._common import iter_maybe
29logger = logging.getLogger(__name__)
32def track(
33 cube: iris.cube.Cube,
34 threshold: float,
35 under_threshold: bool = False,
36 min_size: int = 4,
37 retain_lifetime_on_split: bool = True,
38 tracking_nbhood: int = 5,
39 overlap_threshold: float = 0.3,
40 save_data: bool = False,
41):
42 """Track features between subsequent timesteps.
44 Parameters
45 ----------
46 cube: iris.cube.Cube
47 The cube to identify features in. The cube must be 3D and contain a time coordinate
48 and horizontal coordinates of xy type (not latitude/longitude).
49 threshold: float
50 The threshold value for feature detection.
51 under_threshold: bool, optional
52 If set to True, features are identified where the data is below the threshold.
53 If set to False, features are identified where the data is above the threshold.
54 Default is False.
55 min_size: int, optional
56 The minimum number of contiguous grid points required for a feature to be tracked.
57 Default is 4.
58 retain_lifetime_on_split: bool, optional
59 If set to True, the lifetime of a feature is retained when it splits into
60 multiple features. If set to False, the lifetime is reset when a feature splits.
61 Default is True.
62 tracking_nbhood: int, optional
63 The size of the neighbourhood used for tracking features between timesteps.
64 This dictates the maximum pixel radius from a feature centroid at which new features could
65 reasonably be spawned.
66 Default is 5.
67 overlap_threshold: float, optional
68 The minimum overlap required between features in consecutive timesteps for
69 them to be considered the same feature.
70 Default is 0.3.
71 save_data: bool, optional
72 If set to True, all tracking data is saved to disk for further analysis (including csv
73 and txt files containing feature properties that are not returned in output cubes).
74 Default is False.
76 Returns
77 -------
78 tracking_cubes: iris.cube.CubeList
79 A list of iris cubes containing tracking data, including feature ID, lifetime,
80 and locations of initiating features.
82 Notes
83 -----
84 This operator uses the Simple-Track package to track features between timesteps. Simple-Track is a
85 data-agnostic, threshold-based object tracking algorithm for 2D data. Features are tracked between
86 consecutive frames of data by projecting feature fields onto common timeframes and matching
87 between them based on the degree of overlap. Matched features retain the same identification
88 between all tracked fields, while new features are assigned a unique label.
89 Thus, Simple-Track compiles comprehensive information about feature merging, splitting, accretion,
90 initiation and dissipation.
92 Currently outputs three cubes containing the following data:
93 "feature_id":
94 A 2D field containing the unique label assigned to each feature, which is retained
95 if the feature is tracked across multiple timesteps. This cube can be used as a mask
96 to identify the location of the tracked feature throughout the evaluation period.
97 "feature_lifetime":
98 A 2D field containing the lifetime of each feature in terms of the number of
99 timesteps it has been tracked for. This cube can be used to distinguish between
100 mature and fresh features.
101 "feature_init":
102 A 2D binary field indicating the location of newly initiated features at each timestep.
103 These features are identified as having a lifetime of 1 AND have initiated sufficiently
104 far from other, existing features that they are not considered to have spawned from them.
106 Links
107 ----------
108 .. https://github.com/ParaChute-UK/simple-track
110 Examples
111 --------
112 >>> tracking_cubes = feature.track(threshold=2)
113 >>> lifetime_cube = tracking_cubes.extract_cube("feature_lifetime")
114 # Plot the final timestep of lifetime cube. This will show
115 # the lifetime of features that have been tracked for multiple previous
116 # timesteps, as well as new features that have just been initiated.
117 >>> iplt.pcolormesh(lifetime_cube[-1,:,:],cmap=mpl.cm.bwr)
118 >>> plt.gca().coastlines('10m')
119 >>> plt.clim(-5,5)
120 >>> plt.colorbar()
121 >>> plt.show()
123 """
124 # Setup config
125 tracker_config = {
126 "FEATURE": {
127 "threshold": threshold,
128 "under_threshold": under_threshold,
129 "min_size": min_size,
130 },
131 "TRACKING": {
132 "retain_lifetime_on_split": retain_lifetime_on_split,
133 "overlap_nbhood": tracking_nbhood,
134 "overlap_threshold": overlap_threshold,
135 },
136 "OUTPUT": {
137 "save_data": save_data,
138 "experiment_name": "feature_tracking",
139 "path": f"{os.getcwd()}/tracking_data",
140 },
141 }
142 logger.debug(f"Tracker config: {tracker_config}")
144 # Get cube data into a dict to pass to Tracker
145 times = cube.coord("time").points
146 time_units = cube.coord("time").units
147 times_dt = [time_units.num2pydate(t) for t in times]
148 cube_dict = {
149 time: cube_slice.data
150 for time, cube_slice in zip(times_dt, cube.slices_over("time"), strict=True)
151 }
153 # Run tracking, returning Timeline object
154 timeline = Tracker(tracker_config).run(cube_dict)
155 logger.debug("Tracking completed")
157 # Use input cube as template to make returned cube
158 # By iterating over all cube times, this will ensure all data is present
159 # If a Frame at the given time is not contained in the timeline, error is raised
160 output_type_and_methods = {
161 "lifetime": {
162 "getter": "lifetime_field",
163 "cube_name": "feature_lifetime",
164 },
165 "feature": {
166 "getter": "feature_field",
167 "cube_name": "feature_id",
168 },
169 "init": {
170 "getter": "get_init_field",
171 "cube_name": "feature_init",
172 },
173 }
175 tracking_cubelist = iris.cube.CubeList()
176 for output_type in output_type_and_methods:
177 tracking_data = []
178 for time in times_dt:
179 frame = timeline.get_frame(time)
180 getter = getattr(frame, output_type_and_methods[output_type]["getter"])
181 if callable(getter):
182 tracking_data.append(getter())
183 else:
184 tracking_data.append(getter)
186 # Convert to numpy arrays
187 tracking_data = np.stack(tracking_data, axis=0)
189 # Create cubes
190 tracking_cube = cube.copy(data=tracking_data)
191 tracking_cube.long_name = output_type_and_methods[output_type]["cube_name"]
192 tracking_cube.standard_name = None
193 tracking_cube.var_name = None
194 tracking_cube.units = "1"
195 # Add maximum value of the data to the cube attributes for use in colormap scaling
196 tracking_cube.attributes["max_value"] = np.ma.max(tracking_data)
197 tracking_cubelist.append(tracking_cube)
199 return tracking_cubelist
202def cell_stats(
203 cubes: iris.cube.Cube | iris.cube.CubeList,
204 threshold: float | list[float],
205 under_threshold: bool = False,
206 min_size: int = 4,
207 save_data: bool = False,
208):
209 """Identify features in each timestep and output statistics.
211 Parameters
212 ----------
213 cubes: iris.cube.Cube | iris.cube.CubeList
214 An iris cube (single model) or cubelist (multiple models) containing 2D data to be
215 analysed. Cube must have horizontal coordinates of xy type, not latitude/longitude.
216 The cube must also have a time coordinate, which is used to identify features in
217 each timestep.
218 threshold: float | list[float]
219 The threshold value(s) for feature detection. If a list is provided, each value
220 is used to identify features in the corresponding cube in the cubelist. Therefore,
221 the list should match the number of models.
222 under_threshold: bool, optional
223 If set to True, features are identified where the data is below the threshold.
224 If set to False, features are identified where the data is above the threshold.
225 Default is False.
226 min_size: int, optional
227 The minimum number of contiguous grid points required for a feature to be tracked.
228 Default is 4.
229 save_data: bool, optional
230 If set to True, all tracking data is saved to disk for further analysis (including csv
231 and txt files containing feature properties that are not returned in output cubes).
232 Default is False.
234 Returns
235 -------
236 cell_stats_cubes: iris.cube.CubeList
237 An iris CubeList containing "feature_size", "feature_effective_radius", "feature_mean",
238 and "feature_max" cubes.
240 Notes
241 -----
242 This operator uses the Simple-Track package with tracking disabled to identify features
243 in each timestep and compile cell statistics. Outputs cubes containing feature size (number
244 of grid points), effective radius (in km), mean value within features, and maximum
245 value within features.
247 Links
248 ----------
249 .. https://github.com/ParaChute-UK/simple-track
251 Examples
252 --------
253 >>> cell_stats_cubes = feature.cell_stats(threshold=2)
254 >>> feature_size_cube = cell_stats_cubes.extract_cube("feature_size")
255 >>> plt.hist(feature_size_cube[-1])
256 >>> plt.show()
258 """
259 # Check inputs
260 cubes = iter_maybe(cubes)
262 # Require inputs to have a uniform grid
263 for cube in cubes:
264 _check_uniform_grid(cube)
266 # Setup containing cube list
267 cell_stats_cubelist = iris.cube.CubeList()
269 # If threshold is a list, check that it matches the number of cubes
270 if isinstance(threshold, list):
271 if len(threshold) != len(cubes): 271 ↛ 281line 271 didn't jump to line 281 because the condition on line 271 was always true
272 raise ValueError(
273 f"Length of threshold list ({len(threshold)}) does not match "
274 f"number of cubes ({len(cubes)})."
275 )
276 # else, make it iterable by repeating the same value for each cube
277 else:
278 threshold = [threshold] * len(cubes)
280 # Run tracking on all input data
281 for cube, thresh in zip(cubes, threshold, strict=True):
282 model_name = cube.attributes.get("model_name", None)
283 # Setup config
284 tracker_config = {
285 "FEATURE": {
286 "threshold": thresh,
287 "under_threshold": under_threshold,
288 "min_size": min_size,
289 },
290 "OUTPUT": {
291 "save_data": save_data,
292 "experiment_name": "feature_tracking",
293 "path": f"{os.getcwd()}/{model_name}/cell-stats_data",
294 "skip_tracking": True,
295 },
296 }
297 logger.debug(f"Tracker config: {tracker_config}")
299 # Get cube data into a dict to pass to Tracker
300 times = cube.coord("time").points
301 time_units = cube.coord("time").units
302 times_dt = [time_units.num2pydate(t) for t in times]
303 cube_dict = {
304 time: cube_slice.data
305 for time, cube_slice in zip(times_dt, cube.slices_over("time"), strict=True)
306 }
308 # Run tracking, returning Timeline object
309 timeline = Tracker(tracker_config).run(cube_dict)
310 logger.debug(f"Tracking completed for {model_name}")
312 # Get feature data from each frame of data
313 size_data, mean_data, max_data = _get_cell_stats_arrays_from_timeline(
314 timeline=timeline, expected_frame_times=times_dt
315 )
317 # Get effective radius from feature size, using horizontal coordinate of input cube to estimate grid spacing
318 effective_radius_data, grid_spacing = _get_effective_radius_from_feature_size(
319 size_data=size_data, cube_with_hzntl_coord=cube
320 )
322 # Add grid_spacing as an attribute to the template_cube, so it is copied to
323 # output cubes in following function
324 cube.attributes["grid_spacing"] = grid_spacing
326 # Set output cube properties
327 cube_properties = {
328 "feature_size": {
329 "data": size_data,
330 "long_name": "feature_size",
331 "units": 1,
332 },
333 "feature_mean": {
334 "data": mean_data,
335 "long_name": "feature_mean",
336 "units": 1,
337 },
338 "feature_max": {"data": max_data, "long_name": "feature_max", "units": 1},
339 "feature_effective_radius": {
340 "data": effective_radius_data,
341 "long_name": "feature_effective_radius",
342 "units": "km",
343 },
344 }
346 # Create cubes, add to existing cubelist
347 cell_stats_cubelist.extend(
348 _add_cell_stats_data_to_cubes(
349 data_and_metadata_dict=cube_properties, template_cube=cube
350 )
351 )
353 return cell_stats_cubelist
356def _check_uniform_grid(cube: iris.cube.Cube) -> bool:
357 """Check that the input cube has approximately uniform horizontal grid spacing.
359 Prints warning if cube does not have uniform grid.
361 Parameters
362 ----------
363 cube: iris.cube.Cube
364 An iris cube containing horizontal coordinates.
366 Returns
367 -------
368 bool
369 True if the cube has a uniform grid, False otherwise.
371 Raises
372 ------
373 ValueError
374 If the input cube does not have a uniform grid.
375 """
376 hzntl_coords = [
377 coord
378 for coord in cube.coords()
379 if iris.util.guess_coord_axis(coord) in ["X", "Y"]
380 ]
382 for coord in hzntl_coords:
383 if not iris.util.is_regular(coord):
384 warning_msg = (
385 f"Horizontal coordinate {coord} is not regular. "
386 "Feature statistics calculation may be inaccurate."
387 )
388 logger.warning(warning_msg)
389 print(warning_msg)
390 return False
391 return True
394def _get_cell_stats_arrays_from_timeline(
395 timeline: Timeline, expected_frame_times: list
396) -> list[np.ndarray]:
397 """Extract cell statistics data from a Simple-Track Timeline object.
399 Parameters
400 ----------
401 timeline: Timeline
402 A Simple-Track Timeline object containing tracked features.
404 expected_frame_times: list
405 A list of expected frame times to extract data for.
407 Returns
408 -------
409 size_data: np.ndarray
410 A numpy array containing the size of each feature in grid points.
411 mean_data: np.ndarray
412 A numpy array containing the mean value of each feature.
413 max_data: np.ndarray
414 A numpy array containing the maximum value of each feature.
415 """
416 size_data, mean_data, max_data = [], [], []
417 number_of_features = []
418 for time in expected_frame_times:
419 frame = timeline.get_frame(time)
420 features = frame.features
421 size_data.append([feature.get_size() for feature in features.values()])
422 mean_data.append([feature.mean for feature in features.values()])
423 max_data.append([feature.max for feature in features.values()])
424 number_of_features.append(len(features))
426 # Pad data with NaNs to create arrays of consistent shape (max number of features across
427 # all timesteps)
428 arr_size = max(number_of_features)
430 # Size data is integer, but we need to pad with NaNs (which is a float), so fill
431 # with invalid value first
432 size_data = np.array(
433 [
434 np.pad(sizes, (0, arr_size - len(sizes)), constant_values=-100)
435 for sizes in size_data
436 ],
437 dtype=float,
438 )
439 size_data[size_data == -100] = np.nan
441 # Mean and max data are already float, so can be padded with NaNs directly.
442 mean_data = np.array(
443 [
444 np.pad(means, (0, arr_size - len(means)), constant_values=np.nan)
445 for means in mean_data
446 ]
447 )
448 max_data = np.array(
449 [
450 np.pad(maxs, (0, arr_size - len(maxs)), constant_values=np.nan)
451 for maxs in max_data
452 ]
453 )
455 return size_data, mean_data, max_data
458def _get_effective_radius_from_feature_size(
459 size_data: np.ndarray, cube_with_hzntl_coord: iris.cube.Cube
460) -> np.ndarray:
461 """Convert feature size in grid points to effective radius in km.
463 Parameters
464 ----------
465 size_data: np.ndarray
466 An array containing "feature_size" data, in units of grid points.
467 cube_with_hzntl_coord: iris.cube.Cube
468 An iris cube containing a horizontal coordinate, which is used to
469 estimate the grid spacing for the effective radius calculation.
471 Returns
472 -------
473 effective_radii_data: np.ndarray
474 An array containing "feature_effective_radius" data, in units of km.
476 grid_spacing: float
477 The estimated grid spacing in m, calculated from the horizontal coordinate of the input cube.
479 Notes
480 -----
481 This function assumes that the input cube has a horizontal coordinate system that is regular and
482 that the grid spacing can be estimated from the horizontal coordinates. The effective radius is
483 calculated as the radius of a circle with the same area as the feature size in grid points.
485 """
486 # Guess coord representing horizontal grid (choose first available)
487 hzntl_coord = next(
488 iter(
489 [
490 coord
491 for coord in cube_with_hzntl_coord.coords()
492 if iris.util.guess_coord_axis(coord) in ["X", "Y"]
493 ]
494 )
495 )
497 logger.debug(f"Attempting to convert to effective radius using {hzntl_coord}")
499 # Check coordinate is regular, but only warn if not, this is a naive estimate
500 # and will be inaccurate for irregular grids
501 if not iris.util.is_regular(hzntl_coord): 501 ↛ 502line 501 didn't jump to line 502 because the condition on line 501 was never true
502 logger.warning(
503 f"Horizontal coordinate {hzntl_coord} is not regular. "
504 "Effective radius calculation may be inaccurate."
505 )
507 # Get grid spacing in native coord units (degrees, m, km etc)
508 grid_spacing = iris.util.regular_step(hzntl_coord)
510 if hzntl_coord.units == "m":
511 grid_spacing = grid_spacing / 1000 # Convert to km
513 # If grid spacing is in degrees, convert to km using approximate conversion factor
514 if hzntl_coord.units == "degrees":
515 # Get latitude for better conversion to km
516 lat_coords = [
517 coord
518 for coord in cube_with_hzntl_coord.coords()
519 if iris.util.guess_coord_axis(coord) in ["Y"]
520 ]
521 for coord in lat_coords:
522 if coord.units == "degrees": 522 ↛ 525line 522 didn't jump to line 525 because the condition on line 522 was always true
523 lat_coord_for_conversion = coord
524 else:
525 lat_coord_for_conversion = None
527 # Calculate conversion factor using latitude correction if available,
528 # otherwise use naive 111 km per degree conversion
529 if lat_coord_for_conversion is not None: 529 ↛ 532line 529 didn't jump to line 532 because the condition on line 529 was always true
530 mean_latitude = np.mean(lat_coord_for_conversion.points)
531 else:
532 logger.warning(
533 "No latitude coordinate found for conversion to km. "
534 "Using naive conversion factor of 111 km per degree."
535 )
536 mean_latitude = 0
537 grid_spacing = grid_spacing * 111 * np.cos(np.radians(mean_latitude))
539 effective_radii_data = np.sqrt(size_data * grid_spacing**2 / np.pi)
540 return effective_radii_data, grid_spacing
543def _add_cell_stats_data_to_cubes(
544 data_and_metadata_dict: dict, template_cube: iris.cube.Cube
545) -> iris.cube.CubeList:
546 """Add data to cubes, using template cube for metadata.
548 Parameters
549 ----------
550 data_and_metadata_dict: dict
551 A dictionary containing data and metadata for each cube to be created.
552 The keys are the long names of the cubes, and the values are dictionaries
553 containing the data and units for each cube.
555 template_cube: iris.cube.Cube
556 An iris cube to use as a template for the new cubes. The new cubes will
557 have the same attributes as the template cube.
559 Returns
560 -------
561 cubelist: iris.cube.CubeList
562 A list of iris cubes containing the added data.
564 """
565 cubelist = iris.cube.CubeList()
567 # Construct coordinates for new cubes
568 time_coord = template_cube.coord("time").copy()
569 # To construct feature coordinate, look at the size of dimension 1 for each data
570 arr_size = max(
571 [data_and_metadata_dict[cb]["data"].shape[1] for cb in data_and_metadata_dict]
572 )
573 feature_coord = iris.coords.DimCoord(
574 np.arange(arr_size),
575 long_name="feature_number",
576 var_name="feature_number",
577 units="1",
578 )
579 coords = [time_coord, feature_coord]
580 coords_and_dims = [(coord, i) for i, coord in enumerate(coords)]
582 # Get list of coords to copy from input cube to output cubes
583 copyable_coord_names = [
584 "realization",
585 "hour",
586 "forecast_period",
587 "forecast_reference_time",
588 "model_name",
589 "cset_comparison_base",
590 ]
591 input_cube_coord_names = []
592 for coord in template_cube.coords():
593 input_cube_coord_names.append(coord.standard_name)
594 input_cube_coord_names.append(coord.long_name)
596 coords_to_copy = [
597 coord_name
598 for coord_name in copyable_coord_names
599 if coord_name in input_cube_coord_names
600 ]
602 # Populate cubelist
603 for cb_props in data_and_metadata_dict.values():
604 cell_stats_cube = iris.cube.Cube(
605 data=cb_props["data"],
606 long_name=cb_props["long_name"],
607 units=cb_props["units"],
608 dim_coords_and_dims=coords_and_dims,
609 )
610 # Add other metadata from input cube
611 for coord_name in coords_to_copy:
612 coord = template_cube.coord(coord_name).copy()
613 # Check if this coord represents a dimension of data
614 dims = template_cube.coord_dims(coord)
615 if len(dims) > 0: 615 ↛ 616line 615 didn't jump to line 616 because the condition on line 615 was never true
616 cell_stats_cube.add_aux_coord(coord, dims)
617 else:
618 cell_stats_cube.add_aux_coord(coord)
620 # Copy over attributes
621 cell_stats_cube.attributes = template_cube.attributes
623 # Add to cubelist
624 cubelist.append(cell_stats_cube)
626 return cubelist