diff --git a/docs/src/APIs/Common/Spectra.md b/docs/src/APIs/Common/Spectra.md index 12fbb09bfd4..2c1f5942d02 100644 --- a/docs/src/APIs/Common/Spectra.md +++ b/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 ``` diff --git a/docs/src/APIs/Diagnostics/Diagnostics.md b/docs/src/APIs/Diagnostics/Diagnostics.md index 93338d7a4ab..a77f0c9e77a 100644 --- a/docs/src/APIs/Diagnostics/Diagnostics.md +++ b/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 ``` diff --git a/experiments/AtmosGCM/GCMDriver/GCMDriver.jl b/experiments/AtmosGCM/GCMDriver/GCMDriver.jl index 4701018e748..c6d094a6f71 100755 --- a/experiments/AtmosGCM/GCMDriver/GCMDriver.jl +++ b/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,17 +218,32 @@ 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, @@ -235,7 +251,14 @@ function config_diagnostics(::Type{FT}, driver_config) where {FT} 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, ) diff --git a/experiments/AtmosLES/taylor_green.jl b/experiments/AtmosLES/taylor_green.jl index d5fae8e8a82..3ecad7b8d79 100755 --- a/experiments/AtmosLES/taylor_green.jl +++ b/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, diff --git a/src/Arrays/MPIStateArrays.jl b/src/Arrays/MPIStateArrays.jl index 2015e688e1d..621dc0a23eb 100644 --- a/src/Arrays/MPIStateArrays.jl +++ b/src/Arrays/MPIStateArrays.jl @@ -278,7 +278,7 @@ function Base.similar( ::Type{A}, ::Type{FT}; vars::Type{V} = vars(Q), - nstate = size(Q.data)[2] + nstate = size(Q.data)[2], ) where {A <: AbstractArray, FT <: Number, V} MPIStateArray{FT, V}( Q.mpicomm, @@ -300,7 +300,7 @@ function Base.similar( Q::MPIStateArray{FT}, ::Type{A}; vars::Type{V} = vars(Q), - nstate = size(Q.data)[2] + nstate = size(Q.data)[2], ) where {A <: AbstractArray, FT <: Number, V} similar(Q, A, FT; vars = V, nstate = nstate) end @@ -308,14 +308,14 @@ function Base.similar( Q::MPIStateArray, ::Type{FT}; vars::Type{V} = vars(Q), - nstate = size(Q.data)[2] + nstate = size(Q.data)[2], ) where {FT <: Number, V} similar(Q, typeof(Q.data), FT; vars = V, nstate = nstate) end function Base.similar( Q::MPIStateArray{FT}; vars::Type{V} = vars(Q), - nstate = size(Q.data)[2] + nstate = size(Q.data)[2], ) where {FT, V} similar(Q, FT; vars = V, nstate = nstate) end diff --git a/src/Atmos/Model/AtmosModel.jl b/src/Atmos/Model/AtmosModel.jl index b6fec433214..9687f06919c 100644 --- a/src/Atmos/Model/AtmosModel.jl +++ b/src/Atmos/Model/AtmosModel.jl @@ -530,7 +530,8 @@ function. Contributions from subcomponents are then assembled (pointwise). t::Real, ) ν, D_t, τ = turbulence_tensors(atmos, state, diffusive, aux, t) - ν, D_t, τ = sponge_viscosity_modifier(atmos, atmos.viscoussponge, ν, D_t, τ, aux) + ν, D_t, τ = + sponge_viscosity_modifier(atmos, atmos.viscoussponge, ν, D_t, τ, aux) d_h_tot = -D_t .* diffusive.∇h_tot flux_second_order!(atmos, flux, state, τ, d_h_tot) flux_second_order!(atmos.moisture, flux, state, diffusive, aux, t, D_t) diff --git a/src/Common/Spectra/Spectra.jl b/src/Common/Spectra/Spectra.jl index 0054d36c789..bc3c7a7328d 100644 --- a/src/Common/Spectra/Spectra.jl +++ b/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 diff --git a/src/Common/Spectra/power_spectrum_gcm.jl b/src/Common/Spectra/power_spectrum_gcm.jl new file mode 100644 index 00000000000..7433e955583 --- /dev/null +++ b/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 diff --git a/src/Common/Spectra/power_spectrum.jl b/src/Common/Spectra/power_spectrum_les.jl similarity index 95% rename from src/Common/Spectra/power_spectrum.jl rename to src/Common/Spectra/power_spectrum_les.jl index e69becb5dce..a8276967278 100644 --- a/src/Common/Spectra/power_spectrum.jl +++ b/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) diff --git a/src/Common/Spectra/spherical_helper.jl b/src/Common/Spectra/spherical_helper.jl new file mode 100644 index 00000000000..e8e28514daf --- /dev/null +++ b/src/Common/Spectra/spherical_helper.jl @@ -0,0 +1,257 @@ +# helper functions for spherical harmonics spectra + +""" + SpectralSphericalMesh + +Struct with mesh information. +""" +mutable struct SpectralSphericalMesh + # grid info + num_fourier::Int64 + num_spherical::Int64 + nλ::Int64 + nθ::Int64 + nd::Int64 + Δλ::Float64 + qwg::Array{Float64, 3} + qnm::Array{Float64, 3} # n,m coordinates + wave_numbers::Array{Int64, 2} + + # variables + var_grid::Array{Float64, 3} + var_fourier::Array{ComplexF64, 3} + var_spherical::Array{ComplexF64, 4} + var_spectrum::Array{Float64, 3} + +end +function SpectralSphericalMesh(nθ::Int64, nd::Int64) + nλ = 2nθ + Δλ = 2π / nλ + + num_fourier = Int64((2 * nθ - 1) / 3) # number of truncated zonal wavenumbers (m): minimum truncation given nθ - e.g.: nlat = 32 -> T21 (can change manually for more a severe truncation) + num_spherical = Int64(num_fourier + 1) # number of total wavenumbers (n) + + radius = Float64(6371000) + wave_numbers = compute_wave_numbers(num_fourier, num_spherical) + + qwg = zeros(Float64, num_fourier + 1, num_spherical + 1, nθ) + qnm = zeros(Float64, num_fourier + 1, num_spherical + 2, nθ) + + var_fourier = zeros(Complex{Float64}, nλ, nθ, nd) + var_grid = zeros(Float64, nλ, nθ, nd) + nθ_half = div(nθ, 2) + var_spherical = + zeros(Complex{Float64}, num_fourier + 1, num_spherical + 1, nd, nθ_half) + var_spectrum = zeros(Float64, num_fourier + 1, num_spherical + 1, nd) + + SpectralSphericalMesh( + num_fourier, + num_spherical, + nλ, + nθ, + nd, + Δλ, + qwg, + qnm, + wave_numbers, + var_grid, + var_fourier, + var_spherical, + var_spectrum, + ) +end + +""" + compute_legendre!(num_fourier, num_spherical, sinθ, nθ) + +Normalized associated Legendre polynomials, P_{m,l} = qnm + +# Arguments: +- num_fourier +- num_spherical +- sinθ +- nθ + +# References: +- Ehrendorfer, M. (2011) Spectral Numerical Weather Prediction Models, Appendix B, Society for Industrial and Applied Mathematics +- Winch, D. (2007) Spherical harmonics, in Encyclopedia of Geomagnetism and Paleomagnetism, Eds Gubbins D. and Herrero-Bervera, E., Springer + +# Details (using notation and Eq. references from Ehrendorfer, 2011): + + l=0,1...∞ and m = -l, -l+1, ... l-1, l + + P_{0,0} = 1, such that 1/4π ∫∫YYdS = δ (where Y = spherical harmonics, S = domain surface area) + P_{m,m} = sqrt((2m+1)/2m) cosθ P_{m-1m-1} + P_{m+1,m} = sqrt(2m+3) sinθ P_{m m} + sqrt((l^2-m^2)/(4l^2-1))P_{l,m} = P_{l-1, m} - sqrt(((l-1)^2-m^2)/(4(l-1)^2 - 1))P_{l-2,m} + + THe normalization assures that 1/2 ∫_{-1}^1 P_{l,m}(sinθ) P_{n,m}(sinθ) dsinθ = δ_{n,l} + + Julia index starts with 1, so qnm[m+1,l+1] = P_l^m +""" +function compute_legendre!(num_fourier, num_spherical, sinθ, nθ) + qnm = zeros(Float64, num_fourier + 1, num_spherical + 2, nθ) + + cosθ = sqrt.(1 .- sinθ .^ 2) + ε = zeros(Float64, num_fourier + 1, num_spherical + 2) + + qnm[1, 1, :] .= 1 + for m in 1:num_fourier + qnm[m + 1, m + 1, :] = -sqrt((2m + 1) / (2m)) .* cosθ .* qnm[m, m, :] # Eq. B.20 + qnm[m, m + 1, :] = sqrt(2m + 1) * sinθ .* qnm[m, m, :] # Eq. B.22 + end + qnm[num_fourier + 1, num_fourier + 2, :] = + sqrt(2 * (num_fourier + 2)) * sinθ .* + qnm[num_fourier + 1, num_fourier + 1, :] + + for m in 0:num_fourier + for l in (m + 2):(num_spherical + 1) + ε1 = sqrt(((l - 1)^2 - m^2) ./ (4 * (l - 1)^2 - 1)) + ε2 = sqrt((l^2 - m^2) ./ (4 * l^2 - 1)) + qnm[m + 1, l + 1, :] = + (sinθ .* qnm[m + 1, l, :] - ε1 * qnm[m + 1, l - 1, :]) / ε2 # Eq. B.18 + end + end + + return qnm[:, 1:(num_spherical + 1), :] +end + +""" + compute_gaussian!(n) + +Compute sin(latitude) and the weight factors for Gaussian integration + +# Arguments +- n: number of latitudes + +# References +- Ehrendorfer, M., Spectral Numerical Weather Prediction Models, Appendix B, Society for Industrial and Applied Mathematics, 2011 + +# Details (following notation from Ehrendorfer, 2011): + Pn(x) is an odd function + solve half of the n roots and weightes of Pn(x) # n = 2n_half + P_{-1}(x) = 0 + P_0(x) = 1 + P_1(x) = x + nP_n(x) = (2n-1)xP_{n-1}(x) - (n-1)P_{n-2}(x) + P'_n(x) = n/(x^2-1)(xP_{n}(x) - P_{n-1}(x)) + x -= P_n(x)/P'_{n}() + Initial guess xi^{0} = cos(π(i-0.25)/(n+0.5)) + wi = 2/(1-xi^2)/P_n'(xi)^2 +""" +function compute_gaussian!(n) + itermax = 10000 + tol = 1.0e-15 + + sinθ = zeros(Float64, n) + wts = zeros(Float64, n) + + n_half = Int64(n / 2) + for i in 1:n_half + dp = 0.0 + z = cos(pi * (i - 0.25) / (n + 0.5)) + for iter in 1:itermax + p2 = 0.0 + p1 = 1.0 + + for j in 1:n + p3 = p2 # Pj-2 + p2 = p1 # Pj-1 + p1 = ((2.0 * j - 1.0) * z * p2 - (j - 1.0) * p3) / j #Pj + end + # P'_n + dp = n * (z * p1 - p2) / (z * z - 1.0) + z1 = z + z = z1 - p1 / dp + if (abs(z - z1) <= tol) + break + end + if iter == itermax + @error("Compute_Gaussian! does not converge!") + end + end + + sinθ[i], sinθ[n - i + 1], = -z, z + wts[i] = wts[n - i + 1] = 2.0 / ((1.0 - z * z) * dp * dp) + end + + return sinθ, wts +end + +""" + trans_grid_to_spherical!(mesh::SpectralSphericalMesh, pfield::Array{Float64,2}) + +Transforms a variable on a Gaussian grid (pfield[nλ, nθ]) into the spherical harmonics domain (var_spherical2d[num_fourier+1, num_spherical+1]) + +Here λ = longitude, θ = latitude, η = sinθ, m = zonal wavenumber, n = total wavenumber: +var_spherical2d = F_{m,n} # Output variable in spectral space (Complex{Float64}[num_fourier+1, num_spherical+1]) +qwg = P_{m,n}(η)w(η) # Weighted Legendre polynomials (Float64[num_fourier+1, num_spherical+1, nθ]) +var_fourier2d = g_{m, θ} # Untruncated Fourier transformation (Complex{Float64} [nλ, nθ]) +pfield = F(λ, η) # Input variable on Gaussian grid Float64[nλ, nθ] + +# Arguments +- mesh: struct with mesh information +- pfield: variable on Gaussian grid to be transformed + +# References +- Ehrendorfer, M., Spectral Numerical Weather Prediction Models, Appendix B, Society for Industrial and Applied Mathematics, 2011 +- 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 +""" +function trans_grid_to_spherical!( + mesh::SpectralSphericalMesh, + pfield::Array{Float64, 2}, +) + + num_fourier, num_spherical = mesh.num_fourier, mesh.num_spherical + var_fourier2d, var_spherical2d = + mesh.var_fourier[:, :, 1] * 0.0, mesh.var_spherical[:, :, 1, :] * 0.0 + nλ, nθ, nd = mesh.nλ, mesh.nθ, mesh.nd + + # Retrieve weighted Legendre polynomials + qwg = mesh.qwg # qwg[m,n,nθ] + + # Fourier transformation + for j in 1:nθ + var_fourier2d[:, j] = fft(pfield[:, j], 1) / nλ + end + + # Complete spherical harmonic transformation + @assert(nθ % 2 == 0) + nθ_half = div(nθ, 2) + for m in 1:(num_fourier + 1) + for n in m:num_spherical + var_fourier2d_t = transpose(var_fourier2d[m, :]) # truncates var_fourier(nlon, nhlat) to (nfourier,nlat) + if (n - m) % 2 == 0 + var_spherical2d[m, n, :] .= + ( + var_fourier2d_t[1:nθ_half] .+ + var_fourier2d_t[nθ:-1:(nθ_half + 1)] + ) .* qwg[m, n, 1:nθ_half] ./ 2.0 + else + var_spherical2d[m, n, :] .= + ( + var_fourier2d_t[1:nθ_half] .- + var_fourier2d_t[nθ:-1:(nθ_half + 1)] + ) .* qwg[m, n, 1:nθ_half] ./ 2.0 + end + end + end + + return var_spherical2d +end + +function compute_wave_numbers(num_fourier::Int64, num_spherical::Int64) + """ + See wave_numers[i,j] saves the wave number of this basis + """ + wave_numbers = zeros(Int64, num_fourier + 1, num_spherical + 1) + + for m in 0:num_fourier + for n in m:num_spherical + wave_numbers[m + 1, n + 1] = n + end + end + + return wave_numbers + +end diff --git a/src/Diagnostics/Diagnostics.jl b/src/Diagnostics/Diagnostics.jl index cec1cbd38a7..49c09343096 100644 --- a/src/Diagnostics/Diagnostics.jl +++ b/src/Diagnostics/Diagnostics.jl @@ -12,9 +12,9 @@ export DiagnosticsGroup, setup_atmos_refstate_perturbations, setup_atmos_turbulence_stats, setup_atmos_mass_energy_loss, + setup_atmos_spectra_diagnostics, setup_dump_state_diagnostics, - setup_dump_aux_diagnostics, - setup_dump_spectra_diagnostics + setup_dump_aux_diagnostics using CUDA using Dates diff --git a/src/Diagnostics/atmos_gcm_spectra.jl b/src/Diagnostics/atmos_gcm_spectra.jl new file mode 100644 index 00000000000..d80b9f99f28 --- /dev/null +++ b/src/Diagnostics/atmos_gcm_spectra.jl @@ -0,0 +1,155 @@ +# Spectrum calculator for AtmosGCM + +struct AtmosGCMSpectraDiagnosticsParams <: DiagnosticsGroupParams + nor::Float64 +end + +""" + setup_atmos_spectra_diagnostics( + ::AtmosGCMConfigType, + interval::String, + out_prefix::String; + writer = NetCDFWriter(), + interpol = nothing, + nor = 1.0, + ) + +Create and return a `DiagnosticsGroup` containing a diagnostic that dumps +the spectrum at the specified `interval` after the velocity fields have +been interpolated, into NetCDF files prefixed by `out_prefix`. +""" +function setup_atmos_spectra_diagnostics( + ::AtmosGCMConfigType, + interval::String, + out_prefix::String; + writer = NetCDFWriter(), + interpol = nothing, + nor = 1.0, +) + @assert !isnothing(interpol) + + return DiagnosticsGroup( + "AtmosGCMSpectra", + Diagnostics.atmos_gcm_spectra_init, + Diagnostics.atmos_gcm_spectra_fini, + Diagnostics.atmos_gcm_spectra_collect, + interval, + out_prefix, + writer, + interpol, + AtmosGCMSpectraDiagnosticsParams(nor), + ) +end + +function get_spectra(mpicomm, mpirank, Q, bl, interpol, nor) + FT = eltype(Q) + if array_device(Q) isa CPU + ArrayType = Array + else + ArrayType = CUDA.CuArray + end + + istate = ArrayType{FT}(undef, interpol.Npl, number_states(bl, Prognostic())) + interpolate_local!(interpol, Q.realdata, istate) + + # TODO: get indices here without hard-coding them + _ρu, _ρv, _ρw = 2, 3, 4 + project_cubed_sphere!(interpol, istate, (_ρu, _ρv, _ρw)) + + all_state_data = accumulate_interpolated_data(mpicomm, interpol, istate) + + if mpirank == 0 + var_grid = all_state_data[:, :, :, 3] ./ all_state_data[:, :, :, 1] + dims = dimensions(interpol) + lat = dims["lat"][1] + lon = dims["long"][1] + level = dims["level"][1] .- FT(planet_radius(Settings.param_set)) + + mass_weight = ones(FT, length(level)) # TODO: interpolate on pressure levs and do mass weighted calculations + spectrum_1d, m = power_spectrum_1d( + AtmosGCMConfigType(), + var_grid, + level, + lat, + lon, + mass_weight, + ) + spectrum_2d, m_and_n, _, __ = + power_spectrum_2d(AtmosGCMConfigType(), var_grid, mass_weight) + return spectrum_1d, m, spectrum_2d, m_and_n + end + return nothing, nothing, nothing, nothing +end + +function atmos_gcm_spectra_init(dgngrp, currtime) + Q = Settings.Q + bl = Settings.dg.balance_law + mpicomm = Settings.mpicomm + mpirank = MPI.Comm_rank(mpicomm) + FT = eltype(Q) + interpol = dgngrp.interpol + nor = dgngrp.params.nor + + # get the 1d and 2d spectra and their associated wavenumbers + spectrum_1d, m, spectrum_2d, m_and_n = + get_spectra(mpicomm, mpirank, Q, bl, interpol, nor) + + if mpirank == 0 + # Setup dimensions + interpol_dims = dimensions(interpol) + lat = interpol_dims["lat"] + level = interpol_dims["level"] + level = (level[1] .- FT(planet_radius(Settings.param_set)), level[2]) + num_fourier = ((2 * length(lat[1]) - 1) / 3) # number of truncated zonal wavenumbers + num_spherical = (num_fourier + 1) # number of total wavenumbers (n) + dims = OrderedDict( + "m" => (collect(FT, 1:1:length(lat[1])) .- 1, Dict()), + "m_t" => (collect(FT, 1:1:num_fourier) .- 1, Dict()), + "n" => (collect(FT, 1:1:num_spherical) .- 1, Dict()), + "level" => level, + "lat" => lat, + ) + # Setup variables (NB: 2d spectrum is on the truncated wavenumber grid) + vars = OrderedDict( + "spectrum_1d" => (("m", "lat", "level"), FT, Dict()), + "spectrum_2d" => (("m_t", "n", "level"), FT, Dict()), + ) + + dprefix = @sprintf( + "%s_%s-%s", + dgngrp.out_prefix, + dgngrp.name, + Settings.starttime, + ) + dfilename = joinpath(Settings.output_dir, dprefix) + init_data(dgngrp.writer, dfilename, dims, vars) + end + + return nothing +end + +function atmos_gcm_spectra_collect(dgngrp, currtime) + Q = Settings.Q + bl = Settings.dg.balance_law + mpicomm = Settings.mpicomm + mpirank = MPI.Comm_rank(mpicomm) + FT = eltype(Q) + interpol = dgngrp.interpol + nor = dgngrp.params.nor + + spectrum_1d, _, spectrum_2d, _ = + get_spectra(mpicomm, mpirank, Q, bl, interpol, nor) + + if mpirank == 0 + varvals = OrderedDict( + "spectrum_1d" => spectrum_1d, + "spectrum_2d" => spectrum_2d, + ) + append_data(dgngrp.writer, varvals, currtime) + end + + MPI.Barrier(mpicomm) + return nothing +end + +function atmos_gcm_spectra_fini(dgngrp, currtime) end diff --git a/src/Diagnostics/dump_spectra.jl b/src/Diagnostics/atmos_les_spectra.jl similarity index 70% rename from src/Diagnostics/dump_spectra.jl rename to src/Diagnostics/atmos_les_spectra.jl index d969682f210..1e33e1dfa25 100644 --- a/src/Diagnostics/dump_spectra.jl +++ b/src/Diagnostics/atmos_les_spectra.jl @@ -1,12 +1,12 @@ -# Spectrum calculator +# Spectrum calculator for AtmosLES -struct SpectraDiagnosticsParams <: DiagnosticsGroupParams +struct AtmosLESSpectraDiagnosticsParams <: DiagnosticsGroupParams nor::Float64 end """ - setup_dump_spectra_diagnostics( - ::ClimateMachineConfigType, + setup_atmos_spectra_diagnostics( + ::AtmosLESConfigType, interval::String, out_prefix::String; writer = NetCDFWriter(), @@ -14,13 +14,13 @@ end nor = Inf, ) -Create and return a `DiagnosticsGroup` containing a diagnostic that -dumps the spectrum at the specified -`interval` after the velocity fields have been interpolated, into NetCDF files prefixed by -`out_prefix`. +Create and return a `DiagnosticsGroup` containing a diagnostic for Atmos +LES configurations that dumps the spectrum at the specified `interval` +after the velocity fields have been interpolated, into NetCDF files +prefixed by `out_prefix`. """ -function setup_dump_spectra_diagnostics( - ::ClimateMachineConfigType, +function setup_atmos_spectra_diagnostics( + ::AtmosLESConfigType, interval::String, out_prefix::String, nor::Float64; @@ -30,15 +30,15 @@ function setup_dump_spectra_diagnostics( @assert !isnothing(interpol) return DiagnosticsGroup( - "Spectra", - Diagnostics.dump_spectra_init, - Diagnostics.dump_spectra_fini, - Diagnostics.dump_spectra_collect, + "SpectraLES", + Diagnostics.atmos_les_spectra_init, + Diagnostics.atmos_les_spectra_fini, + Diagnostics.atmos_les_spectra_collect, interval, out_prefix, writer, interpol, - SpectraDiagnosticsParams(nor), + AtmosLESSpectraDiagnosticsParams(nor), ) end @@ -53,13 +53,13 @@ function get_spectrum(mpicomm, mpirank, Q, bl, interpol, nor) w = all_state_data[:, :, :, 4] ./ all_state_data[:, :, :, 1] x1 = Array(interpol.x1g) d = length(x1) - s, k = power_spectrum(u, v, w, x1[d] - x1[1], d, nor) + s, k = power_spectrum_3d(AtmosLESConfigType(), u, v, w, x1[d] - x1[1], d, nor) return s, k end return nothing, nothing end -function dump_spectra_init(dgngrp, currtime) +function atmos_les_spectra_init(dgngrp, currtime) Q = Settings.Q bl = Settings.dg.balance_law mpicomm = Settings.mpicomm @@ -86,7 +86,7 @@ function dump_spectra_init(dgngrp, currtime) return nothing end -function dump_spectra_collect(dgngrp, currtime) +function atmos_les_spectra_collect(dgngrp, currtime) Q = Settings.Q bl = Settings.dg.balance_law mpicomm = Settings.mpicomm @@ -106,4 +106,4 @@ function dump_spectra_collect(dgngrp, currtime) return nothing end -function dump_spectra_fini(dgngrp, currtime) end +function atmos_les_spectra_fini(dgngrp, currtime) end diff --git a/src/Diagnostics/groups.jl b/src/Diagnostics/groups.jl index ff82d7b70aa..dc8ecea8e21 100644 --- a/src/Diagnostics/groups.jl +++ b/src/Diagnostics/groups.jl @@ -85,7 +85,8 @@ include("atmos_les_default_perturbations.jl") include("atmos_refstate_perturbations.jl") include("atmos_turbulence_stats.jl") include("atmos_mass_energy_loss.jl") +include("atmos_les_spectra.jl") +include("atmos_gcm_spectra.jl") include("dump_init.jl") include("dump_state.jl") include("dump_aux.jl") -include("dump_spectra.jl") diff --git a/src/Driver/diagnostics_configs.jl b/src/Driver/diagnostics_configs.jl index 1321f35422d..b78184b9515 100644 --- a/src/Driver/diagnostics_configs.jl +++ b/src/Driver/diagnostics_configs.jl @@ -24,7 +24,8 @@ end ) Creates an `InterpolationTopology` (either an `InterpolationBrick` or an -`InterpolationCubedSphere`) to be used with a `DiagnosticsGroup`. +`InterpolationCubedSphere`) to be used with a `DiagnosticsGroup`. The axes +are set up based on `boundaries` and `resolution`. """ function InterpolationConfiguration( driver_config::DriverConfiguration, @@ -64,7 +65,6 @@ function InterpolationConfiguration( FT(_planet_radius + info.domain_height), nelem = info.nelem_vert, ) - axes = ( collect(range( boundaries[1, 1], @@ -94,3 +94,48 @@ function InterpolationConfiguration( @error "Cannot set up interpolation for this topology." end end + +""" + InterpolationConfiguration( + driver_config::DriverConfiguration, + boundaries::Array, + axes, + ) + +Creates an `InterpolationTopology` (either an `InterpolationBrick` or an +`InterpolationCubedSphere`) to be used with a `DiagnosticsGroup`. The axes +are directly specified in lat/lon/lvl order. +""" +function InterpolationConfiguration( + driver_config::DriverConfiguration, + boundaries::Array, + axes, +) + param_set = driver_config.bl.param_set + grid = driver_config.grid + if isa(grid.topology, StackedBrickTopology) + return InterpolationBrick(grid, boundaries, axes[1], axes[2], axes[3]) + + elseif isa(grid.topology, StackedCubedSphereTopology) + + FT = eltype(grid) + _planet_radius::FT = planet_radius(param_set) + info = driver_config.config_info + vert_range = grid1d( + _planet_radius, + FT(_planet_radius + info.domain_height), + nelem = info.nelem_vert, + ) + + return InterpolationCubedSphere( + grid, + collect(vert_range), + info.nelem_horz, + axes[1], + axes[2], + axes[3], + ) + else + @error "Cannot set up interpolation for this topology." + end +end diff --git a/src/Numerics/Mesh/Metrics.jl b/src/Numerics/Mesh/Metrics.jl index fda70c6782e..bf029831ae5 100644 --- a/src/Numerics/Mesh/Metrics.jl +++ b/src/Numerics/Mesh/Metrics.jl @@ -379,27 +379,28 @@ function computemetric!( end for j in 1:Nq, i in 1:Nq - n1[i, j, 1, e] = -J[ 1, i, j, e] * ξ1x1[ 1, i, j, e] - n1[i, j, 2, e] = J[Nq, i, j, e] * ξ1x1[Nq, i, j, e] - n1[i, j, 3, e] = -J[ i, 1, j, e] * ξ2x1[ i, 1, j, e] - n1[i, j, 4, e] = J[ i,Nq, j, e] * ξ2x1[ i,Nq, j, e] - n1[i, j, 5, e] = -J[ i, j, 1, e] * ξ3x1[ i, j, 1, e] - n1[i, j, 6, e] = J[ i, j,Nq, e] * ξ3x1[ i, j,Nq, e] - n2[i, j, 1, e] = -J[ 1, i, j, e] * ξ1x2[ 1, i, j, e] - n2[i, j, 2, e] = J[Nq, i, j, e] * ξ1x2[Nq, i, j, e] - n2[i, j, 3, e] = -J[ i, 1, j, e] * ξ2x2[ i, 1, j, e] - n2[i, j, 4, e] = J[ i,Nq, j, e] * ξ2x2[ i,Nq, j, e] - n2[i, j, 5, e] = -J[ i, j, 1, e] * ξ3x2[ i, j, 1, e] - n2[i, j, 6, e] = J[ i, j,Nq, e] * ξ3x2[ i, j,Nq, e] - n3[i, j, 1, e] = -J[ 1, i, j, e] * ξ1x3[ 1, i, j, e] - n3[i, j, 2, e] = J[Nq, i, j, e] * ξ1x3[Nq, i, j, e] - n3[i, j, 3, e] = -J[ i, 1, j, e] * ξ2x3[ i, 1, j, e] - n3[i, j, 4, e] = J[ i,Nq, j, e] * ξ2x3[ i,Nq, j, e] - n3[i, j, 5, e] = -J[ i, j, 1, e] * ξ3x3[ i, j, 1, e] - n3[i, j, 6, e] = J[ i, j,Nq, e] * ξ3x3[ i, j,Nq, e] - - for n = 1:nface - sJ[i, j, n, e] = hypot(n1[i, j, n, e], n2[i, j, n, e], n3[i, j, n, e]) + n1[i, j, 1, e] = -J[1, i, j, e] * ξ1x1[1, i, j, e] + n1[i, j, 2, e] = J[Nq, i, j, e] * ξ1x1[Nq, i, j, e] + n1[i, j, 3, e] = -J[i, 1, j, e] * ξ2x1[i, 1, j, e] + n1[i, j, 4, e] = J[i, Nq, j, e] * ξ2x1[i, Nq, j, e] + n1[i, j, 5, e] = -J[i, j, 1, e] * ξ3x1[i, j, 1, e] + n1[i, j, 6, e] = J[i, j, Nq, e] * ξ3x1[i, j, Nq, e] + n2[i, j, 1, e] = -J[1, i, j, e] * ξ1x2[1, i, j, e] + n2[i, j, 2, e] = J[Nq, i, j, e] * ξ1x2[Nq, i, j, e] + n2[i, j, 3, e] = -J[i, 1, j, e] * ξ2x2[i, 1, j, e] + n2[i, j, 4, e] = J[i, Nq, j, e] * ξ2x2[i, Nq, j, e] + n2[i, j, 5, e] = -J[i, j, 1, e] * ξ3x2[i, j, 1, e] + n2[i, j, 6, e] = J[i, j, Nq, e] * ξ3x2[i, j, Nq, e] + n3[i, j, 1, e] = -J[1, i, j, e] * ξ1x3[1, i, j, e] + n3[i, j, 2, e] = J[Nq, i, j, e] * ξ1x3[Nq, i, j, e] + n3[i, j, 3, e] = -J[i, 1, j, e] * ξ2x3[i, 1, j, e] + n3[i, j, 4, e] = J[i, Nq, j, e] * ξ2x3[i, Nq, j, e] + n3[i, j, 5, e] = -J[i, j, 1, e] * ξ3x3[i, j, 1, e] + n3[i, j, 6, e] = J[i, j, Nq, e] * ξ3x3[i, j, Nq, e] + + for n in 1:nface + sJ[i, j, n, e] = + hypot(n1[i, j, n, e], n2[i, j, n, e], n3[i, j, n, e]) n1[i, j, n, e] /= sJ[i, j, n, e] n2[i, j, n, e] /= sJ[i, j, n, e] n3[i, j, n, e] /= sJ[i, j, n, e] diff --git a/test/Atmos/EDMF/helper_funcs/subdomain_thermo_states.jl b/test/Atmos/EDMF/helper_funcs/subdomain_thermo_states.jl index a850f2510c3..c5b87bd51a5 100644 --- a/test/Atmos/EDMF/helper_funcs/subdomain_thermo_states.jl +++ b/test/Atmos/EDMF/helper_funcs/subdomain_thermo_states.jl @@ -133,12 +133,7 @@ function new_thermo_state_up( θ_liq_up = ρaθ_liq_up / ρa_up q_tot_up = ρaq_tot_up / ρa_up - PhaseEquil_pθq( - m.param_set, - p, - θ_liq_up, - q_tot_up, - ) + PhaseEquil_pθq(m.param_set, p, θ_liq_up, q_tot_up) end return ts_up end @@ -171,12 +166,7 @@ function new_thermo_state_en( @print("q_tot_en = ", q_tot_en, "\n") error("Environment q_tot_en out-of-bounds in new_thermo_state_en") end - ts_en = PhaseEquil_pθq( - m.param_set, - p, - θ_liq_en, - q_tot_en, - ) + ts_en = PhaseEquil_pθq(m.param_set, p, θ_liq_en, q_tot_en) return ts_en end diff --git a/test/Common/Spectra/runtests.jl b/test/Common/Spectra/runtests.jl new file mode 100644 index 00000000000..d58b528e0a4 --- /dev/null +++ b/test/Common/Spectra/runtests.jl @@ -0,0 +1,100 @@ +using Test +using FFTW + +using ClimateMachine.ConfigTypes +using ClimateMachine.Spectra +using ClimateMachine.Spectra: + compute_gaussian!, + compute_legendre!, + SpectralSphericalMesh, + trans_grid_to_spherical!, + compute_wave_numbers + +include("spherical_helper_test.jl") + +@testset "power_spectrum_1d (GCM)" begin + FT = Float64 + # -- TEST 1: power_spectrum_1d(AtmosGCMConfigType(), var_grid, z, lat, lon, weight) + nlats = 32 + + # Setup grid + sinθ, wts = compute_gaussian!(nlats) + yarray = asin.(sinθ) .* 180 / π + xarray = 180.0 ./ nlats * collect(FT, 1:1:(2nlats))[:] .- 180.0 + z = 1 + + # Setup variable + mass_weight = ones(Float64, length(z)) + var_grid = + 1.0 * reshape( + sin.(xarray / xarray[end] * 5.0 * 2π) .* (yarray .* 0.0 .+ 1.0)', + length(xarray), + length(yarray), + 1, + ) + + 1.0 * reshape( + sin.(xarray / xarray[end] * 10.0 * 2π) .* (yarray .* 0.0 .+ 1.0)', + length(xarray), + length(yarray), + 1, + ) + nm_spectrum, wave_numbers = power_spectrum_1d( + AtmosGCMConfigType(), + var_grid, + z, + yarray, + xarray, + mass_weight, + ) + + nm_spectrum_ = nm_spectrum[:, 10, 1] + var_grid_ = var_grid[:, 10, 1] + sum_spec = sum((nm_spectrum_ .* conj(nm_spectrum_))) + sum_grid = sum(0.5 .* (var_grid_ .^ 2)) / length(var_grid_) + + sum_res = (sum_spec - sum_grid) / sum_grid + + @test sum_res < 0.1 +end + +@testset "power_spectrum_2d (GCM)" begin + # -- TEST 2: power_spectrum_2d + # Setup grid + FT = Float64 + nlats = 32 + sinθ, wts = compute_gaussian!(nlats) + cosθ = sqrt.(1 .- sinθ .^ 2) + yarray = asin.(sinθ) .* 180 / π + xarray = 180.0 ./ nlats * collect(FT, 1:1:(2nlats))[:] .- 180.0 + z = 1 + + # Setup variable: use an example analytical P_nm function + P_32 = sqrt(105 / 8) * (sinθ .- sinθ .^ 3) + var_grid = + 1.0 * reshape( + sin.(xarray / xarray[end] * 3.0 * π) .* P_32', + length(xarray), + length(yarray), + 1, + ) + + mass_weight = ones(Float64, z) + spectrum, wave_numbers, spherical, mesh = + power_spectrum_2d(AtmosGCMConfigType(), var_grid, mass_weight) + + # Grid to spherical to grid reconstruction + reconstruction = trans_spherical_to_grid!(mesh, spherical) + + sum_spec = sum((spectrum)) + sum_grid = + sum(0.5 .* var_grid[:, :, 1] .^ 2 * reshape(cosθ, (length(cosθ), 1))) + sum_reco = sum( + 0.5 .* reconstruction[:, :, 1] .^ 2 * reshape(cosθ, (length(cosθ), 1)), + ) + + sum_res_1 = (sum_spec - sum_grid) / sum_grid + sum_res_2 = (sum_reco - sum_grid) / sum_grid + + @test abs(sum_res_1) < 0.5 + @test abs(sum_res_2) < 0.5 +end diff --git a/test/Common/Spectra/spherical_helper_test.jl b/test/Common/Spectra/spherical_helper_test.jl new file mode 100644 index 00000000000..54341499911 --- /dev/null +++ b/test/Common/Spectra/spherical_helper_test.jl @@ -0,0 +1,76 @@ +# additional helper functions for spherical harmonics spectra + +""" + TransSphericalToGrid!(mesh, snm ) + +Transforms a variable expressed in spherical harmonics (var_spherical[num_fourier+1, num_spherical+1]) onto a Gaussian grid (pfield[nλ, nθ]) + +[THIS IS USED FOR TESTING ONLY] + + With F_{m,n} = (-1)^m F_{-m,n}* + P_{m,n} = (-1)^m P_{-m,n} + + F(λ, η) = ∑_{m= -N}^{N} ∑_{n=|m|}^{N} F_{m,n} P_{m,n}(η) e^{imλ} + = ∑_{m= 0}^{N} ∑_{n=m}^{N} F_{m,n} P_{m,n} e^{imλ} + ∑_{m= 1}^{N} ∑_{n=m}^{N} F_{-m,n} P_{-m,n} e^{-imλ} + + Here η = sinθ, N = num_fourier, and denote + ! extra coeffients in snm n > N are not used. + + ∑_{n=m}^{N} F_{m,n} P_{m,n} = g_{m}(η) m = 1, ... N + ∑_{n=m}^{N} F_{m,n} P_{m,n}/2.0 = g_{m}(η) m = 0 + + We have + + F(λ, η) = ∑_{m= 0}^{N} g_{m}(η) e^{imλ} + ∑_{m= 0}^{N} g_{m}(η)* e^{-imλ} + = 2real{ ∑_{m= 0}^{N} g_{m}(η) e^{imλ} } + + snm = F_{m,n} # Complex{Float64} [num_fourier+1, num_spherical+1] + qnm = P_{m,n,η} # Float64[num_fourier+1, num_spherical+1, nθ] + fourier_g = g_{m, η} # Complex{Float64} nλ×nθ with padded 0s fourier_g[num_fourier+2, :] == 0.0 + pfiled = F(λ, η) # Float64[nλ, nθ] + + ! use all spherical harmonic modes + + +# Arguments +- mesh: struct with mesh information +- snm: spherical variable + +# References +- Ehrendorfer, M., Spectral Numerical Weather Prediction Models, Appendix B, Society for Industrial and Applied Mathematics, 2011 +""" +function trans_spherical_to_grid!(mesh, snm) + num_fourier, num_spherical = mesh.num_fourier, mesh.num_spherical + nλ, nθ, nd = mesh.nλ, mesh.nθ, mesh.nd + + qnm = mesh.qnm + + fourier_g = mesh.var_fourier .* 0.0 + fourier_s = mesh.var_fourier .* 0.0 + + @assert(nθ % 2 == 0) + nθ_half = div(nθ, 2) + for m in 1:(num_fourier + 1) + for n in m:num_spherical + snm_t = transpose(snm[m, n, :, 1:nθ_half]) #snm[m,n, :] is complex number + if (n - m) % 2 == 0 + fourier_s[m, 1:nθ_half, :] .+= + qnm[m, n, 1:nθ_half] .* sum(snm_t, dims = 1) #even function part + else + fourier_s[m, (nθ_half + 1):nθ, :] .+= + qnm[m, n, 1:nθ_half] .* sum(snm_t, dims = 1) #odd function part + end + end + end + fourier_g[:, 1:nθ_half, :] .= + fourier_s[:, 1:nθ_half, :] .+ fourier_s[:, (nθ_half + 1):nθ, :] + fourier_g[:, nθ:-1:(nθ_half + 1), :] .= + fourier_s[:, 1:nθ_half, :] .- fourier_s[:, (nθ_half + 1):nθ, :] # this got ignored... + + fourier_g[1, :, :] ./= 2.0 + pfield = zeros(Float64, nλ, nθ, nd) + for j in 1:nθ + pfield[:, j, :] .= 2.0 * nλ * real.(ifft(fourier_g[:, j, :], 1)) #fourier for the first dimension + end + return pfield +end diff --git a/test/Common/Spectra/standalone_visual_test.jl b/test/Common/Spectra/standalone_visual_test.jl new file mode 100644 index 00000000000..0f69652da92 --- /dev/null +++ b/test/Common/Spectra/standalone_visual_test.jl @@ -0,0 +1,98 @@ +# Standalone test file that tests spectra visually +#using Plots # uncomment when using the plotting code below + +using ClimateMachine.Spectra: + compute_gaussian!, + compute_legendre!, + SpectralSphericalMesh, + trans_grid_to_spherical!, + compute_wave_numbers +using FFTW + + +include("spherical_helper_test.jl") + +FT = Float64 +# -- TEST 1: power_spectrum_gcm_1d(AtmosGCMConfigType(), var_grid, z, lat, lon, weight) +nlats = 32 + +# Setup grid +sinθ, wts = compute_gaussian!(nlats) +yarray = asin.(sinθ) .* 180 / π +xarray = 180.0 ./ nlats * collect(FT, 1:1:(2nlats))[:] .- 180.0 +z = 1 + +# Setup variable +mass_weight = ones(Float64, length(z)); +var_grid = + 1.0 * reshape( + sin.(xarray / xarray[end] * 5.0 * 2π) .* (yarray .* 0.0 .+ 1.0)', + length(xarray), + length(yarray), + 1, + ) + + 1.0 * reshape( + sin.(xarray / xarray[end] * 10.0 * 2π) .* (yarray .* 0.0 .+ 1.0)', + length(xarray), + length(yarray), + 1, + ) +nm_spectrum, wave_numbers = power_spectrum_gcm_1d( + AtmosGCMConfigType(), + var_grid, + z, + yarray, + xarray, + mass_weight, +); + + +# Check visually +plot(wave_numbers[:, 16, 1], nm_spectrum[:, 16, 1], xlims = (0, 20)) +contourf(var_grid[:, :, 1]) +contourf(nm_spectrum[2:20, :, 1]) + +# -- TEST 2: power_spectrum_gcm_2d +# Setup grid +sinθ, wts = compute_gaussian!(nlats) +yarray = asin.(sinθ) .* 180 / π +xarray = 180.0 ./ nlats * collect(FT, 1:1:(2nlats))[:] .- 180.0 +z = 1 + +# Setup variable: use an example analytical P_nm function +P_32 = sqrt(105 / 8) * (sinθ .- sinθ .^ 3) +var_grid = + 1.0 * reshape( + sin.(xarray / xarray[end] * 3.0 * π) .* P_32', + length(xarray), + length(yarray), + 1, + ) + +mass_weight = ones(Float64, z); +spectrum, wave_numbers, spherical, mesh = + power_spectrum_gcm_2d(AtmosGCMConfigType(), var_grid, mass_weight) + +# Grid to spherical to grid reconstruction +reconstruction = trans_spherical_to_grid!(mesh, spherical) + +# Check visually +contourf(var_grid[:, :, 1]) +contourf(reconstruction[:, :, 1]) +contourf(var_grid[:, :, 1] .- reconstruction[:, :, 1]) + +# Spectrum +contourf( + collect(0:1:(mesh.num_fourier - 1))[:], + collect(0:1:(mesh.num_spherical - 1))[:], + (spectrum[:, :, 1])', + xlabel = "m", + ylabel = "n", +) + +# Check magnitude +println(sum(spectrum)) +println(sum( + 0.5 .* var_grid[:, :, 1] .^ 2 * + reshape(sqrt.(1 .- sinθ .^ 2), (length(sinθ), 1)), +)) diff --git a/test/Numerics/DGMethods/Euler/acousticwave_1d_imex.jl b/test/Numerics/DGMethods/Euler/acousticwave_1d_imex.jl index 77bc19603ad..42f5d0bb29f 100644 --- a/test/Numerics/DGMethods/Euler/acousticwave_1d_imex.jl +++ b/test/Numerics/DGMethods/Euler/acousticwave_1d_imex.jl @@ -16,11 +16,7 @@ using ClimateMachine.VTK: writevtk, writepvtu using ClimateMachine.GenericCallbacks: EveryXWallTimeSeconds, EveryXSimulationSteps using ClimateMachine.Thermodynamics: - air_density, - soundspeed_air, - internal_energy, - PhaseDry_pT, - PhasePartition + air_density, soundspeed_air, internal_energy, PhaseDry_pT, PhasePartition using ClimateMachine.TemperatureProfiles: IsothermalProfile using ClimateMachine.Atmos: AtmosModel, diff --git a/test/Numerics/DGMethods/Euler/acousticwave_mrigark.jl b/test/Numerics/DGMethods/Euler/acousticwave_mrigark.jl index b97a5f46fa6..9d8ecd3978f 100644 --- a/test/Numerics/DGMethods/Euler/acousticwave_mrigark.jl +++ b/test/Numerics/DGMethods/Euler/acousticwave_mrigark.jl @@ -20,11 +20,7 @@ using ClimateMachine.VTK: writevtk, writepvtu using ClimateMachine.GenericCallbacks: EveryXWallTimeSeconds, EveryXSimulationSteps using ClimateMachine.Thermodynamics: - air_density, - soundspeed_air, - internal_energy, - PhaseDry_pT, - PhasePartition + air_density, soundspeed_air, internal_energy, PhaseDry_pT, PhasePartition using ClimateMachine.TemperatureProfiles: IsothermalProfile using ClimateMachine.Atmos: AtmosModel, diff --git a/tutorials/Atmos/dry_rayleigh_benard.jl b/tutorials/Atmos/dry_rayleigh_benard.jl index 8869fb598be..4783daec345 100644 --- a/tutorials/Atmos/dry_rayleigh_benard.jl +++ b/tutorials/Atmos/dry_rayleigh_benard.jl @@ -38,8 +38,7 @@ using ClimateMachine.Diagnostics using ClimateMachine.GenericCallbacks using ClimateMachine.ODESolvers using ClimateMachine.Mesh.Filters -using ClimateMachine.Thermodynamics: - PhaseEquil_pTq, internal_energy +using ClimateMachine.Thermodynamics: PhaseEquil_pTq, internal_energy using ClimateMachine.TurbulenceClosures using ClimateMachine.VariableTemplates