Coverage for src/CSET/operators/power_spectrum.py: 91%
157 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-09-08 11:47 +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 print("CUBE in _power_spectrum ", cube)
315 if cube.ndim == 2:
316 cube_3d = cube.data[np.newaxis, :, :]
317 logger.debug("Adding in new axis for a 2 dimensional cube.")
318 elif cube.ndim == 3:
319 cube_3d = cube.data
320 else:
321 raise ValueError(
322 f"Cube is {cube.ndim} dimensional. Cube should be 2 or 3 dimensional."
323 )
325 # Calculate spectrum
326 ps_array = _DCT_ps(cube_3d)
328 # Make wavenumber comparable between models with different domain sizes and
329 # resolutions. Try to find appropriate spatial coordinates.
330 coord_pairs = (
331 ("projection_x_coordinate", "projection_y_coordinate"),
332 ("grid_latitude", "grid_longitude"),
333 ("latitude", "longitude"),
334 )
335 for x_coord_name, y_coord_name in coord_pairs:
336 try:
337 # Try projection coordinates first (most common for limited area models)
338 x_coord = cube.coord(x_coord_name)
339 y_coord = cube.coord(y_coord_name)
340 except iris.exceptions.CoordinateNotFoundError:
341 continue
342 logger.debug(
343 "Using %s and %s coordinates for grid spacing calculation.",
344 x_coord_name,
345 y_coord_name,
346 )
347 break # Break out of loop if we found usable coords.
348 else:
349 # Raise error if no usable coords found.
350 raise ValueError(
351 "Could not find appropriate spatial coordinates. "
352 "Expected one of: 'projection_x_coordinate'/'projection_y_coordinate', "
353 "'grid_latitude'/'grid_longitude', or 'latitude'/'longitude'."
354 )
356 # Calculate grid spacing.
357 dx = np.abs(np.diff(x_coord.points).mean())
358 dy = np.abs(np.diff(y_coord.points).mean())
359 if "latitude" in x_coord.name():
360 # Convert from degrees to meters. x is lat, y is lon.
361 R_earth = 6371000 # meters
362 lat_mid = np.mean(x_coord.points)
363 dx = dx * np.pi / 180 * R_earth * np.cos(lat_mid * np.pi / 180)
364 dy = dy * np.pi / 180 * R_earth
365 domain_size_km = ((dx * cube_3d.shape[2]) + (dy * cube_3d.shape[1])) / 2 / 1000
367 # Convert wavenumber into physically meaningful wavenumber coordinate in
368 # cycles per km rather than wavenumber per index k.
369 ps_len = ps_array.shape[1]
370 k_indices = np.arange(1, ps_len + 1)
371 physical_wavenumbers = k_indices / domain_size_km # cycles/km
373 # Create a new DimCoord with physical wavenumber
374 physical_wavenumbers_coord = iris.coords.DimCoord(
375 physical_wavenumbers, long_name="physical_wavenumber", units="km-1"
376 )
378 # Calculate wavelength and add as auxiliary coordinate
379 wavelengths = domain_size_km / k_indices # km
380 wavelength_coord = iris.coords.AuxCoord(
381 wavelengths, long_name="wavelength", units="km"
382 )
384 # Ensure power spectrum output is 2D: (time, frequency)
385 if ps_array.ndim == 1: 385 ↛ 386line 385 didn't jump to line 386 because the condition on line 385 was never true
386 ps_array = ps_array[np.newaxis, :]
388 # Prepare time coordinate
389 numeric_time = time_coord.units.date2num(time_points)
390 numeric_time = np.atleast_1d(numeric_time)
392 # Make time coord length match the number of spectra
393 if len(numeric_time) != ps_array.shape[0]: 393 ↛ 394line 393 didn't jump to line 394 because the condition on line 393 was never true
394 numeric_time = np.repeat(numeric_time[0], ps_array.shape[0])
396 new_time_coord = iris.coords.DimCoord(
397 numeric_time,
398 standard_name="time",
399 units=time_coord.units,
400 )
402 # Create output cube with physical coordinates
403 ps_cube = iris.cube.Cube(
404 ps_array,
405 dim_coords_and_dims=[
406 (new_time_coord, 0),
407 (physical_wavenumbers_coord, 1),
408 ],
409 long_name="power_spectral_density",
410 )
412 # Add wavelength as auxiliary coordinate
413 # Realization coordinate is added in _calculate_power_spectrum
414 ps_cube.add_aux_coord(wavelength_coord, data_dims=1)
416 return ps_cube
419def _DCT_ps(y_3d):
420 """Calculate power spectra for regional domains.
422 Parameters
423 ----------
424 y_3d: 3D array
425 3 dimensional array to calculate spectrum for.
426 (2D field data with 3rd dimension of time)
428 Returns
429 -------
430 ps_array:
431 Array of power spectra values calculated for input field (for each time)
432 """
433 Nt, Ny, Nx = y_3d.shape
435 # Max coefficient
436 Nmin = min(Nx - 1, Ny - 1)
438 # Create alpha matrix (of wavenumbers)
439 alpha_matrix = _create_alpha_matrix(Ny, Nx)
441 # Prepare output array
442 ps_array = np.zeros((Nt, Nmin))
444 # Loop over time to get spectrum for each time.
445 for t in range(Nt):
446 y_2d = y_3d[t]
448 # Apply 2D DCT to transform y_3d[t] from physical space to spectral space.
449 # fkk is a 2D array of DCT coefficients, representing the amplitudes of
450 # cosine basis functions at different spatial frequencies.
452 # DCT transform and normalise spectrum to allow comparison between models.
453 fkk = fft.dctn(y_2d, norm="ortho")
455 # calculate variance (energy) of spectral coefficient at each wavenumber pair (k_x, k_y)
456 # as the square of the DCT coefficient, normalised by the total number of grid points (Nx * Ny).
457 sigma_2 = fkk**2 / Nx / Ny
459 # Group ellipses of alphas into the same wavenumber k/Nmin
460 for k in range(1, Nmin + 1):
461 # Define the bounds of the current normalised wavenumber magnitude of bin k
462 alpha = k / Nmin
463 alpha_p1 = (k + 1) / Nmin
465 # Sum up elements matching in bin k and divide by bin size
466 mask_k = np.where((alpha_matrix >= alpha) & (alpha_matrix < alpha_p1))
467 n_coeffs = len(mask_k[0]) # number of coefficients in bin k
468 if n_coeffs > 0: 468 ↛ 473line 468 didn't jump to line 473 because the condition on line 468 was always true
469 ps_array[t, k - 1] = (
470 np.sum(sigma_2[mask_k]) / n_coeffs
471 ) # average power in bin k
472 else:
473 ps_array[t, k - 1] = 0.0
475 return ps_array
478def _create_alpha_matrix(Ny, Nx):
479 """Construct an array of 2D wavenumbers from 2D wavenumber pair.
481 Parameters
482 ----------
483 Ny, Nx:
484 Dimensions of the 2D field for which the power spectra is calculated. Used to
485 create the array of 2D wavenumbers. Each Ny, Nx pair is associated with a
486 single-scale parameter.
488 Returns
489 -------
490 alpha_matrix:
491 normalisation of 2D wavenumber axes, transforming the spectral domain into
492 an elliptic coordinate system.
494 """
495 # Create x_indices: each row is [1, 2, ..., Nx]
496 x_indices = np.tile(np.arange(1, Nx + 1), (Ny, 1))
498 # Create y_indices: each column is [1, 2, ..., Ny]
499 y_indices = np.tile(np.arange(1, Ny + 1).reshape(Ny, 1), (1, Nx))
501 # Compute alpha_matrix
502 alpha_matrix = np.sqrt((x_indices**2) / Nx**2 + (y_indices**2) / Ny**2)
504 return alpha_matrix
507def _coord_dimension(cube, coord_name):
508 """Return the single dimension associated with a coordinate.
510 Parameters
511 ----------
512 cube : iris.cube.Cube
513 Cube containing the coordinate.
514 coord_name : str
515 Name of the coordinate for which to retrieve the associated
516 dimension.
518 Returns
519 -------
520 int
521 Index of the dimension associated with the coordinate.
523 Raises
524 ------
525 ValueError
526 Raised if the coordinate is not associated with exactly one
527 dimension.
529 """
530 coord_dims = cube.coord_dims(coord_name)
532 if len(coord_dims) != 1: 532 ↛ 535line 532 didn't jump to line 535 because the condition on line 532 was always true
533 raise ValueError(f"Expected {coord_name} to be a one-dimensional coordinate.")
535 return coord_dims[0]