Coverage for src/CSET/operators/_spectra.py: 100%

1# © Crown copyright, Met Office (2022-2026) 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. 

14 

15"""Functions to support calculation of power spectra.""" 

16 

17import numpy as np 

18import scipy.fft as fft 

19 

20 

21def DCT_ps(y_3d): 

22 """Calculate power spectra for regional domains. 

23 

24 Parameters 

25 ---------- 

26 y_3d: 3D array 

27 3 dimensional array to calculate spectrum for. 

28 (2D field data with 3rd dimension of time) 

29 

30 Returns: ps_array 

31 Array of power spectra values calculated for input field (for each time) 

32 

33 Method for regional domains: 

34 Calculate power spectra over limited area domain using Discrete Cosine Transform (DCT) 

35 as described in Denis et al 2002 [Denis_etal_2002]_. 

36 

37 References 

38 ---------- 

39 .. [Denis_etal_2002] Bertrand Denis, Jean Côté and René Laprise (2002) 

40 "Spectral Decomposition of Two-Dimensional Atmospheric Fields on 

41 Limited-Area Domains Using the Discrete Cosine Transform (DCT)" 

42 Monthly Weather Review, Vol. 130, 1812-1828 

43 doi: https://doi.org/10.1175/1520-0493(2002)130<1812:SDOTDA>2.0.CO;2 

44 """ 

45 Nt, Ny, Nx = y_3d.shape 

46 

47 # Max coefficient 

48 Nmin = min(Nx - 1, Ny - 1) 

49 

50 # Create alpha matrix (of wavenumbers) 

51 alpha_matrix = create_alpha_matrix(Ny, Nx) 

52 

53 # Prepare output array 

54 ps_array = np.zeros((Nt, Nmin)) 

55 

56 # Loop over time to get spectrum for each time. 

57 for t in range(Nt): 

58 y_2d = y_3d[t] 

59 

60 # Apply 2D DCT to transform y_3d[t] from physical space to spectral space. 

61 # fkk is a 2D array of DCT coefficients, representing the amplitudes of 

62 # cosine basis functions at different spatial frequencies. 

63 

64 # normalise spectrum to allow comparison between models. 

65 fkk = fft.dctn(y_2d, norm="ortho") 

66 

67 # Normalise fkk 

68 fkk = fkk / np.sqrt(Ny * Nx) 

69 

70 # calculate variance of spectral coefficient 

71 sigma_2 = fkk**2 / Nx / Ny 

72 

73 # Group ellipses of alphas into the same wavenumber k/Nmin 

74 for k in range(1, Nmin + 1): 

75 alpha = k / Nmin 

76 alpha_p1 = (k + 1) / Nmin 

77 

78 # Sum up elements matching k 

79 mask_k = np.where((alpha_matrix >= alpha) & (alpha_matrix < alpha_p1)) 

80 ps_array[t, k - 1] = np.sum(sigma_2[mask_k]) 

81 

82 return ps_array 

83 

84 

85def create_alpha_matrix(Ny, Nx): 

86 """Construct an array of 2D wavenumbers from 2D wavenumber pair. 

87 

88 Parameters 

89 ---------- 

90 Ny, Nx: 

91 Dimensions of the 2D field for which the power spectra is calculated. Used to 

92 create the array of 2D wavenumbers. Each Ny, Nx pair is associated with a 

93 single-scale parameter. 

94 

95 Returns: alpha_matrix 

96 normalisation of 2D wavenumber axes, transforming the spectral domain into 

97 an elliptic coordinate system. 

98 

99 """ 

100 # Create x_indices: each row is [1, 2, ..., Nx] 

101 x_indices = np.tile(np.arange(1, Nx + 1), (Ny, 1)) 

102 

103 # Create y_indices: each column is [1, 2, ..., Ny] 

104 y_indices = np.tile(np.arange(1, Ny + 1).reshape(Ny, 1), (1, Nx)) 

105 

106 # Compute alpha_matrix 

107 alpha_matrix = np.sqrt((x_indices**2) / Nx**2 + (y_indices**2) / Ny**2) 

108 

109 return alpha_matrix