Coverage for src/CSET/operators/ageofair.py: 94%
142 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-07 15:12 +0000
1# © Crown copyright, Met Office (2022-2024) 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.
15"""
16Age of air operator.
18The age of air diagnostic provides a qualtitative view of how old air is within
19the domain, by calculating a back trajectory at each grid point at each lead time
20to determine when air entered through the lateral boundary. This is useful for
21diagnosing how quickly air ventilates the domain, depending on its size and the
22prevailing meteorology.
24The diagnostic uses the u, v and w components of wind, along with geopotential height to
25perform the back trajectory. Data is first regridded to 0.5 degrees.
27Note: the code here does not consider sub-grid transport, and only uses the postprocessed
28velocity fields and geopotential height. Its applicability is for large-scale flow O(1000 km),
29and not small scale flow where mixing is likely to play a larger role.
30"""
32import datetime
33import logging
34import multiprocessing
35import os
36import tempfile
37from functools import partial
38from math import asin, cos, radians, sin, sqrt
40import numpy as np
41from iris.cube import Cube
42from scipy.ndimage import gaussian_filter
44from CSET.operators._utils import get_cube_yxcoordname
46logger = logging.getLogger(__name__)
49def _calc_dist(coord_1, coord_2):
50 """Calculate distance between two coordinate tuples.
52 Arguments
53 ----------
54 coord_1: tuple
55 A tuple containing (latitude, longitude) coordinate floats
56 coord_2: tuple
57 A tuple containing (latitude, longitude) coordinate floats
59 Returns
60 -------
61 distance: float
62 Distance between the two coordinate points in meters
64 Notes
65 -----
66 The function uses the Haversine approximation to calculate distance in metres.
68 """
69 # Approximate radius of earth in m
70 # Source: https://nssdc.gsfc.nasa.gov/planetary/factsheet/earthfact.html
71 radius = 6378000
73 # Extract coordinates and convert to radians
74 lat1 = radians(coord_1[0])
75 lon1 = radians(coord_1[1])
76 lat2 = radians(coord_2[0])
77 lon2 = radians(coord_2[1])
79 # Find out delta latitude, longitude
80 dlon = lon2 - lon1
81 dlat = lat2 - lat1
83 # Compute distance
84 a = sin(dlat / 2) ** 2 + cos(lat1) * cos(lat2) * sin(dlon / 2) ** 2
85 c = 2 * asin(sqrt(a))
86 distance = radius * c
88 return distance
91def _aoa_core(
92 x_arr: np.ndarray,
93 y_arr: np.ndarray,
94 z_arr: np.ndarray,
95 g_arr: np.ndarray,
96 lats: np.ndarray,
97 lons: np.ndarray,
98 dt: int,
99 plev_idx: int,
100 timeunit: str,
101 cyclic: bool,
102 tmpdir: str,
103 lon_pnt: int,
104):
105 """AOA multiprocessing core.
107 Runs the core age of air code on a specific longitude point (all latitudes) for
108 parallelisation.
110 Arguments
111 ---------
112 x_arr: np.ndarray
113 A numpy array containing x wind data.
114 y_arr: np.ndarray
115 A numpy array containing y wind data.
116 z_arr: np.ndarray
117 A numpy array containing w wind data.
118 g_arr: np.ndarray
119 A numpy array containing geopotential height data.
120 lats: np.ndarray
121 A numpy array containing latitude points.
122 lons: np.ndarray
123 A numpy array containing longitude points.
124 dt: int
125 Gap between time intervals
126 plev_idx: int
127 Index of pressure level requested to run back trajectories on.
128 timeunit: str
129 Units of time, currently only accepts 'hour'
130 cyclic: bool
131 Whether to wrap at east/west boundaries. See compute_ageofair for a fuller description.
132 tmpdir: str
133 Path to store intermediate data
134 lon_pnt: int
135 Longitude point to extract and run back trajectories on for parallelisation.
136 """
137 # Initialise empty array to store age of air for this latitude strip.
138 ageofair_local = np.zeros((x_arr.shape[0], x_arr.shape[2]))
139 logger.debug("Working on %s", lon_pnt)
141 # Ignore leadtime 0 as this is trivial.
142 for leadtime in range(1, x_arr.shape[0]):
143 # Initialise leadtime slice with current leadtime.
144 ageofair_local[leadtime, :] = leadtime * dt
145 for lat_pnt in range(x_arr.shape[2]):
146 # Gridpoint initialised as within LAM by construction.
147 outside_lam = False
149 # If final column, look at dist from prev column, otherwise look at next column.
150 if lon_pnt == len(lons) - 1:
151 ew_spacing = _calc_dist(
152 (lats[lat_pnt], lons[lon_pnt]), (lats[lat_pnt], lons[lon_pnt - 1])
153 )
154 else:
155 ew_spacing = _calc_dist(
156 (lats[lat_pnt], lons[lon_pnt]), (lats[lat_pnt], lons[lon_pnt + 1])
157 )
159 # If final row, look at dist from row column, otherwise look at next row.
160 if lat_pnt == len(lats) - 1:
161 ns_spacing = _calc_dist(
162 (lats[lat_pnt], lons[lon_pnt]), (lats[lat_pnt - 1], lons[lon_pnt])
163 )
164 else:
165 ns_spacing = _calc_dist(
166 (lats[lat_pnt], lons[lon_pnt]), (lats[lat_pnt + 1], lons[lon_pnt])
167 )
169 # Go through past timeslices
170 for n in range(leadtime):
171 # First step back, so we use i,j coords to find out parcel location
172 # in terms of array point
173 if n == 0:
174 x = lon_pnt
175 y = lat_pnt
176 z = plev_idx
178 # Only seek preceding wind if its inside domain.
179 if not outside_lam:
180 # Get vector profile at current time - nearest whole gridpoint.
181 u = x_arr[leadtime - n, int(z), int(y), int(x)]
182 v = y_arr[leadtime - n, int(z), int(y), int(x)]
183 w = z_arr[leadtime - n, int(z), int(y), int(x)]
184 g = g_arr[leadtime - n, int(z), int(y), int(x)]
186 # First, compute horizontal displacement using inverse of horizontal vector
187 # Convert m/s to m/[samplingrate]h, then m -> model gridpoints
188 if timeunit == "hour": 188 ↛ 194line 188 didn't jump to line 194 because the condition on line 188 was always true
189 du = ((u * 60 * 60 * dt) / ew_spacing) * -1.0
190 dv = ((v * 60 * 60 * dt) / ns_spacing) * -1.0
191 dz = (w * 60 * 60 * dt) * -1.0
193 # Get column of geopot height.
194 g_col = g_arr[(leadtime - n), :, int(y), int(x)]
196 # New geopotential height of parcel - store 'capacity' between timesteps as vertical motions smaller.
197 if n == 0:
198 new_g = g + dz
199 pre_g = new_g
200 else:
201 new_g = pre_g + dz
203 # Calculate which geopot level is closest to new geopot level.
204 z = np.argmin(np.abs(g_col - new_g))
206 # Update x,y location based on displacement. Z already updated
207 x = x + du
208 y = y + dv
210 # If it is now outside domain, then save age and don't process further with outside LAM flag.
211 # Support cyclic domains like K-SCALE, where x coord out of domain gets moved through dateline.
212 if cyclic:
213 if (
214 x < 0
215 ): # as for example -0.3 would still be in domain, but x_arr.shape-0.3 would result in index error
216 x = x_arr.shape[3] + x # wrap back around dateline
217 elif x >= x_arr.shape[3]:
218 x = x_arr.shape[3] - x
219 else:
220 if x < 0 or x >= x_arr.shape[3]:
221 ageofair_local[leadtime, lat_pnt] = n * dt
222 outside_lam = True
224 if y < 0 or y >= x_arr.shape[2]:
225 ageofair_local[leadtime, lat_pnt] = n * dt
226 outside_lam = True
228 # Save 3d array containing age of air
229 np.save(tmpdir + f"/aoa_frag_{lon_pnt:04d}.npy", ageofair_local)
232def compute_ageofair(
233 XWIND: Cube,
234 YWIND: Cube,
235 WWIND: Cube,
236 GEOPOT: Cube,
237 plev: int,
238 cyclic: bool = False,
239 multicore=True,
240):
241 """Compute back trajectories for a given forecast.
243 This allows us to determine when air entered through the boundaries. This will run on all available
244 lead-times, and thus return an age of air cube of shape ntime, nlat, nlon. It supports multiprocesing,
245 by iterating over longitude, or if set as None, will run on a single core, which is easier for debugging.
246 This function supports ensembles, where it will check if realization dimension exists and if so, loop
247 over this axis.
249 Arguments
250 ----------
251 XWIND: Cube
252 An iris cube containing the x component of wind on pressure levels, on a 0p5 degree grid.
253 Requires 4 dimensions, ordered time, pressure, latitude and longitude. Must contain at
254 least 2 time points to compute back trajectory.
255 YWIND: Cube
256 An iris cube in the same format as XWIND.
257 WWIND: Cube
258 An iris cube in the same format as XWIND.
259 GEOPOT: Cube
260 An iris cube in the same format as XWIND.
261 plev: int
262 The pressure level of which to compute the back trajectory on. The function will search to
263 see if this exists and if not, will raise an exception.
264 cyclic: bool
265 If cyclic is True, then the code will assume no east/west boundary and if a back trajectory
266 reaches the boundary, it will emerge out of the other side. This option is useful for large
267 domains such as the K-SCALE tropical channel, where there are only north/south boundaries in
268 the domain.
269 multicore: bool
270 If true, split up age of air diagnostic to use multiple cores (defaults to number of cores available to the process), otherwise run
271 using a single process, which is easier to debug if developing the code.
273 Returns
274 -------
275 ageofair_cube: Cube
276 An iris cube of the age of air data, with 3 dimensions (time, latitude, longitude).
278 Notes
279 -----
280 The age of air diagnostic was used in Warner et al. (2023) [Warneretal2023]_ to identify the relative
281 role of spin-up from initial conditions and lateral boundary conditions over tropical Africa to explore
282 the impact of new data assimilation techniques. A further paper is currently in review ([Warneretal2024]_)
283 which applies the diagnostic more widely to the Australian ACCESS convection-permitting models.
285 """
286 # Set up temporary directory to store intermediate age of air slices.
287 tmpdir = tempfile.TemporaryDirectory(dir=os.getenv("CYLC_TASK_WORK_DIR"))
288 logger.info("Made tmpdir %s", tmpdir.name)
290 # Check that all cubes are of same size (will catch different dimension orders too).
291 if not XWIND.shape == YWIND.shape == WWIND.shape == GEOPOT.shape:
292 raise ValueError("Cubes are not the same shape")
294 # Get time units and assign for later
295 if str(XWIND.coord("time").units).startswith("hours since "):
296 timeunit = "hour"
297 else:
298 raise NotImplementedError("Unsupported time base")
300 # Make data non-lazy to speed up code.
301 logger.info("Making data non-lazy...")
302 x_arr = XWIND.data
303 y_arr = YWIND.data
304 z_arr = WWIND.data
305 g_arr = GEOPOT.data
307 # Get coord points
308 lat_name, lon_name = get_cube_yxcoordname(XWIND)
309 lats = XWIND.coord(lat_name).points
310 lons = XWIND.coord(lon_name).points
311 time = XWIND.coord("time").points
313 # Get time spacing of cube to determine whether the spacing in time is the
314 # same throughout the cube. If not, then not supported.
315 dt = XWIND.coord("time").points[1:] - XWIND.coord("time").points[:-1]
316 if np.all(dt == dt[0]):
317 dt = dt[0]
318 else:
319 raise NotImplementedError("Time intervals are not consistent")
321 # Some logic to determine which index each axis is, and check for ensembles.
322 dimension_mapping = {}
323 for coord in XWIND.dim_coords:
324 dim_index = XWIND.coord_dims(coord.name())[0]
325 dimension_mapping[coord.name()] = dim_index
327 if "realization" in dimension_mapping:
328 ensemble_mode = True
329 if dimension_mapping != {
330 "realization": 0,
331 "time": 1,
332 "pressure": 2,
333 lat_name: 3,
334 lon_name: 4,
335 }:
336 raise ValueError(
337 f"Dimension mapping not correct, ordered {dimension_mapping}"
338 )
339 else:
340 ensemble_mode = False
341 if dimension_mapping != {"time": 0, "pressure": 1, lat_name: 2, lon_name: 3}:
342 raise ValueError(
343 f"Dimension mapping not correct, ordered {dimension_mapping}"
344 )
346 # Smooth vertical velocity to 2sigma (standard for 0.5 degree).
347 logger.info("Smoothing vertical velocity...")
348 if ensemble_mode:
349 z_arr = gaussian_filter(z_arr, 2, mode="nearest", axes=(3, 4))
350 else:
351 z_arr = gaussian_filter(z_arr, 2, mode="nearest", axes=(2, 3))
353 # Get array index for user specified pressure level.
354 if plev not in XWIND.coord("pressure").points:
355 raise IndexError(f"Can't find plev {plev} in {XWIND.coord('pressure').points}")
357 # Find corresponding pressure level index
358 plev_idx = np.where(XWIND.coord("pressure").points == plev)[0][0]
360 # Initialise cube containing age of air.
361 if ensemble_mode:
362 ageofair_cube = Cube(
363 np.zeros(
364 (
365 len(XWIND.coord("realization").points),
366 len(time),
367 len(lats),
368 len(lons),
369 )
370 ),
371 long_name="age_of_air",
372 units="hours",
373 dim_coords_and_dims=[
374 (XWIND.coord("realization"), 0),
375 (XWIND.coord("time"), 1),
376 (XWIND.coord(lat_name), 2),
377 (XWIND.coord(lon_name), 3),
378 ],
379 )
380 else:
381 ageofair_cube = Cube(
382 np.zeros((len(time), len(lats), len(lons))),
383 long_name="age_of_air",
384 units="hours",
385 dim_coords_and_dims=[
386 (XWIND.coord("time"), 0),
387 (XWIND.coord(lat_name), 1),
388 (XWIND.coord(lon_name), 2),
389 ],
390 )
392 # Unix API for getting set of usable CPUs.
393 # See https://docs.python.org/3/library/os.html#os.cpu_count
394 if multicore: 394 ↛ 395line 394 didn't jump to line 395 because the condition on line 394 was never true
395 num_usable_cores = len(os.sched_getaffinity(0))
396 # Use "spawn" method to avoid warnings before the default is changed in
397 # python 3.14. See the (not very good) warning here:
398 # https://docs.python.org/3/library/multiprocessing.html#contexts-and-start-methods
399 mp_context = multiprocessing.get_context("spawn")
400 pool = mp_context.Pool(num_usable_cores)
402 logger.info("STARTING AOA DIAG...")
403 start = datetime.datetime.now()
405 # Main call for calculating age of air diagnostic
406 if ensemble_mode:
407 for e in range(len(XWIND.coord("realization").points)):
408 logger.info(f"Working on member {e}")
410 # Multiprocessing on each longitude slice
411 func = partial(
412 _aoa_core,
413 np.copy(x_arr[e, :, :, :, :]),
414 np.copy(y_arr[e, :, :, :, :]),
415 np.copy(z_arr[e, :, :, :, :]),
416 np.copy(g_arr[e, :, :, :, :]),
417 lats,
418 lons,
419 dt,
420 plev_idx,
421 timeunit,
422 cyclic,
423 tmpdir.name,
424 )
425 if multicore: 425 ↛ 426line 425 didn't jump to line 426 because the condition on line 425 was never true
426 pool.map(func, range(XWIND.shape[4]))
427 else:
428 # Convert to list to ensure everything is processed.
429 list(map(func, range(XWIND.shape[4])))
431 for i in range(XWIND.shape[4]):
432 file = f"{tmpdir.name}/aoa_frag_{i:04}.npy"
433 ageofair_cube.data[e, :, :, i] = np.load(file)
435 else:
436 # Multiprocessing on each longitude slice
437 func = partial(
438 _aoa_core,
439 np.copy(x_arr),
440 np.copy(y_arr),
441 np.copy(z_arr),
442 np.copy(g_arr),
443 lats,
444 lons,
445 dt,
446 plev_idx,
447 timeunit,
448 cyclic,
449 tmpdir.name,
450 )
451 if multicore: 451 ↛ 452line 451 didn't jump to line 452 because the condition on line 451 was never true
452 pool.map(func, range(XWIND.shape[3]))
453 else:
454 # Convert to list to ensure everything is processed.
455 list(map(func, range(XWIND.shape[3])))
457 for i in range(XWIND.shape[3]):
458 file = f"{tmpdir.name}/aoa_frag_{i:04}.npy"
459 ageofair_cube.data[:, :, i] = np.load(file)
461 if multicore: 461 ↛ 463line 461 didn't jump to line 463 because the condition on line 461 was never true
462 # Wait for tasks to finish then clean up worker processes.
463 pool.terminate()
464 pool.join()
466 # Verbose for time taken to run, and collate tmp ndarrays into final cube, and return
467 logger.info(
468 "AOA DIAG DONE, took %s s",
469 (datetime.datetime.now() - start).total_seconds(),
470 )
472 # Clean tmpdir
473 tmpdir.cleanup()
475 return ageofair_cube