Coverage for src/CSET/operators/power_spectrum.py: 96%
116 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-12 16:03 +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(cubes: iris.cube.Cube | iris.cube.CubeList):
32 """Wrap power spectrum code.
34 This function is a wrapper that handles power spectrum
35 calculations for both single cubes and cube lists and includes ensembles.
37 The input cube is split up into a cube for each model,
38 time and realization and a power spectrum calculated for each before
39 combining into one cube ahead of plotting. This is done to retain the
40 model_name attribute correctly for different models.
42 In case of a CubeList (Multiple models and ensembles): It iterates through
43 each cube and calculates an individual power spectrum. In case of a
44 single cube (one model) it directly calculates the power spectrum.
46 Method for regional domains:
47 Calculate power spectra over limited area domain using Discrete Cosine Transform (DCT)
48 as described in Denis et al 2002 [Denis_etal_2002]_.
50 Parameters
51 ----------
52 cubes: Cube | CubeList
53 Field over which to calculate a power spectrum.
55 Returns
56 -------
57 Cube | CubeList:
58 CubeList of power spectra.
59 """
60 out = iris.cube.CubeList()
61 for cube in iter_maybe(cubes):
62 model = cube.attributes.get("model_name")
64 if cube.coords("realization") and cube.coord_dims("forecast_reference_time"):
65 members = []
67 for frt_cube in cube.slices_over("forecast_reference_time"):
68 frt = frt_cube.coord("forecast_reference_time").points[0]
70 for member in frt_cube.slices_over("realization"):
71 realiz = member.coord("realization").points[0]
72 members.append((member, realiz, frt))
74 elif cube.coords("realization"):
75 members = []
77 for member in cube.slices_over("realization"):
78 realiz = member.coord("realization").points[0]
79 members.append((member, realiz, None))
81 elif cube.coords("forecast_reference_time"):
82 members = []
84 for frt_cube in cube.slices_over("forecast_reference_time"):
85 frt = frt_cube.coord("forecast_reference_time").points[0]
86 members.append((frt_cube, None, frt))
88 else:
89 members = [(cube, None, None)]
91 member_power_spectra = iris.cube.CubeList()
93 for member, realiz, frt in members:
94 ps = _power_spectrum(member)
96 if model:
97 ps.attributes["model_name"] = model
99 if realiz is not None:
100 ps.add_aux_coord(
101 iris.coords.AuxCoord(
102 realiz,
103 long_name="realization",
104 units="1",
105 )
106 )
107 ps = iris.util.new_axis(ps, "realization")
109 if frt is not None:
110 # attach forecast_reference_time to the time dimension
112 time_dim = ps.coord_dims("time")[0]
114 ps.add_aux_coord(
115 iris.coords.AuxCoord(
116 frt,
117 long_name="forecast_reference_time",
118 units=member.coord("forecast_reference_time").units,
119 ),
120 data_dims=(time_dim,),
121 )
123 member_power_spectra.append(ps)
125 combined_cube = member_power_spectra.concatenate_cube()
126 combined_cube = iris.util.squeeze(combined_cube)
127 out.append(combined_cube)
129 if len(out) == 1:
130 return out[0]
131 else:
132 return out
135def _power_spectrum(cube: iris.cube.Cube) -> iris.cube.Cube:
136 """Calculate power spectrum for a single cube for 1 vertical level at 1 time.
138 Parameters
139 ----------
140 cube: Cube
141 Data to plot as power spectrum.
142 The cubes should cover the same phenomenon i.e. all cubes contain temperature data.
143 We do not support different data such as temperature and humidity in the same CubeList
144 for plotting.
146 Returns
147 -------
148 iris.cube.Cube
149 The power spectrum of the data.
150 To be plotted and aggregation performed after.
152 Raises
153 ------
154 ValueError
155 If the cube doesn't have the right dimensions.
156 TypeError
157 If the cube isn't a Cube.
158 """
159 # Extract time coordinate and convert to datetime
160 time_coord = cube.coord("time")
161 time_points = time_coord.units.num2date(time_coord.points)
163 if cube.ndim == 2:
164 cube_3d = cube.data[np.newaxis, :, :]
165 logger.debug("Adding in new axis for a 2 dimensional cube.")
166 elif cube.ndim == 3:
167 cube_3d = cube.data
168 else:
169 raise ValueError(
170 f"Cube is {cube.ndim} dimensional. Cube should be 2 or 3 dimensional."
171 )
173 # Calculate spectrum
174 ps_array = _DCT_ps(cube_3d)
176 # Make wavenumber comparable between models with different domain sizes and
177 # resolutions. Try to find appropriate spatial coordinates.
178 coord_pairs = (
179 ("projection_x_coordinate", "projection_y_coordinate"),
180 ("grid_latitude", "grid_longitude"),
181 ("latitude", "longitude"),
182 )
183 for x_coord_name, y_coord_name in coord_pairs:
184 try:
185 # Try projection coordinates first (most common for limited area models)
186 x_coord = cube.coord(x_coord_name)
187 y_coord = cube.coord(y_coord_name)
188 except iris.exceptions.CoordinateNotFoundError:
189 continue
190 logger.debug(
191 "Using %s and %s coordinates for grid spacing calculation.",
192 x_coord_name,
193 y_coord_name,
194 )
195 break # Break out of loop if we found usable coords.
196 else:
197 # Raise error if no usable coords found.
198 raise ValueError(
199 "Could not find appropriate spatial coordinates. "
200 "Expected one of: 'projection_x_coordinate'/'projection_y_coordinate', "
201 "'grid_latitude'/'grid_longitude', or 'latitude'/'longitude'."
202 )
204 # Calculate grid spacing.
205 dx = np.abs(np.diff(x_coord.points).mean())
206 dy = np.abs(np.diff(y_coord.points).mean())
207 if "latitude" in x_coord.name():
208 # Convert from degrees to meters. x is lat, y is lon.
209 R_earth = 6371000 # meters
210 lat_mid = np.mean(x_coord.points)
211 dx = dx * np.pi / 180 * R_earth * np.cos(lat_mid * np.pi / 180)
212 dy = dy * np.pi / 180 * R_earth
213 domain_size_km = ((dx * cube_3d.shape[2]) + (dy * cube_3d.shape[1])) / 2 / 1000
215 # Convert wavenumber into physically meaningful wavenumber coordinate in
216 # cycles per km rather than wavenumber per index k.
217 ps_len = ps_array.shape[1]
218 k_indices = np.arange(1, ps_len + 1)
219 physical_wavenumbers = k_indices / domain_size_km # cycles/km
221 # Create a new DimCoord with physical wavenumber
222 physical_wavenumbers_coord = iris.coords.DimCoord(
223 physical_wavenumbers, long_name="physical_wavenumber", units="km-1"
224 )
226 # Calculate wavelength and add as auxiliary coordinate
227 wavelengths = domain_size_km / k_indices # km
228 wavelength_coord = iris.coords.AuxCoord(
229 wavelengths, long_name="wavelength", units="km"
230 )
232 # Ensure power spectrum output is 2D: (time, frequency)
233 if ps_array.ndim == 1: 233 ↛ 234line 233 didn't jump to line 234 because the condition on line 233 was never true
234 ps_array = ps_array[np.newaxis, :]
236 # Prepare time coordinate
237 numeric_time = time_coord.units.date2num(time_points)
238 numeric_time = np.atleast_1d(numeric_time)
240 # Make time coord length match the number of spectra
241 if len(numeric_time) != ps_array.shape[0]: 241 ↛ 242line 241 didn't jump to line 242 because the condition on line 241 was never true
242 numeric_time = np.repeat(numeric_time[0], ps_array.shape[0])
244 new_time_coord = iris.coords.DimCoord(
245 numeric_time,
246 standard_name="time",
247 units=time_coord.units,
248 )
250 # Create output cube with physical coordinates
251 ps_cube = iris.cube.Cube(
252 ps_array,
253 dim_coords_and_dims=[
254 (new_time_coord, 0),
255 (physical_wavenumbers_coord, 1),
256 ],
257 long_name="power_spectral_density",
258 )
260 # Add wavelength as auxiliary coordinate
261 # Realization coordinate is added in _calculate_power_spectrum
262 ps_cube.add_aux_coord(wavelength_coord, data_dims=1)
264 return ps_cube
267def _DCT_ps(y_3d):
268 """Calculate power spectra for regional domains.
270 Parameters
271 ----------
272 y_3d: 3D array
273 3 dimensional array to calculate spectrum for.
274 (2D field data with 3rd dimension of time)
276 Returns
277 -------
278 ps_array:
279 Array of power spectra values calculated for input field (for each time)
280 """
281 Nt, Ny, Nx = y_3d.shape
283 # Max coefficient
284 Nmin = min(Nx - 1, Ny - 1)
286 # Create alpha matrix (of wavenumbers)
287 alpha_matrix = _create_alpha_matrix(Ny, Nx)
289 # Prepare output array
290 ps_array = np.zeros((Nt, Nmin))
292 # Loop over time to get spectrum for each time.
293 for t in range(Nt):
294 y_2d = y_3d[t]
296 # Apply 2D DCT to transform y_3d[t] from physical space to spectral space.
297 # fkk is a 2D array of DCT coefficients, representing the amplitudes of
298 # cosine basis functions at different spatial frequencies.
300 # DCT transform and normalise spectrum to allow comparison between models.
301 fkk = fft.dctn(y_2d, norm="ortho")
303 # calculate variance (energy) of spectral coefficient at each wavenumber pair (k_x, k_y)
304 # as the square of the DCT coefficient, normalised by the total number of grid points (Nx * Ny).
305 sigma_2 = fkk**2 / Nx / Ny
307 # Group ellipses of alphas into the same wavenumber k/Nmin
308 for k in range(1, Nmin + 1):
309 # Define the bounds of the current normalised wavenumber magnitude of bin k
310 alpha = k / Nmin
311 alpha_p1 = (k + 1) / Nmin
313 # Sum up elements matching in bin k and divide by bin size
314 mask_k = np.where((alpha_matrix >= alpha) & (alpha_matrix < alpha_p1))
315 n_coeffs = len(mask_k[0]) # number of coefficients in bin k
316 if n_coeffs > 0: 316 ↛ 321line 316 didn't jump to line 321 because the condition on line 316 was always true
317 ps_array[t, k - 1] = (
318 np.sum(sigma_2[mask_k]) / n_coeffs
319 ) # average power in bin k
320 else:
321 ps_array[t, k - 1] = 0.0
323 return ps_array
326def _create_alpha_matrix(Ny, Nx):
327 """Construct an array of 2D wavenumbers from 2D wavenumber pair.
329 Parameters
330 ----------
331 Ny, Nx:
332 Dimensions of the 2D field for which the power spectra is calculated. Used to
333 create the array of 2D wavenumbers. Each Ny, Nx pair is associated with a
334 single-scale parameter.
336 Returns
337 -------
338 alpha_matrix:
339 normalisation of 2D wavenumber axes, transforming the spectral domain into
340 an elliptic coordinate system.
342 """
343 # Create x_indices: each row is [1, 2, ..., Nx]
344 x_indices = np.tile(np.arange(1, Nx + 1), (Ny, 1))
346 # Create y_indices: each column is [1, 2, ..., Ny]
347 y_indices = np.tile(np.arange(1, Ny + 1).reshape(Ny, 1), (1, Nx))
349 # Compute alpha_matrix
350 alpha_matrix = np.sqrt((x_indices**2) / Nx**2 + (y_indices**2) / Ny**2)
352 return alpha_matrix