From f08c69ba9ffddf24f125f034faee3aa3f1e1330c Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Thu, 22 Dec 2022 17:28:51 +1100 Subject: [PATCH] add warning for using MultiLayerQG w nlayer=1 --- src/multilayerqg.jl | 173 +++++++++++++++++++++++--------------------- 1 file changed, 90 insertions(+), 83 deletions(-) diff --git a/src/multilayerqg.jl b/src/multilayerqg.jl index 79959d2a..94a8889e 100644 --- a/src/multilayerqg.jl +++ b/src/multilayerqg.jl @@ -121,16 +121,23 @@ function Problem(nlayers::Int, # number of fluid layers To use MultiLayerQG with 3 fluid layers or more we suggest, for now, to restrict running on CPU.""" end - + + if nlayers == 1 + @warn """MultiLayerQG module does work for single-layer configuration but may not be as + optimized. We suggest using SingleLayerQG module for single-layer QG simulation unless + you have reasons to use MultiLayerQG in a single-layer configuration, e.g., you want to + compare solutions with varying number of fluid layers.""" + end + # topographic PV eta === nothing && (eta = zeros(dev, T, (nx, ny))) - + grid = TwoDGrid(dev; nx, Lx, ny, Ly, aliased_fraction, T) - + params = Params(nlayers, g, f₀, β, ρ, H, U, eta, μ, ν, nν, grid; calcFq) - + vars = calcFq == nothingfunction ? DecayingVars(grid, params) : (stochastic ? StochasticForcedVars(grid, params) : ForcedVars(grid, params)) - + equation = linear ? LinearEquation(params, grid) : Equation(params, grid) FourierFlows.Problem(equation, stepper, dt, grid, vars, params) @@ -152,15 +159,15 @@ struct Params{T, Aphys3D, Aphys2D, Aphys1D, Atrans4D, Trfft} <: AbstractParams "constant planetary vorticity" f₀ :: T "planetary vorticity ``y``-gradient" - β :: T + β :: T "array with density of each fluid layer" ρ :: Aphys3D "array with rest height of each fluid layer" - H :: Aphys3D + H :: Aphys3D "array with imposed constant zonal flow ``U(y)`` in each fluid layer" - U :: Aphys3D + U :: Aphys3D "array containing the topographic PV" - eta :: Aphys2D + eta :: Aphys2D "linear bottom drag coefficient" μ :: T "small-scale (hyper)-viscosity coefficient" @@ -290,7 +297,7 @@ function convert_U_to_U3D(dev, nlayers, grid, U::Number) return A(U_3D) end -function Params(nlayers, g, f₀, β, ρ, H, U, eta, μ, ν, nν, grid; calcFq=nothingfunction, effort=FFTW.MEASURE) where TU +function Params(nlayers, g, f₀, β, ρ, H, U, eta, μ, ν, nν, grid; calcFq=nothingfunction, effort=FFTW.MEASURE) dev = grid.device T = eltype(grid) A = device_array(dev) @@ -298,7 +305,7 @@ function Params(nlayers, g, f₀, β, ρ, H, U, eta, μ, ν, nν, grid; calcFq=n ny, nx = grid.ny , grid.nx nkr, nl = grid.nkr, grid.nl kr, l = grid.kr , grid.l - + U = convert_U_to_U3D(dev, nlayers, grid, U) Uyy = real.(ifft(-l.^2 .* fft(U))) @@ -314,30 +321,30 @@ function Params(nlayers, g, f₀, β, ρ, H, U, eta, μ, ν, nν, grid; calcFq=n Qy = zeros(dev, T, (nx, ny, nlayers)) Qy = T(β) .- Uyy # T(β) is needed to ensure that Qy remains same type as U @views @. Qy[:, :, nlayers] += etay - + rfftplanlayered = plan_flows_rfft(A{T, 3}(undef, grid.nx, grid.ny, nlayers), [1, 2]; flags=effort) - + if nlayers==1 return SingleLayerParams(T(β), U, eta, T(μ), T(ν), nν, calcFq, Qx, Qy, rfftplanlayered) - + else # if nlayers≥2 - + ρ = reshape(T.(ρ), (1, 1, nlayers)) H = reshape(T.(H), (1, 1, nlayers)) - g′ = T(g) * (ρ[2:nlayers] - ρ[1:nlayers-1]) ./ ρ[2:nlayers] # reduced gravity at each interface + g′ = T(g) * (ρ[2:nlayers] - ρ[1:nlayers-1]) ./ ρ[2:nlayers] # reduced gravity at each interface Fm = @. T(f₀^2 / (g′ * H[2:nlayers])) Fp = @. T(f₀^2 / (g′ * H[1:nlayers-1])) typeofSkl = SArray{Tuple{nlayers, nlayers}, T, 2, nlayers^2} # StaticArrays of type T and dims = (nlayers x nlayers) - + S = Array{typeofSkl, 2}(undef, (nkr, nl)) calcS!(S, Fp, Fm, nlayers, grid) S⁻¹ = Array{typeofSkl, 2}(undef, (nkr, nl)) calcS⁻¹!(S⁻¹, Fp, Fm, nlayers, grid) - + S, S⁻¹, Fp, Fm = A(S), A(S⁻¹), A(Fp), A(Fm) # convert to appropriate ArrayType CUDA.@allowscalar @views Qy[:, :, 1] = @. Qy[:, :, 1] - Fp[1] * (U[:, :, 2] - U[:, :, 1]) @@ -365,7 +372,7 @@ numberoflayers(::TwoLayerParams) = 2 """ hyperviscosity(params, grid) -Return the linear operator `L` that corresponds to (hyper)-viscosity of order ``n_ν`` with +Return the linear operator `L` that corresponds to (hyper)-viscosity of order ``n_ν`` with coefficient ``ν`` for ``n`` fluid layers. ```math L_j = - ν |𝐤|^{2 n_ν}, \\ j = 1, ...,n . @@ -378,37 +385,37 @@ function hyperviscosity(params, grid) L = device_array(dev){T}(undef, (grid.nkr, grid.nl, numberoflayers(params))) @. L = - params.ν * grid.Krsq^params.nν @views @. L[1, 1, :] = 0 - + return L end """ LinearEquation(params, grid) -Return the equation for a multi-layer quasi-geostrophic problem with `params` and `grid`. -The linear opeartor ``L`` includes only (hyper)-viscosity and is computed via +Return the equation for a multi-layer quasi-geostrophic problem with `params` and `grid`. +The linear opeartor ``L`` includes only (hyper)-viscosity and is computed via `hyperviscosity(params, grid)`. The nonlinear term is computed via function `calcNlinear!`. """ function LinearEquation(params, grid) L = hyperviscosity(params, grid) - + return FourierFlows.Equation(L, calcNlinear!, grid) end - + """ Equation(params, grid) -Return the equation for a multi-layer quasi-geostrophic problem with `params` and `grid`. -The linear opeartor ``L`` includes only (hyper)-viscosity and is computed via +Return the equation for a multi-layer quasi-geostrophic problem with `params` and `grid`. +The linear opeartor ``L`` includes only (hyper)-viscosity and is computed via `hyperviscosity(params, grid)`. The nonlinear term is computed via function `calcN!`. """ function Equation(params, grid) L = hyperviscosity(params, grid) - + return FourierFlows.Equation(L, calcN!, grid) end @@ -460,10 +467,10 @@ function DecayingVars(grid, params) Dev = typeof(grid.device) T = eltype(grid) nlayers = numberoflayers(params) - + @devzeros Dev T (grid.nx, grid.ny, nlayers) q ψ u v @devzeros Dev Complex{T} (grid.nkr, grid.nl, nlayers) qh ψh uh vh - + return Vars(q, ψ, u, v, qh, ψh, uh, vh, nothing, nothing) end @@ -476,10 +483,10 @@ function ForcedVars(grid, params) Dev = typeof(grid.device) T = eltype(grid) nlayers = numberoflayers(params) - + @devzeros Dev T (grid.nx, grid.ny, nlayers) q ψ u v @devzeros Dev Complex{T} (grid.nkr, grid.nl, nlayers) qh ψh uh vh Fqh - + return Vars(q, ψ, u, v, qh, ψh, uh, vh, Fqh, nothing) end @@ -492,10 +499,10 @@ function StochasticForcedVars(grid, params) Dev = typeof(grid.device) T = eltype(grid) nlayers = numberoflayers(params) - + @devzeros Dev T (grid.nx, grid.ny, nlayers) q ψ u v @devzeros Dev Complex{T} (grid.nkr, grid.nl, nlayers) qh ψh uh vh Fqh prevsol - + return Vars(q, ψ, u, v, qh, ψh, uh, vh, Fqh, prevsol) end @@ -516,14 +523,14 @@ invtransform!(var, varh, params::AbstractParams) = ldiv!(var, params.rfftplan, v """ pvfromstreamfunction!(qh, ψh, params, grid) -Obtain the Fourier transform of the PV from the streamfunction `ψh` in each layer using +Obtain the Fourier transform of the PV from the streamfunction `ψh` in each layer using `qh = params.S * ψh`. """ function pvfromstreamfunction!(qh, ψh, params, grid) for j=1:grid.nl, i=1:grid.nkr - CUDA.@allowscalar @views qh[i, j, :] .= params.S[i, j] * ψh[i, j, :] + CUDA.@allowscalar @views qh[i, j, :] .= params.S[i, j] * ψh[i, j, :] end - + return nothing end @@ -535,7 +542,7 @@ case of a single fluid layer configuration. In this case, ``q̂ = - k² ψ̂``. """ function pvfromstreamfunction!(qh, ψh, params::SingleLayerParams, grid) @. qh = -grid.Krsq * ψh - + return nothing end @@ -558,12 +565,12 @@ on the GPU.) """ function pvfromstreamfunction!(qh, ψh, params::TwoLayerParams, grid) f₀, g′, H₁, H₂ = params.f₀, params.g′, params.H[1], params.H[2] - + ψ1h, ψ2h = view(ψh, :, :, 1), view(ψh, :, :, 2) @views @. qh[:, :, 1] = - grid.Krsq * ψ1h + f₀^2 / (g′ * H₁) * (ψ2h - ψ1h) @views @. qh[:, :, 2] = - grid.Krsq * ψ2h + f₀^2 / (g′ * H₂) * (ψ1h - ψ2h) - + return nothing end @@ -577,7 +584,7 @@ function streamfunctionfrompv!(ψh, qh, params, grid) for j=1:grid.nl, i=1:grid.nkr CUDA.@allowscalar @views ψh[i, j, :] .= params.S⁻¹[i, j] * qh[i, j, :] end - + return nothing end @@ -589,7 +596,7 @@ case of a single fluid layer configuration. In this case, ``ψ̂ = - k⁻² q̂` """ function streamfunctionfrompv!(ψh, qh, params::SingleLayerParams, grid) @. ψh = -grid.invKrsq * qh - + return nothing end @@ -614,12 +621,12 @@ on the GPU.) """ function streamfunctionfrompv!(ψh, qh, params::TwoLayerParams, grid) f₀, g′, H₁, H₂ = params.f₀, params.g′, params.H[1], params.H[2] - + q1h, q2h = view(qh, :, :, 1), view(qh, :, :, 2) @views @. ψh[:, :, 1] = - grid.Krsq * q1h - f₀^2 / g′ * (q1h / H₂ + q2h / H₁) @views @. ψh[:, :, 2] = - grid.Krsq * q2h - f₀^2 / g′ * (q1h / H₂ + q2h / H₁) - + for j in 1:2 @views @. ψh[:, :, j] *= grid.invKrsq / (grid.Krsq + f₀^2 / g′ * (H₁ + H₂) / (H₁ * H₂)) end @@ -630,39 +637,39 @@ end """ calcS!(S, Fp, Fm, nlayers, grid) -Construct the array ``𝕊``, which consists of `nlayer` x `nlayer` static arrays ``𝕊_𝐤`` that +Construct the array ``𝕊``, which consists of `nlayer` x `nlayer` static arrays ``𝕊_𝐤`` that relate the ``q̂_j``'s and ``ψ̂_j``'s for every wavenumber: ``q̂_𝐤 = 𝕊_𝐤 ψ̂_𝐤``. """ function calcS!(S, Fp, Fm, nlayers, grid) F = Matrix(Tridiagonal(Fm, -([Fp; 0] + [0; Fm]), Fp)) - + for n=1:grid.nl, m=1:grid.nkr k² = CUDA.@allowscalar grid.Krsq[m, n] Skl = SMatrix{nlayers, nlayers}(- k² * I + F) S[m, n] = Skl end - + return nothing end """ calcS⁻¹!(S, Fp, Fm, nlayers, grid) -Construct the array ``𝕊⁻¹``, which consists of `nlayer` x `nlayer` static arrays ``(𝕊_𝐤)⁻¹`` +Construct the array ``𝕊⁻¹``, which consists of `nlayer` x `nlayer` static arrays ``(𝕊_𝐤)⁻¹`` that relate the ``q̂_j``'s and ``ψ̂_j``'s for every wavenumber: ``ψ̂_𝐤 = (𝕊_𝐤)⁻¹ q̂_𝐤``. """ function calcS⁻¹!(S⁻¹, Fp, Fm, nlayers, grid) F = Matrix(Tridiagonal(Fm, -([Fp; 0] + [0; Fm]), Fp)) - + for n=1:grid.nl, m=1:grid.nkr k² = CUDA.@allowscalar grid.Krsq[m, n] == 0 ? 1 : grid.Krsq[m, n] Skl = - k² * I + F S⁻¹[m, n] = SMatrix{nlayers, nlayers}(I / Skl) end - + T = eltype(grid) S⁻¹[1, 1] = SMatrix{nlayers, nlayers}(zeros(T, (nlayers, nlayers))) - + return nothing end @@ -683,15 +690,15 @@ N_j = - \\widehat{𝖩(ψ_j, q_j)} - \\widehat{U_j ∂_x Q_j} - \\widehat{U_j """ function calcN!(N, sol, t, clock, vars, params, grid) nlayers = numberoflayers(params) - + dealias!(sol, grid) - + calcN_advection!(N, sol, vars, params, grid) - + @views @. N[:, :, nlayers] += params.μ * grid.Krsq * vars.ψh[:, :, nlayers] # bottom linear drag - + addforcing!(N, sol, t, clock, vars, params, grid) - + return nothing end @@ -701,17 +708,17 @@ end Compute the nonlinear term of the linearized equations: ```math -N_j = - \\widehat{U_j ∂_x Q_j} - \\widehat{U_j ∂_x q_j} + \\widehat{(∂_y ψ_j)(∂_x Q_j)} +N_j = - \\widehat{U_j ∂_x Q_j} - \\widehat{U_j ∂_x q_j} + \\widehat{(∂_y ψ_j)(∂_x Q_j)} - \\widehat{(∂_x ψ_j)(∂_y Q_j)} + δ_{j, n} μ |𝐤|^2 ψ̂_n + F̂_j . ``` """ function calcNlinear!(N, sol, t, clock, vars, params, grid) nlayers = numberoflayers(params) - + calcN_linearadvection!(N, sol, vars, params, grid) @views @. N[:, :, nlayers] += params.μ * grid.Krsq * vars.ψh[:, :, nlayers] # bottom linear drag addforcing!(N, sol, t, clock, vars, params, grid) - + return nothing end @@ -735,21 +742,21 @@ function calcN_advection!(N, sol, vars, params, grid) invtransform!(vars.u, vars.uh, params) @. vars.u += params.U # add the imposed zonal flow U - + uQx, uQxh = vars.q, vars.uh # use vars.q and vars.uh as scratch variables @. uQx = vars.u * params.Qx # (U+u)*∂Q/∂x fwdtransform!(uQxh, uQx, params) @. N = - uQxh # -\hat{(U+u)*∂Q/∂x} invtransform!(vars.v, vars.vh, params) - + vQy, vQyh = vars.q, vars.vh # use vars.q and vars.vh as scratch variables @. vQy = vars.v * params.Qy # v*∂Q/∂y fwdtransform!(vQyh, vQy, params) @. N -= vQyh # -\hat{v*∂Q/∂y} invtransform!(vars.q, vars.qh, params) - + uq , vq = vars.u , vars.v # use vars.u and vars.v as scratch variables uqh, vqh = vars.uh, vars.vh # use vars.uh and vars.vh as scratch variables @. uq *= vars.q # (U+u)*q @@ -790,7 +797,7 @@ function calcN_linearadvection!(N, sol, vars, params, grid) @. N = - uQxh # -\hat{(U+u)*∂Q/∂x} invtransform!(vars.v, vars.vh, params) - + vQy, vQyh = vars.q, vars.vh # use vars.q and vars.vh as scratch variables @. vQy = vars.v * params.Qy # v*∂Q/∂y @@ -798,7 +805,7 @@ function calcN_linearadvection!(N, sol, vars, params, grid) @. N -= vQyh # -\hat{v*∂Q/∂y} invtransform!(vars.q, vars.qh, params) - + @. vars.u = params.U Uq , Uqh = vars.u , vars.uh # use vars.u and vars.uh as scratch variables @. Uq *= vars.q # U*q @@ -813,8 +820,8 @@ end """ addforcing!(N, sol, t, clock, vars, params, grid) - -When the problem includes forcing, calculate the forcing term ``F̂`` for each layer and add + +When the problem includes forcing, calculate the forcing term ``F̂`` for each layer and add it to the nonlinear term ``N``. """ addforcing!(N, sol, t, clock, vars::Vars, params, grid) = nothing @@ -822,7 +829,7 @@ addforcing!(N, sol, t, clock, vars::Vars, params, grid) = nothing function addforcing!(N, sol, t, clock, vars::ForcedVars, params, grid) params.calcFq!(vars.Fqh, sol, t, clock, vars, params, grid) @. N += vars.Fqh - + return nothing end @@ -849,7 +856,7 @@ function updatevars!(vars, params, grid, sol) invtransform!(vars.ψ, deepcopy(vars.ψh), params) invtransform!(vars.u, deepcopy(vars.uh), params) invtransform!(vars.v, deepcopy(vars.vh), params) - + return nothing end @@ -868,7 +875,7 @@ function set_q!(sol, params, vars, grid, q) @. vars.qh[1, 1, :] = 0 @. sol = vars.qh updatevars!(vars, params, grid, sol) - + return nothing end @@ -877,7 +884,7 @@ function set_q!(sol, params::SingleLayerParams, vars, grid, q::AbstractArray{T, q_3D = vars.q @views q_3D[:, :, 1] = A(q) set_q!(sol, params, vars, grid, q_3D) - + return nothing end @@ -896,9 +903,9 @@ function set_ψ!(sol, params, vars, grid, ψ) fwdtransform!(vars.ψh, A(ψ), params) pvfromstreamfunction!(vars.qh, vars.ψh, params, grid) invtransform!(vars.q, vars.qh, params) - + set_q!(sol, params, vars, grid, vars.q) - + return nothing end @@ -906,10 +913,10 @@ function set_ψ!(sol, params::SingleLayerParams, vars, grid, ψ::AbstractArray{T A = typeof(vars.ψ[:, :, 1]) ψ_3D = vars.ψ @views ψ_3D[:, :, 1] = A(ψ) - + set_ψ!(sol, params, vars, grid, ψ_3D) - - return nothing + + return nothing end set_ψ!(prob, ψ) = set_ψ!(prob.sol, prob.params, prob.vars, prob.grid, ψ) @@ -945,7 +952,7 @@ function energies(vars, params, grid, sol) abs²∇𝐮h = vars.uh # use vars.uh as scratch variable @. abs²∇𝐮h = grid.Krsq * abs2(vars.ψh) - + for j = 1:nlayers CUDA.@allowscalar KE[j] = 1 / (2 * grid.Lx * grid.Ly) * parsevalsum(abs²∇𝐮h[:, :, j], grid) * params.H[j] / sum(params.H) end @@ -960,10 +967,10 @@ end function energies(vars, params::TwoLayerParams, grid, sol) nlayers = numberoflayers(params) KE, PE = zeros(nlayers), zeros(nlayers-1) - + @. vars.qh = sol streamfunctionfrompv!(vars.ψh, vars.qh, params, grid) - + abs²∇𝐮h = vars.uh # use vars.uh as scratch variable @. abs²∇𝐮h = grid.Krsq * abs2(vars.ψh) @@ -974,7 +981,7 @@ function energies(vars, params::TwoLayerParams, grid, sol) end PE = @views 1 / (2 * grid.Lx * grid.Ly * sum(params.H)) * params.f₀^2 / params.g′ * parsevalsum(abs2.(ψ2h .- ψ1h), grid) - + return KE, PE end @@ -984,7 +991,7 @@ function energies(vars, params::SingleLayerParams, grid, sol) abs²∇𝐮h = vars.uh # use vars.uh as scratch variable @. abs²∇𝐮h = grid.Krsq * abs2(vars.ψh) - + return 1 / (2 * grid.Lx * grid.Ly) * parsevalsum(abs²∇𝐮h, grid) end @@ -1002,11 +1009,11 @@ verticalfluxes``_{3/2},...,``verticalfluxes``_{n-1/2}``, where ``n`` is the tota The lateral eddy fluxes within the ``j``-th fluid layer are ```math -\\textrm{lateralfluxes}_j = \\frac{H_j}{H} \\int U_j v_j ∂_y u_j +\\textrm{lateralfluxes}_j = \\frac{H_j}{H} \\int U_j v_j ∂_y u_j \\frac{𝖽x 𝖽y}{L_x L_y} , \\ j = 1, ..., n , ``` -while the vertical eddy fluxes at the ``j+1/2``-th fluid interface (i.e., interface between +while the vertical eddy fluxes at the ``j+1/2``-th fluid interface (i.e., interface between the ``j``-th and ``(j+1)``-th fluid layer) are ```math @@ -1016,7 +1023,7 @@ v_{j+1} ψ_{j} \\frac{𝖽x 𝖽y}{L_x L_y} , \\ j = 1, ..., n-1. """ function fluxes(vars, params, grid, sol) nlayers = numberoflayers(params) - + lateralfluxes, verticalfluxes = zeros(nlayers), zeros(nlayers-1) updatevars!(vars, params, grid, sol) @@ -1058,7 +1065,7 @@ function fluxes(vars, params::TwoLayerParams, grid, sol) U₁, U₂ = view(params.U, :, :, 1), view(params.U, :, :, 2) ψ₁ = view(vars.ψ, :, :, 1) v₂ = view(vars.v, :, :, 2) - + verticalfluxes = sum(@views @. params.f₀^2 / params.g′ * (U₁ - U₂) * v₂ * ψ₁; dims=(1, 2)) verticalfluxes *= grid.dx * grid.dy / (grid.Lx * grid.Ly * sum(params.H))