Add 1d and 2d spectra calculators for Atmos GCM

Along with test cases and a diagnostics group. Refactor old spectra
diagnostics group for Atmos LES.

Co-authored-by: Jia <jiahe.gatech@gmail.com>
Co-authored-by: K Pamnany <kpamnany@gmail.com>

docs/src/APIs/Common/Spectra.md

@@ -7,5 +7,7 @@ CurrentModule = ClimateMachine.Spectra
 ## Methods
 
 ```@docs
-power_spectrum
+power_spectrum_1d
+power_spectrum_2d
+power_spectrum_3d
 ```

docs/src/APIs/Diagnostics/Diagnostics.md

@@ -27,7 +27,7 @@ Diagnostics.setup_atmos_default_perturbations
 Diagnostics.setup_atmos_refstate_perturbations
 Diagnostics.setup_atmos_turbulence_stats
 Diagnostics.setup_atmos_mass_energy_loss
+Diagnostics.setup_atmos_spectra_diagnostics
 Diagnostics.setup_dump_state_diagnostics
 Diagnostics.setup_dump_aux_diagnostics
-Diagnostics.setup_dump_spectra_diagnostics
 ```

experiments/AtmosGCM/GCMDriver/GCMDriver.jl

@@ -39,6 +39,7 @@ using ClimateMachine.TemperatureProfiles
 using ClimateMachine.Thermodynamics
 using ClimateMachine.TurbulenceClosures
 using ClimateMachine.VariableTemplates
+using ClimateMachine.Spectra: compute_gaussian!
 
 using CLIMAParameters
 using CLIMAParameters.Planet
@@ -217,25 +218,47 @@ function config_diagnostics(::Type{FT}, driver_config) where {FT}
     _planet_radius = planet_radius(param_set)::FT
 
     info = driver_config.config_info
+
+    # Setup diagnostic grid(s)
+    nlats = 32
+
+    sinθ, wts = compute_gaussian!(nlats)
+    lats = asin.(sinθ) .* 180 / π
+    lons = 180.0 ./ nlats * collect(FT, 1:1:(2nlats))[:] .- 180.0
+
     boundaries = [
-        FT(-90.0) FT(-180.0) _planet_radius
-        FT(90.0) FT(180.0) FT(_planet_radius + info.domain_height)
+        FT(lats[1]) FT(lons[1]) _planet_radius
+        FT(lats[end]) FT(lons[end]) FT(_planet_radius + info.domain_height)
     ]
-    resolution = (FT(1), FT(1), FT(1000)) # in (deg, deg, m)
+
+    lvls = collect(range(
+        boundaries[1, 3],
+        boundaries[2, 3],
+        step = FT(1000), # in m
+    ))
+
     interpol = ClimateMachine.InterpolationConfiguration(
         driver_config,
         boundaries,
-        resolution,
+        [lats, lons, lvls],
     )
 
+    # Setup diagnostics group(s)
     dgngrp = setup_atmos_default_diagnostics(
         AtmosGCMConfigType(),
         interval,
         driver_config.name,
         interpol = interpol,
     )
 
-    return ClimateMachine.DiagnosticsConfiguration([dgngrp])
+    ds_dgngrp = setup_atmos_spectra_diagnostics(
+        AtmosGCMConfigType(),
+        interval,
+        driver_config.name,
+        interpol = interpol,
+    )
+
+    return ClimateMachine.DiagnosticsConfiguration([dgngrp, ds_dgngrp])
 end
 
 # Entry point
@@ -359,10 +382,16 @@ function main()
         nothing
     end
 
+    check_cons = (
+        ClimateMachine.ConservationCheck("ρ", "100steps", FT(0.0001)),
+        ClimateMachine.ConservationCheck("ρe", "100steps", FT(0.0025)),
+    )
+
     # Run the model
     result = ClimateMachine.invoke!(
         solver_config;
         diagnostics_config = dgn_config,
+        check_cons = check_cons,
         user_callbacks = (cbtmarfilter, cbfilter),
         check_euclidean_distance = true,
     )

experiments/AtmosLES/taylor_green.jl

@@ -168,7 +168,7 @@ function config_diagnostics(
         boundaries,
         resolution,
     )
-    ds_dgngrp = setup_dump_spectra_diagnostics(
+    ds_dgngrp = setup_atmos_spectra_diagnostics(
         AtmosLESConfigType(),
         "0.06ssecs",
         driver_config.name,

src/Common/Spectra/Spectra.jl

@@ -1,18 +1,15 @@
 module Spectra
 
-export power_spectrum
+export power_spectrum_1d, power_spectrum_2d, power_spectrum_3d
 
+using FFTW
 using MPI
-using ClimateMachine
-using ..BalanceLaws
+
 using ..ConfigTypes
-using ..DGMethods
-using ..Mesh.Interpolation
 using ..Mesh.Grids
 using ..MPIStateArrays
-using ..VariableTemplates
-using ..Writers
 
-include("power_spectrum.jl")
+include("power_spectrum_les.jl")
+include("power_spectrum_gcm.jl")
 
 end

src/Common/Spectra/power_spectrum_gcm.jl

@@ -0,0 +1,110 @@
+include("spherical_helper.jl")
+
+"""
+    power_spectrum_1d(::AtmosGCMConfigType, var_grid, z, lat, lon, weight)
+
+Calculates the 1D (zonal) power spectra using the fourier transform at each latitude and level from a 3D velocity field.
+The snapshots of these spectra should be averaged to obtain a time-average.
+The input velocities must be interpolated to a Gaussian grid.
+
+# References
+- A. Wiin-Nielsen (1967) On the annual variation and spectral distribution of atmospheric energy, Tellus, 19:4, 540-559, DOI: 10.3402/tellusa.v19i4.9822
+- Koshyk, J. N., and K. Hamilton, 2001: The Horizontal Kinetic Energy Spectrum and Spectral Budget Simulated by a High-Resolution Troposphere–Stratosphere–Mesosphere GCM. J. Atmos. Sci., 58, 329–348, https://doi.org/10.1175/1520-0469(2001)058<0329:THKESA>2.0.CO;2.
+
+# Arguments
+- var_grid: variable (typically u or v) on a Gausian (lon, lat, z) grid to be transformed
+- z: vertical coordinate (height or pressure)
+- lat: latitude
+- lon: longitude
+- mass_weight: for mass-weighted calculations
+"""
+function power_spectrum_1d(::AtmosGCMConfigType, var_grid, z, lat, lon, weight)
+    num_lev = length(z)
+    num_lat = length(lat)
+    num_lon = length(lon)
+    num_fourier = Int64(num_lon)
+
+    # get number of positive Fourier coefficients incl. 0
+    if mod(num_lon, 2) == 0 # even
+        num_pfourier = div(num_lon, 2)
+    else # odd
+        num_pfourier = div(num_lon, 2) + 1
+    end
+
+    zon_spectrum = zeros(Float64, num_pfourier, num_lat, num_lev)
+    freqs = zeros(Float64, num_pfourier, num_lat, num_lev)
+
+    for k in 1:num_lev
+        for j in 1:num_lat
+            # compute fft frequencies for each latitude
+            x = lon ./ 180.0 .* π
+            dx = (lon[2] - lon[1]) ./ 180.0 .* π
+
+            freqs_ = fftfreq(num_fourier, 1.0 / dx) # 0,+ve freq,-ve freqs (lowest to highest)
+            freqs[:, j, k] = freqs_[1:num_pfourier] .* 2.0 .* π
+
+            # compute the fourier coefficients for all latitudes
+            fourier = fft(var_grid[:, j, k]) # e.g. vcos_grid, ucos_grid
+            fourier = (fourier / num_fourier)
+
+            # convert to energy spectra
+            zon_spectrum[1, j, k] =
+                zon_spectrum[1, j, k] +
+                0.5 * weight[k] * fourier[1] .* conj(fourier[1])
+
+            for m in 2:num_pfourier
+                zon_spectrum[m, j, k] =
+                    zon_spectrum[m, j, k] +
+                    2.0 * weight[k] * fourier[m] * conj(fourier[m]) # factor 2 for neg freq contribution
+            end
+        end
+    end
+    return zon_spectrum, freqs
+end
+
+"""
+    power_spectrum_2d(::AtmosGCMConfigType, var_grid, mass_weight)
+
+- transform variable on grid to the 2d spectral space using fft on latitude circles
+(as for the 1D spectrum) and Legendre polynomials for meridians, and calculate spectra
+
+# Arguments
+- var_grid: variable (typically u or v) on a Gausian (lon, lat, z) grid to be transformed
+- mass_weight: weight for mass-weighted calculations
+
+# References
+Baer, F., 1972: An Alternate Scale Representation of Atmospheric Energy Spectra. J. Atmos. Sci., 29, 649–664, https://doi.org/10.1175/1520-0469(1972)029<0649:AASROA>2.0.CO;2
+"""
+function power_spectrum_2d(::AtmosGCMConfigType, var_grid, mass_weight)
+    #  initialize spherical mesh variables
+    nθ, nd = (size(var_grid, 2), size(var_grid, 3))
+    mesh = SpectralSphericalMesh(nθ, nd)
+    var_spectrum = mesh.var_spectrum
+    var_spherical = mesh.var_spherical
+
+    sinθ, wts = compute_gaussian!(mesh.nθ) # latitude weights using Gaussian quadrature, to orthogonalize Legendre polynomials upon summation
+    mesh.qnm =
+        compute_legendre!(mesh.num_fourier, mesh.num_spherical, sinθ, mesh.nθ) #  normalized associated Legendre polynomials
+
+    for k in 1:(mesh.nd)
+        # apply Gaussian quadrature weights
+        for i in 1:(mesh.nθ)
+            mesh.qwg[:, :, i] .= mesh.qnm[:, :, i] * wts[i] * mass_weight[k]
+        end
+
+        # Transform variable using spherical harmonics
+        var_spherical[:, :, k, :] =
+            trans_grid_to_spherical!(mesh, var_grid[:, :, k]) # var_spherical[m,n,k,sinθ]
+
+        # Calculate energy spectra
+        var_spectrum[:, :, k] =
+            sum(var_spherical[:, :, k, :], dims = 3) .*
+            conj(sum(var_spherical[:, :, k, :], dims = 3))  # var_spectrum[m,n,k]
+    end
+    return var_spectrum[2:end, 2:end, :] .* nθ .^ 2,
+    mesh.wave_numbers,
+    var_spherical,
+    mesh
+end
+
+# TODO: enable mass weighted and vertically integrated calculations

src/Common/Spectra/power_spectrum.jl → src/Common/Spectra/power_spectrum_les.jl

@@ -1,7 +1,5 @@
-using FFTW
-
 """
-    power_spectrum(u, v, w, L, dim, nor)
+    power_spectrum_3d(::AtmosLESConfigType, u, v, w, L, dim, nor)
 
 Calculates the Powerspectrum from the 3D velocity fields `u`, `v`, `w`. Inputs
 need to be equi-spaced and the domain is assumed to be the same size and
@@ -12,7 +10,7 @@ have the same number of points in all directions.
 - dim: number of points
 - nor: normalization factor
 """
-function power_spectrum(u, v, w, L, dim, nor)
+function power_spectrum_3d(::AtmosLESConfigType, u, v, w, L, dim, nor)
     u_fft = fft(u)
     v_fft = fft(v)
     w_fft = fft(w)