Coverage for src/CSET/operators/power_spectrum.py: 91%
156 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-09-11 09:22 +0000
1# © Crown copyright, Met Office (2022-2025) 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"""Operators for calculating power spectra."""
17import logging
19import iris
20import iris.coords
21import iris.cube
22import iris.exceptions
23import numpy as np
24from scipy import fft
26from CSET._common import iter_maybe
28logger = logging.getLogger(__name__)
31def calculate_power_spectrum(
32 cubes: iris.cube.Cube | iris.cube.CubeList,
33):
34 """Wrap power spectrum code.
36 This function is a wrapper that handles power spectrum
37 calculations for both single cubes and cube lists and includes
38 ensembles.
40 The input cube is split into a cube for each model, time,
41 forecast_reference_time and realization. A power spectrum is
42 calculated for each before combining them into one cube ahead of
43 plotting. Attributes (model_name, realization,
44 forecast_reference_time) are retained from original cube.
46 In the case of a CubeList containing multiple models, multiple
47 cases or ensembles, the function iterates through each cube and
48 calculates an individual power spectrum.
50 In the case of a single cube, it directly calculates the power
51 spectrum.
53 [Denis_etal_2002]_
55 Parameters
56 ----------
57 cubes : iris.cube.Cube | iris.cube.CubeList
58 Field over which to calculate a power spectrum.
60 Returns
61 -------
62 iris.cube.Cube | iris.cube.CubeList
63 Power-spectrum cube, or a CubeList for multiple models.
64 """
65 out = iris.cube.CubeList()
67 for input_cube in iter_maybe(cubes):
68 model = input_cube.attributes.get("model_name")
70 # Check whether data has realization and/or
71 # forecast_reference_time coordinates
72 has_realization = bool(input_cube.coords("realization"))
73 has_frt = bool(input_cube.coords("forecast_reference_time"))
75 # Build a list containing:
76 # (cube_slice, realization_value, frt_value)
77 #
78 # Retain realization and frt values and restore to
79 # power-spectrum cubes later.
80 if has_realization and has_frt:
81 # Both realization and forecast_reference_time coords
82 members = []
84 for frt_cube in input_cube.slices_over("forecast_reference_time"):
85 frt = frt_cube.coord("forecast_reference_time").points[0]
87 for member in frt_cube.slices_over("realization"):
88 realiz = member.coord("realization").points[0]
90 members.append((member, realiz, frt))
92 elif has_realization: 92 ↛ 94line 92 didn't jump to line 94 because the condition on line 92 was never true
93 # Only realization coord
94 members = []
96 for member in input_cube.slices_over("realization"):
97 realiz = member.coord("realization").points[0]
99 members.append((member, realiz, None))
101 elif has_frt:
102 # Only forecast_reference_time coord.
103 members = []
105 for frt_cube in input_cube.slices_over("forecast_reference_time"):
106 frt = frt_cube.coord("forecast_reference_time").points[0]
108 members.append((frt_cube, None, frt))
110 else:
111 # Neither realization nor forecast_reference_time coords.
112 members = [(input_cube, None, None)]
114 member_power_spectra = iris.cube.CubeList()
116 # Calculate the spectrum separately for every member/FRT
117 # combination.
118 for member, realiz, frt in members:
119 # Calculate power spectrum
120 ps = _power_spectrum(member)
122 # Attach model name if available
123 if model:
124 ps.attributes["model_name"] = model
126 # Add the correct realization from the parent cube.
127 if realiz is not None:
128 ps.add_aux_coord(
129 iris.coords.AuxCoord(
130 realiz,
131 long_name="realization",
132 units="1",
133 )
134 )
136 ps = iris.util.new_axis(
137 ps,
138 "realization",
139 )
141 # Add the forecast_reference_time from the parent cube.
142 if frt is not None: 142 ↛ 157line 142 didn't jump to line 157 because the condition on line 142 was always true
143 ps.add_aux_coord(
144 iris.coords.AuxCoord(
145 frt,
146 standard_name=("forecast_reference_time"),
147 units=member.coord("forecast_reference_time").units,
148 )
149 )
151 # Promote to dimension coordinate.
152 ps = iris.util.new_axis(
153 ps,
154 "forecast_reference_time",
155 )
157 member_power_spectra.append(ps)
159 # If both realization and FRT vary. Concatenate in stages:
160 #
161 # 1. Concatenate realizations within each FRT.
162 # 2. Concatenate the resulting cubes over FRT.
163 if has_realization and has_frt:
164 # Both realization and forecast_reference_time coords
165 frt_power_spectra = iris.cube.CubeList()
167 frt_values = np.unique(
168 [
169 ps_cube.coord("forecast_reference_time").points[0]
170 for ps_cube in member_power_spectra
171 ]
172 )
174 for frt in frt_values:
175 cubes_for_frt = iris.cube.CubeList(
176 [
177 ps_cube
178 for ps_cube in member_power_spectra
179 if (ps_cube.coord("forecast_reference_time").points[0] == frt)
180 ]
181 )
183 # Within one FRT, realization is the coordinate
184 # that varies.
185 frt_cube = cubes_for_frt.concatenate_cube()
187 frt_power_spectra.append(frt_cube)
189 # If there is only one FRT, no second concatenation is
190 # required.
191 if len(frt_power_spectra) == 1:
192 combined_cube = frt_power_spectra[0]
194 else:
195 # There are multiple FRTs.
196 #
197 # If every FRT has one time point, time and FRT vary
198 # together. Make time an AuxCoord attached to the FRT
199 # dimension so Iris has one concatenation dimension.
200 one_time_per_frt = all(
201 frt_cube.coord("time").shape == (1,)
202 for frt_cube in frt_power_spectra
203 )
205 if one_time_per_frt:
206 cubes_for_frt_concat = iris.cube.CubeList()
208 for frt_cube in frt_power_spectra:
209 frt_cube = frt_cube.copy()
211 time_coord = frt_cube.coord("time").copy()
213 time_dims = frt_cube.coord_dims("time")
215 time_dim = time_dims[0]
217 # Select the only point on the time
218 # dimension. This removes that dimension but
219 # initially leaves time as a scalar coord.
220 index = [slice(None)] * frt_cube.ndim
221 index[time_dim] = 0
223 frt_cube = frt_cube[tuple(index)]
225 frt_cube.remove_coord("time")
227 frt_dims = frt_cube.coord_dims("forecast_reference_time")
229 frt_dim = frt_dims[0]
231 # Attach the one time point to the FRT
232 # dimension.
233 frt_cube.add_aux_coord(
234 time_coord,
235 data_dims=(frt_dim,),
236 )
238 cubes_for_frt_concat.append(frt_cube)
240 combined_cube = cubes_for_frt_concat.concatenate_cube()
242 else:
243 # If all FRT cubes have the same time coordinate,
244 # only FRT varies and normal concatenation should
245 # work.
246 first_time = frt_power_spectra[0].coord("time")
248 matching_times = all(
249 frt_cube.coord("time") == first_time
250 for frt_cube in frt_power_spectra[1:]
251 )
253 if not matching_times: 253 ↛ 262line 253 didn't jump to line 262 because the condition on line 253 was always true
254 raise ValueError(
255 "Cannot combine power spectra: "
256 "multiple forecast reference times "
257 "have different multi-point time "
258 "coordinates."
259 )
261 # Combine individual cubes into single cube.
262 combined_cube = frt_power_spectra.concatenate_cube()
264 else:
265 # Only one of realization or FRT varies, or neither
266 # exists. In those cases only one concatenation axis is
267 # required.
268 # Combine individual cubes into single cube.
269 if len(member_power_spectra) == 1: 269 ↛ 272line 269 didn't jump to line 272 because the condition on line 269 was always true
270 combined_cube = member_power_spectra[0]
271 else:
272 combined_cube = member_power_spectra.concatenate_cube()
274 combined_cube = iris.util.squeeze(combined_cube)
276 out.append(combined_cube)
278 # Directly return cube if only one.
279 if len(out) == 1:
280 return out[0]
282 return out
285def _power_spectrum(cube: iris.cube.Cube) -> iris.cube.Cube:
286 """Calculate power spectrum for a single cube for 1 vertical level at 1 time.
288 Parameters
289 ----------
290 cube: Cube
291 Data to plot as power spectrum.
292 The cubes should cover the same phenomenon i.e. all cubes contain temperature data.
293 We do not support different data such as temperature and humidity in the same CubeList
294 for plotting.
296 Returns
297 -------
298 iris.cube.Cube
299 The power spectrum of the data.
300 To be plotted and aggregation performed after.
302 Raises
303 ------
304 ValueError
305 If the cube doesn't have the right dimensions.
306 TypeError
307 If the cube isn't a Cube.
308 """
309 # Extract time coordinate and convert to datetime
310 time_coord = cube.coord("time")
311 time_points = time_coord.units.num2date(time_coord.points)
313 if cube.ndim == 2:
314 cube_3d = cube.data[np.newaxis, :, :]
315 logger.debug("Adding in new axis for a 2 dimensional cube.")
316 elif cube.ndim == 3:
317 cube_3d = cube.data
318 else:
319 raise ValueError(
320 f"Cube is {cube.ndim} dimensional. Cube should be 2 or 3 dimensional."
321 )
323 # Calculate spectrum
324 ps_array = _DCT_ps(cube_3d)
326 # Make wavenumber comparable between models with different domain sizes and
327 # resolutions. Try to find appropriate spatial coordinates.
328 coord_pairs = (
329 ("projection_x_coordinate", "projection_y_coordinate"),
330 ("grid_latitude", "grid_longitude"),
331 ("latitude", "longitude"),
332 )
333 for x_coord_name, y_coord_name in coord_pairs:
334 try:
335 # Try projection coordinates first (most common for limited area models)
336 x_coord = cube.coord(x_coord_name)
337 y_coord = cube.coord(y_coord_name)
338 except iris.exceptions.CoordinateNotFoundError:
339 continue
340 logger.debug(
341 "Using %s and %s coordinates for grid spacing calculation.",
342 x_coord_name,
343 y_coord_name,
344 )
345 break # Break out of loop if we found usable coords.
346 else:
347 # Raise error if no usable coords found.
348 raise ValueError(
349 "Could not find appropriate spatial coordinates. "
350 "Expected one of: 'projection_x_coordinate'/'projection_y_coordinate', "
351 "'grid_latitude'/'grid_longitude', or 'latitude'/'longitude'."
352 )
354 # Calculate grid spacing.
355 dx = np.abs(np.diff(x_coord.points).mean())
356 dy = np.abs(np.diff(y_coord.points).mean())
357 if "latitude" in x_coord.name():
358 # Convert from degrees to meters. x is lat, y is lon.
359 R_earth = 6371000 # meters
360 lat_mid = np.mean(x_coord.points)
361 dx = dx * np.pi / 180 * R_earth * np.cos(lat_mid * np.pi / 180)
362 dy = dy * np.pi / 180 * R_earth
363 domain_size_km = ((dx * cube_3d.shape[2]) + (dy * cube_3d.shape[1])) / 2 / 1000
365 # Convert wavenumber into physically meaningful wavenumber coordinate in
366 # cycles per km rather than wavenumber per index k.
367 ps_len = ps_array.shape[1]
368 k_indices = np.arange(1, ps_len + 1)
369 physical_wavenumbers = k_indices / domain_size_km # cycles/km
371 # Create a new DimCoord with physical wavenumber
372 physical_wavenumbers_coord = iris.coords.DimCoord(
373 physical_wavenumbers, long_name="physical_wavenumber", units="km-1"
374 )
376 # Calculate wavelength and add as auxiliary coordinate
377 wavelengths = domain_size_km / k_indices # km
378 wavelength_coord = iris.coords.AuxCoord(
379 wavelengths, long_name="wavelength", units="km"
380 )
382 # Ensure power spectrum output is 2D: (time, frequency)
383 if ps_array.ndim == 1: 383 ↛ 384line 383 didn't jump to line 384 because the condition on line 383 was never true
384 ps_array = ps_array[np.newaxis, :]
386 # Prepare time coordinate
387 numeric_time = time_coord.units.date2num(time_points)
388 numeric_time = np.atleast_1d(numeric_time)
390 # Make time coord length match the number of spectra
391 if len(numeric_time) != ps_array.shape[0]: 391 ↛ 392line 391 didn't jump to line 392 because the condition on line 391 was never true
392 numeric_time = np.repeat(numeric_time[0], ps_array.shape[0])
394 new_time_coord = iris.coords.DimCoord(
395 numeric_time,
396 standard_name="time",
397 units=time_coord.units,
398 )
400 # Create output cube with physical coordinates
401 ps_cube = iris.cube.Cube(
402 ps_array,
403 dim_coords_and_dims=[
404 (new_time_coord, 0),
405 (physical_wavenumbers_coord, 1),
406 ],
407 long_name="power_spectral_density",
408 )
410 # Add wavelength as auxiliary coordinate
411 # Realization coordinate is added in _calculate_power_spectrum
412 ps_cube.add_aux_coord(wavelength_coord, data_dims=1)
414 return ps_cube
417def _DCT_ps(y_3d):
418 """Calculate power spectra for regional domains.
420 Parameters
421 ----------
422 y_3d: 3D array
423 3 dimensional array to calculate spectrum for.
424 (2D field data with 3rd dimension of time)
426 Returns
427 -------
428 ps_array:
429 Array of power spectra values calculated for input field (for each time)
430 """
431 Nt, Ny, Nx = y_3d.shape
433 # Max coefficient
434 Nmin = min(Nx - 1, Ny - 1)
436 # Create alpha matrix (of wavenumbers)
437 alpha_matrix = _create_alpha_matrix(Ny, Nx)
439 # Prepare output array
440 ps_array = np.zeros((Nt, Nmin))
442 # Loop over time to get spectrum for each time.
443 for t in range(Nt):
444 y_2d = y_3d[t]
446 # Apply 2D DCT to transform y_3d[t] from physical space to spectral space.
447 # fkk is a 2D array of DCT coefficients, representing the amplitudes of
448 # cosine basis functions at different spatial frequencies.
450 # DCT transform and normalise spectrum to allow comparison between models.
451 fkk = fft.dctn(y_2d, norm="ortho")
453 # calculate variance (energy) of spectral coefficient at each wavenumber pair (k_x, k_y)
454 # as the square of the DCT coefficient, normalised by the total number of grid points (Nx * Ny).
455 sigma_2 = fkk**2 / Nx / Ny
457 # Group ellipses of alphas into the same wavenumber k/Nmin
458 for k in range(1, Nmin + 1):
459 # Define the bounds of the current normalised wavenumber magnitude of bin k
460 alpha = k / Nmin
461 alpha_p1 = (k + 1) / Nmin
463 # Sum up elements matching in bin k and divide by bin size
464 mask_k = np.where((alpha_matrix >= alpha) & (alpha_matrix < alpha_p1))
465 n_coeffs = len(mask_k[0]) # number of coefficients in bin k
466 if n_coeffs > 0: 466 ↛ 471line 466 didn't jump to line 471 because the condition on line 466 was always true
467 ps_array[t, k - 1] = (
468 np.sum(sigma_2[mask_k]) / n_coeffs
469 ) # average power in bin k
470 else:
471 ps_array[t, k - 1] = 0.0
473 return ps_array
476def _create_alpha_matrix(Ny, Nx):
477 """Construct an array of 2D wavenumbers from 2D wavenumber pair.
479 Parameters
480 ----------
481 Ny, Nx:
482 Dimensions of the 2D field for which the power spectra is calculated. Used to
483 create the array of 2D wavenumbers. Each Ny, Nx pair is associated with a
484 single-scale parameter.
486 Returns
487 -------
488 alpha_matrix:
489 normalisation of 2D wavenumber axes, transforming the spectral domain into
490 an elliptic coordinate system.
492 """
493 # Create x_indices: each row is [1, 2, ..., Nx]
494 x_indices = np.tile(np.arange(1, Nx + 1), (Ny, 1))
496 # Create y_indices: each column is [1, 2, ..., Ny]
497 y_indices = np.tile(np.arange(1, Ny + 1).reshape(Ny, 1), (1, Nx))
499 # Compute alpha_matrix
500 alpha_matrix = np.sqrt((x_indices**2) / Nx**2 + (y_indices**2) / Ny**2)
502 return alpha_matrix
505def _coord_dimension(cube, coord_name):
506 """Return the single dimension associated with a coordinate.
508 Parameters
509 ----------
510 cube : iris.cube.Cube
511 Cube containing the coordinate.
512 coord_name : str
513 Name of the coordinate for which to retrieve the associated
514 dimension.
516 Returns
517 -------
518 int
519 Index of the dimension associated with the coordinate.
521 Raises
522 ------
523 ValueError
524 Raised if the coordinate is not associated with exactly one
525 dimension.
527 """
528 coord_dims = cube.coord_dims(coord_name)
530 if len(coord_dims) != 1: 530 ↛ 533line 530 didn't jump to line 533 because the condition on line 530 was always true
531 raise ValueError(f"Expected {coord_name} to be a one-dimensional coordinate.")
533 return coord_dims[0]