From 10f9e87064c0a55e7af2758abc5ac6204575681b Mon Sep 17 00:00:00 2001 From: Navid Constantinou Date: Thu, 28 May 2020 15:50:40 +1000 Subject: [PATCH 1/3] tests for BarotropicQGQL run on both CPU/GPU --- examples/barotropicqgql_betaforced.jl | 2 +- src/barotropicqgql.jl | 422 ++++++++++++-------------- test/runtests.jl | 43 ++- test/test_barotropicqgql.jl | 144 ++++----- 4 files changed, 287 insertions(+), 324 deletions(-) diff --git a/examples/barotropicqgql_betaforced.jl b/examples/barotropicqgql_betaforced.jl index 94f11f42..70010474 100644 --- a/examples/barotropicqgql_betaforced.jl +++ b/examples/barotropicqgql_betaforced.jl @@ -50,7 +50,7 @@ forcing_bandwidth = 1.5 # the width of the forcing spectrum gr = TwoDGrid(nx, Lx) -k = [ gr.kr[i] for i=1:gr.nkr, j=1:gr.nl] # a 2D grid with the zonal wavenumber +k = [ gr.kr[i] for i=1:gr.nkr, j=1:gr.nl ] # a 2D grid with the zonal wavenumber forcing_spectrum = @. exp( -(sqrt(gr.Krsq)-forcing_wavenumber)^2 / (2forcing_bandwidth^2) ) @. forcing_spectrum[ gr.Krsq < (2π/Lx*2)^2 ] = 0 diff --git a/src/barotropicqgql.jl b/src/barotropicqgql.jl index b80a3b6e..4b32af0c 100644 --- a/src/barotropicqgql.jl +++ b/src/barotropicqgql.jl @@ -22,7 +22,6 @@ using FourierFlows: parsevalsum, parsevalsum2 import FFTW: rfft abstract type BarotropicQGQLVars <: AbstractVars end -abstract type BarotropicQGQLForcedVars <: BarotropicQGQLVars end const physicalvars = [:zeta, :psi, :u, :v, :uzeta, :vzeta, :U, :Zeta, :Psi] const transformvars = [ Symbol(var, :h) for var in physicalvars ] @@ -32,11 +31,11 @@ const stochforcedvars = [:prevsol] nothingfunction(args...) = nothing """ - Problem(; parameters...) + Problem(dev=CPU(); parameters...) -Construct a BarotropicQGQL turbulence problem. +Construct a BarotropicQGQL turbulence problem on device `dev`. """ -function Problem(; +function Problem(dev::Device=CPU(); # Numerical parameters nx = 256, Lx = 2π, @@ -44,13 +43,12 @@ function Problem(; Ly = Lx, dt = 0.01, # Physical parameters - f0 = 1.0, - beta = 0.0, + β = 0.0, eta = nothing, # Drag and/or hyper-/hypo-viscosity - nu = 0.0, - nnu = 1, - mu = 0.0, + ν = 0.0, + nν = 1, + μ = 0.0, # Timestepper and eqn options stepper = "RK4", calcF = nothingfunction, @@ -58,27 +56,21 @@ function Problem(; T = Float64) # the grid - gr = TwoDGrid(nx, Lx, ny, Ly; T=T) - x, y = gridpoints(gr) + grid = TwoDGrid(dev, nx, Lx, ny, Ly; T=T) + x, y = gridpoints(grid) # topographic PV if eta==nothing - eta = 0*x - etah = rfft(eta) + eta = zeros(dev, T, (nx, ny)) end - if typeof(eta)!=Array{T, 2} #this is true if eta was passes in Problem as a function - pr = Params(gr, f0, beta, eta, mu, nu, nnu, calcF) - else - pr = Params{T}(f0, beta, eta, rfft(eta), mu, nu, nnu, calcF) - end - vs = calcF == nothingfunction ? Vars(gr) : (stochastic ? StochasticForcedVars(gr) : ForcedVars(gr)) - eq = BarotropicQGQL.Equation(pr, gr) - FourierFlows.Problem(eq, stepper, dt, gr, vs, pr) -end + params = !(typeof(eta)<:ArrayType(dev)) ? Params(grid, β, eta, μ, ν, nν, calcF) : Params(β, eta, rfft(eta), μ, ν, nν, calcF) + + vars = calcF == nothingfunction ? Vars(dev, grid) : (stochastic ? StochasticForcedVars(dev, grid) : ForcedVars(dev, grid)) -InitialValueProblem(; kwargs...) = Problem(; kwargs...) -ForcedProblem(; kwargs...) = Problem(; kwargs...) + equation = BarotropicQGQL.Equation(params, grid) + FourierFlows.Problem(equation, stepper, dt, grid, vars, params, dev) +end # ---------- @@ -86,31 +78,30 @@ ForcedProblem(; kwargs...) = Problem(; kwargs...) # ---------- """ - Params(g::TwoDGrid, f0, beta, FU, eta, mu, nu, nnu) + Params -Returns the params for an unforced two-dimensional barotropic QG problem. +A struct that contains all parameter values for a two-dimensional barotropic QG QL problem. """ -struct Params{T} <: AbstractParams - f0::T # Constant planetary vorticity - beta::T # Planetary vorticity y-gradient - eta::Array{T, 2} # Topographic PV - etah::Array{Complex{T}, 2} # FFT of Topographic PV - mu::T # Linear drag - nu::T # Viscosity coefficient - nnu::Int # Hyperviscous order (nnu=1 is plain old viscosity) - calcF!::Function # Function that calculates the forcing on QGPV q +struct Params{T, Aphys, Atrans} <: AbstractParams + β :: T # Planetary vorticity y-gradient + eta :: Aphys # Topographic PV + etah :: Atrans # FFT of Topographic PV + μ :: T # Linear drag + ν :: T # Viscosity coefficient + nν :: Int # Hyperviscous order (nν=1 is plain old viscosity) + calcF! :: Function # Function that calculates the forcing on QGPV q end """ - Params(g::TwoDGrid, f0, beta, eta::Function, mu, nu, nnu) + Params(g::TwoDGrid, β, eta::Function, μ, ν, nν, calcF) -Constructor for Params that accepts a generating function for the topographic PV. +Constructor for `params` that accepts a generating function for the topographic PV. """ -function Params(g::AbstractGrid{T}, f0, beta, eta::Function, mu, nu, nnu, calcF) where T - x, y = gridpoints(g) - etagrid = eta(x, y) - etah = rfft(etagrid) - Params{T}(f0, beta, etagrid, etah, mu, nu, nnu, calcF) +function Params(grid::AbstractGrid{T, A}, β, eta::Function, μ, ν, nν, calcF) where {T, A} + x, y = gridpoints(grid) + eta_on_grid = A([eta(grid.x[i], grid.y[j]) for i=1:grid.nx, j=1:grid.ny]) + etah_on_grid = rfft(eta_on_grid) + Params(β, eta_on_grid, etah_on_grid, μ, ν, nν, calcF) end @@ -119,14 +110,14 @@ end # --------- """ - Equation(p, g) + Equation(params, grid) -Returns the equation for two-dimensional barotropic QG problem with params p and grid g. +Returns the equation for two-dimensional barotropic QG QL problem with parameters `params` and on `grid`. """ -function Equation(p::Params, g::AbstractGrid{T}) where T - L = @. -p.mu - p.nu*g.Krsq^p.nnu + im*p.beta*g.kr*g.invKrsq +function Equation(params::Params, grid::AbstractGrid) + L = @. -params.μ - params.ν * grid.Krsq^params.nν + im*params.β * grid.kr * grid.invKrsq L[1, 1] = 0 - FourierFlows.Equation(L, calcN!, g) + FourierFlows.Equation(L, calcN!, grid) end @@ -135,98 +126,69 @@ end # ---- """ - Vars(g) - -Returns the vars for unforced two-dimensional barotropic QG problem with grid g. + Vars +A struct that contains all variables for a two-dimensional barotropic QG QL problem. """ -mutable struct Vars{T} <: BarotropicQGQLVars - u::Array{T,2} - v::Array{T,2} - U::Array{T,2} - uzeta::Array{T,2} - vzeta::Array{T,2} - zeta::Array{T,2} - Zeta::Array{T,2} - psi::Array{T,2} - Psi::Array{T,2} - Nz::Array{Complex{T},2} - NZ::Array{Complex{T},2} - uh::Array{Complex{T},2} - vh::Array{Complex{T},2} - Uh::Array{Complex{T},2} - zetah::Array{Complex{T},2} - Zetah::Array{Complex{T},2} - psih::Array{Complex{T},2} - Psih::Array{Complex{T},2} +mutable struct Vars{Aphys, Atrans, F, P} <: BarotropicQGQLVars + u :: Aphys + v :: Aphys + U :: Aphys + uzeta :: Aphys + vzeta :: Aphys + zeta :: Aphys + Zeta :: Aphys + psi :: Aphys + Psi :: Aphys + Nz :: Atrans + NZ :: Atrans + uh :: Atrans + vh :: Atrans + Uh :: Atrans + zetah :: Atrans + Zetah :: Atrans + psih :: Atrans + Psih :: Atrans + Fh :: F + prevsol :: P end -function Vars(g::AbstractGrid{T}) where T - @createarrays T (g.nx, g.ny) u v U uzeta vzeta zeta Zeta psi Psi - @createarrays Complex{T} (g.nkr, g.nl) N NZ uh vh Uh zetah Zetah psih Psih - Vars(u, v, U, uzeta, vzeta, zeta, Zeta, psi, Psi, N, NZ, uh, vh, Uh, zetah, Zetah, psih, Psih) -end +const ForcedVars = Vars{<:AbstractArray, <:AbstractArray, <:AbstractArray, Nothing} +const StochasticForcedVars = Vars{<:AbstractArray, <:AbstractArray, <:AbstractArray, <:AbstractArray} """ - ForcedVars(g) + Vars(dev, grid) -Returns the vars for forced two-dimensional barotropic QG problem with grid g. +Returns the vars for unforced two-dimensional barotropic QG problem with `grid`. """ -mutable struct ForcedVars{T} <: BarotropicQGQLForcedVars - u::Array{T,2} - v::Array{T,2} - U::Array{T,2} - uzeta::Array{T,2} - vzeta::Array{T,2} - zeta::Array{T,2} - Zeta::Array{T,2} - psi::Array{T,2} - Psi::Array{T,2} - Nz::Array{Complex{T},2} - NZ::Array{Complex{T},2} - uh::Array{Complex{T},2} - vh::Array{Complex{T},2} - Uh::Array{Complex{T},2} - zetah::Array{Complex{T},2} - Zetah::Array{Complex{T},2} - psih::Array{Complex{T},2} - Psih::Array{Complex{T},2} - Fh::Array{Complex{T},2} -end - -function ForcedVars(g::AbstractGrid{T}) where T - v = Vars(g) - Fh = zeros(Complex{T}, (g.nkr, g.nl)) - ForcedVars(getfield.(Ref(v), fieldnames(typeof(v)))..., Fh) +function Vars(dev::Dev, grid::AbstractGrid) where Dev + T = eltype(grid) + @devzeros Dev T (grid.nx, grid.ny) u v U uzeta vzeta zeta Zeta psi Psi + @devzeros Dev Complex{T} (grid.nkr, grid.nl) N NZ uh vh Uh zetah Zetah psih Psih + Vars(u, v, U, uzeta, vzeta, zeta, Zeta, psi, Psi, N, NZ, uh, vh, Uh, zetah, Zetah, psih, Psih, nothing, nothing) end +""" + ForcedVars(dev, grid) -mutable struct StochasticForcedVars{T} <: BarotropicQGQLForcedVars - u::Array{T,2} - v::Array{T,2} - U::Array{T,2} - uzeta::Array{T,2} - vzeta::Array{T,2} - zeta::Array{T,2} - Zeta::Array{T,2} - psi::Array{T,2} - Psi::Array{T,2} - Nz::Array{Complex{T},2} - NZ::Array{Complex{T},2} - uh::Array{Complex{T},2} - vh::Array{Complex{T},2} - Uh::Array{Complex{T},2} - zetah::Array{Complex{T},2} - Zetah::Array{Complex{T},2} - psih::Array{Complex{T},2} - Psih::Array{Complex{T},2} - Fh::Array{Complex{T},2} - prevsol::Array{Complex{T},2} +Returns the `vars` for forced two-dimensional barotropic QG QL problem on device `dev` and `grid`. +""" +function ForcedVars(dev::Dev, grid::AbstractGrid) where Dev + T = eltype(grid) + @devzeros Dev T (grid.nx, grid.ny) u v U uzeta vzeta zeta Zeta psi Psi + @devzeros Dev Complex{T} (grid.nkr, grid.nl) N NZ uh vh Uh zetah Zetah psih Psih Fh + Vars(u, v, U, uzeta, vzeta, zeta, Zeta, psi, Psi, N, NZ, uh, vh, Uh, zetah, Zetah, psih, Psih, Fh, nothing) end -function StochasticForcedVars(g::AbstractGrid{T}) where T - v = ForcedVars(g) - prevsol = zeros(Complex{T}, (g.nkr, g.nl)) - StochasticForcedVars(getfield.(Ref(v), fieldnames(typeof(v)))..., prevsol) +""" + StochasticForcedVars(dev, grid) + +Returns the `vars` for stochastically forced two-dimensional barotropic QG QL problem on device `dev` and `grid`. +""" +function StochasticForcedVars(dev::Dev, grid::AbstractGrid) where Dev + T = eltype(grid) + @devzeros Dev T (grid.nx, grid.ny) u v U uzeta vzeta zeta Zeta psi Psi + @devzeros Dev Complex{T} (grid.nkr, grid.nl) N NZ uh vh Uh zetah Zetah psih Psih Fh prevsol + Vars(u, v, U, uzeta, vzeta, zeta, Zeta, psi, Psi, N, NZ, uh, vh, Uh, zetah, Zetah, psih, Psih, Fh, prevsol) end @@ -234,67 +196,66 @@ end # Solvers # ------- -function calcN_advection!(N, sol, t, cl, v, p, g) - # Note that U = sol[1, 1]. For all other elements ζ = sol - Kr = [ g.kr[i] for i=1:g.nkr, j=1:g.nl] - @. v.zetah = sol - @. v.zetah[Kr .== 0] = 0 - @. v.Zetah = sol - @. v.Zetah[abs.(Kr) .> 0] = 0 - - @. v.uh = im * g.l * g.invKrsq * v.zetah - @. v.vh = -im * g.kr * g.invKrsq * v.zetah - @. v.Uh = im * g.l * g.invKrsq * v.Zetah - - ldiv!(v.zeta, g.rfftplan, v.zetah) - ldiv!(v.u, g.rfftplan, v.uh) - ldiv!(v.v, g.rfftplan, v.vh) - - ldiv!(v.Zeta, g.rfftplan, v.Zetah) - ldiv!(v.U, g.rfftplan, v.Uh) - - @. v.uzeta = v.u*v.zeta # u*ζ - @. v.vzeta = v.v*v.zeta # v*ζ - - mul!(v.uh, g.rfftplan, v.uzeta) # \hat{u*q} - @. v.NZ = -im*g.kr*v.uh # -∂[u*q]/∂x - mul!(v.vh, g.rfftplan, v.vzeta) # \hat{v*q} - @. v.NZ += - im*g.l*v.vh # -∂[v*q]/∂y - @. v.NZ[abs.(Kr) .> 0] = 0 - - @. v.U = v.U*v.zeta # U*ζ - @. v.u = v.u*v.Zeta # u*Ζ - @. v.v = v.v*v.Zeta # v*Ζ - - mul!(v.uh, g.rfftplan, v.U + v.u) # \hat{U*ζ + u*Ζ} - @. v.Nz = -im*g.kr*v.uh # -∂[U*ζ + u*Ζ]/∂x - mul!(v.vh, g.rfftplan, v.v) # \hat{v*Z} - @. v.Nz += - im*g.l*v.vh # -∂[v*Z]/∂y - @. v.Nz[abs.(Kr) .== 0] = 0 - - @. N = v.NZ + v.Nz +function calcN_advection!(N, sol, t, clock, vars, params, grid) + Kr = [ grid.kr[i] for i=1:grid.nkr, j=1:grid.nl] + @. vars.zetah = sol + @. vars.zetah[Kr .== 0] = 0 + @. vars.Zetah = sol + @. vars.Zetah[abs.(Kr) .> 0] = 0 + + @. vars.uh = im * grid.l * grid.invKrsq * vars.zetah + @. vars.vh = -im * grid.kr * grid.invKrsq * vars.zetah + @. vars.Uh = im * grid.l * grid.invKrsq * vars.Zetah + + ldiv!(vars.zeta, grid.rfftplan, vars.zetah) + ldiv!(vars.u, grid.rfftplan, vars.uh) + ldiv!(vars.v, grid.rfftplan, vars.vh) + + ldiv!(vars.Zeta, grid.rfftplan, vars.Zetah) + ldiv!(vars.U, grid.rfftplan, vars.Uh) + + @. vars.uzeta = vars.u * vars.zeta # u*ζ + @. vars.vzeta = vars.v * vars.zeta # v*ζ + + mul!(vars.uh, grid.rfftplan, vars.uzeta) # \hat{u*q} + @. vars.NZ = -im * grid.kr * vars.uh # -∂[u*q]/∂x + mul!(vars.vh, grid.rfftplan, vars.vzeta) # \hat{v*q} + @. vars.NZ += - im * grid.l * vars.vh # -∂[v*q]/∂y + @. vars.NZ[abs.(Kr) .> 0] = 0 + + @. vars.U = vars.U * vars.zeta # U*ζ + @. vars.u = vars.u * vars.Zeta # u*Ζ + @. vars.v = vars.v * vars.Zeta # v*Ζ + + mul!(vars.uh, grid.rfftplan, vars.U + vars.u) # \hat{U*ζ + u*Ζ} + @. vars.Nz = -im * grid.kr*vars.uh # -∂[U*ζ + u*Ζ]/∂x + mul!(vars.vh, grid.rfftplan, vars.v) # \hat{v*Z} + @. vars.Nz += - im * grid.l*vars.vh # -∂[v*Z]/∂y + @. vars.Nz[abs.(Kr) .== 0] = 0 + + @. N = vars.NZ + vars.Nz end -function calcN!(N, sol, t, cl, v, p, g) - calcN_advection!(N, sol, t, cl, v, p, g) - addforcing!(N, sol, t, cl, v, p, g) +function calcN!(N, sol, t, clock, vars, params, grid) + calcN_advection!(N, sol, t, clock, vars, params, grid) + addforcing!(N, sol, t, clock, vars, params, grid) nothing end -addforcing!(N, sol, t, cl, v::Vars, p, g) = nothing +addforcing!(N, sol, t, cl, vars::Vars, params, grid) = nothing -function addforcing!(N, sol, t, cl, v::ForcedVars, p, g) - p.calcF!(v.Fh, sol, t, cl, v, p, g) - @. N += v.Fh +function addforcing!(N, sol, t, clock, vars::ForcedVars, params, grid) + params.calcF!(vars.Fh, sol, t, clock, vars, params, grid) + @. N += vars.Fh nothing end -function addforcing!(N, sol, t, cl, v::StochasticForcedVars, p, g) - if t == cl.t # not a substep - @. v.prevsol = sol # sol at previous time-step is needed to compute budgets for stochastic forcing - p.calcF!(v.Fh, sol, t, cl, v, p, g) +function addforcing!(N, sol, t, clock, vars::StochasticForcedVars, params, grid) + if t == clock.t # not a substep + @. vars.prevsol = sol # sol at previous time-step is needed to compute budgets for stochastic forcing + params.calcF!(vars.Fh, sol, t, clock, vars, params, grid) end - @. N += v.Fh + @. N += vars.Fh nothing end @@ -304,31 +265,32 @@ end # ---------------- """ - updatevars!(v, s, g) + updatevars!(sol, vars, params, grid) + updatevars!(prob) -Update the vars in v on the grid g with the solution in sol. +Update the `vars` of a problem `prob` that has `grid` and `params` with the solution in `sol`. """ -function updatevars!(sol, v, p, g) - Kr = [ g.kr[i] for i=1:g.nkr, j=1:g.nl] +function updatevars!(sol, vars, params, grid) + Kr = [ grid.kr[i] for i=1:grid.nkr, j=1:grid.nl] sol[1, 1] = 0 - @. v.zetah = sol - @. v.zetah[Kr .== 0] = 0 - @. v.Zetah = sol - @. v.Zetah[abs.(Kr) .> 0] = 0 - - @. v.Psih = -v.Zetah * g.invKrsq - @. v.psih = -v.zetah * g.invKrsq - @. v.uh = -im * g.l * v.psih - @. v.vh = im * g.kr * v.psih - @. v.Uh = im * g.l * v.Zetah * g.invKrsq - - ldiv!(v.zeta, g.rfftplan, deepcopy(v.zetah)) - ldiv!(v.Zeta, g.rfftplan, deepcopy(v.Zetah)) - ldiv!(v.psi, g.rfftplan, v.psih) - ldiv!(v.Psi, g.rfftplan, v.Psih) - ldiv!(v.u, g.rfftplan, deepcopy(v.uh)) - ldiv!(v.v, g.rfftplan, deepcopy(v.vh)) - ldiv!(v.U, g.rfftplan, deepcopy(v.Uh)) + @. vars.zetah = sol + @. vars.zetah[Kr .== 0] = 0 + @. vars.Zetah = sol + @. vars.Zetah[abs.(Kr) .> 0] = 0 + + @. vars.Psih = -vars.Zetah * grid.invKrsq + @. vars.psih = -vars.zetah * grid.invKrsq + @. vars.uh = -im * grid.l * vars.psih + @. vars.vh = im * grid.kr * vars.psih + @. vars.Uh = im * grid.l * vars.Zetah * grid.invKrsq + + ldiv!(vars.zeta, grid.rfftplan, deepcopy(vars.zetah)) + ldiv!(vars.Zeta, grid.rfftplan, deepcopy(vars.Zetah)) + ldiv!(vars.psi, grid.rfftplan, vars.psih) + ldiv!(vars.Psi, grid.rfftplan, vars.Psih) + ldiv!(vars.u, grid.rfftplan, deepcopy(vars.uh)) + ldiv!(vars.v, grid.rfftplan, deepcopy(vars.vh)) + ldiv!(vars.U, grid.rfftplan, deepcopy(vars.Uh)) nothing end @@ -340,14 +302,14 @@ updatevars!(prob) = updatevars!(prob.sol, prob.vars, prob.params, prob.grid) set_zeta!(sol, v, g, zeta) Set the solution sol as the transform of zeta and update variables v -on the grid g. +on the `grid`. """ -function set_zeta!(sol, v, p, g, zeta) - mul!(v.zetah, g.rfftplan, zeta) - v.zetah[1, 1] = 0.0 - @. sol = v.zetah +function set_zeta!(sol, vars, params, grid, zeta) + mul!(vars.zetah, grid.rfftplan, zeta) + vars.zetah[1, 1] = 0.0 + @. sol = vars.zetah - updatevars!(sol, v, p, g) + updatevars!(sol, vars, params, grid) nothing end @@ -359,9 +321,9 @@ set_zeta!(prob, zeta) = set_zeta!(prob.sol, prob.vars, prob.params, prob.grid, z Returns the domain-averaged kinetic energy of sol. """ -function energy(sol, g::AbstractGrid) - 0.5*(parsevalsum2(g.kr.*g.invKrsq.*sol, g) - + parsevalsum2(g.l.*g.invKrsq.*sol, g))/(g.Lx*g.Ly) +function energy(sol, grid::AbstractGrid) + return 0.5 * (parsevalsum2(grid.kr .* grid.invKrsq .* sol, grid) + + parsevalsum2(grid.l .* grid.invKrsq .* sol, grid)) / (grid.Lx * grid.Ly) end energy(prob) = energy(prob.sol, prob.grid) @@ -371,10 +333,10 @@ energy(prob) = energy(prob.sol, prob.grid) Returns the domain-averaged enstrophy of sol. """ -function enstrophy(sol, g::AbstractGrid, v::AbstractVars) - @. v.uh = sol - v.uh[1, 1] = 0 - 0.5*parsevalsum2(v.uh, g)/(g.Lx*g.Ly) +function enstrophy(sol, grid::AbstractGrid, vars::AbstractVars) + @. vars.uh = sol + vars.uh[1, 1] = 0 + return 0.5 * parsevalsum2(vars.uh, grid) / (grid.Lx * grid.Ly) end enstrophy(prob) = enstrophy(prob.sol, prob.grid, prob.vars) @@ -383,12 +345,12 @@ enstrophy(prob) = enstrophy(prob.sol, prob.grid, prob.vars) dissipation(prob) dissipation(sol, v, p, g) -Returns the domain-averaged dissipation rate. nnu must be >= 1. +Returns the domain-averaged dissipation rate. nν must be >= 1. """ -@inline function dissipation(sol, v, p, g) - @. v.uh = g.Krsq^(p.nnu-1) * abs2(sol) - v.uh[1, 1] = 0 - p.nu/(g.Lx*g.Ly)*parsevalsum(v.uh, g) +@inline function dissipation(sol, vars, params, grid) + @. vars.uh = grid.Krsq^(params.nν-1) * abs(sol)^2 + vars.uh[1, 1] = 0 + params.ν / (grid.Lx * grid.Ly) * parsevalsum(vars.uh, grid) end @inline dissipation(prob) = dissipation(prob.sol, prob.vars, prob.params, prob.grid) @@ -399,15 +361,15 @@ end Returns the domain-averaged rate of work of energy by the forcing Fh. """ -@inline function work(sol, v::ForcedVars, g) - @. v.uh = g.invKrsq * sol * conj(v.Fh) - 1/(g.Lx*g.Ly)*parsevalsum(v.uh, g) +@inline function work(sol, vars::ForcedVars, grid) + @. vars.uh = grid.invKrsq * sol * conj(vars.Fh) + 1/(grid.Lx * grid.Ly)*parsevalsum(vars.uh, grid) end -@inline function work(sol, v::StochasticForcedVars, g) - @. v.uh = g.invKrsq * (v.prevsol + sol)/2.0 * conj(v.Fh) # Stratonovich - # @. v.uh = g.invKrsq * v.prevsol * conj(v.Fh) # Ito - 1/(g.Lx*g.Ly)*parsevalsum(v.uh, g) +@inline function work(sol, vars::StochasticForcedVars, grid) + @. vars.uh = grid.invKrsq * (vars.prevsol + sol)/2.0 * conj(vars.Fh) # Stratonovich + # @. vars.uh = grid.invKrsq * vars.prevsol * conj(vars.Fh) # Ito + 1/(grid.Lx * grid.Ly) * parsevalsum(vars.uh, grid) end @inline work(prob) = work(prob.sol, prob.vars, prob.grid) @@ -416,13 +378,13 @@ end drag(prob) drag(sol, v, p, g) -Returns the extraction of domain-averaged energy by drag mu. +Returns the extraction of domain-averaged energy by drag μ. """ @inline function drag(prob) - sol, v, p, g = prob.sol, prob.vars, prob.params, prob.grid - @. v.uh = g.invKrsq * abs2(sol) - v.uh[1, 1] = 0 - p.mu/(g.Lx*g.Ly)*parsevalsum(v.uh, g) + sol, vars, params, grid = prob.sol, prob.vars, prob.params, prob.grid + @. vars.uh = grid.invKrsq * abs(sol)^2 + vars.uh[1, 1] = 0 + params.μ / (grid.Lx * grid.Ly) * parsevalsum(vars.uh, grid) end end # module diff --git a/test/runtests.jl b/test/runtests.jl index 9ec6f4a6..bdf54b8a 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -54,7 +54,7 @@ for dev in devices @test test_twodnavierstokes_problemtype(Float32) @test TwoDNavierStokes.nothingfunction() == nothing end - + @testset "BarotropicQG" begin include("test_barotropicqg.jl") @@ -76,6 +76,25 @@ for dev in devices @test BarotropicQG.nothingfunction() == nothing end + @testset "BarotropicQGQL" begin + include("test_barotropicqgql.jl") + + @test test_bqgql_rossbywave("ETDRK4", 1e-2, 20, dev) + @test test_bqgql_rossbywave("FilteredETDRK4", 1e-2, 20, dev) + @test test_bqgql_rossbywave("RK4", 1e-2, 20, dev) + @test test_bqgql_rossbywave("FilteredRK4", 1e-2, 20, dev) + @test test_bqgql_rossbywave("AB3", 1e-3, 200, dev) + @test test_bqgql_rossbywave("FilteredAB3", 1e-3, 200, dev) + @test test_bqgql_rossbywave("ForwardEuler", 1e-4, 2000, dev) + @test test_bqgql_rossbywave("FilteredForwardEuler", 1e-4, 2000, dev) + @test test_bqgql_deterministicforcingbudgets(dev) + @test test_bqgql_stochasticforcingbudgets(dev) + @test test_bqgql_advection(0.0005, "ForwardEuler", dev) + @test test_bqgql_energyenstrophy(dev) + @test test_bqgql_problemtype(dev, Float32) + @test BarotropicQGQL.nothingfunction() == nothing + end + @testset "MultilayerQG" begin include("test_multilayerqg.jl") @@ -97,28 +116,6 @@ for dev in devices end - -println("rest of tests only on CPU") - -@testset "BarotropicQGQL" begin - include("test_barotropicqgql.jl") - - @test test_bqgql_rossbywave("ETDRK4", 1e-2, 20) - @test test_bqgql_rossbywave("FilteredETDRK4", 1e-2, 20) - @test test_bqgql_rossbywave("RK4", 1e-2, 20) - @test test_bqgql_rossbywave("FilteredRK4", 1e-2, 20) - @test test_bqgql_rossbywave("AB3", 1e-3, 200) - @test test_bqgql_rossbywave("FilteredAB3", 1e-3, 200) - @test test_bqgql_rossbywave("ForwardEuler", 1e-4, 2000) - @test test_bqgql_rossbywave("FilteredForwardEuler", 1e-4, 2000) - @test test_bqgql_deterministicforcingbudgets() - @test test_bqgql_stochasticforcingbudgets() - @test test_bqgql_advection(0.0005, "ForwardEuler") - @test test_bqgql_energyenstrophy() - @test test_bqgql_problemtype(Float32) - @test BarotropicQGQL.nothingfunction() == nothing -end - end # time println("Total test time: ", testtime) diff --git a/test/test_barotropicqgql.jl b/test/test_barotropicqgql.jl index 3bf84a0d..8b7567aa 100644 --- a/test/test_barotropicqgql.jl +++ b/test/test_barotropicqgql.jl @@ -3,32 +3,33 @@ Evolvesa a Rossby wave and compares with the analytic solution. """ -function test_bqgql_rossbywave(stepper, dt, nsteps) - nx = 64 - beta = 2.0 - Lx = 2π - mu = 0.0 - nu = 0.0 +function test_bqgql_rossbywave(stepper, dt, nsteps, dev::Device=CPU()) + nx = 64 + β = 2.0 + Lx = 2π + μ = 0.0 + ν = 0.0 + T = Float64 # the following if statement is called so that all the cases of # Problem() fuction are tested if stepper=="ForwardEuler" - eta = zeros(nx, nx) + eta = zeros(dev, T, (nx, nx)) else eta(x, y) = 0*x end - prob = BarotropicQGQL.InitialValueProblem(nx=nx, Lx=Lx, eta=eta, beta=beta, mu=mu, nu=nu, stepper=stepper, dt=dt) + prob = BarotropicQGQL.Problem(dev; nx=nx, Lx=Lx, eta=eta, β=β, μ=μ, ν=ν, stepper=stepper, dt=dt) sol, cl, v, p, g = prob.sol, prob.clock, prob.vars, prob.params, prob.grid x, y = gridpoints(g) # the Rossby wave initial condition ampl = 1e-2 - kwave = 3.0*2π/g.Lx - lwave = 2.0*2π/g.Ly - ω = -p.beta*kwave/(kwave^2.0 + lwave^2.0) - ζ0 = @. ampl*cos(kwave*x)*cos(lwave*y) + kwave = 3*2π/g.Lx + lwave = 2*2π/g.Ly + ω = - p.β * kwave / (kwave^2 + lwave^2) + ζ0 = @. ampl * cos(kwave*x) * cos(lwave*y) ζ0h = rfft(ζ0) BarotropicQGQL.set_zeta!(prob, ζ0) @@ -39,31 +40,31 @@ function test_bqgql_rossbywave(stepper, dt, nsteps) ζ_theory = @. ampl*cos(kwave*(x - ω/kwave*cl.t))*cos(lwave*y) - isapprox(ζ_theory, v.zeta, rtol=g.nx*g.ny*nsteps*1e-12) + return isapprox(ζ_theory, v.zeta, rtol=g.nx*g.ny*nsteps*1e-12) end """ - test_stochasticforcingbudgets(; kwargs...) + test_stochasticforcingbudgets(dev; kwargs...) -Tests if the energy budgets are closed for BarotropicQG with stochastic forcing. +Tests if the energy budget is closed for BarotropicQG problem with stochastic forcing. """ -function test_bqgql_stochasticforcingbudgets(; n=256, dt=0.01, L=2π, nu=1e-7, nnu=2, mu=1e-1) +function test_bqgql_stochasticforcingbudgets(dev::Device=CPU(); n=256, dt=0.01, L=2π, ν=1e-7, nν=2, μ=1e-1, T=Float64) n, L = 256, 2π - nu, nnu = 1e-7, 2 - mu = 1e-1 - dt, tf = 0.005, 0.1/mu + ν, nν = 1e-7, 2 + μ = 1e-1 + dt, tf = 0.005, 0.1/μ nt = round(Int, tf/dt) - ns = 1 + + gr = TwoDGrid(dev, n, L) + x, y = gridpoints(gr) # Forcing kf, dkf = 12.0, 2.0 ε = 0.1 - gr = TwoDGrid(n, L) - x, y = gridpoints(gr) - - Kr = [ gr.kr[i] for i=1:gr.nkr, j=1:gr.nl] + + Kr = ArrayType(dev)([ gr.kr[i] for i=1:gr.nkr, j=1:gr.nl]) - forcingcovariancespectrum = zero(gr.Krsq) + forcingcovariancespectrum = zeros(dev, T, (gr.nkr, gr.nl)) @. forcingcovariancespectrum = exp.(-(sqrt(gr.Krsq)-kf)^2/(2*dkf^2)) @. forcingcovariancespectrum[gr.Krsq .< 2.0^2 ] = 0 @. forcingcovariancespectrum[gr.Krsq .> 20.0^2 ] = 0 @@ -74,13 +75,13 @@ function test_bqgql_stochasticforcingbudgets(; n=256, dt=0.01, L=2π, nu=1e-7, n Random.seed!(1234) function calcF!(F, sol, t, cl, v, p, g) - eta = exp.(2π*im*rand(Float64, size(sol)))/sqrt(cl.dt) + eta = ArrayType(dev)(exp.(2π*im*rand(T, size(sol)))/sqrt(cl.dt)) eta[1, 1] = 0 @. F = eta*sqrt(forcingcovariancespectrum) - nothing + return nothing end - prob = BarotropicQGQL.ForcedProblem(nx=n, Lx=L, nu=nu, nnu=nnu, mu=mu, dt=dt, + prob = BarotropicQGQL.Problem(dev; nx=n, Lx=L, ν=ν, nν=nν, μ=μ, dt=dt, stepper="RK4", calcF=calcF!, stochastic=true) sol, cl, v, p, g = prob.sol, prob.clock, prob.vars, prob.params, prob.grid @@ -93,7 +94,6 @@ function test_bqgql_stochasticforcingbudgets(; n=256, dt=0.01, L=2π, nu=1e-7, n diags = [E, D, W, R] # Step forward - stepforward!(prob, diags, round(Int, nt)) BarotropicQGQL.updatevars!(prob) @@ -102,7 +102,7 @@ function test_bqgql_stochasticforcingbudgets(; n=256, dt=0.01, L=2π, nu=1e-7, n E, D, W, R = diags - t = round(mu*cl.t, digits=2) + t = round(μ*cl.t, digits=2) i₀ = 1 dEdt = (E[(i₀+1):E.i] - E[i₀:E.i-1])/cl.dt @@ -117,36 +117,36 @@ function test_bqgql_stochasticforcingbudgets(; n=256, dt=0.01, L=2π, nu=1e-7, n residual = dEdt - total - # println(mean(abs.(residual))) - isapprox(mean(abs.(residual)), 0, atol=1e-4) + return isapprox(mean(abs.(residual)), 0, atol=1e-4) end """ - test_stochasticforcingbudgets(; kwargs...) + test_bqgql_deterministicforcingbudgets(dev ; kwargs...) -Tests if the energy budgets are closed for BarotropicQG with stochastic forcing. +Tests if the energy budget is closed for BarotropicQGQL problem with deterministic forcing. """ -function test_bqgql_deterministicforcingbudgets(; n=256, dt=0.01, L=2π, nu=1e-7, nnu=2, mu=1e-1) +function test_bqgql_deterministicforcingbudgets(dev::Device=CPU(); n=256, dt=0.01, L=2π, ν=1e-7, nν=2, μ=1e-1) n, L = 256, 2π - nu, nnu = 1e-7, 2 - mu = 1e-1 - dt, tf = 0.005, 0.1/mu + ν, nν = 1e-7, 2 + μ = 1e-1 + dt, tf = 0.005, 0.1/μ nt = round(Int, tf/dt) - ns = 1 + gr = TwoDGrid(dev, n, L) + x, y = gridpoints(gr) + k0, l0 = 2π/gr.Lx, 2π/gr.Ly # Forcing = 0.01cos(4x)cos(5y)cos(2t) - gr = TwoDGrid(n, L) - x, y = gridpoints(gr) - f = @. 0.01*cos(4*gr.kr[2]*x)*cos(5*gr.l[2]*y) + f = @. 0.01 * cos(4k0*x) * cos(5l0*y) fh = rfft(f) function calcF!(Fh, sol, t, cl, v, p, g) - @. Fh = fh*cos(2*t) - nothing + cos2t = cos(2*t) + @. Fh = fh*cos2t + return nothing end - prob = BarotropicQGQL.ForcedProblem(nx=n, Lx=L, nu=nu, nnu=nnu, mu=mu, dt=dt, + prob = BarotropicQGQL.Problem(dev; nx=n, Lx=L, ν=ν, nν=nν, μ=μ, dt=dt, stepper="RK4", calcF=calcF!, stochastic=false) sol, cl, v, p, g = prob.sol, prob.clock, prob.vars, prob.params, prob.grid @@ -160,15 +160,13 @@ function test_bqgql_deterministicforcingbudgets(; n=256, dt=0.01, L=2π, nu=1e-7 # Step forward - stepforward!(prob, diags, round(Int, nt)) + stepforward!(prob, diags, nt) BarotropicQGQL.updatevars!(prob) - cfl = cl.dt*maximum([maximum(v.v)/g.dx, maximum(v.u)/g.dy]) - E, D, W, R = diags - t = round(mu*cl.t, digits=2) + t = round(μ*cl.t, digits=2) i₀ = 1 dEdt = (E[(i₀+1):E.i] - E[i₀:E.i-1])/cl.dt @@ -181,62 +179,66 @@ function test_bqgql_deterministicforcingbudgets(; n=256, dt=0.01, L=2π, nu=1e-7 residual = dEdt - total # println(mean(abs.(residual))) - isapprox(mean(abs.(residual)), 0, atol=1e-8) + return isapprox(mean(abs.(residual)), 0, atol=1e-8) end """ - test_bqgql_nonlinearadvection(dt, stepper; kwargs...) + test_bqgql_nonlinearadvection(dt, stepper, dev; kwargs...) Tests the advection term in the TwoDNavierStokes module by timestepping a test problem with timestep dt and timestepper identified by the string stepper. The test problem is derived by picking a solution ζf (with associated streamfunction ψf) for which the advection term J(ψf, ζf) is non-zero. Next, a -forcing Ff is derived according to Ff = ∂ζf/∂t + J(ψf, ζf) - nuΔζf. One solution +forcing Ff is derived according to Ff = ∂ζf/∂t + J(ψf, ζf) - νΔζf. One solution to the vorticity equation forced by this Ff is then ζf. (This solution may not be realized, at least at long times, if it is unstable.) """ -function test_bqgql_advection(dt, stepper; n=128, L=2π, nu=1e-2, nnu=1, mu=0.0) +function test_bqgql_advection(dt, stepper, dev::Device=CPU(); n=128, L=2π, ν=1e-2, nν=1, μ=0.0) n, L = 128, 2π - nu, nnu = 1e-2, 1 - mu = 0.0 + ν, nν = 1e-2, 1 + μ = 0.0 tf = 1.0 nt = round(Int, tf/dt) - gr = TwoDGrid(n, L) + gr = TwoDGrid(dev, n, L) x, y = gridpoints(gr) psif = @. cos(3y) + sin(2x)*cos(2y) + 2sin(x)*cos(3y) qf = @. - 9cos(3y) - 8sin(2x)*cos(2y) - 20sin(x)*cos(3y) - Ff = @. nu*( 81cos(3y) + 200cos(3y)*sin(x) + 64cos(2y)*sin(2x) ) - + Ff = @. ν*( 81cos(3y) + 200cos(3y)*sin(x) + 64cos(2y)*sin(2x) ) - 3sin(3y)*(-16cos(2x)*cos(2y) - 20cos(x)*cos(3y)) - 27sin(3y)*(2cos(2x)*cos(2y) + 2cos(x)*cos(3y)) + 0*(-8cos(x)*cos(3y)*sin(2x)*sin(2y) + 24*cos(2x)*cos(2y)*sin(x)*sin(3y)) - Ffh = -rfft(Ff) # Forcing function calcF!(Fh, sol, t, cl, v, p, g) Fh .= Ffh - nothing + return nothing end - prob = BarotropicQGQL.ForcedProblem(nx=n, Lx=L, nu=nu, nnu=nnu, mu=mu, dt=dt, stepper=stepper, calcF=calcF!) + prob = BarotropicQGQL.Problem(dev; nx=n, Lx=L, ν=ν, nν=nν, μ=μ, dt=dt, stepper=stepper, calcF=calcF!) sol, cl, v, p, g = prob.sol, prob.clock, prob.vars, prob.params, prob.grid BarotropicQGQL.set_zeta!(prob, qf) # Step forward stepforward!(prob, round(Int, nt)) BarotropicQGQL.updatevars!(prob) - isapprox(v.zeta+v.Zeta, qf, rtol=1e-13) + return isapprox(v.zeta+v.Zeta, qf, rtol=1e-13) end -function test_bqgql_energyenstrophy() +""" + test_bqgql_energyenstrophy(dev) + +Tests the energy and enstrophy function for a BarotropicQGQL problem. +""" +function test_bqgql_energyenstrophy(dev::Device=CPU()) nx, Lx = 64, 2π ny, Ly = 64, 3π - g = TwoDGrid(nx, Lx, ny, Ly) - k0, l0 = g.k[2], g.l[2] # fundamental wavenumbers + g = TwoDGrid(dev, nx, Lx, ny, Ly) + k0, l0 = 2π/g.Lx, 2π/g.Ly # fundamental wavenumbers x, y = gridpoints(g) energy_calc = 29/9 @@ -246,7 +248,7 @@ function test_bqgql_energyenstrophy() psi0 = @. sin(2k0*x)*cos(2l0*y) + 2sin(k0*x)*cos(3l0*y) zeta0 = @. -((2k0)^2+(2l0)^2)*sin(2k0*x)*cos(2l0*y) - (k0^2+(3l0)^2)*2sin(k0*x)*cos(3l0*y) - prob = BarotropicQGQL.InitialValueProblem(nx=nx, Lx=Lx, ny=ny, Ly=Ly, eta=eta, stepper="ForwardEuler") + prob = BarotropicQGQL.Problem(dev; nx=nx, Lx=Lx, ny=ny, Ly=Ly, eta=eta, stepper="ForwardEuler") sol, cl, v, p, g = prob.sol, prob.clock, prob.vars, prob.params, prob.grid BarotropicQGQL.set_zeta!(prob, zeta0) @@ -255,11 +257,13 @@ function test_bqgql_energyenstrophy() energyzeta0 = BarotropicQGQL.energy(prob) enstrophyzeta0 = BarotropicQGQL.enstrophy(prob) - isapprox(energyzeta0, energy_calc, rtol=1e-13) && isapprox(enstrophyzeta0, enstrophy_calc, rtol=1e-13) && BarotropicQGQL.addforcing!(prob.timestepper.N, sol, cl.t, cl, v, p, g)==nothing + return isapprox(energyzeta0, energy_calc, rtol=1e-13) && isapprox(enstrophyzeta0, enstrophy_calc, rtol=1e-13) && BarotropicQGQL.addforcing!(prob.timestepper.N, sol, cl.t, cl, v, p, g)==nothing end -function test_bqgql_problemtype(T=Float32) - prob = BarotropicQGQL.Problem(T=T) - - (typeof(prob.sol)==Array{Complex{T},2} && typeof(prob.grid.Lx)==T && eltype(prob.grid.x)==T && typeof(prob.vars.u)==Array{T,2}) +function test_bqgql_problemtype(dev, T) + prob = BarotropicQGQL.Problem(dev; T=T) + + A = ArrayType(dev) + + return (typeof(prob.sol)<:A{Complex{T}, 2} && typeof(prob.grid.Lx)==T && eltype(prob.grid.x)==T && typeof(prob.vars.u)<:A{T, 2}) end From 7d6a4c89386004e6404ad3197b09ffd223885a8c Mon Sep 17 00:00:00 2001 From: Navid Constantinou Date: Thu, 28 May 2020 16:10:46 +1000 Subject: [PATCH 2/3] =?UTF-8?q?beta=20->=20=CE=B2,=20mu=20->=20=CE=BC?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- examples/barotropicqgql_betaforced.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/barotropicqgql_betaforced.jl b/examples/barotropicqgql_betaforced.jl index 70010474..d55f9e12 100644 --- a/examples/barotropicqgql_betaforced.jl +++ b/examples/barotropicqgql_betaforced.jl @@ -77,7 +77,7 @@ nothing # hide # a viscosity coefficient ν leads to the module's default value: ν=0. In this # example numerical instability due to accumulation of enstrophy in high wavenumbers # is taken care with the `FilteredTimestepper` we picked. -prob = BarotropicQGQL.Problem(nx=nx, Lx=Lx, beta=β, mu=μ, dt=dt, stepper=stepper, +prob = BarotropicQGQL.Problem(nx=nx, Lx=Lx, β=β, μ=μ, dt=dt, stepper=stepper, calcF=calcF!, stochastic=true) nothing # hide From 7dfb7014c555ea2ddfcc504cdacfc82e8d1696c0 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Fri, 29 May 2020 09:28:41 +1000 Subject: [PATCH 3/3] simplifies eta and params/vars style within Problem() --- src/barotropicqgql.jl | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/src/barotropicqgql.jl b/src/barotropicqgql.jl index 4b32af0c..40ceece0 100644 --- a/src/barotropicqgql.jl +++ b/src/barotropicqgql.jl @@ -60,15 +60,19 @@ function Problem(dev::Device=CPU(); x, y = gridpoints(grid) # topographic PV - if eta==nothing - eta = zeros(dev, T, (nx, ny)) - end + eta === nothing && ( eta = zeros(dev, T, (nx, ny)) ) - params = !(typeof(eta)<:ArrayType(dev)) ? Params(grid, β, eta, μ, ν, nν, calcF) : Params(β, eta, rfft(eta), μ, ν, nν, calcF) + params = !(typeof(eta)<:ArrayType(dev)) ? + Params(grid, β, eta, μ, ν, nν, calcF) : + Params(β, eta, rfft(eta), μ, ν, nν, calcF) - vars = calcF == nothingfunction ? Vars(dev, grid) : (stochastic ? StochasticForcedVars(dev, grid) : ForcedVars(dev, grid)) + vars = calcF == nothingfunction ? + Vars(dev, grid) : stochastic ? + StochasticForcedVars(dev, grid) : + ForcedVars(dev, grid) equation = BarotropicQGQL.Equation(params, grid) + FourierFlows.Problem(equation, stepper, dt, grid, vars, params, dev) end