Coverage for src/CSET/operators/power_spectrum.py: 87%
155 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-18 10:30 +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 Parameters
54 ----------
55 cubes : iris.cube.Cube | iris.cube.CubeList
56 Field over which to calculate a power spectrum.
58 Returns
59 -------
60 iris.cube.Cube | iris.cube.CubeList
61 Power-spectrum cube, or a CubeList for multiple models.
62 """
63 out = iris.cube.CubeList()
65 for input_cube in iter_maybe(cubes):
66 model = input_cube.attributes.get("model_name")
68 # Check whether data has realization and/or
69 # forecast_reference_time coordinates
70 has_realization = bool(input_cube.coords("realization"))
71 has_frt = bool(input_cube.coords("forecast_reference_time"))
73 # Build a list containing:
74 # (cube_slice, realization_value, frt_value)
75 #
76 # Retain realization and frt values and restore to
77 # power-spectrum cubes later.
78 if has_realization and has_frt:
79 # Both realization and forecast_reference_time coords
80 members = []
82 for frt_cube in input_cube.slices_over("forecast_reference_time"):
83 frt = frt_cube.coord("forecast_reference_time").points[0]
85 for member in frt_cube.slices_over("realization"):
86 realiz = member.coord("realization").points[0]
88 members.append((member, realiz, frt))
90 elif has_realization: 90 ↛ 92line 90 didn't jump to line 92 because the condition on line 90 was never true
91 # Only realization coord
92 members = []
94 for member in input_cube.slices_over("realization"):
95 realiz = member.coord("realization").points[0]
97 members.append((member, realiz, None))
99 elif has_frt:
100 # Only forecast_reference_time coord.
101 members = []
103 for frt_cube in input_cube.slices_over("forecast_reference_time"):
104 frt = frt_cube.coord("forecast_reference_time").points[0]
106 members.append((frt_cube, None, frt))
108 else:
109 # Neither realization nor forecast_reference_time coords.
110 members = [(input_cube, None, None)]
112 member_power_spectra = iris.cube.CubeList()
114 # Calculate the spectrum separately for every member/FRT
115 # combination.
116 for member, realiz, frt in members:
117 # Calculate power spectrum
118 ps = _power_spectrum(member)
120 # Attach model name if available
121 if model:
122 ps.attributes["model_name"] = model
124 # Add the correct realization from the parent cube.
125 if realiz is not None:
126 ps.add_aux_coord(
127 iris.coords.AuxCoord(
128 realiz,
129 long_name="realization",
130 units="1",
131 )
132 )
134 ps = iris.util.new_axis(
135 ps,
136 "realization",
137 )
139 # Add the forecast_reference_time from the parent cube.
140 if frt is not None: 140 ↛ 155line 140 didn't jump to line 155 because the condition on line 140 was always true
141 ps.add_aux_coord(
142 iris.coords.AuxCoord(
143 frt,
144 standard_name=("forecast_reference_time"),
145 units=member.coord("forecast_reference_time").units,
146 )
147 )
149 # Promote to dimension coordinate.
150 ps = iris.util.new_axis(
151 ps,
152 "forecast_reference_time",
153 )
155 member_power_spectra.append(ps)
157 # If both realization and FRT vary. Concatenate in stages:
158 #
159 # 1. Concatenate realizations within each FRT.
160 # 2. Concatenate the resulting cubes over FRT.
161 if has_realization and has_frt:
162 # Both realization and forecast_reference_time coords
163 frt_power_spectra = iris.cube.CubeList()
165 frt_values = np.unique(
166 [
167 ps_cube.coord("forecast_reference_time").points[0]
168 for ps_cube in member_power_spectra
169 ]
170 )
172 for frt in frt_values:
173 cubes_for_frt = iris.cube.CubeList(
174 [
175 ps_cube
176 for ps_cube in member_power_spectra
177 if (ps_cube.coord("forecast_reference_time").points[0] == frt)
178 ]
179 )
181 # Within one FRT, realization is the coordinate
182 # that varies.
183 frt_cube = cubes_for_frt.concatenate_cube()
185 frt_power_spectra.append(frt_cube)
187 # If there is only one FRT, no second concatenation is
188 # required.
189 if len(frt_power_spectra) == 1:
190 combined_cube = frt_power_spectra[0]
192 else:
193 # There are multiple FRTs.
194 #
195 # If every FRT has one time point, time and FRT vary
196 # together. Make time an AuxCoord attached to the FRT
197 # dimension so Iris has one concatenation dimension.
198 one_time_per_frt = all(
199 frt_cube.coord("time").shape == (1,)
200 for frt_cube in frt_power_spectra
201 )
203 if one_time_per_frt: 203 ↛ 257line 203 didn't jump to line 257 because the condition on line 203 was always true
204 cubes_for_frt_concat = iris.cube.CubeList()
206 for frt_cube in frt_power_spectra:
207 frt_cube = frt_cube.copy()
209 time_coord = frt_cube.coord("time").copy()
211 time_dims = frt_cube.coord_dims("time")
213 if len(time_dims) != 1: 213 ↛ 214line 213 didn't jump to line 214 because the condition on line 213 was never true
214 raise ValueError(
215 "Expected time to be a one-dimensional coordinate."
216 )
218 time_dim = time_dims[0]
220 # Select the only point on the time
221 # dimension. This removes that dimension but
222 # initially leaves time as a scalar coord.
223 index = [slice(None)] * frt_cube.ndim
224 index[time_dim] = 0
226 frt_cube = frt_cube[tuple(index)]
228 frt_cube.remove_coord("time")
230 frt_dims = frt_cube.coord_dims("forecast_reference_time")
232 if len(frt_dims) != 1: 232 ↛ 233line 232 didn't jump to line 233 because the condition on line 232 was never true
233 raise ValueError(
234 "Expected "
235 "forecast_reference_time to "
236 "be a one-dimensional "
237 "coordinate."
238 )
240 frt_dim = frt_dims[0]
242 # Attach the one time point to the FRT
243 # dimension.
244 frt_cube.add_aux_coord(
245 time_coord,
246 data_dims=(frt_dim,),
247 )
249 cubes_for_frt_concat.append(frt_cube)
251 combined_cube = cubes_for_frt_concat.concatenate_cube()
253 else:
254 # If all FRT cubes have the same time coordinate,
255 # only FRT varies and normal concatenation should
256 # work.
257 first_time = frt_power_spectra[0].coord("time")
259 matching_times = all(
260 frt_cube.coord("time") == first_time
261 for frt_cube in frt_power_spectra[1:]
262 )
264 if not matching_times:
265 raise ValueError(
266 "Cannot combine power spectra: "
267 "multiple forecast reference times "
268 "have different multi-point time "
269 "coordinates."
270 )
272 # Combine individual cubes into single cube.
273 combined_cube = frt_power_spectra.concatenate_cube()
275 else:
276 # Only one of realization or FRT varies, or neither
277 # exists. In those cases only one concatenation axis is
278 # required.
279 # Combine individual cubes into single cube.
280 if len(member_power_spectra) == 1: 280 ↛ 283line 280 didn't jump to line 283 because the condition on line 280 was always true
281 combined_cube = member_power_spectra[0]
282 else:
283 combined_cube = member_power_spectra.concatenate_cube()
285 combined_cube = iris.util.squeeze(combined_cube)
287 out.append(combined_cube)
289 # Directly return cube if only one.
290 if len(out) == 1:
291 return out[0]
293 return out
296def _power_spectrum(cube: iris.cube.Cube) -> iris.cube.Cube:
297 """Calculate power spectrum for a single cube for 1 vertical level at 1 time.
299 Parameters
300 ----------
301 cube: Cube
302 Data to plot as power spectrum.
303 The cubes should cover the same phenomenon i.e. all cubes contain temperature data.
304 We do not support different data such as temperature and humidity in the same CubeList
305 for plotting.
307 Returns
308 -------
309 iris.cube.Cube
310 The power spectrum of the data.
311 To be plotted and aggregation performed after.
313 Raises
314 ------
315 ValueError
316 If the cube doesn't have the right dimensions.
317 TypeError
318 If the cube isn't a Cube.
319 """
320 # Extract time coordinate and convert to datetime
321 time_coord = cube.coord("time")
322 time_points = time_coord.units.num2date(time_coord.points)
324 if cube.ndim == 2:
325 cube_3d = cube.data[np.newaxis, :, :]
326 logger.debug("Adding in new axis for a 2 dimensional cube.")
327 elif cube.ndim == 3:
328 cube_3d = cube.data
329 else:
330 raise ValueError(
331 f"Cube is {cube.ndim} dimensional. Cube should be 2 or 3 dimensional."
332 )
334 # Calculate spectrum
335 ps_array = _DCT_ps(cube_3d)
337 # Make wavenumber comparable between models with different domain sizes and
338 # resolutions. Try to find appropriate spatial coordinates.
339 coord_pairs = (
340 ("projection_x_coordinate", "projection_y_coordinate"),
341 ("grid_latitude", "grid_longitude"),
342 ("latitude", "longitude"),
343 )
344 for x_coord_name, y_coord_name in coord_pairs:
345 try:
346 # Try projection coordinates first (most common for limited area models)
347 x_coord = cube.coord(x_coord_name)
348 y_coord = cube.coord(y_coord_name)
349 except iris.exceptions.CoordinateNotFoundError:
350 continue
351 logger.debug(
352 "Using %s and %s coordinates for grid spacing calculation.",
353 x_coord_name,
354 y_coord_name,
355 )
356 break # Break out of loop if we found usable coords.
357 else:
358 # Raise error if no usable coords found.
359 raise ValueError(
360 "Could not find appropriate spatial coordinates. "
361 "Expected one of: 'projection_x_coordinate'/'projection_y_coordinate', "
362 "'grid_latitude'/'grid_longitude', or 'latitude'/'longitude'."
363 )
365 # Calculate grid spacing.
366 dx = np.abs(np.diff(x_coord.points).mean())
367 dy = np.abs(np.diff(y_coord.points).mean())
368 if "latitude" in x_coord.name():
369 # Convert from degrees to meters. x is lat, y is lon.
370 R_earth = 6371000 # meters
371 lat_mid = np.mean(x_coord.points)
372 dx = dx * np.pi / 180 * R_earth * np.cos(lat_mid * np.pi / 180)
373 dy = dy * np.pi / 180 * R_earth
374 domain_size_km = ((dx * cube_3d.shape[2]) + (dy * cube_3d.shape[1])) / 2 / 1000
376 # Convert wavenumber into physically meaningful wavenumber coordinate in
377 # cycles per km rather than wavenumber per index k.
378 ps_len = ps_array.shape[1]
379 k_indices = np.arange(1, ps_len + 1)
380 physical_wavenumbers = k_indices / domain_size_km # cycles/km
382 # Create a new DimCoord with physical wavenumber
383 physical_wavenumbers_coord = iris.coords.DimCoord(
384 physical_wavenumbers, long_name="physical_wavenumber", units="km-1"
385 )
387 # Calculate wavelength and add as auxiliary coordinate
388 wavelengths = domain_size_km / k_indices # km
389 wavelength_coord = iris.coords.AuxCoord(
390 wavelengths, long_name="wavelength", units="km"
391 )
393 # Ensure power spectrum output is 2D: (time, frequency)
394 if ps_array.ndim == 1: 394 ↛ 395line 394 didn't jump to line 395 because the condition on line 394 was never true
395 ps_array = ps_array[np.newaxis, :]
397 # Prepare time coordinate
398 numeric_time = time_coord.units.date2num(time_points)
399 numeric_time = np.atleast_1d(numeric_time)
401 # Make time coord length match the number of spectra
402 if len(numeric_time) != ps_array.shape[0]: 402 ↛ 403line 402 didn't jump to line 403 because the condition on line 402 was never true
403 numeric_time = np.repeat(numeric_time[0], ps_array.shape[0])
405 new_time_coord = iris.coords.DimCoord(
406 numeric_time,
407 standard_name="time",
408 units=time_coord.units,
409 )
411 # Create output cube with physical coordinates
412 ps_cube = iris.cube.Cube(
413 ps_array,
414 dim_coords_and_dims=[
415 (new_time_coord, 0),
416 (physical_wavenumbers_coord, 1),
417 ],
418 long_name="power_spectral_density",
419 )
421 # Add wavelength as auxiliary coordinate
422 # Realization coordinate is added in _calculate_power_spectrum
423 ps_cube.add_aux_coord(wavelength_coord, data_dims=1)
425 return ps_cube
428def _DCT_ps(y_3d):
429 """Calculate power spectra for regional domains.
431 Parameters
432 ----------
433 y_3d: 3D array
434 3 dimensional array to calculate spectrum for.
435 (2D field data with 3rd dimension of time)
437 Returns
438 -------
439 ps_array:
440 Array of power spectra values calculated for input field (for each time)
441 """
442 Nt, Ny, Nx = y_3d.shape
444 # Max coefficient
445 Nmin = min(Nx - 1, Ny - 1)
447 # Create alpha matrix (of wavenumbers)
448 alpha_matrix = _create_alpha_matrix(Ny, Nx)
450 # Prepare output array
451 ps_array = np.zeros((Nt, Nmin))
453 # Loop over time to get spectrum for each time.
454 for t in range(Nt):
455 y_2d = y_3d[t]
457 # Apply 2D DCT to transform y_3d[t] from physical space to spectral space.
458 # fkk is a 2D array of DCT coefficients, representing the amplitudes of
459 # cosine basis functions at different spatial frequencies.
461 # DCT transform and normalise spectrum to allow comparison between models.
462 fkk = fft.dctn(y_2d, norm="ortho")
464 # calculate variance (energy) of spectral coefficient at each wavenumber pair (k_x, k_y)
465 # as the square of the DCT coefficient, normalised by the total number of grid points (Nx * Ny).
466 sigma_2 = fkk**2 / Nx / Ny
468 # Group ellipses of alphas into the same wavenumber k/Nmin
469 for k in range(1, Nmin + 1):
470 # Define the bounds of the current normalised wavenumber magnitude of bin k
471 alpha = k / Nmin
472 alpha_p1 = (k + 1) / Nmin
474 # Sum up elements matching in bin k and divide by bin size
475 mask_k = np.where((alpha_matrix >= alpha) & (alpha_matrix < alpha_p1))
476 n_coeffs = len(mask_k[0]) # number of coefficients in bin k
477 if n_coeffs > 0: 477 ↛ 482line 477 didn't jump to line 482 because the condition on line 477 was always true
478 ps_array[t, k - 1] = (
479 np.sum(sigma_2[mask_k]) / n_coeffs
480 ) # average power in bin k
481 else:
482 ps_array[t, k - 1] = 0.0
484 return ps_array
487def _create_alpha_matrix(Ny, Nx):
488 """Construct an array of 2D wavenumbers from 2D wavenumber pair.
490 Parameters
491 ----------
492 Ny, Nx:
493 Dimensions of the 2D field for which the power spectra is calculated. Used to
494 create the array of 2D wavenumbers. Each Ny, Nx pair is associated with a
495 single-scale parameter.
497 Returns
498 -------
499 alpha_matrix:
500 normalisation of 2D wavenumber axes, transforming the spectral domain into
501 an elliptic coordinate system.
503 """
504 # Create x_indices: each row is [1, 2, ..., Nx]
505 x_indices = np.tile(np.arange(1, Nx + 1), (Ny, 1))
507 # Create y_indices: each column is [1, 2, ..., Ny]
508 y_indices = np.tile(np.arange(1, Ny + 1).reshape(Ny, 1), (1, Nx))
510 # Compute alpha_matrix
511 alpha_matrix = np.sqrt((x_indices**2) / Nx**2 + (y_indices**2) / Ny**2)
513 return alpha_matrix