Coverage for src/CSET/operators/feature.py: 93%
137 statements
« prev ^ index » next coverage.py v7.16.1, created at 2026-09-17 15:05 +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 threshold: float
49 The threshold value for feature detection.
50 under_threshold: bool, optional
51 If set to True, features are identified where the data is below the threshold.
52 If set to False, features are identified where the data is above the threshold.
53 Default is False.
54 min_size: int, optional
55 The minimum number of contiguous grid points required for a feature to be tracked.
56 Default is 4.
57 retain_lifetime_on_split: bool, optional
58 If set to True, the lifetime of a feature is retained when it splits into
59 multiple features. If set to False, the lifetime is reset when a feature splits.
60 Default is True.
61 tracking_nbhood: int, optional
62 The size of the neighbourhood used for tracking features between timesteps.
63 This dictates the maximum pixel radius from a feature centroid at which new features could
64 reasonably be spawned.
65 Default is 5.
66 overlap_threshold: float, optional
67 The minimum overlap required between features in consecutive timesteps for
68 them to be considered the same feature.
69 Default is 0.3.
70 save_data: bool, optional
71 If set to True, all tracking data is saved to disk for further analysis (including csv
72 and txt files containing feature properties that are not returned in output cubes).
73 Default is False.
75 Returns
76 -------
77 tracking_cubes: iris.cube.CubeList
78 A list of iris cubes containing tracking data, including feature ID, lifetime,
79 and locations of initiating features.
81 Notes
82 -----
83 This operator uses the Simple-Track package to track features between timesteps. Simple-Track is a
84 data-agnostic, threshold-based object tracking algorithm for 2D data. Features are tracked between
85 consecutive frames of data by projecting feature fields onto common timeframes and matching
86 between them based on the degree of overlap. Matched features retain the same identification
87 between all tracked fields, while new features are assigned a unique label.
88 Thus, Simple-Track compiles comprehensive information about feature merging, splitting, accretion,
89 initiation and dissipation.
91 Currently outputs three cubes containing the following data:
92 "feature_id":
93 A 2D field containing the unique label assigned to each feature, which is retained
94 if the feature is tracked across multiple timesteps. This cube can be used as a mask
95 to identify the location of the tracked feature throughout the evaluation period.
96 "feature_lifetime":
97 A 2D field containing the lifetime of each feature in terms of the number of
98 timesteps it has been tracked for. This cube can be used to distinguish between
99 mature and fresh features.
100 "feature_init":
101 A 2D binary field indicating the location of newly initiated features at each timestep.
102 These features are identified as having a lifetime of 1 AND have initiated sufficiently
103 far from other, existing features that they are not considered to have spawned from them.
105 Links
106 ----------
107 .. https://github.com/ParaChute-UK/simple-track
109 Examples
110 --------
111 >>> tracking_cubes = feature.track(threshold=2)
112 >>> lifetime_cube = tracking_cubes.extract_cube("feature_lifetime")
113 # Plot the final timestep of lifetime cube. This will show
114 # the lifetime of features that have been tracked for multiple previous
115 # timesteps, as well as new features that have just been initiated.
116 >>> iplt.pcolormesh(lifetime_cube[-1,:,:],cmap=mpl.cm.bwr)
117 >>> plt.gca().coastlines('10m')
118 >>> plt.clim(-5,5)
119 >>> plt.colorbar()
120 >>> plt.show()
122 """
123 # Setup config
124 tracker_config = {
125 "FEATURE": {
126 "threshold": threshold,
127 "under_threshold": under_threshold,
128 "min_size": min_size,
129 },
130 "TRACKING": {
131 "retain_lifetime_on_split": retain_lifetime_on_split,
132 "overlap_nbhood": tracking_nbhood,
133 "overlap_threshold": overlap_threshold,
134 },
135 "OUTPUT": {
136 "save_data": save_data,
137 "experiment_name": "feature_tracking",
138 "path": f"{os.getcwd()}/tracking_data",
139 },
140 }
141 logger.debug(f"Tracker config: {tracker_config}")
143 # Get cube data into a dict to pass to Tracker
144 times = cube.coord("time").points
145 time_units = cube.coord("time").units
146 times_dt = [time_units.num2pydate(t) for t in times]
147 cube_dict = {
148 time: cube_slice.data
149 for time, cube_slice in zip(times_dt, cube.slices_over("time"), strict=True)
150 }
152 # Run tracking, returning Timeline object
153 timeline = Tracker(tracker_config).run(cube_dict)
154 logger.debug("Tracking completed")
156 # Use input cube as template to make returned cube
157 # By iterating over all cube times, this will ensure all data is present
158 # If a Frame at the given time is not contained in the timeline, error is raised
159 output_type_and_methods = {
160 "lifetime": {
161 "getter": "lifetime_field",
162 "cube_name": "feature_lifetime",
163 },
164 "feature": {
165 "getter": "feature_field",
166 "cube_name": "feature_id",
167 },
168 "init": {
169 "getter": "get_init_field",
170 "cube_name": "feature_init",
171 },
172 }
174 tracking_cubelist = iris.cube.CubeList()
175 for output_type in output_type_and_methods:
176 tracking_data = []
177 for time in times_dt:
178 frame = timeline.get_frame(time)
179 getter = getattr(frame, output_type_and_methods[output_type]["getter"])
180 if callable(getter):
181 tracking_data.append(getter())
182 else:
183 tracking_data.append(getter)
185 # Convert to numpy arrays
186 tracking_data = np.stack(tracking_data, axis=0)
188 # Create cubes
189 tracking_cube = cube.copy(data=tracking_data)
190 tracking_cube.long_name = output_type_and_methods[output_type]["cube_name"]
191 tracking_cube.standard_name = None
192 tracking_cube.var_name = None
193 tracking_cube.units = "1"
194 # Add maximum value of the data to the cube attributes for use in colormap scaling
195 tracking_cube.attributes["max_value"] = np.ma.max(tracking_data)
196 tracking_cubelist.append(tracking_cube)
198 return tracking_cubelist
201def cell_stats(
202 cubes: iris.cube.Cube | iris.cube.CubeList,
203 threshold: float | list[float],
204 under_threshold: bool = False,
205 min_size: int = 4,
206 save_data: bool = False,
207):
208 """Identify features in each timestep and output statistics.
210 Parameters
211 ----------
212 cubes: iris.cube.Cube | iris.cube.CubeList
213 An iris cube (single model) or cubelist (multiple models) containing 2D data to be
214 analysed. Cube must have horizontal coordinates on a regular grid.
215 The cube must also have a time coordinate, which is used to identify features in
216 each timestep.
217 threshold: float | list[float]
218 The threshold value(s) for feature detection. If a list is provided, each value
219 is used to identify features in the corresponding cube in the cubelist. Therefore,
220 the list should match the number of models.
221 under_threshold: bool, optional
222 If set to True, features are identified where the data is below the threshold.
223 If set to False, features are identified where the data is above the threshold.
224 Default is False.
225 min_size: int, optional
226 The minimum number of contiguous grid points required for a feature to be tracked.
227 Default is 4.
228 save_data: bool, optional
229 If set to True, all tracking data is saved to disk for further analysis (including csv
230 and txt files containing feature properties that are not returned in output cubes).
231 Default is False.
233 Returns
234 -------
235 cell_stats_cubelist: iris.cube.CubeList
236 An iris CubeList containing "feature_size", "feature_effective_diameter", "feature_mean",
237 and "feature_max" cubes.
239 Notes
240 -----
241 This operator uses the Simple-Track package with tracking disabled to identify features
242 in each timestep and compile cell statistics. Outputs cubes containing feature size (number
243 of grid points), effective diameter (in km), mean value within features, and maximum
244 value within features.
246 Links
247 ----------
248 .. https://github.com/ParaChute-UK/simple-track
250 Examples
251 --------
252 >>> cell_stats_cubes = feature.cell_stats(threshold=2)
253 >>> feature_size_cube = cell_stats_cubes.extract_cube("feature_size")
254 >>> plt.hist(feature_size_cube[-1])
255 >>> plt.show()
257 """
258 # Check inputs
259 cubes = iter_maybe(cubes)
261 # Require inputs to have a uniform grid
262 for cube in cubes:
263 _check_uniform_grid(cube)
265 # Setup containing cube list
266 cell_stats_cubelist = iris.cube.CubeList()
268 # If threshold is a list, check that it matches the number of cubes
269 if isinstance(threshold, list):
270 if len(threshold) != len(cubes): 270 ↛ 280line 270 didn't jump to line 280 because the condition on line 270 was always true
271 raise ValueError(
272 f"Length of threshold list ({len(threshold)}) does not match "
273 f"number of cubes ({len(cubes)})."
274 )
275 # else, make it iterable by repeating the same value for each cube
276 else:
277 threshold = [threshold] * len(cubes)
279 # Run tracking on all input data
280 for cube, thresh in zip(cubes, threshold, strict=True):
281 model_name = cube.attributes.get("model_name", None)
282 # Setup config
283 tracker_config = {
284 "FEATURE": {
285 "threshold": thresh,
286 "under_threshold": under_threshold,
287 "min_size": min_size,
288 },
289 "OUTPUT": {
290 "save_data": save_data,
291 "experiment_name": "feature_tracking",
292 "path": f"{os.getcwd()}/{model_name}/cell-stats_data",
293 "skip_tracking": True,
294 },
295 }
296 logger.debug(f"Tracker config: {tracker_config}")
298 # Get cube data into a dict to pass to Tracker
299 times = cube.coord("time").points
300 time_units = cube.coord("time").units
301 times_dt = [time_units.num2pydate(t) for t in times]
302 cube_dict = {
303 time: cube_slice.data
304 for time, cube_slice in zip(times_dt, cube.slices_over("time"), strict=True)
305 }
307 # Run tracking, returning Timeline object
308 timeline = Tracker(tracker_config).run(cube_dict)
309 logger.debug(f"Tracking completed for {model_name}")
311 # Get feature data from each frame of data
312 size_data, mean_data, max_data = _get_cell_stats_arrays_from_timeline(
313 timeline=timeline, expected_frame_times=times_dt
314 )
316 # Get effective diameter from feature size, using horizontal coordinate of input cube to estimate grid spacing
317 effective_diameter_data, grid_spacing = (
318 _get_effective_diameter_from_feature_size(
319 size_data=size_data, cube_with_hzntl_coord=cube
320 )
321 )
323 # Add grid_spacing as an attribute to the template_cube, so it is copied to
324 # output cubes in following function
325 cube.attributes["grid_spacing"] = grid_spacing
327 # Set output cube properties
328 cube_properties = {
329 "feature_size": {
330 "data": size_data,
331 "long_name": "feature_size",
332 "units": 1,
333 },
334 "feature_mean": {
335 "data": mean_data,
336 "long_name": "feature_mean",
337 "units": 1,
338 },
339 "feature_max": {"data": max_data, "long_name": "feature_max", "units": 1},
340 "feature_effective_diameter": {
341 "data": effective_diameter_data,
342 "long_name": "feature_effective_diameter",
343 "units": "km",
344 },
345 }
347 # Create cubes, add to existing cubelist
348 cell_stats_cubelist.extend(
349 _add_cell_stats_data_to_cubes(
350 data_and_metadata_dict=cube_properties, template_cube=cube
351 )
352 )
354 return cell_stats_cubelist
357def _check_uniform_grid(cube: iris.cube.Cube) -> bool:
358 """Check that the input cube has approximately uniform horizontal grid spacing.
360 Prints warning if cube does not have uniform grid.
362 Parameters
363 ----------
364 cube: iris.cube.Cube
365 An iris cube containing horizontal coordinates.
367 Returns
368 -------
369 bool
370 True if the cube has a uniform grid, False otherwise.
372 Raises
373 ------
374 ValueError
375 If the input cube does not have a uniform grid.
376 """
377 hzntl_coords = [
378 coord
379 for coord in cube.coords()
380 if iris.util.guess_coord_axis(coord) in ["X", "Y"]
381 ]
383 for coord in hzntl_coords:
384 if not iris.util.is_regular(coord):
385 warning_msg = (
386 f"Horizontal coordinate {coord} is not regular. "
387 "Feature statistics calculation may be inaccurate."
388 )
389 logger.warning(warning_msg)
390 print(warning_msg)
391 return False
392 return True
395def _get_cell_stats_arrays_from_timeline(
396 timeline: Timeline, expected_frame_times: list
397) -> list[np.ndarray]:
398 """Extract cell statistics data from a Simple-Track Timeline object.
400 Parameters
401 ----------
402 timeline: Timeline
403 A Simple-Track Timeline object containing tracked features.
405 expected_frame_times: list
406 A list of expected frame times to extract data for.
408 Returns
409 -------
410 size_data: np.ndarray
411 A numpy array containing the size of each feature in grid points.
412 mean_data: np.ndarray
413 A numpy array containing the mean value of each feature.
414 max_data: np.ndarray
415 A numpy array containing the maximum value of each feature.
416 """
417 size_data, mean_data, max_data = [], [], []
418 number_of_features = []
419 for time in expected_frame_times:
420 frame = timeline.get_frame(time)
421 features = frame.features
422 size_data.append([feature.get_size() for feature in features.values()])
423 mean_data.append([feature.mean for feature in features.values()])
424 max_data.append([feature.max for feature in features.values()])
425 number_of_features.append(len(features))
427 # Pad data with NaNs to create arrays of consistent shape (max number of features across
428 # all timesteps)
429 arr_size = max(number_of_features)
431 # Size data is integer, but we need to pad with NaNs (which is a float), so fill
432 # with invalid value first
433 size_data = np.array(
434 [
435 np.pad(sizes, (0, arr_size - len(sizes)), constant_values=-100)
436 for sizes in size_data
437 ],
438 dtype=float,
439 )
440 size_data[size_data == -100] = np.nan
442 # Mean and max data are already float, so can be padded with NaNs directly.
443 mean_data = np.array(
444 [
445 np.pad(means, (0, arr_size - len(means)), constant_values=np.nan)
446 for means in mean_data
447 ]
448 )
449 max_data = np.array(
450 [
451 np.pad(maxs, (0, arr_size - len(maxs)), constant_values=np.nan)
452 for maxs in max_data
453 ]
454 )
456 return size_data, mean_data, max_data
459def _get_effective_diameter_from_feature_size(
460 size_data: np.ndarray, cube_with_hzntl_coord: iris.cube.Cube
461) -> np.ndarray:
462 """Convert feature size in grid points to effective diameter in km.
464 Parameters
465 ----------
466 size_data: np.ndarray
467 An array containing "feature_size" data, in units of grid points.
468 cube_with_hzntl_coord: iris.cube.Cube
469 An iris cube containing a horizontal coordinate, which is used to
470 estimate the grid spacing for the effective radius calculation.
472 Returns
473 -------
474 effective_diameters_data: np.ndarray
475 An array containing "feature_effective_diameter" data, in units of km.
477 grid_spacing: float
478 The estimated grid spacing in m, calculated from the horizontal coordinate of the input cube.
480 Notes
481 -----
482 This function assumes that the input cube has a horizontal coordinate system that is regular and
483 that the grid spacing can be estimated from the horizontal coordinates. The effective diameter is
484 calculated as the diameter of a circle with the same area as the feature size in grid points.
486 """
487 # Guess coord representing horizontal grid (choose first available)
488 try:
489 hzntl_coord = next(
490 iter(
491 [
492 coord
493 for coord in cube_with_hzntl_coord.coords()
494 if iris.util.guess_coord_axis(coord) in ["X", "Y"]
495 ]
496 )
497 )
498 except StopIteration:
499 raise ValueError(
500 "No horizontal coordinate found in input cube. "
501 "Cannot calculate effective diameter."
502 ) from None
504 logger.debug(f"Attempting to convert to effective diameter using {hzntl_coord}")
506 # Check coordinate is regular, but only warn if not, this is a naive estimate
507 # and will be inaccurate for irregular grids
508 # Additionally, the current function only checks for regularity in one horizontal
509 # coordinate. Again, this is a bit of naive estimate.
510 if not iris.util.is_regular(hzntl_coord): 510 ↛ 511line 510 didn't jump to line 511 because the condition on line 510 was never true
511 logger.warning(
512 f"Horizontal coordinate {hzntl_coord} is not regular. "
513 "Effective diameter calculation may be inaccurate."
514 )
516 # Get grid spacing in native coord units (degrees, m, km etc)
517 grid_spacing = iris.util.regular_step(hzntl_coord)
519 if hzntl_coord.units == "m":
520 grid_spacing = grid_spacing / 1000 # Convert to km
522 # If grid spacing is in degrees, convert to km using approximate conversion factor
523 if hzntl_coord.units == "degrees":
524 # Get latitude for better conversion to km
525 lat_coords = [
526 coord
527 for coord in cube_with_hzntl_coord.coords()
528 if iris.util.guess_coord_axis(coord) in ["Y"]
529 ]
530 for coord in lat_coords:
531 if coord.units == "degrees": 531 ↛ 534line 531 didn't jump to line 534 because the condition on line 531 was always true
532 lat_coord_for_conversion = coord
533 else:
534 lat_coord_for_conversion = None
536 # Calculate conversion factor using latitude correction if available,
537 # otherwise use naive 111 km per degree conversion
538 if lat_coord_for_conversion is not None: 538 ↛ 541line 538 didn't jump to line 541 because the condition on line 538 was always true
539 mean_latitude = np.mean(lat_coord_for_conversion.points)
540 else:
541 logger.warning(
542 "No latitude coordinate found for conversion to km. "
543 "Using naive conversion factor of 111 km per degree."
544 )
545 mean_latitude = 0
546 grid_spacing = grid_spacing * 111 * np.cos(np.radians(mean_latitude))
548 effective_diameter_data = np.sqrt(size_data * grid_spacing**2 / np.pi) * 2
549 return effective_diameter_data, grid_spacing
552def _add_cell_stats_data_to_cubes(
553 data_and_metadata_dict: dict, template_cube: iris.cube.Cube
554) -> iris.cube.CubeList:
555 """Add data to cubes, using template cube for metadata.
557 Parameters
558 ----------
559 data_and_metadata_dict: dict
560 A dictionary containing data and metadata for each cube to be created.
561 The keys are the long names of the cubes, and the values are dictionaries
562 containing the data and units for each cube.
564 template_cube: iris.cube.Cube
565 An iris cube to use as a template for the new cubes. The new cubes will
566 have the same attributes as the template cube.
568 Returns
569 -------
570 cubelist: iris.cube.CubeList
571 A list of iris cubes containing the added data.
573 """
574 cubelist = iris.cube.CubeList()
576 # Construct coordinates for new cubes
577 time_coord = template_cube.coord("time").copy()
578 # To construct feature coordinate, look at the size of dimension 1 for each data
579 arr_size = max(
580 [data_and_metadata_dict[cb]["data"].shape[1] for cb in data_and_metadata_dict]
581 )
582 feature_coord = iris.coords.DimCoord(
583 np.arange(arr_size),
584 long_name="feature_number",
585 var_name="feature_number",
586 units="1",
587 )
588 coords = [time_coord, feature_coord]
589 coords_and_dims = [(coord, i) for i, coord in enumerate(coords)]
591 # Get list of coords to copy from input cube to output cubes
592 copyable_coord_names = [
593 "realization",
594 "hour",
595 "forecast_period",
596 "forecast_reference_time",
597 "model_name",
598 "cset_comparison_base",
599 ]
600 input_cube_coord_names = []
601 for coord in template_cube.coords():
602 input_cube_coord_names.append(coord.standard_name)
603 input_cube_coord_names.append(coord.long_name)
605 coords_to_copy = [
606 coord_name
607 for coord_name in copyable_coord_names
608 if coord_name in input_cube_coord_names
609 ]
611 # Populate cubelist
612 for cb_props in data_and_metadata_dict.values():
613 cell_stats_cube = iris.cube.Cube(
614 data=cb_props["data"],
615 long_name=cb_props["long_name"],
616 units=cb_props["units"],
617 dim_coords_and_dims=coords_and_dims,
618 )
619 # Add other metadata from input cube
620 for coord_name in coords_to_copy:
621 coord = template_cube.coord(coord_name).copy()
622 # Check if this coord represents a dimension of data
623 dims = template_cube.coord_dims(coord)
624 if len(dims) > 0: 624 ↛ 625line 624 didn't jump to line 625 because the condition on line 624 was never true
625 cell_stats_cube.add_aux_coord(coord, dims)
626 else:
627 cell_stats_cube.add_aux_coord(coord)
629 # Copy over attributes
630 cell_stats_cube.attributes = template_cube.attributes
632 # Add to cubelist
633 cubelist.append(cell_stats_cube)
635 return cubelist