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