Coverage for src/CSET/operators/power_spectrum.py: 95%
98 statements
« prev ^ index » next coverage.py v7.15.3, created at 2026-08-04 08:32 +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 Parameters
47 ----------
48 cubes: Cube | CubeList
49 Field over which to calculate a power spectrum.
51 Returns
52 -------
53 Cube | CubeList:
54 CubeList of power spectra.
55 """
56 out = iris.cube.CubeList()
57 for cube in iter_maybe(cubes):
58 model = cube.attributes.get("model_name")
60 # Check whether data has a realization coord.
61 if cube.coords("realization"):
62 members_and_realizations = [
63 (member, int(member.coord("realization").points[0]))
64 for member in cube.slices_over("realization")
65 ]
66 else:
67 members_and_realizations = [(cube, None)]
69 # Loop over each realization.
70 member_power_spectra = iris.cube.CubeList()
71 for member, realiz in members_and_realizations:
72 # Calculate power spectrum.
73 ps = _power_spectrum(member)
74 # Attach model name if available.
75 if model:
76 ps.attributes["model_name"] = model
77 # Add the correct realization from the parent cube.
78 if realiz is not None:
79 ps.add_aux_coord(
80 iris.coords.AuxCoord(realiz, long_name="realization", units="1")
81 )
82 # Promote to dimension coordinate.
83 ps = iris.util.new_axis(ps, "realization")
84 member_power_spectra.append(ps)
86 # Merge the individual realization cubes into a single cube, then
87 # squeeze off length 1 realization coordinates.
88 combined_cube = member_power_spectra.concatenate_cube()
89 combined_cube = iris.util.squeeze(combined_cube)
90 out.append(combined_cube)
92 # Directly return cube if we only have one.
93 if len(out) == 1:
94 return out[0]
95 else:
96 return out
99def _power_spectrum(cube: iris.cube.Cube) -> iris.cube.Cube:
100 """Calculate power spectrum for a single cube for 1 vertical level at 1 time.
102 Parameters
103 ----------
104 cube: Cube
105 Data to plot as power spectrum.
106 The cubes should cover the same phenomenon i.e. all cubes contain temperature data.
107 We do not support different data such as temperature and humidity in the same CubeList
108 for plotting.
110 Returns
111 -------
112 iris.cube.Cube
113 The power spectrum of the data.
114 To be plotted and aggregation performed after.
116 Raises
117 ------
118 ValueError
119 If the cube doesn't have the right dimensions.
120 TypeError
121 If the cube isn't a Cube.
122 """
123 # Extract time coordinate and convert to datetime
124 time_coord = cube.coord("time")
125 time_points = time_coord.units.num2date(time_coord.points)
127 if cube.ndim == 2:
128 cube_3d = cube.data[np.newaxis, :, :]
129 logger.debug("Adding in new axis for a 2 dimensional cube.")
130 elif cube.ndim == 3:
131 cube_3d = cube.data
132 else:
133 raise ValueError(
134 f"Cube is {cube.ndim} dimensional. Cube should be 2 or 3 dimensional."
135 )
137 # Calculate spectrum
138 ps_array = _DCT_ps(cube_3d)
140 # Make wavenumber comparable between models with different domain sizes and
141 # resolutions. Try to find appropriate spatial coordinates.
142 coord_pairs = (
143 ("projection_x_coordinate", "projection_y_coordinate"),
144 ("grid_latitude", "grid_longitude"),
145 ("latitude", "longitude"),
146 )
147 for x_coord_name, y_coord_name in coord_pairs:
148 try:
149 # Try projection coordinates first (most common for limited area models)
150 x_coord = cube.coord(x_coord_name)
151 y_coord = cube.coord(y_coord_name)
152 except iris.exceptions.CoordinateNotFoundError:
153 continue
154 logger.debug(
155 "Using %s and %s coordinates for grid spacing calculation.",
156 x_coord_name,
157 y_coord_name,
158 )
159 break # Break out of loop if we found usable coords.
160 else:
161 # Raise error if no usable coords found.
162 raise ValueError(
163 "Could not find appropriate spatial coordinates. "
164 "Expected one of: 'projection_x_coordinate'/'projection_y_coordinate', "
165 "'grid_latitude'/'grid_longitude', or 'latitude'/'longitude'."
166 )
168 # Calculate grid spacing.
169 dx = np.abs(np.diff(x_coord.points).mean())
170 dy = np.abs(np.diff(y_coord.points).mean())
171 if "latitude" in x_coord.name():
172 # Convert from degrees to meters. x is lat, y is lon.
173 R_earth = 6371000 # meters
174 lat_mid = np.mean(x_coord.points)
175 dx = dx * np.pi / 180 * R_earth * np.cos(lat_mid * np.pi / 180)
176 dy = dy * np.pi / 180 * R_earth
177 domain_size_km = ((dx * cube_3d.shape[2]) + (dy * cube_3d.shape[1])) / 2 / 1000
179 # Convert wavenumber into physically meaningful wavenumber coordinate in
180 # cycles per km rather than wavenumber per index k.
181 ps_len = ps_array.shape[1]
182 k_indices = np.arange(1, ps_len + 1)
183 physical_wavenumbers = k_indices / domain_size_km # cycles/km
185 # Create a new DimCoord with physical wavenumber
186 physical_wavenumbers_coord = iris.coords.DimCoord(
187 physical_wavenumbers, long_name="physical_wavenumber", units="km-1"
188 )
190 # Calculate wavelength and add as auxiliary coordinate
191 wavelengths = domain_size_km / k_indices # km
192 wavelength_coord = iris.coords.AuxCoord(
193 wavelengths, long_name="wavelength", units="km"
194 )
196 # Ensure power spectrum output is 2D: (time, frequency)
197 if ps_array.ndim == 1: 197 ↛ 198line 197 didn't jump to line 198 because the condition on line 197 was never true
198 ps_array = ps_array[np.newaxis, :]
200 # Prepare time coordinate
201 numeric_time = time_coord.units.date2num(time_points)
202 numeric_time = np.atleast_1d(numeric_time)
204 # Make time coord length match the number of spectra
205 if len(numeric_time) != ps_array.shape[0]: 205 ↛ 206line 205 didn't jump to line 206 because the condition on line 205 was never true
206 numeric_time = np.repeat(numeric_time[0], ps_array.shape[0])
208 new_time_coord = iris.coords.DimCoord(
209 numeric_time,
210 standard_name="time",
211 units=time_coord.units,
212 )
214 # Create output cube with physical coordinates
215 ps_cube = iris.cube.Cube(
216 ps_array,
217 dim_coords_and_dims=[
218 (new_time_coord, 0),
219 (physical_wavenumbers_coord, 1),
220 ],
221 long_name="power_spectral_density",
222 )
224 # Add wavelength as auxiliary coordinate
225 # Realization coordinate is added in _calculate_power_spectrum
226 ps_cube.add_aux_coord(wavelength_coord, data_dims=1)
228 return ps_cube
231def _DCT_ps(y_3d):
232 """Calculate power spectra for regional domains.
234 Parameters
235 ----------
236 y_3d: 3D array
237 3 dimensional array to calculate spectrum for.
238 (2D field data with 3rd dimension of time)
240 Returns
241 -------
242 ps_array:
243 Array of power spectra values calculated for input field (for each time)
245 Method for regional domains:
246 Calculate power spectra over limited area domain using Discrete Cosine Transform (DCT)
247 as described in Denis et al 2002 [Denis_etal_2002]_.
249 References
250 ----------
251 .. [Denis_etal_2002] Bertrand Denis, Jean Côté and René Laprise (2002)
252 "Spectral Decomposition of Two-Dimensional Atmospheric Fields on
253 Limited-Area Domains Using the Discrete Cosine Transform (DCT)"
254 Monthly Weather Review, Vol. 130, 1812-1828
255 doi: https://doi.org/10.1175/1520-0493(2002)130<1812:SDOTDA>2.0.CO;2
256 """
257 Nt, Ny, Nx = y_3d.shape
259 # Max coefficient
260 Nmin = min(Nx - 1, Ny - 1)
262 # Create alpha matrix (of wavenumbers)
263 alpha_matrix = _create_alpha_matrix(Ny, Nx)
265 # Prepare output array
266 ps_array = np.zeros((Nt, Nmin))
268 # Loop over time to get spectrum for each time.
269 for t in range(Nt):
270 y_2d = y_3d[t]
272 # Apply 2D DCT to transform y_3d[t] from physical space to spectral space.
273 # fkk is a 2D array of DCT coefficients, representing the amplitudes of
274 # cosine basis functions at different spatial frequencies.
276 # DCT transform and normalise spectrum to allow comparison between models.
277 fkk = fft.dctn(y_2d, norm="ortho")
279 # calculate variance (energy) of spectral coefficient at each wavenumber pair (k_x, k_y)
280 # as the square of the DCT coefficient, normalised by the total number of grid points (Nx * Ny).
281 sigma_2 = fkk**2 / Nx / Ny
283 # Group ellipses of alphas into the same wavenumber k/Nmin
284 for k in range(1, Nmin + 1):
285 # Define the bounds of the current normalised wavenumber magnitude of bin k
286 alpha = k / Nmin
287 alpha_p1 = (k + 1) / Nmin
289 # Sum up elements matching in bin k and divide by bin size
290 mask_k = np.where((alpha_matrix >= alpha) & (alpha_matrix < alpha_p1))
291 n_coeffs = len(mask_k[0]) # number of coefficients in bin k
292 if n_coeffs > 0: 292 ↛ 297line 292 didn't jump to line 297 because the condition on line 292 was always true
293 ps_array[t, k - 1] = (
294 np.sum(sigma_2[mask_k]) / n_coeffs
295 ) # average power in bin k
296 else:
297 ps_array[t, k - 1] = 0.0
299 return ps_array
302def _create_alpha_matrix(Ny, Nx):
303 """Construct an array of 2D wavenumbers from 2D wavenumber pair.
305 Parameters
306 ----------
307 Ny, Nx:
308 Dimensions of the 2D field for which the power spectra is calculated. Used to
309 create the array of 2D wavenumbers. Each Ny, Nx pair is associated with a
310 single-scale parameter.
312 Returns
313 -------
314 alpha_matrix:
315 normalisation of 2D wavenumber axes, transforming the spectral domain into
316 an elliptic coordinate system.
318 """
319 # Create x_indices: each row is [1, 2, ..., Nx]
320 x_indices = np.tile(np.arange(1, Nx + 1), (Ny, 1))
322 # Create y_indices: each column is [1, 2, ..., Ny]
323 y_indices = np.tile(np.arange(1, Ny + 1).reshape(Ny, 1), (1, Nx))
325 # Compute alpha_matrix
326 alpha_matrix = np.sqrt((x_indices**2) / Nx**2 + (y_indices**2) / Ny**2)
328 return alpha_matrix