Coverage for src/CSET/operators/ageofair.py: 94%
142 statements
« prev ^ index » next coverage.py v7.15.3, created at 2026-08-04 08:32 +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 References
286 ----------
287 .. [Warneretal2023] Warner, J.L., Petch, J., Short, C., Bain, C., 2023. Assessing the impact of an NWP warm-start
288 system on model spin-up over tropical Africa. QJ, 149( 751), pp.621-636. doi:10.1002/qj.4429
289 .. [Warneretal2024] Diagnosing lateral boundary spin-up in regional models using an age of air diagnostic
290 James L. Warner, Charmaine N. Franklin, Belinda Roux, Shaun Cooper, Susan Rennie, Vinod
291 Kumar.
292 Submitted for Quarterly Journal of the Royal Meteorological Society.
294 """
295 # Set up temporary directory to store intermediate age of air slices.
296 tmpdir = tempfile.TemporaryDirectory(dir=os.getenv("CYLC_TASK_WORK_DIR"))
297 logger.info("Made tmpdir %s", tmpdir.name)
299 # Check that all cubes are of same size (will catch different dimension orders too).
300 if not XWIND.shape == YWIND.shape == WWIND.shape == GEOPOT.shape:
301 raise ValueError("Cubes are not the same shape")
303 # Get time units and assign for later
304 if str(XWIND.coord("time").units).startswith("hours since "):
305 timeunit = "hour"
306 else:
307 raise NotImplementedError("Unsupported time base")
309 # Make data non-lazy to speed up code.
310 logger.info("Making data non-lazy...")
311 x_arr = XWIND.data
312 y_arr = YWIND.data
313 z_arr = WWIND.data
314 g_arr = GEOPOT.data
316 # Get coord points
317 lat_name, lon_name = get_cube_yxcoordname(XWIND)
318 lats = XWIND.coord(lat_name).points
319 lons = XWIND.coord(lon_name).points
320 time = XWIND.coord("time").points
322 # Get time spacing of cube to determine whether the spacing in time is the
323 # same throughout the cube. If not, then not supported.
324 dt = XWIND.coord("time").points[1:] - XWIND.coord("time").points[:-1]
325 if np.all(dt == dt[0]):
326 dt = dt[0]
327 else:
328 raise NotImplementedError("Time intervals are not consistent")
330 # Some logic to determine which index each axis is, and check for ensembles.
331 dimension_mapping = {}
332 for coord in XWIND.dim_coords:
333 dim_index = XWIND.coord_dims(coord.name())[0]
334 dimension_mapping[coord.name()] = dim_index
336 if "realization" in dimension_mapping:
337 ensemble_mode = True
338 if dimension_mapping != {
339 "realization": 0,
340 "time": 1,
341 "pressure": 2,
342 lat_name: 3,
343 lon_name: 4,
344 }:
345 raise ValueError(
346 f"Dimension mapping not correct, ordered {dimension_mapping}"
347 )
348 else:
349 ensemble_mode = False
350 if dimension_mapping != {"time": 0, "pressure": 1, lat_name: 2, lon_name: 3}:
351 raise ValueError(
352 f"Dimension mapping not correct, ordered {dimension_mapping}"
353 )
355 # Smooth vertical velocity to 2sigma (standard for 0.5 degree).
356 logger.info("Smoothing vertical velocity...")
357 if ensemble_mode:
358 z_arr = gaussian_filter(z_arr, 2, mode="nearest", axes=(3, 4))
359 else:
360 z_arr = gaussian_filter(z_arr, 2, mode="nearest", axes=(2, 3))
362 # Get array index for user specified pressure level.
363 if plev not in XWIND.coord("pressure").points:
364 raise IndexError(f"Can't find plev {plev} in {XWIND.coord('pressure').points}")
366 # Find corresponding pressure level index
367 plev_idx = np.where(XWIND.coord("pressure").points == plev)[0][0]
369 # Initialise cube containing age of air.
370 if ensemble_mode:
371 ageofair_cube = Cube(
372 np.zeros(
373 (
374 len(XWIND.coord("realization").points),
375 len(time),
376 len(lats),
377 len(lons),
378 )
379 ),
380 long_name="age_of_air",
381 units="hours",
382 dim_coords_and_dims=[
383 (XWIND.coord("realization"), 0),
384 (XWIND.coord("time"), 1),
385 (XWIND.coord(lat_name), 2),
386 (XWIND.coord(lon_name), 3),
387 ],
388 )
389 else:
390 ageofair_cube = Cube(
391 np.zeros((len(time), len(lats), len(lons))),
392 long_name="age_of_air",
393 units="hours",
394 dim_coords_and_dims=[
395 (XWIND.coord("time"), 0),
396 (XWIND.coord(lat_name), 1),
397 (XWIND.coord(lon_name), 2),
398 ],
399 )
401 # Unix API for getting set of usable CPUs.
402 # See https://docs.python.org/3/library/os.html#os.cpu_count
403 if multicore: 403 ↛ 404line 403 didn't jump to line 404 because the condition on line 403 was never true
404 num_usable_cores = len(os.sched_getaffinity(0))
405 # Use "spawn" method to avoid warnings before the default is changed in
406 # python 3.14. See the (not very good) warning here:
407 # https://docs.python.org/3/library/multiprocessing.html#contexts-and-start-methods
408 mp_context = multiprocessing.get_context("spawn")
409 pool = mp_context.Pool(num_usable_cores)
411 logger.info("STARTING AOA DIAG...")
412 start = datetime.datetime.now()
414 # Main call for calculating age of air diagnostic
415 if ensemble_mode:
416 for e in range(len(XWIND.coord("realization").points)):
417 logger.info(f"Working on member {e}")
419 # Multiprocessing on each longitude slice
420 func = partial(
421 _aoa_core,
422 np.copy(x_arr[e, :, :, :, :]),
423 np.copy(y_arr[e, :, :, :, :]),
424 np.copy(z_arr[e, :, :, :, :]),
425 np.copy(g_arr[e, :, :, :, :]),
426 lats,
427 lons,
428 dt,
429 plev_idx,
430 timeunit,
431 cyclic,
432 tmpdir.name,
433 )
434 if multicore: 434 ↛ 435line 434 didn't jump to line 435 because the condition on line 434 was never true
435 pool.map(func, range(XWIND.shape[4]))
436 else:
437 # Convert to list to ensure everything is processed.
438 list(map(func, range(XWIND.shape[4])))
440 for i in range(XWIND.shape[4]):
441 file = f"{tmpdir.name}/aoa_frag_{i:04}.npy"
442 ageofair_cube.data[e, :, :, i] = np.load(file)
444 else:
445 # Multiprocessing on each longitude slice
446 func = partial(
447 _aoa_core,
448 np.copy(x_arr),
449 np.copy(y_arr),
450 np.copy(z_arr),
451 np.copy(g_arr),
452 lats,
453 lons,
454 dt,
455 plev_idx,
456 timeunit,
457 cyclic,
458 tmpdir.name,
459 )
460 if multicore: 460 ↛ 461line 460 didn't jump to line 461 because the condition on line 460 was never true
461 pool.map(func, range(XWIND.shape[3]))
462 else:
463 # Convert to list to ensure everything is processed.
464 list(map(func, range(XWIND.shape[3])))
466 for i in range(XWIND.shape[3]):
467 file = f"{tmpdir.name}/aoa_frag_{i:04}.npy"
468 ageofair_cube.data[:, :, i] = np.load(file)
470 if multicore: 470 ↛ 472line 470 didn't jump to line 472 because the condition on line 470 was never true
471 # Wait for tasks to finish then clean up worker processes.
472 pool.terminate()
473 pool.join()
475 # Verbose for time taken to run, and collate tmp ndarrays into final cube, and return
476 logger.info(
477 "AOA DIAG DONE, took %s s",
478 (datetime.datetime.now() - start).total_seconds(),
479 )
481 # Clean tmpdir
482 tmpdir.cleanup()
484 return ageofair_cube