From 8a5b1e44b1e12e33b4f7488c74facf517dc4ba60 Mon Sep 17 00:00:00 2001 From: Brodie Pearson Date: Wed, 7 Apr 2021 13:35:41 -0700 Subject: [PATCH 01/29] Add hypo-viscosity to SQG scheme --- src/surfaceqg.jl | 116 +++++++++++++++++++++++++++------------- src/twodnavierstokes.jl | 74 ++++++++++++------------- 2 files changed, 117 insertions(+), 73 deletions(-) diff --git a/src/surfaceqg.jl b/src/surfaceqg.jl index 39e13898..a057b50d 100644 --- a/src/surfaceqg.jl +++ b/src/surfaceqg.jl @@ -7,7 +7,8 @@ export kinetic_energy, buoyancy_variance, - buoyancy_dissipation, + buoyancy_dissipation_hyperviscosity, + buoyancy_dissipation_hypoviscosity, buoyancy_work, buoyancy_advection, kinetic_energy_advection @@ -36,9 +37,11 @@ function Problem(dev::Device=CPU(); ny = nx, Ly = Lx, dt = 0.01, - # Hyper-viscosity parameters + # Drag and/or hyper-/hypo-viscosity parameters ν = 0, nν = 1, + μ = 0, + nμ = 0, # Timestepper and equation options stepper = "RK4", calcF = nothingfunction, @@ -47,7 +50,7 @@ function Problem(dev::Device=CPU(); grid = TwoDGrid(dev, nx, Lx, ny, Ly; T=T) - params = Params{T}(ν, nν, calcF) + params = Params{T}(ν, nν, μ, nμ, calcF) vars = calcF == nothingfunction ? DecayingVars(dev, grid) : (stochastic ? StochasticForcedVars(dev, grid) : ForcedVars(dev, grid)) @@ -62,7 +65,7 @@ end # ---------- """ - Params{T}(ν, nν, calcF!) + Params{T}(ν, nν, μ, nμ, calcF!) A struct containing the parameters for Surface QG dynamics. Included are: @@ -73,6 +76,10 @@ struct Params{T} <: AbstractParams ν :: T "buoyancy (hyper)-viscosity order" nν :: Int + "buoyancy (hypo)-viscosity coefficient" + μ :: T + "buoyancy (hypo)-viscosity order" + nμ :: Int "function that calculates the Fourier transform of the forcing, ``F̂``" calcF! :: Function end @@ -87,21 +94,22 @@ Params(ν, nν) = Params(ν, nν, nothingfunction) """ Equation(params, grid) -Return the `equation` for surface QG dynamics with `params` and `grid`. The linear -opeartor ``L`` includes (hyper)-viscosity of order ``n_ν`` with coefficient ``ν``, +Return the `equation` for surface QG dynamics with `params` and `grid`. The linear +opeartor ``L`` includes (hyper)-viscosity of order ``n_ν`` with coefficient ``ν`` and +hypo-viscocity of order ``n_μ`` with coefficient ``μ``, ```math -L = - ν |𝐤|^{2 n_ν} . +L = - ν |𝐤|^{2 n_ν} - μ |𝐤|^{2 n_μ} . ``` -Plain old viscocity corresponds to ``n_ν=1``. +Plain old viscocity corresponds to ``n_ν=1`` while ``n_μ=0`` corresponds to linear drag. The nonlinear term is computed via function `calcN!()`. """ function Equation(params::Params, grid::AbstractGrid) - L = @. - params.ν * grid.Krsq^params.nν + L = @. - params.ν * grid.Krsq^params.nν - params.μ * grid.Krsq^params.nμ CUDA.@allowscalar L[1, 1] = 0 - + return FourierFlows.Equation(L, calcN!, grid) end @@ -151,7 +159,7 @@ function DecayingVars(::Dev, grid::AbstractGrid) where Dev T = eltype(grid) @devzeros Dev T (grid.nx, grid.ny) b u v @devzeros Dev Complex{T} (grid.nkr, grid.nl) bh uh vh - + return Vars(b, u, v, bh, uh, vh, nothing, nothing) end @@ -164,7 +172,7 @@ function ForcedVars(dev::Dev, grid) where Dev T = eltype(grid) @devzeros Dev T (grid.nx, grid.ny) b u v @devzeros Dev Complex{T} (grid.nkr, grid.nl) bh uh vh Fh - + return Vars(b, u, v, bh, uh, vh, Fh, nothing) end @@ -175,10 +183,10 @@ Return the `vars` for stochastically forced surface QG dynamics on device `dev` """ function StochasticForcedVars(dev::Dev, grid) where Dev T = eltype(grid) - + @devzeros Dev T (grid.nx, grid.ny) b u v @devzeros Dev Complex{T} (grid.nkr, grid.nl) bh uh vh Fh prevsol - + return Vars(b, u, v, bh, uh, vh, Fh, prevsol) end @@ -190,7 +198,7 @@ end """ calcN_advection(N, sol, t, clock, vars, params, grid) -Calculate the Fourier transform of the advection term, ``- 𝖩(ψ, b)`` in conservative +Calculate the Fourier transform of the advection term, ``- 𝖩(ψ, b)`` in conservative form, i.e., ``- ∂_x[(∂_y ψ)b] - ∂_y[(∂_x ψ)b]`` and store it in `N`: ```math @@ -205,10 +213,10 @@ function calcN_advection!(N, sol, t, clock, vars, params, grid) ldiv!(vars.u, grid.rfftplan, vars.uh) ldiv!(vars.v, grid.rfftplan, vars.vh) ldiv!(vars.b, grid.rfftplan, vars.bh) - + ub, ubh = vars.u, vars.uh # use vars.u, vars.uh as scratch variables vb, vbh = vars.v, vars.vh # use vars.v, vars.vh as scratch variables - + @. ub *= vars.b # u*b @. vb *= vars.b # v*b @@ -216,7 +224,7 @@ function calcN_advection!(N, sol, t, clock, vars, params, grid) mul!(vbh, grid.rfftplan, vb) # \hat{v*b} @. N = - im * grid.kr * ubh - im * grid.l * vbh - + return nothing end @@ -232,14 +240,14 @@ N = - \\widehat{𝖩(ψ, b)} + F̂ . 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) - + return nothing end """ addforcing!(N, sol, t, clock, vars, params, grid) -When the problem includes forcing, calculate the forcing term ``F̂`` and add it to the +When the problem includes forcing, calculate the forcing term ``F̂`` and add it to the nonlinear term ``N``. """ addforcing!(N, sol, t, clock, vars::Vars, params, grid) = nothing @@ -247,7 +255,7 @@ addforcing!(N, sol, t, clock, vars::Vars, params, grid) = nothing function addforcing!(N, sol, t, clock, vars::ForcedVars, params, grid) params.calcF!(vars.Fh, sol, t, clock, vars, params, grid) @. N += vars.Fh - + return nothing end @@ -256,9 +264,9 @@ function addforcing!(N, sol, t, clock, vars::StochasticForcedVars, params, grid) @. 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 += vars.Fh - + return nothing end @@ -274,15 +282,15 @@ Update variables in `vars` with solution in `sol`. """ function updatevars!(prob) vars, grid, sol = prob.vars, prob.grid, prob.sol - + @. vars.bh = sol @. vars.uh = im * grid.l * sqrt(grid.invKrsq) * sol @. vars.vh = - im * grid.kr * sqrt(grid.invKrsq) * sol - + ldiv!(vars.b, grid.rfftplan, deepcopy(vars.bh)) ldiv!(vars.u, grid.rfftplan, deepcopy(vars.uh)) ldiv!(vars.v, grid.rfftplan, deepcopy(vars.vh)) - + return nothing end @@ -294,9 +302,9 @@ Set the solution `sol` as the transform of `b` and update all variables. function set_b!(prob, b) mul!(prob.sol, prob.grid.rfftplan, b) CUDA.@allowscalar prob.sol[1, 1] = 0 # zero domain average - + updatevars!(prob) - + return nothing end @@ -314,10 +322,10 @@ In SQG, this is identical to half the domain-averaged surface buoyancy variance. ψh = vars.uh # use vars.uh as scratch variable kinetic_energyh = vars.bh # use vars.bh as scratch variable - + @. ψh = sqrt(grid.invKrsq) * sol @. kinetic_energyh = 1 / 2 * grid.Krsq * abs2(ψh) - + return 1 / (grid.Lx * grid.Ly) * parsevalsum(kinetic_energyh, grid) end @@ -328,17 +336,17 @@ Return the buoyancy variance, ```math \\int b² \\frac{𝖽x 𝖽y}{L_x L_y} = \\sum_{𝐤} |b̂|² . ``` -In SQG, this is identical to the velocity variance (i.e., twice the domain-averaged kinetic +In SQG, this is identical to the velocity variance (i.e., twice the domain-averaged kinetic energy). """ @inline function buoyancy_variance(prob) sol, grid = prob.sol, prob.grid - + return 1 / (grid.Lx * grid.Ly) * parsevalsum(abs2.(sol), grid) end """ - buoyancy_dissipation(prob) + buoyancy_dissipation_hyperviscosity(prob) Return the domain-averaged dissipation rate of surface buoyancy variance due to small scale (hyper)-viscosity, @@ -351,13 +359,49 @@ In SQG, this is identical to twice the rate of kinetic energy dissipation @inline function buoyancy_dissipation(prob) sol, vars, params, grid = prob.sol, prob.vars, prob.params, prob.grid buoyancy_dissipationh = vars.uh # use vars.uh as scratch variable - + @. buoyancy_dissipationh = 2 * params.ν * grid.Krsq^params.nν * abs2(sol) CUDA.@allowscalar buoyancy_dissipationh[1, 1] = 0 - + + return 1 / (grid.Lx * grid.Ly) * parsevalsum(buoyancy_dissipationh, grid) +end + +""" + buoyancy_dissipation(prob, ξ, nξ) + +Return the domain-averaged dissipation rate of surface buoyancy variance due +to large scale (hypo)-viscosity, +```math +2 ξ (-1)^{n_ξ+1} \\int b ∇^{2n_ξ} b \\frac{𝖽x 𝖽y}{L_x L_y} = - 2 ν \\sum_{𝐤} |𝐤|^{2n_ξ} |b̂|² , +``` +where ``ν`` the (hyper)-viscosity coefficient ``ν`` and ``nν`` the (hyper)-viscosity order. +In SQG, this is identical to twice the rate of kinetic energy dissipation +""" +@inline function buoyancy_dissipation(prob, ξ, nξ) + sol, vars, params, grid = prob.sol, prob.vars, prob.params, prob.grid + buoyancy_dissipationh = vars.uh # use vars.uh as scratch variable + + @. buoyancy_dissipationh = 2 * ξ * grid.Krsq^nξ * abs2(sol) + CUDA.@allowscalar buoyancy_dissipationh[1, 1] = 0 + return 1 / (grid.Lx * grid.Ly) * parsevalsum(buoyancy_dissipationh, grid) end +""" + buoyancy_dissipation_hyperviscosity(prob, ξ, νξ) + +Return the domain-averaged buoyancy dissipation rate done by the ``ν`` (hyper)-viscosity. +""" +buoyancy_dissipation_hyperviscosity(prob) = buoyancy_dissipation(prob, prob.params.ν, prob.params.nν) + +""" + buoyancy_dissipation_hypoviscosity(prob, ξ, νξ) + +Return the domain-averaged buoyancy dissipation rate done by the ``μ`` (hypo)-viscosity. +""" +buoyancy_dissipation_hypoviscosity(prob) = buoyancy_dissipation(prob, prob.params.μ, prob.params.nμ) + + """ buoyancy_work(prob) buoyancy_work(sol, vars, grid) @@ -376,7 +420,7 @@ end @inline function buoyancy_work(sol, vars::StochasticForcedVars, grid) buoyancy_workh = vars.uh # use vars.uh as scratch variable - + @. buoyancy_workh = (vars.prevsol + sol) * conj(vars.Fh) # Stratonovich return 1 / (grid.Lx * grid.Ly) * parsevalsum(buoyancy_workh, grid) end diff --git a/src/twodnavierstokes.jl b/src/twodnavierstokes.jl index 926b2512..7f4f54c1 100644 --- a/src/twodnavierstokes.jl +++ b/src/twodnavierstokes.jl @@ -96,7 +96,7 @@ Params(ν, nν) = Params(ν, nν, typeof(ν)(0), 0, nothingfunction) Equation(params, grid) Return the `equation` for two-dimensional Navier-Stokes with `params` and `grid`. The linear -operator ``L`` includes (hyper)-viscosity of order ``n_ν`` with coefficient ``ν`` and +operator ``L`` includes (hyper)-viscosity of order ``n_ν`` with coefficient ``ν`` and hypo-viscocity of order ``n_μ`` with coefficient ``μ``, ```math @@ -110,7 +110,7 @@ The nonlinear term is computed via the function `calcN!`. function Equation(params::Params, grid::AbstractGrid) L = @. - params.ν * grid.Krsq^params.nν - params.μ * grid.Krsq^params.nμ CUDA.@allowscalar L[1, 1] = 0 - + return FourierFlows.Equation(L, calcN!, grid) end @@ -154,15 +154,15 @@ const StochasticForcedVars = Vars{<:AbstractArray, <:AbstractArray, <:AbstractAr """ DecayingVars(dev, grid) -Return the `vars` for unforced two-dimensional Navier-Stokes problem on device `dev` and +Return the `vars` for unforced two-dimensional Navier-Stokes problem on device `dev` and with `grid`. """ function DecayingVars(::Dev, grid::AbstractGrid) where Dev T = eltype(grid) - + @devzeros Dev T (grid.nx, grid.ny) ζ u v @devzeros Dev Complex{T} (grid.nkr, grid.nl) ζh uh vh - + return Vars(ζ, u, v, ζh, uh, vh, nothing, nothing) end @@ -173,25 +173,25 @@ Return the `vars` for forced two-dimensional Navier-Stokes on device `dev` and w """ function ForcedVars(dev::Dev, grid::AbstractGrid) where Dev T = eltype(grid) - + @devzeros Dev T (grid.nx, grid.ny) ζ u v @devzeros Dev Complex{T} (grid.nkr, grid.nl) ζh uh vh Fh - + return Vars(ζ, u, v, ζh, uh, vh, Fh, nothing) end """ StochasticForcedVars(dev, grid) -Return the `vars` for stochastically forced two-dimensional Navier-Stokes on device `dev` and +Return the `vars` for stochastically forced two-dimensional Navier-Stokes on device `dev` and with `grid`. """ function StochasticForcedVars(dev::Dev, grid::AbstractGrid) where Dev T = eltype(grid) - + @devzeros Dev T (grid.nx, grid.ny) ζ u v @devzeros Dev Complex{T} (grid.nkr, grid.nl) ζh uh vh Fh prevsol - + return Vars(ζ, u, v, ζh, uh, vh, Fh, prevsol) end @@ -203,7 +203,7 @@ end """ calcN_advection!(N, sol, t, clock, vars, params, grid) -Calculate the Fourier transform of the advection term, ``- 𝖩(ψ, ζ)`` in conservative form, +Calculate the Fourier transform of the advection term, ``- 𝖩(ψ, ζ)`` in conservative form, i.e., ``- ∂_x[(∂_y ψ)ζ] - ∂_y[(∂_x ψ)ζ]`` and store it in `N`: ```math @@ -218,19 +218,19 @@ function calcN_advection!(N, sol, t, clock, vars, params, grid) ldiv!(vars.u, grid.rfftplan, vars.uh) ldiv!(vars.v, grid.rfftplan, vars.vh) ldiv!(vars.ζ, grid.rfftplan, vars.ζh) - + uζ = vars.u # use vars.u as scratch variable @. uζ *= vars.ζ # u*ζ vζ = vars.v # use vars.v as scratch variable @. vζ *= vars.ζ # v*ζ - + uζh = vars.uh # use vars.uh as scratch variable mul!(uζh, grid.rfftplan, uζ) # \hat{u*ζ} vζh = vars.vh # use vars.vh as scratch variable mul!(vζh, grid.rfftplan, vζ) # \hat{v*ζ} @. N = - im * grid.kr * uζh - im * grid.l * vζh - + return nothing end @@ -245,25 +245,25 @@ N = - \\widehat{𝖩(ψ, ζ)} + F̂ . """ 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) - + return nothing end """ addforcing!(N, sol, t, clock, vars, params, grid) -When the problem includes forcing, calculate the forcing term ``F̂`` and add it to the +When the problem includes forcing, calculate the forcing term ``F̂`` and add it to the nonlinear term ``N``. """ addforcing!(N, sol, t, clock, vars::DecayingVars, params, grid) = nothing function addforcing!(N, sol, t, clock, vars::ForcedVars, params, grid) params.calcF!(vars.Fh, sol, t, clock, vars, params, grid) - + @. N += vars.Fh - + return nothing end function addforcing!(N, sol, t, clock, vars::StochasticForcedVars, params, grid) @@ -272,7 +272,7 @@ function addforcing!(N, sol, t, clock, vars::StochasticForcedVars, params, grid) params.calcF!(vars.Fh, sol, t, clock, vars, params, grid) end @. N += vars.Fh - + return nothing end @@ -288,15 +288,15 @@ Update variables in `vars` with solution in `sol`. """ function updatevars!(prob) vars, grid, sol = prob.vars, prob.grid, prob.sol - + @. vars.ζh = sol @. vars.uh = im * grid.l * grid.invKrsq * sol @. vars.vh = - im * grid.kr * grid.invKrsq * sol - + ldiv!(vars.ζ, grid.rfftplan, deepcopy(vars.ζh)) # deepcopy() since inverse real-fft destroys its input ldiv!(vars.u, grid.rfftplan, deepcopy(vars.uh)) # deepcopy() since inverse real-fft destroys its input ldiv!(vars.v, grid.rfftplan, deepcopy(vars.vh)) # deepcopy() since inverse real-fft destroys its input - + return nothing end @@ -307,18 +307,18 @@ Set the solution `sol` as the transform of `ζ` and then update variables in `va """ function set_ζ!(prob, ζ) mul!(prob.sol, prob.grid.rfftplan, ζ) - + CUDA.@allowscalar prob.sol[1, 1] = 0 # enforce zero domain average - + updatevars!(prob) - + return nothing end """ energy(prob) -Return the domain-averaged kinetic energy. Since ``u² + v² = |{\\bf ∇} ψ|²``, the domain-averaged +Return the domain-averaged kinetic energy. Since ``u² + v² = |{\\bf ∇} ψ|²``, the domain-averaged kinetic energy is ```math @@ -328,7 +328,7 @@ kinetic energy is @inline function energy(prob) sol, vars, grid = prob.sol, prob.vars, prob.grid energyh = vars.uh # use vars.uh as scratch variable - + @. energyh = 1 / 2 * grid.invKrsq * abs2(sol) return 1 / (grid.Lx * grid.Ly) * parsevalsum(energyh, grid) end @@ -348,20 +348,20 @@ Returns the domain-averaged enstrophy, end """ - energy_dissipation(prob, ξ, νξ) + energy_dissipation(prob, ξ, nξ) Return the domain-averaged energy dissipation rate done by the viscous term, ```math - ξ (-1)^{n_ξ+1} \\int ψ ∇^{2n_ξ} ζ \\frac{𝖽x 𝖽y}{L_x L_y} = - ξ \\sum_{𝐤} |𝐤|^{2(n_ξ-1)} |ζ̂|² , ``` -where ``ξ`` and ``nξ`` could be either the (hyper)-viscosity coefficient ``ν`` and its order +where ``ξ`` and ``nξ`` could be either the (hyper)-viscosity coefficient ``ν`` and its order ``n_ν``, or the hypo-viscocity coefficient ``μ`` and its order ``n_μ``. """ @inline function energy_dissipation(prob, ξ, nξ) sol, vars, grid = prob.sol, prob.vars, prob.grid energy_dissipationh = vars.uh # use vars.uh as scratch variable - + @. energy_dissipationh = - ξ * grid.Krsq^(nξ - 1) * abs2(sol) CUDA.@allowscalar energy_dissipationh[1, 1] = 0 return 1 / (grid.Lx * grid.Ly) * parsevalsum(energy_dissipationh, grid) @@ -390,13 +390,13 @@ Return the domain-averaged enstrophy dissipation rate done by the viscous term, ξ (-1)^{n_ξ+1} \\int ζ ∇^{2n_ξ} ζ \\frac{𝖽x 𝖽y}{L_x L_y} = - ξ \\sum_{𝐤} |𝐤|^{2n_ξ} |ζ̂|² , ``` -where ``ξ`` and ``nξ`` could be either the (hyper)-viscosity coefficient ``ν`` and its order +where ``ξ`` and ``nξ`` could be either the (hyper)-viscosity coefficient ``ν`` and its order ``n_ν``, or the hypo-viscocity coefficient ``μ`` and its order ``n_μ``. """ @inline function enstrophy_dissipation(prob, ξ, nξ) sol, vars, grid = prob.sol, prob.vars, prob.grid enstrophy_dissipationh = vars.uh # use vars.uh as scratch variable - + @. enstrophy_dissipationh = - ξ * grid.Krsq^nξ * abs2(sol) CUDA.@allowscalar enstrophy_dissipationh[1, 1] = 0 return 1 / (grid.Lx * grid.Ly) * parsevalsum(enstrophy_dissipationh, grid) @@ -428,14 +428,14 @@ Return the domain-averaged rate of work of energy by the forcing ``F``, """ @inline function energy_work(sol, vars::ForcedVars, grid) energy_workh = vars.uh # use vars.uh as scratch variable - + @. energy_workh = grid.invKrsq * sol * conj(vars.Fh) return 1 / (grid.Lx * grid.Ly) * parsevalsum(energy_workh, grid) end @inline function energy_work(sol, vars::StochasticForcedVars, grid) energy_workh = vars.uh # use vars.uh as scratch variable - + @. energy_workh = grid.invKrsq * (vars.prevsol + sol) / 2 * conj(vars.Fh) return 1 / (grid.Lx * grid.Ly) * parsevalsum(energy_workh, grid) end @@ -454,14 +454,14 @@ Return the domain-averaged rate of work of enstrophy by the forcing ``F``, """ @inline function enstrophy_work(sol, vars::ForcedVars, grid) enstrophy_workh = vars.uh # use vars.uh as scratch variable - + @. enstrophy_workh = sol * conj(vars.Fh) return 1 / (grid.Lx * grid.Ly) * parsevalsum(enstrophy_workh, grid) end @inline function enstrophy_work(sol, vars::StochasticForcedVars, grid) enstrophy_workh = vars.uh # use vars.uh as scratch variable - + @. enstrophy_workh = (vars.prevsol + sol) / 2 * conj(vars.Fh) return 1 / (grid.Lx * grid.Ly) * parsevalsum(enstrophy_workh, grid) end From 33a7abcfccd98ac4a4b701ad38c39c43caca4f32 Mon Sep 17 00:00:00 2001 From: Brodie Pearson Date: Fri, 31 Dec 2021 14:20:39 -0800 Subject: [PATCH 02/29] Create twodnavierstokeswithtracer.jl --- src/twodnavierstokeswithtracer.jl | 509 ++++++++++++++++++++++++++++++ 1 file changed, 509 insertions(+) create mode 100644 src/twodnavierstokeswithtracer.jl diff --git a/src/twodnavierstokeswithtracer.jl b/src/twodnavierstokeswithtracer.jl new file mode 100644 index 00000000..b4826a84 --- /dev/null +++ b/src/twodnavierstokeswithtracer.jl @@ -0,0 +1,509 @@ +module TwoDNavierStokes + +export + Problem, + set_ζ!, + updatevars!, + + energy, + energy_dissipation_hyperviscosity, + energy_dissipation_hypoviscosity, + energy_work, + enstrophy, + enstrophy_dissipation_hyperviscosity, + enstrophy_dissipation_hypoviscosity, + enstrophy_work + +using + CUDA, + Reexport, + DocStringExtensions + +@reexport using FourierFlows + +using LinearAlgebra: mul!, ldiv! +using FourierFlows: parsevalsum + +nothingfunction(args...) = nothing + +""" + Problem(dev::Device=CPU(); + nx = 256, + ny = nx, + Lx = 2π, + Ly = Lx, + ν = 0, + nν = 1, + μ = 0, + nμ = 0, + dt = 0.01, + stepper = "RK4", + calcF = nothingfunction, + stochastic = false, + aliased_fraction = 1/3, + T = Float64) + +Construct a two-dimensional Navier-Stokes `problem` on device `dev`. + +Keyword arguments +================= + - `dev`: (required) `CPU()` or `GPU()`; computer architecture used to time-step `problem`. + - `nx`: Number of grid points in ``x``-domain. + - `ny`: Number of grid points in ``y``-domain. + - `Lx`: Extent of the ``x``-domain. + - `Ly`: Extent of the ``y``-domain. + - `ν`: Small-scale (hyper)-viscosity coefficient. + - `nν`: (Hyper)-viscosity order, `nν```≥ 1``. + - `μ`: Large-scale (hypo)-viscosity coefficient. + - `nμ`: (Hypo)-viscosity order, `nμ```≤ 0``. + - `dt`: Time-step. + - `stepper`: Time-stepping method. + - `calcF`: Function that calculates the Fourier transform of the forcing, ``F̂``. + - `stochastic`: `true` or `false`; boolean denoting whether `calcF` is temporally stochastic. + - `aliased_fraction`: the fraction of high-wavenumbers that are zero-ed out by `dealias!()`. + - `T`: `Float32` or `Float64`; floating point type used for `problem` data. +""" +function Problem(dev::Device=CPU(); + # Numerical parameters + nx = 256, + ny = nx, + Lx = 2π, + Ly = Lx, + # Drag and/or hyper-/hypo-viscosity + ν = 0, + nν = 1, + μ = 0, + nμ = 0, + # Timestepper and equation options + dt = 0.01, + stepper = "RK4", + calcF = nothingfunction, + stochastic = false, + # Float type and dealiasing + aliased_fraction = 1/3, + T = Float64) + + grid = TwoDGrid(dev, nx, Lx, ny, Ly; aliased_fraction=aliased_fraction, T=T) + + params = Params{T}(ν, nν, μ, nμ, calcF) + + vars = calcF == nothingfunction ? DecayingVars(dev, grid) : (stochastic ? StochasticForcedVars(dev, grid) : ForcedVars(dev, grid)) + + equation = Equation(params, grid) + + return FourierFlows.Problem(equation, stepper, dt, grid, vars, params, dev) +end + + +# ---------- +# Parameters +# ---------- + +""" + Params{T}(ν, nν, μ, nμ, calcF!) + +A struct containing the parameters for the two-dimensional Navier-Stokes. Included are: + +$(TYPEDFIELDS) +""" +struct Params{T} <: AbstractParams + "small-scale (hyper)-viscosity coefficient" + ν :: T + "(hyper)-viscosity order, `nν```≥ 1``" + nν :: Int + "large-scale (hypo)-viscosity coefficient" + μ :: T + "(hypo)-viscosity order, `nμ```≤ 0``" + nμ :: Int + "function that calculates the Fourier transform of the forcing, ``F̂``" + calcF! :: Function +end + +Params(ν, nν) = Params(ν, nν, typeof(ν)(0), 0, nothingfunction) + + +# --------- +# Equations +# --------- + +""" + Equation(params, grid) + +Return the `equation` for two-dimensional Navier-Stokes with `params` and `grid`. The linear +operator ``L`` includes (hyper)-viscosity of order ``n_ν`` with coefficient ``ν`` and +hypo-viscocity of order ``n_μ`` with coefficient ``μ``, + +```math +L = - ν |𝐤|^{2 n_ν} - μ |𝐤|^{2 n_μ} . +``` + +Plain old viscocity corresponds to ``n_ν=1`` while ``n_μ=0`` corresponds to linear drag. + +The nonlinear term is computed via the function `calcN!`. +""" +function Equation(params::Params, grid::AbstractGrid) + L = @. - params.ν * grid.Krsq^params.nν - params.μ * grid.Krsq^params.nμ + CUDA.@allowscalar L[1, 1] = 0 + + return FourierFlows.Equation(L, calcN!, grid) +end + + +# ---- +# Vars +# ---- + +abstract type TwoDNavierStokesVars <: AbstractVars end + +""" + Vars{Aphys, Atrans, F, P}(ζ, u, v, ζh, uh, vh, Fh, prevsol) + +The variables for two-dimensional Navier-Stokes: + +$(FIELDS) +""" +struct Vars{Aphys, Atrans, F, P} <: TwoDNavierStokesVars + "relative vorticity" + ζ :: Aphys + "x-component of velocity" + u :: Aphys + "y-component of velocity" + v :: Aphys + "Fourier transform of relative vorticity" + ζh :: Atrans + "Fourier transform of x-component of velocity" + uh :: Atrans + "Fourier transform of y-component of velocity" + vh :: Atrans + "Fourier transform of forcing" + Fh :: F + "`sol` at previous time-step" + prevsol :: P +end + +const DecayingVars = Vars{<:AbstractArray, <:AbstractArray, Nothing, Nothing} +const ForcedVars = Vars{<:AbstractArray, <:AbstractArray, <:AbstractArray, Nothing} +const StochasticForcedVars = Vars{<:AbstractArray, <:AbstractArray, <:AbstractArray, <:AbstractArray} + +""" + DecayingVars(dev, grid) + +Return the `vars` for unforced two-dimensional Navier-Stokes problem on device `dev` and +with `grid`. +""" +function DecayingVars(::Dev, grid::AbstractGrid) where Dev + T = eltype(grid) + + @devzeros Dev T (grid.nx, grid.ny) ζ u v + @devzeros Dev Complex{T} (grid.nkr, grid.nl) ζh uh vh + + return Vars(ζ, u, v, ζh, uh, vh, nothing, nothing) +end + +""" + ForcedVars(dev, grid) + +Return the `vars` for forced two-dimensional Navier-Stokes on device `dev` and with `grid`. +""" +function ForcedVars(dev::Dev, grid::AbstractGrid) where Dev + T = eltype(grid) + + @devzeros Dev T (grid.nx, grid.ny) ζ u v + @devzeros Dev Complex{T} (grid.nkr, grid.nl) ζh uh vh Fh + + return Vars(ζ, u, v, ζh, uh, vh, Fh, nothing) +end + +""" + StochasticForcedVars(dev, grid) + +Return the `vars` for stochastically forced two-dimensional Navier-Stokes on device `dev` and +with `grid`. +""" +function StochasticForcedVars(dev::Dev, grid::AbstractGrid) where Dev + T = eltype(grid) + + @devzeros Dev T (grid.nx, grid.ny) ζ u v + @devzeros Dev Complex{T} (grid.nkr, grid.nl) ζh uh vh Fh prevsol + + return Vars(ζ, u, v, ζh, uh, vh, Fh, prevsol) +end + + +# ------- +# Solvers +# ------- + +""" + calcN_advection!(N, sol, t, clock, vars, params, grid) + +Calculate the Fourier transform of the advection term, ``- 𝖩(ψ, ζ)`` in conservative form, +i.e., ``- ∂_x[(∂_y ψ)ζ] - ∂_y[(∂_x ψ)ζ]`` and store it in `N`: + +```math +N = - \\widehat{𝖩(ψ, ζ)} = - i k_x \\widehat{u ζ} - i k_y \\widehat{v ζ} . +``` +""" +function calcN_advection!(N, sol, t, clock, vars, params, grid) + @. vars.uh = im * grid.l * grid.invKrsq * sol + @. vars.vh = - im * grid.kr * grid.invKrsq * sol + @. vars.ζh = sol + + ldiv!(vars.u, grid.rfftplan, vars.uh) + ldiv!(vars.v, grid.rfftplan, vars.vh) + ldiv!(vars.ζ, grid.rfftplan, vars.ζh) + + uζ = vars.u # use vars.u as scratch variable + @. uζ *= vars.ζ # u*ζ + vζ = vars.v # use vars.v as scratch variable + @. vζ *= vars.ζ # v*ζ + + uζh = vars.uh # use vars.uh as scratch variable + mul!(uζh, grid.rfftplan, uζ) # \hat{u*ζ} + vζh = vars.vh # use vars.vh as scratch variable + mul!(vζh, grid.rfftplan, vζ) # \hat{v*ζ} + + @. N = - im * grid.kr * uζh - im * grid.l * vζh + + return nothing +end + +""" + calcN!(N, sol, t, clock, vars, params, grid) + +Calculate the nonlinear term, that is the advection term and the forcing, + +```math +N = - \\widehat{𝖩(ψ, ζ)} + F̂ . +``` +""" +function calcN!(N, sol, t, clock, vars, params, grid) + dealias!(sol, grid) + + calcN_advection!(N, sol, t, clock, vars, params, grid) + + addforcing!(N, sol, t, clock, vars, params, grid) + + return nothing +end + +""" + addforcing!(N, sol, t, clock, vars, params, grid) + +When the problem includes forcing, calculate the forcing term ``F̂`` and add it to the +nonlinear term ``N``. +""" +addforcing!(N, sol, t, clock, vars::DecayingVars, params, grid) = nothing + +function addforcing!(N, sol, t, clock, vars::ForcedVars, params, grid) + params.calcF!(vars.Fh, sol, t, clock, vars, params, grid) + + @. N += vars.Fh + + return nothing +end +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 += vars.Fh + + return nothing +end + + +# ---------------- +# Helper functions +# ---------------- + +""" + updatevars!(prob) + +Update variables in `vars` with solution in `sol`. +""" +function updatevars!(prob) + vars, grid, sol = prob.vars, prob.grid, prob.sol + + dealias!(sol, grid) + + @. vars.ζh = sol + @. vars.uh = im * grid.l * grid.invKrsq * sol + @. vars.vh = - im * grid.kr * grid.invKrsq * sol + + ldiv!(vars.ζ, grid.rfftplan, deepcopy(vars.ζh)) # deepcopy() since inverse real-fft destroys its input + ldiv!(vars.u, grid.rfftplan, deepcopy(vars.uh)) # deepcopy() since inverse real-fft destroys its input + ldiv!(vars.v, grid.rfftplan, deepcopy(vars.vh)) # deepcopy() since inverse real-fft destroys its input + + return nothing +end + +""" + set_ζ!(prob, ζ) + +Set the solution `sol` as the transform of `ζ` and then update variables in `vars`. +""" +function set_ζ!(prob, ζ) + mul!(prob.sol, prob.grid.rfftplan, ζ) + + CUDA.@allowscalar prob.sol[1, 1] = 0 # enforce zero domain average + + updatevars!(prob) + + return nothing +end + +""" + energy(prob) + +Return the domain-averaged kinetic energy. Since ``u² + v² = |{\\bf ∇} ψ|²``, the domain-averaged +kinetic energy is + +```math +\\int \\frac1{2} |{\\bf ∇} ψ|² \\frac{𝖽x 𝖽y}{L_x L_y} = \\sum_{𝐤} \\frac1{2} |𝐤|² |ψ̂|² . +``` +""" +@inline function energy(prob) + sol, vars, grid = prob.sol, prob.vars, prob.grid + energyh = vars.uh # use vars.uh as scratch variable + + @. energyh = 1 / 2 * grid.invKrsq * abs2(sol) + return 1 / (grid.Lx * grid.Ly) * parsevalsum(energyh, grid) +end + +""" + enstrophy(prob) + +Returns the domain-averaged enstrophy, + +```math +\\int \\frac1{2} ζ² \\frac{𝖽x 𝖽y}{L_x L_y} = \\sum_{𝐤} \\frac1{2} |ζ̂|² . +``` +""" +@inline function enstrophy(prob) + sol, grid = prob.sol, prob.grid + return 1 / (2 * grid.Lx * grid.Ly) * parsevalsum(abs2.(sol), grid) +end + +""" + energy_dissipation(prob, ξ, nξ) + +Return the domain-averaged energy dissipation rate done by the viscous term, + +```math +- ξ (-1)^{n_ξ+1} \\int ψ ∇^{2n_ξ} ζ \\frac{𝖽x 𝖽y}{L_x L_y} = - ξ \\sum_{𝐤} |𝐤|^{2(n_ξ-1)} |ζ̂|² , +``` +where ``ξ`` and ``nξ`` could be either the (hyper)-viscosity coefficient ``ν`` and its order +``n_ν``, or the hypo-viscocity coefficient ``μ`` and its order ``n_μ``. +""" +@inline function energy_dissipation(prob, ξ, nξ) + sol, vars, grid = prob.sol, prob.vars, prob.grid + energy_dissipationh = vars.uh # use vars.uh as scratch variable + + @. energy_dissipationh = - ξ * grid.Krsq^(nξ - 1) * abs2(sol) + CUDA.@allowscalar energy_dissipationh[1, 1] = 0 + return 1 / (grid.Lx * grid.Ly) * parsevalsum(energy_dissipationh, grid) +end + +""" + energy_dissipation_hyperviscosity(prob, ξ, νξ) + +Return the domain-averaged energy dissipation rate done by the ``ν`` (hyper)-viscosity. +""" +energy_dissipation_hyperviscosity(prob) = energy_dissipation(prob, prob.params.ν, prob.params.nν) + +""" + energy_dissipation_hypoviscosity(prob, ξ, νξ) + +Return the domain-averaged energy dissipation rate done by the ``μ`` (hypo)-viscosity. +""" +energy_dissipation_hypoviscosity(prob) = energy_dissipation(prob, prob.params.μ, prob.params.nμ) + +""" + enstrophy_dissipation(prob, ξ, νξ) + +Return the domain-averaged enstrophy dissipation rate done by the viscous term, + +```math +ξ (-1)^{n_ξ+1} \\int ζ ∇^{2n_ξ} ζ \\frac{𝖽x 𝖽y}{L_x L_y} = - ξ \\sum_{𝐤} |𝐤|^{2n_ξ} |ζ̂|² , +``` + +where ``ξ`` and ``nξ`` could be either the (hyper)-viscosity coefficient ``ν`` and its order +``n_ν``, or the hypo-viscocity coefficient ``μ`` and its order ``n_μ``. +""" +@inline function enstrophy_dissipation(prob, ξ, nξ) + sol, vars, grid = prob.sol, prob.vars, prob.grid + enstrophy_dissipationh = vars.uh # use vars.uh as scratch variable + + @. enstrophy_dissipationh = - ξ * grid.Krsq^nξ * abs2(sol) + CUDA.@allowscalar enstrophy_dissipationh[1, 1] = 0 + return 1 / (grid.Lx * grid.Ly) * parsevalsum(enstrophy_dissipationh, grid) +end + +""" + enstrophy_dissipation_hyperviscosity(prob, ξ, νξ) + +Return the domain-averaged enstrophy dissipation rate done by the ``ν`` (hyper)-viscosity. +""" +enstrophy_dissipation_hyperviscosity(prob) = enstrophy_dissipation(prob, prob.params.ν, prob.params.nν) + +""" + enstrophy_dissipation_hypoviscosity(prob, ξ, νξ) + +Return the domain-averaged enstrophy dissipation rate done by the ``μ`` (hypo)-viscosity. +""" +enstrophy_dissipation_hypoviscosity(prob) = enstrophy_dissipation(prob, prob.params.μ, prob.params.nμ) + +""" + energy_work(prob) + energy_work(sol, vars, grid) + +Return the domain-averaged rate of work of energy by the forcing ``F``, + +```math +- \\int ψ F \\frac{𝖽x 𝖽y}{L_x L_y} = - \\sum_{𝐤} ψ̂ F̂^* . +``` +""" +@inline function energy_work(sol, vars::ForcedVars, grid) + energy_workh = vars.uh # use vars.uh as scratch variable + + @. energy_workh = grid.invKrsq * sol * conj(vars.Fh) + return 1 / (grid.Lx * grid.Ly) * parsevalsum(energy_workh, grid) +end + +@inline function energy_work(sol, vars::StochasticForcedVars, grid) + energy_workh = vars.uh # use vars.uh as scratch variable + + @. energy_workh = grid.invKrsq * (vars.prevsol + sol) / 2 * conj(vars.Fh) + return 1 / (grid.Lx * grid.Ly) * parsevalsum(energy_workh, grid) +end + +@inline energy_work(prob) = energy_work(prob.sol, prob.vars, prob.grid) + +""" + enstrophy_work(prob) + enstrophy_work(sol, vars, grid) + +Return the domain-averaged rate of work of enstrophy by the forcing ``F``, + +```math +\\int ζ F \\frac{𝖽x 𝖽y}{L_x L_y} = \\sum_{𝐤} ζ̂ F̂^* . +``` +""" +@inline function enstrophy_work(sol, vars::ForcedVars, grid) + enstrophy_workh = vars.uh # use vars.uh as scratch variable + + @. enstrophy_workh = sol * conj(vars.Fh) + return 1 / (grid.Lx * grid.Ly) * parsevalsum(enstrophy_workh, grid) +end + +@inline function enstrophy_work(sol, vars::StochasticForcedVars, grid) + enstrophy_workh = vars.uh # use vars.uh as scratch variable + + @. enstrophy_workh = (vars.prevsol + sol) / 2 * conj(vars.Fh) + return 1 / (grid.Lx * grid.Ly) * parsevalsum(enstrophy_workh, grid) +end + +@inline enstrophy_work(prob) = enstrophy_work(prob.sol, prob.vars, prob.grid) + +end # module From 523da705a2a24457c339de5f0d346bdc0a230cbd Mon Sep 17 00:00:00 2001 From: Brodie Pearson Date: Fri, 31 Dec 2021 14:48:35 -0800 Subject: [PATCH 03/29] Adds tracer to new 2D dynamics module --- src/twodnavierstokeswithtracer.jl | 224 +++++++++++++++++++----------- 1 file changed, 146 insertions(+), 78 deletions(-) diff --git a/src/twodnavierstokeswithtracer.jl b/src/twodnavierstokeswithtracer.jl index b4826a84..9509b3c6 100644 --- a/src/twodnavierstokeswithtracer.jl +++ b/src/twodnavierstokeswithtracer.jl @@ -1,6 +1,8 @@ -module TwoDNavierStokes +module TwoDNavierStokesTracer export + fwdtransform!, + invtransform!, Problem, set_ζ!, updatevars!, @@ -15,6 +17,7 @@ export enstrophy_work using + FFTW, CUDA, Reexport, DocStringExtensions @@ -22,12 +25,13 @@ using @reexport using FourierFlows using LinearAlgebra: mul!, ldiv! -using FourierFlows: parsevalsum +using FourierFlows: parsevalsum, plan_flows_rfft nothingfunction(args...) = nothing """ - Problem(dev::Device=CPU(); + Problem(ntracers::Int, + dev::Device=CPU(); nx = 256, ny = nx, Lx = 2π, @@ -36,6 +40,8 @@ nothingfunction(args...) = nothing nν = 1, μ = 0, nμ = 0, + κ = 0, + nκ = 1, dt = 0.01, stepper = "RK4", calcF = nothingfunction, @@ -43,10 +49,11 @@ nothingfunction(args...) = nothing aliased_fraction = 1/3, T = Float64) -Construct a two-dimensional Navier-Stokes `problem` on device `dev`. +Construct a two-dimensional Navier-Stokes with tracers `problem` on device `dev`. Keyword arguments ================= + - `ntracers`: (required) Number of tracers. - `dev`: (required) `CPU()` or `GPU()`; computer architecture used to time-step `problem`. - `nx`: Number of grid points in ``x``-domain. - `ny`: Number of grid points in ``y``-domain. @@ -56,6 +63,8 @@ Keyword arguments - `nν`: (Hyper)-viscosity order, `nν```≥ 1``. - `μ`: Large-scale (hypo)-viscosity coefficient. - `nμ`: (Hypo)-viscosity order, `nμ```≤ 0``. + - `κ`: Small-scale diffusivity coefficient. + - `nκ`: Diffusivity order, `nκ```≥ 1``. - `dt`: Time-step. - `stepper`: Time-stepping method. - `calcF`: Function that calculates the Fourier transform of the forcing, ``F̂``. @@ -63,7 +72,8 @@ Keyword arguments - `aliased_fraction`: the fraction of high-wavenumbers that are zero-ed out by `dealias!()`. - `T`: `Float32` or `Float64`; floating point type used for `problem` data. """ -function Problem(dev::Device=CPU(); +function Problem(ntracers::Int, # number of tracers + dev::Device=CPU(); # Numerical parameters nx = 256, ny = nx, @@ -74,6 +84,8 @@ function Problem(dev::Device=CPU(); nν = 1, μ = 0, nμ = 0, + κ = 0.0, + nκ = 0, # Timestepper and equation options dt = 0.01, stepper = "RK4", @@ -85,11 +97,11 @@ function Problem(dev::Device=CPU(); grid = TwoDGrid(dev, nx, Lx, ny, Ly; aliased_fraction=aliased_fraction, T=T) - params = Params{T}(ν, nν, μ, nμ, calcF) + params = Params{T}(ν, nν, μ, nμ, κ, nκ, calcF) - vars = calcF == nothingfunction ? DecayingVars(dev, grid) : (stochastic ? StochasticForcedVars(dev, grid) : ForcedVars(dev, grid)) + vars = calcF == nothingfunction ? DecayingVars(dev, grid, params) : (stochastic ? StochasticForcedVars(dev, grid, params) : ForcedVars(dev, grid, params)) - equation = Equation(params, grid) + equation = Equation(dev, params, grid) return FourierFlows.Problem(equation, stepper, dt, grid, vars, params, dev) end @@ -100,13 +112,15 @@ end # ---------- """ - Params{T}(ν, nν, μ, nμ, calcF!) + Params{T}(ν, nν, μ, nμ, κ, nκ calcF!) -A struct containing the parameters for the two-dimensional Navier-Stokes. Included are: +A struct containing the parameters for the two-dimensional Navier-Stokes with tracers. Included are: $(TYPEDFIELDS) """ -struct Params{T} <: AbstractParams +struct Params{T, Trfft} <: AbstractParams + "number of tracers" + ntracers :: Int "small-scale (hyper)-viscosity coefficient" ν :: T "(hyper)-viscosity order, `nν```≥ 1``" @@ -115,46 +129,83 @@ struct Params{T} <: AbstractParams μ :: T "(hypo)-viscosity order, `nμ```≤ 0``" nμ :: Int + "tracer diffusivity coefficient" + κ :: T + "tracer diffusivity order" + nκ :: Int "function that calculates the Fourier transform of the forcing, ``F̂``" - calcF! :: Function + calcF! :: Function + "rfft plan for FFTs (Derived parameter)" + rfftplan :: Trfft end -Params(ν, nν) = Params(ν, nν, typeof(ν)(0), 0, nothingfunction) +function Params(ntracers, ν, nν, μ, nμ, κ, nκ, β, grid; calcF=nothingfunction, effort=FFTW.MEASURE, dev::Device=CPU()) where TU + T = eltype(grid) + A = ArrayType(dev) + ny, nx = grid.ny , grid.nx + nkr, nl = grid.nkr, grid.nl + kr, l = grid.kr , grid.l + + rfftplanlayered = plan_flows_rfft(A{T, 3}(undef, grid.nx, grid.ny, ntracers + 1), [1, 2]; flags=effort) -# --------- -# Equations -# --------- + return Params(ntracers, ν, nν, μ, nμ, κ, nκ, β, calcF, rfftplanlayered) +end """ - Equation(params, grid) - -Return the `equation` for two-dimensional Navier-Stokes with `params` and `grid`. The linear -operator ``L`` includes (hyper)-viscosity of order ``n_ν`` with coefficient ``ν`` and -hypo-viscocity of order ``n_μ`` with coefficient ``μ``, + fwdtransform!(varh, var, params) +Compute the Fourier transform of `var` and store it in `varh`. +""" +fwdtransform!(varh, var, params::AbstractParams) = mul!(varh, params.rfftplan, var) -```math -L = - ν |𝐤|^{2 n_ν} - μ |𝐤|^{2 n_μ} . -``` +""" + invtransform!(var, varh, params) +Compute the inverse Fourier transform of `varh` and store it in `var`. +""" +invtransform!(var, varh, params::AbstractParams) = ldiv!(var, params.rfftplan, varh) -Plain old viscocity corresponds to ``n_ν=1`` while ``n_μ=0`` corresponds to linear drag. -The nonlinear term is computed via the function `calcN!`. -""" -function Equation(params::Params, grid::AbstractGrid) - L = @. - params.ν * grid.Krsq^params.nν - params.μ * grid.Krsq^params.nμ - CUDA.@allowscalar L[1, 1] = 0 +# --------- +# Equations +# --------- +"Create a variable for number of layers. First layer describes motion, the other layers are tracers." +numberoflayers(params) = params.ntracers + 1 + +""" + Equation(dev, params, grid) + Return the `equation` for two-dimensional dynamics and the `equations`` for tracer evolution under these dynamics, + using `params` and `grid`. The first layer of the linear operator ``L`` describes the dynamics and + includes (hyper)-viscosity of order ``n_ν`` with coefficient ``ν`` and hypo-viscocity of order ``n_μ`` + with coefficient ``μ``. The second layer onwards of ``L`` describe the tracer diffusion of order ``n_κ`` + with coefficient ``κ``. + + ```math + L[:, :, 1] = - ν |𝐤|^{2 n_ν} - μ |𝐤|^{2 n_μ} . + ``` + ```math + L[:, :, 2:ntracers+1] = - κ |𝐤|^{2 n_κ}. + ``` + Plain old viscocity corresponds to ``n_ν=1`` while ``n_μ=0`` corresponds to linear drag. + The nonlinear term is computed via function `calcN!()`. +""" +function Equation(dev, params, grid) + #L = hyperviscosity(dev, params, grid) + nlayers = numberoflayers(params) + T = eltype(grid) + L = ArrayType(dev){T}(undef, (grid.nkr, grid.nl, nlayers)) + @views @. L[:,:,1] = - params.ν * grid.Krsq^params.nν - params.μ * grid.Krsq^params.nμ - params.β * im * grid.kr * sqrt(grid.invKrsq) + @views @. L[:,:,2:nlayers] = - params.κ * grid.Krsq^params.nκ + # Need to add diffusivities for different layers + @views @. L[1, 1, :] = 0 + return FourierFlows.Equation(L, calcN!, grid) end - # ---- # Vars # ---- -abstract type TwoDNavierStokesVars <: AbstractVars end - """ Vars{Aphys, Atrans, F, P}(ζ, u, v, ζh, uh, vh, Fh, prevsol) @@ -162,15 +213,15 @@ The variables for two-dimensional Navier-Stokes: $(FIELDS) """ -struct Vars{Aphys, Atrans, F, P} <: TwoDNavierStokesVars - "relative vorticity" - ζ :: Aphys +struct Vars{Aphys3D, Aphys, Atrans3D, Atrans, F, P} <: AbstractVars + "relative vorticity ([:, :, 1] layer) and tracers (other layers)" + ζ :: Aphys3D "x-component of velocity" u :: Aphys "y-component of velocity" v :: Aphys - "Fourier transform of relative vorticity" - ζh :: Atrans + "Fourier transform of relative vorticity ([:, :, 1] layer) and tracers (other layers)" + ζh :: Atrans3D "Fourier transform of x-component of velocity" uh :: Atrans "Fourier transform of y-component of velocity" @@ -192,10 +243,13 @@ Return the `vars` for unforced two-dimensional Navier-Stokes problem on device ` with `grid`. """ function DecayingVars(::Dev, grid::AbstractGrid) where Dev + nlayers = numberoflayers(params) T = eltype(grid) - @devzeros Dev T (grid.nx, grid.ny) ζ u v - @devzeros Dev Complex{T} (grid.nkr, grid.nl) ζh uh vh + @devzeros Dev T (grid.nx, grid.ny, nlayers) ζ + @devzeros Dev T (grid.nx, grid.ny) u v + @devzeros Dev Complex{T} (grid.nkr, grid.nl, nlayers) ζh + @devzeros Dev Complex{T} (grid.nkr, grid.nl) uh vh return Vars(ζ, u, v, ζh, uh, vh, nothing, nothing) end @@ -206,10 +260,13 @@ end Return the `vars` for forced two-dimensional Navier-Stokes on device `dev` and with `grid`. """ function ForcedVars(dev::Dev, grid::AbstractGrid) where Dev + nlayers = numberoflayers(params) T = eltype(grid) - @devzeros Dev T (grid.nx, grid.ny) ζ u v - @devzeros Dev Complex{T} (grid.nkr, grid.nl) ζh uh vh Fh + @devzeros Dev T (grid.nx, grid.ny, nlayers) ζ + @devzeros Dev T (grid.nx, grid.ny) u v + @devzeros Dev Complex{T} (grid.nkr, grid.nl, nlayers) ζh + @devzeros Dev Complex{T} (grid.nkr, grid.nl) uh vh Fh return Vars(ζ, u, v, ζh, uh, vh, Fh, nothing) end @@ -221,10 +278,13 @@ Return the `vars` for stochastically forced two-dimensional Navier-Stokes on dev with `grid`. """ function StochasticForcedVars(dev::Dev, grid::AbstractGrid) where Dev + nlayers = numberoflayers(params) T = eltype(grid) - @devzeros Dev T (grid.nx, grid.ny) ζ u v - @devzeros Dev Complex{T} (grid.nkr, grid.nl) ζh uh vh Fh prevsol + @devzeros Dev T (grid.nx, grid.ny, nlayers) ζ + @devzeros Dev T (grid.nx, grid.ny) u v + @devzeros Dev Complex{T} (grid.nkr, grid.nl, nlayers) ζh + @devzeros Dev Complex{T} (grid.nkr, grid.nl) uh vh Fh prevsol return Vars(ζ, u, v, ζh, uh, vh, Fh, prevsol) end @@ -243,27 +303,33 @@ i.e., ``- ∂_x[(∂_y ψ)ζ] - ∂_y[(∂_x ψ)ζ]`` and store it in `N`: ```math N = - \\widehat{𝖩(ψ, ζ)} = - i k_x \\widehat{u ζ} - i k_y \\widehat{v ζ} . ``` + +Note the first layer ``N[:, :, 1]`` of this advection term corresponds to advection of vorticity ``ζ``. +The subsequent layers of ``N[:, :, 2:ntracers+1]`` correspond to advection of the each tracer. """ function calcN_advection!(N, sol, t, clock, vars, params, grid) - @. vars.uh = im * grid.l * grid.invKrsq * sol - @. vars.vh = - im * grid.kr * grid.invKrsq * sol + nlayers = numberoflayers(params) @. vars.ζh = sol - ldiv!(vars.u, grid.rfftplan, vars.uh) - ldiv!(vars.v, grid.rfftplan, vars.vh) - ldiv!(vars.ζ, grid.rfftplan, vars.ζh) + for j in 1:nlayers + @. vars.uh = im * grid.l * sqrt(grid.invKrsq) * sol[:, :, 1] + @. vars.vh = - im * grid.kr * sqrt(grid.invKrsq) * sol[:, :, 1] + ldiv!(vars.u, grid.rfftplan, vars.ζh[:, :, j]) + vars.ζ[:, : ,j] = vars.u + ldiv!(vars.u, grid.rfftplan, vars.uh) + ldiv!(vars.v, grid.rfftplan, vars.vh) + + uζ, uζh = vars.u, vars.uh # use vars.u, vars.uh as scratch variables + vζ, vζh = vars.v, vars.vh # use vars.v, vars.vh as scratch variables - uζ = vars.u # use vars.u as scratch variable - @. uζ *= vars.ζ # u*ζ - vζ = vars.v # use vars.v as scratch variable - @. vζ *= vars.ζ # v*ζ + @. uζ *= vars.ζ[:, :, j] # u*ζ [note u is 2D array and ζ is 3D array] + @. vζ *= vars.ζ[:, :, j] # v*ζ [note v is 2D array and ζ is 3D array] - uζh = vars.uh # use vars.uh as scratch variable - mul!(uζh, grid.rfftplan, uζ) # \hat{u*ζ} - vζh = vars.vh # use vars.vh as scratch variable - mul!(vζh, grid.rfftplan, vζ) # \hat{v*ζ} + mul!(uζh, grid.rfftplan, uζ) # \hat{u*ζ} + mul!(vζh, grid.rfftplan, vζ) # \hat{v*ζ} - @. N = - im * grid.kr * uζh - im * grid.l * vζh + @views @. N[:, :, j] = - im * grid.kr * uζh - im * grid.l * vζh + end return nothing end @@ -298,16 +364,16 @@ addforcing!(N, sol, t, clock, vars::DecayingVars, params, grid) = nothing function addforcing!(N, sol, t, clock, vars::ForcedVars, params, grid) params.calcF!(vars.Fh, sol, t, clock, vars, params, grid) - @. N += vars.Fh + @views @. N[:, :, 1] += vars.Fh return nothing end 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 + @. vars.prevsol = sol[:, :, 1] # 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 += vars.Fh + @views @. N[:, :, 1] += vars.Fh return nothing end @@ -328,10 +394,10 @@ function updatevars!(prob) dealias!(sol, grid) @. vars.ζh = sol - @. vars.uh = im * grid.l * grid.invKrsq * sol - @. vars.vh = - im * grid.kr * grid.invKrsq * sol - - ldiv!(vars.ζ, grid.rfftplan, deepcopy(vars.ζh)) # deepcopy() since inverse real-fft destroys its input + @. vars.uh = im * grid.l * grid.invKrsq * sol[:,:,1] + @. vars.vh = - im * grid.kr * grid.invKrsq * sol[:,:,1] + + invtransform!(vars.ζ, deepcopy(vars.ζh), prob.params) # deepcopy() since inverse real-fft destroys its input ldiv!(vars.u, grid.rfftplan, deepcopy(vars.uh)) # deepcopy() since inverse real-fft destroys its input ldiv!(vars.v, grid.rfftplan, deepcopy(vars.vh)) # deepcopy() since inverse real-fft destroys its input @@ -339,14 +405,16 @@ function updatevars!(prob) end """ - set_ζ!(prob, ζ) + set_ζ_and_tracers!(prob, ζ) -Set the solution `sol` as the transform of `ζ` and then update variables in `vars`. +Set the first solution layer `sol[:,:,1]` as the transform of `ζ` and lower layers as transform of `tracers`. +Then update variables in `vars`. """ -function set_ζ!(prob, ζ) +function set_ζ_and_tracers!(prob, ζ, tracers) mul!(prob.sol, prob.grid.rfftplan, ζ) - - CUDA.@allowscalar prob.sol[1, 1] = 0 # enforce zero domain average + full_ζ = cat(ζ, tracers, dims=3) # append b and tracers for use in sol + fwdtransform!(prob.sol, full_ζ, prob.params) + @views CUDA.@allowscalar prob.sol[1, 1, 1] = 0 # zero domain average updatevars!(prob) @@ -367,7 +435,7 @@ kinetic energy is sol, vars, grid = prob.sol, prob.vars, prob.grid energyh = vars.uh # use vars.uh as scratch variable - @. energyh = 1 / 2 * grid.invKrsq * abs2(sol) + @. energyh = 1 / 2 * grid.invKrsq * abs2(sol[:,:,1]) return 1 / (grid.Lx * grid.Ly) * parsevalsum(energyh, grid) end @@ -382,7 +450,7 @@ Returns the domain-averaged enstrophy, """ @inline function enstrophy(prob) sol, grid = prob.sol, prob.grid - return 1 / (2 * grid.Lx * grid.Ly) * parsevalsum(abs2.(sol), grid) + return 1 / (2 * grid.Lx * grid.Ly) * parsevalsum(abs2.(sol[:,:,1]), grid) end """ @@ -400,7 +468,7 @@ where ``ξ`` and ``nξ`` could be either the (hyper)-viscosity coefficient ``ν` sol, vars, grid = prob.sol, prob.vars, prob.grid energy_dissipationh = vars.uh # use vars.uh as scratch variable - @. energy_dissipationh = - ξ * grid.Krsq^(nξ - 1) * abs2(sol) + @. energy_dissipationh = - ξ * grid.Krsq^(nξ - 1) * abs2(sol[:,:,1]) CUDA.@allowscalar energy_dissipationh[1, 1] = 0 return 1 / (grid.Lx * grid.Ly) * parsevalsum(energy_dissipationh, grid) end @@ -435,7 +503,7 @@ where ``ξ`` and ``nξ`` could be either the (hyper)-viscosity coefficient ``ν` sol, vars, grid = prob.sol, prob.vars, prob.grid enstrophy_dissipationh = vars.uh # use vars.uh as scratch variable - @. enstrophy_dissipationh = - ξ * grid.Krsq^nξ * abs2(sol) + @. enstrophy_dissipationh = - ξ * grid.Krsq^nξ * abs2(sol[:,:,1]) CUDA.@allowscalar enstrophy_dissipationh[1, 1] = 0 return 1 / (grid.Lx * grid.Ly) * parsevalsum(enstrophy_dissipationh, grid) end @@ -467,14 +535,14 @@ Return the domain-averaged rate of work of energy by the forcing ``F``, @inline function energy_work(sol, vars::ForcedVars, grid) energy_workh = vars.uh # use vars.uh as scratch variable - @. energy_workh = grid.invKrsq * sol * conj(vars.Fh) + @. energy_workh = grid.invKrsq * sol[:,:,1] * conj(vars.Fh) return 1 / (grid.Lx * grid.Ly) * parsevalsum(energy_workh, grid) end @inline function energy_work(sol, vars::StochasticForcedVars, grid) energy_workh = vars.uh # use vars.uh as scratch variable - @. energy_workh = grid.invKrsq * (vars.prevsol + sol) / 2 * conj(vars.Fh) + @. energy_workh = grid.invKrsq * (vars.prevsol + sol[:,:,1]) / 2 * conj(vars.Fh) return 1 / (grid.Lx * grid.Ly) * parsevalsum(energy_workh, grid) end @@ -493,14 +561,14 @@ Return the domain-averaged rate of work of enstrophy by the forcing ``F``, @inline function enstrophy_work(sol, vars::ForcedVars, grid) enstrophy_workh = vars.uh # use vars.uh as scratch variable - @. enstrophy_workh = sol * conj(vars.Fh) + @. enstrophy_workh = sol[:,:,1] * conj(vars.Fh) return 1 / (grid.Lx * grid.Ly) * parsevalsum(enstrophy_workh, grid) end @inline function enstrophy_work(sol, vars::StochasticForcedVars, grid) enstrophy_workh = vars.uh # use vars.uh as scratch variable - @. enstrophy_workh = (vars.prevsol + sol) / 2 * conj(vars.Fh) + @. enstrophy_workh = (vars.prevsol + sol[:,:,1]) / 2 * conj(vars.Fh) return 1 / (grid.Lx * grid.Ly) * parsevalsum(enstrophy_workh, grid) end From 98fb239d899e357b9f69d046ab5f7ae74910d497 Mon Sep 17 00:00:00 2001 From: Brodie Pearson Date: Fri, 31 Dec 2021 14:48:40 -0800 Subject: [PATCH 04/29] Update GeophysicalFlows.jl --- src/GeophysicalFlows.jl | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/GeophysicalFlows.jl b/src/GeophysicalFlows.jl index b3401da9..6c4971f1 100644 --- a/src/GeophysicalFlows.jl +++ b/src/GeophysicalFlows.jl @@ -16,12 +16,14 @@ using include("utils.jl") include("twodnavierstokes.jl") +include("twodnavierstokeswithtracer.jl") include("singlelayerqg.jl") include("multilayerqg.jl") include("surfaceqg.jl") include("barotropicqgql.jl") @reexport using GeophysicalFlows.TwoDNavierStokes +@reexport using GeophysicalFlows.TwoDNavierStokesTracer @reexport using GeophysicalFlows.SingleLayerQG @reexport using GeophysicalFlows.MultiLayerQG @reexport using GeophysicalFlows.SurfaceQG From 25b0e205c15dd10eb48fbec84fdb87ff30e24b2b Mon Sep 17 00:00:00 2001 From: Brodie Pearson Date: Fri, 31 Dec 2021 14:58:23 -0800 Subject: [PATCH 05/29] Fixes Params typo --- src/twodnavierstokeswithtracer.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/twodnavierstokeswithtracer.jl b/src/twodnavierstokeswithtracer.jl index 9509b3c6..9887ed6e 100644 --- a/src/twodnavierstokeswithtracer.jl +++ b/src/twodnavierstokeswithtracer.jl @@ -97,7 +97,7 @@ function Problem(ntracers::Int, # number of tracers grid = TwoDGrid(dev, nx, Lx, ny, Ly; aliased_fraction=aliased_fraction, T=T) - params = Params{T}(ν, nν, μ, nμ, κ, nκ, calcF) + params = Params(ntracers, ν, nν, μ, nμ, κ, nκ, β, grid; calcF, dev) vars = calcF == nothingfunction ? DecayingVars(dev, grid, params) : (stochastic ? StochasticForcedVars(dev, grid, params) : ForcedVars(dev, grid, params)) From e1e1d55d89a2f1392da104d871e7e1b067d6a739 Mon Sep 17 00:00:00 2001 From: Brodie Pearson Date: Fri, 31 Dec 2021 15:07:13 -0800 Subject: [PATCH 06/29] Removes beta term --- src/twodnavierstokeswithtracer.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/twodnavierstokeswithtracer.jl b/src/twodnavierstokeswithtracer.jl index 9887ed6e..d2e4c38f 100644 --- a/src/twodnavierstokeswithtracer.jl +++ b/src/twodnavierstokeswithtracer.jl @@ -97,7 +97,7 @@ function Problem(ntracers::Int, # number of tracers grid = TwoDGrid(dev, nx, Lx, ny, Ly; aliased_fraction=aliased_fraction, T=T) - params = Params(ntracers, ν, nν, μ, nμ, κ, nκ, β, grid; calcF, dev) + params = Params(ntracers, ν, nν, μ, nμ, κ, nκ, grid; calcF, dev) vars = calcF == nothingfunction ? DecayingVars(dev, grid, params) : (stochastic ? StochasticForcedVars(dev, grid, params) : ForcedVars(dev, grid, params)) From aa8baa30069148679fc389237272cf238db497c8 Mon Sep 17 00:00:00 2001 From: Brodie Pearson Date: Fri, 31 Dec 2021 15:07:18 -0800 Subject: [PATCH 07/29] Update Project.toml --- examples/Project.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/examples/Project.toml b/examples/Project.toml index 1669ce76..3e974ded 100644 --- a/examples/Project.toml +++ b/examples/Project.toml @@ -1,5 +1,6 @@ [deps] CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" +FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341" GeophysicalFlows = "44ee3b1c-bc02-53fa-8355-8e347616e15e" JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819" Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80" From cb1fe5c376f54dcdb8348d1543d10ef375e9447a Mon Sep 17 00:00:00 2001 From: Brodie Pearson Date: Fri, 31 Dec 2021 15:08:43 -0800 Subject: [PATCH 08/29] Removes more betas --- src/twodnavierstokeswithtracer.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/twodnavierstokeswithtracer.jl b/src/twodnavierstokeswithtracer.jl index d2e4c38f..36b9899e 100644 --- a/src/twodnavierstokeswithtracer.jl +++ b/src/twodnavierstokeswithtracer.jl @@ -139,7 +139,7 @@ struct Params{T, Trfft} <: AbstractParams rfftplan :: Trfft end -function Params(ntracers, ν, nν, μ, nμ, κ, nκ, β, grid; calcF=nothingfunction, effort=FFTW.MEASURE, dev::Device=CPU()) where TU +function Params(ntracers, ν, nν, μ, nμ, κ, nκ, grid; calcF=nothingfunction, effort=FFTW.MEASURE, dev::Device=CPU()) where TU T = eltype(grid) A = ArrayType(dev) @@ -149,7 +149,7 @@ function Params(ntracers, ν, nν, μ, nμ, κ, nκ, β, grid; calcF=nothingfunc rfftplanlayered = plan_flows_rfft(A{T, 3}(undef, grid.nx, grid.ny, ntracers + 1), [1, 2]; flags=effort) - return Params(ntracers, ν, nν, μ, nμ, κ, nκ, β, calcF, rfftplanlayered) + return Params(ntracers, ν, nν, μ, nμ, κ, nκ, calcF, rfftplanlayered) end """ From 5b66e9b06d7afd9ce4554a828d53730aeca10561 Mon Sep 17 00:00:00 2001 From: Brodie Pearson Date: Fri, 31 Dec 2021 15:14:15 -0800 Subject: [PATCH 09/29] Allow non-integer diffusivities --- src/twodnavierstokeswithtracer.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/twodnavierstokeswithtracer.jl b/src/twodnavierstokeswithtracer.jl index 36b9899e..b4c1e838 100644 --- a/src/twodnavierstokeswithtracer.jl +++ b/src/twodnavierstokeswithtracer.jl @@ -80,9 +80,9 @@ function Problem(ntracers::Int, # number of tracers Lx = 2π, Ly = Lx, # Drag and/or hyper-/hypo-viscosity - ν = 0, + ν = 0.0, nν = 1, - μ = 0, + μ = 0.0, nμ = 0, κ = 0.0, nκ = 0, From 7ee2473cb9bb16b80066f097aa4545c0d12f5009 Mon Sep 17 00:00:00 2001 From: Brodie Pearson Date: Fri, 31 Dec 2021 15:15:55 -0800 Subject: [PATCH 10/29] Allow Vars to use params --- src/twodnavierstokeswithtracer.jl | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/twodnavierstokeswithtracer.jl b/src/twodnavierstokeswithtracer.jl index b4c1e838..f7d4ef5f 100644 --- a/src/twodnavierstokeswithtracer.jl +++ b/src/twodnavierstokeswithtracer.jl @@ -242,7 +242,7 @@ const StochasticForcedVars = Vars{<:AbstractArray, <:AbstractArray, <:AbstractAr Return the `vars` for unforced two-dimensional Navier-Stokes problem on device `dev` and with `grid`. """ -function DecayingVars(::Dev, grid::AbstractGrid) where Dev +function DecayingVars(::Dev, grid::AbstractGrid, params) where Dev nlayers = numberoflayers(params) T = eltype(grid) @@ -259,7 +259,7 @@ end Return the `vars` for forced two-dimensional Navier-Stokes on device `dev` and with `grid`. """ -function ForcedVars(dev::Dev, grid::AbstractGrid) where Dev +function ForcedVars(dev::Dev, grid::AbstractGrid, params) where Dev nlayers = numberoflayers(params) T = eltype(grid) @@ -277,7 +277,7 @@ end Return the `vars` for stochastically forced two-dimensional Navier-Stokes on device `dev` and with `grid`. """ -function StochasticForcedVars(dev::Dev, grid::AbstractGrid) where Dev +function StochasticForcedVars(dev::Dev, grid::AbstractGrid, params) where Dev nlayers = numberoflayers(params) T = eltype(grid) From acb8e72a430547ff5a18d5b7f69f2ae708f9ca73 Mon Sep 17 00:00:00 2001 From: Brodie Pearson Date: Fri, 31 Dec 2021 15:17:19 -0800 Subject: [PATCH 11/29] Remove final beta --- src/twodnavierstokeswithtracer.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/twodnavierstokeswithtracer.jl b/src/twodnavierstokeswithtracer.jl index f7d4ef5f..7d8997a9 100644 --- a/src/twodnavierstokeswithtracer.jl +++ b/src/twodnavierstokeswithtracer.jl @@ -194,7 +194,7 @@ function Equation(dev, params, grid) nlayers = numberoflayers(params) T = eltype(grid) L = ArrayType(dev){T}(undef, (grid.nkr, grid.nl, nlayers)) - @views @. L[:,:,1] = - params.ν * grid.Krsq^params.nν - params.μ * grid.Krsq^params.nμ - params.β * im * grid.kr * sqrt(grid.invKrsq) + @views @. L[:,:,1] = - params.ν * grid.Krsq^params.nν - params.μ * grid.Krsq^params.nμ @views @. L[:,:,2:nlayers] = - params.κ * grid.Krsq^params.nκ # Need to add diffusivities for different layers @views @. L[1, 1, :] = 0 From a5fbb8919573207a6fb5453a9d6ba16dcbcbb356 Mon Sep 17 00:00:00 2001 From: Brodie Pearson Date: Fri, 31 Dec 2021 15:24:10 -0800 Subject: [PATCH 12/29] Typo in initial condition setter --- src/twodnavierstokeswithtracer.jl | 1 - 1 file changed, 1 deletion(-) diff --git a/src/twodnavierstokeswithtracer.jl b/src/twodnavierstokeswithtracer.jl index 7d8997a9..9b240844 100644 --- a/src/twodnavierstokeswithtracer.jl +++ b/src/twodnavierstokeswithtracer.jl @@ -411,7 +411,6 @@ Set the first solution layer `sol[:,:,1]` as the transform of `ζ` and lower lay Then update variables in `vars`. """ function set_ζ_and_tracers!(prob, ζ, tracers) - mul!(prob.sol, prob.grid.rfftplan, ζ) full_ζ = cat(ζ, tracers, dims=3) # append b and tracers for use in sol fwdtransform!(prob.sol, full_ζ, prob.params) @views CUDA.@allowscalar prob.sol[1, 1, 1] = 0 # zero domain average From bc79601f46d771da1e4f13eb572f36de53c6ce65 Mon Sep 17 00:00:00 2001 From: Brodie Pearson Date: Fri, 31 Dec 2021 15:32:49 -0800 Subject: [PATCH 13/29] Remove Stochastic due to bug --- src/twodnavierstokeswithtracer.jl | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/src/twodnavierstokeswithtracer.jl b/src/twodnavierstokeswithtracer.jl index 9b240844..cc40ca79 100644 --- a/src/twodnavierstokeswithtracer.jl +++ b/src/twodnavierstokeswithtracer.jl @@ -369,11 +369,11 @@ function addforcing!(N, sol, t, clock, vars::ForcedVars, params, grid) return nothing end function addforcing!(N, sol, t, clock, vars::StochasticForcedVars, params, grid) - if t == clock.t # not a substep - @. vars.prevsol = sol[:, :, 1] # sol at previous time-step is needed to compute budgets for stochastic forcing - params.calcF!(vars.Fh, sol, t, clock, vars, params, grid) - end - @views @. N[:, :, 1] += vars.Fh + #if t == clock.t # not a substep + # @. vars.prevsol = sol[:, :, 1] # sol at previous time-step is needed to compute budgets for stochastic forcing + # params.calcF!(vars.Fh, sol, t, clock, vars, params, grid) + #end + #@views @. N[:, :, 1] += vars.Fh return nothing end From 02625e5c43452dcdb0792d48d7d30e46632e7eed Mon Sep 17 00:00:00 2001 From: Brodie Pearson Date: Fri, 31 Dec 2021 15:41:07 -0800 Subject: [PATCH 14/29] Update twodnavierstokeswithtracer.jl --- src/twodnavierstokeswithtracer.jl | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/src/twodnavierstokeswithtracer.jl b/src/twodnavierstokeswithtracer.jl index cc40ca79..6ec0ab6e 100644 --- a/src/twodnavierstokeswithtracer.jl +++ b/src/twodnavierstokeswithtracer.jl @@ -283,8 +283,8 @@ function StochasticForcedVars(dev::Dev, grid::AbstractGrid, params) where Dev @devzeros Dev T (grid.nx, grid.ny, nlayers) ζ @devzeros Dev T (grid.nx, grid.ny) u v - @devzeros Dev Complex{T} (grid.nkr, grid.nl, nlayers) ζh - @devzeros Dev Complex{T} (grid.nkr, grid.nl) uh vh Fh prevsol + @devzeros Dev Complex{T} (grid.nkr, grid.nl, nlayers) ζh prevsol + @devzeros Dev Complex{T} (grid.nkr, grid.nl) uh vh Fh return Vars(ζ, u, v, ζh, uh, vh, Fh, prevsol) end @@ -359,12 +359,13 @@ end When the problem includes forcing, calculate the forcing term ``F̂`` and add it to the nonlinear term ``N``. """ -addforcing!(N, sol, t, clock, vars::DecayingVars, params, grid) = nothing - -function addforcing!(N, sol, t, clock, vars::ForcedVars, params, grid) - params.calcF!(vars.Fh, sol, t, clock, vars, params, grid) +#addforcing!(N, sol, t, clock, vars::DecayingVars, params, grid) = nothing - @views @. N[:, :, 1] += vars.Fh +function addforcing!(N, sol, t, clock, vars::Vars, params, grid) + if !isnothing(vars.Fh) # Ignore if there is no forcing + params.calcF!(vars.Fh, sol, t, clock, vars, params, grid) + @views @. N[:, :, 1] += vars.Fh + end return nothing end From 20d99f7f2d827e1e3ddcae05cea40dd1340277e7 Mon Sep 17 00:00:00 2001 From: Brodie Pearson Date: Fri, 31 Dec 2021 15:49:09 -0800 Subject: [PATCH 15/29] Fixed typo in calcN! (from SQG conversion) --- src/twodnavierstokeswithtracer.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/twodnavierstokeswithtracer.jl b/src/twodnavierstokeswithtracer.jl index 6ec0ab6e..c4e813f1 100644 --- a/src/twodnavierstokeswithtracer.jl +++ b/src/twodnavierstokeswithtracer.jl @@ -312,8 +312,8 @@ function calcN_advection!(N, sol, t, clock, vars, params, grid) @. vars.ζh = sol for j in 1:nlayers - @. vars.uh = im * grid.l * sqrt(grid.invKrsq) * sol[:, :, 1] - @. vars.vh = - im * grid.kr * sqrt(grid.invKrsq) * sol[:, :, 1] + @. vars.uh = im * grid.l * grid.invKrsq * sol[:, :, 1] + @. vars.vh = - im * grid.kr * grid.invKrsq * sol[:, :, 1] ldiv!(vars.u, grid.rfftplan, vars.ζh[:, :, j]) vars.ζ[:, : ,j] = vars.u ldiv!(vars.u, grid.rfftplan, vars.uh) From 118862208f8661d251f7275a0e020739f8349e9a Mon Sep 17 00:00:00 2001 From: Brodie Pearson Date: Fri, 31 Dec 2021 16:01:32 -0800 Subject: [PATCH 16/29] Add example with tracers in 2D flow --- examples/twodnavierstokes_decaying_tracer.jl | 259 +++++++++++++++++++ 1 file changed, 259 insertions(+) create mode 100644 examples/twodnavierstokes_decaying_tracer.jl diff --git a/examples/twodnavierstokes_decaying_tracer.jl b/examples/twodnavierstokes_decaying_tracer.jl new file mode 100644 index 00000000..75c3394b --- /dev/null +++ b/examples/twodnavierstokes_decaying_tracer.jl @@ -0,0 +1,259 @@ +# # 2D decaying turbulence with tracers +# +# A simulation of decaying two-dimensional turbulence with tracers (in this case, 3 tracers with different initial conditions). +# +# ## Install dependencies +# +# First let's make sure we have all required packages installed. + +# ```julia +# using Pkg +# pkg"add GeophysicalFlows, Printf, Random, Plots" +# ``` + +# ## Let's begin +# Let's load `GeophysicalFlows.jl` and some other needed packages. +# +using GeophysicalFlows, Printf, Random, Plots + +using Random: seed! +using GeophysicalFlows: peakedisotropicspectrum + + +# ## Choosing a device: CPU or GPU + +dev = CPU() # Device (CPU/GPU) +nothing # hide + + +# ## Numerical, domain, and simulation parameters +# +# First, we pick some numerical and physical parameters for our model. + +n, L = 128, 2π # grid resolution and domain length +nothing # hide + +# Then we pick the time-stepper parameters + dt = 1e-2 # timestep +nsteps = 4000 # total number of steps + nsubs = 20 # number of steps between each plot + ntracers = 3 # number of tracers +nothing # hide + + +# ## Problem setup +# We initialize a `Problem` by providing a set of keyword arguments. The +# `stepper` keyword defines the time-stepper to be used. +prob = TwoDNavierStokesTracer.Problem(ntracers, dev; nx=n, Lx=L, ny=n, Ly=L, dt=dt, stepper="FilteredRK4") +nothing # hide + +# Next we define some shortcuts for convenience. +sol, clock, vars, grid = prob.sol, prob.clock, prob.vars, prob.grid +x, y = grid.x, grid.y +nothing # hide + + +# ## Setting initial conditions + +# Our initial condition tries to reproduce the initial condition used by McWilliams (_JFM_, 1984). +seed!(1234) +k₀, E₀ = 6, 0.5 +ζ₀ = peakedisotropicspectrum(grid, k₀, E₀, mask=prob.timestepper.filter[:,:,1]) + +# Our initial condition for tracers (``c``). +# For the first tracer we use a gaussian centered at ``(x, y) = (L_x/5, 0)``. +gaussian(x, y, σ) = exp(-(x^2 + y^2) / (2σ^2)) + +amplitude, spread = 0.5, 0.15 +for j in 1:ntracers + if j>1 + global c₀ = cat(c₀, 0.0.*ζ₀, dims=3) + else + global c₀ = 0.0.*ζ₀ + end +end +c₀[:,:,1] = [amplitude * gaussian(x[i] - 0.1 * grid.Lx, y[j], spread) for i=1:grid.nx, j=1:grid.ny] +# For the second tracer we use a wider gaussian centered at a different location +c₀[:,:,2] = [amplitude * gaussian(x[i] - 0.2 * grid.Lx, y[j] + 0.1 * grid.Ly, 4*spread) for i=1:grid.nx, j=1:grid.ny] +# For the third tracer we use a straight band in the x-direction that is an eighth as wide as the domain +width = 1/8 +[c₀[i,j,3] = amplitude for i=1:grid.nx, j=Int( (1/2) * grid.ny ):Int( (1/2 + width) * grid.ny )] + +TwoDNavierStokesTracer.set_ζ_and_tracers!(prob, ζ₀, c₀) +nothing # hide + +# Let's plot the initial vorticity field. Note that when plotting, we decorate the variable +# to be plotted with `Array()` to make sure it is brought back on the CPU when `vars` live on +# the GPU. +heatmap(x, y, Array(vars.ζ[:,:,1]'), + aspectratio = 1, + c = :balance, + clim = (-40, 40), + xlims = (-L/2, L/2), + ylims = (-L/2, L/2), + xticks = -3:3, + yticks = -3:3, + xlabel = "x", + ylabel = "y", + title = "initial vorticity", + framestyle = :box) + + +# ## Diagnostics + +# Create Diagnostics -- `energy` and `enstrophy` functions are imported at the top. +E = Diagnostic(TwoDNavierStokesTracer.energy, prob; nsteps=nsteps) +Z = Diagnostic(TwoDNavierStokesTracer.enstrophy, prob; nsteps=nsteps) +diags = [E, Z] # A list of Diagnostics types passed to "stepforward!" will be updated every timestep. +nothing # hide + + +# ## Output + +# We choose folder for outputing `.jld2` files and snapshots (`.png` files). +filepath = "." +plotpath = "./plots_decayingTwoDNavierStokesTracer" +plotname = "snapshots" +filename = joinpath(filepath, "decayingTwoDNavierStokesTracer.jld2") +nothing # hide + +# Do some basic file management +if isfile(filename); rm(filename); end +if !isdir(plotpath); mkdir(plotpath); end +nothing # hide + +# And then create Output +get_sol(prob) = prob.sol # extracts the Fourier-transformed solution +get_u(prob) = irfft(im * grid.l .* grid.invKrsq .* sol, grid.nx) +out = Output(prob, filename, (:sol, get_sol), (:u, get_u)) +saveproblem(out) +nothing # hide + + +# ## Visualizing the simulation + +# We initialize a plot with the vorticity field and the time-series of +# energy and enstrophy diagnostics. + +p1 = heatmap(x, y, Array(vars.ζ[:,:,1]'), + aspectratio = 1, + c = :balance, + clim = (-40, 40), + xlims = (-L/2, L/2), + ylims = (-L/2, L/2), + xticks = -3:3, + yticks = -3:3, + xlabel = "x", + ylabel = "y", + title = "vorticity, t=" * @sprintf("%.2f", clock.t), + framestyle = :box) + +p2 = plot(2, # this means "a plot with two series" + label = ["energy E(t)/E(0)" "enstrophy Z(t)/Z(0)"], + legend = :right, + linewidth = 2, + alpha = 0.7, + xlabel = "t", + xlims = (0, 41), + ylims = (0, 1.1)) + +ptimeϕ₁ = heatmap(x, y, Array(vars.ζ[:,:,2]'), + aspectratio = 1, + c = :balance, + clim = (-0.5, 0.5), + xlims = (-L/2, L/2), + ylims = (-L/2, L/2), + xticks = -3:3, + yticks = -3:3, + xlabel = "x", + ylabel = "y", + title = "ϕ₁(x, y, t=" * @sprintf("%.2f", clock.t) * ")", + framestyle = :box) + + ptimeϕ₂ = heatmap(x, y, Array(vars.ζ[:,:,3]'), + aspectratio = 1, + c = :balance, + clim = (-0.5, 0.5), + xlims = (-L/2, L/2), + ylims = (-L/2, L/2), + xticks = -3:3, + yticks = -3:3, + xlabel = "x", + ylabel = "y", + title = "ϕ₂(x, y, t=" * @sprintf("%.2f", clock.t) * ")", + framestyle = :box) + + ptimeϕ₃ = heatmap(x, y, Array(vars.ζ[:,:,4]'), + aspectratio = 1, + c = :balance, + clim = (-0.5, 0.5), + xlims = (-L/2, L/2), + ylims = (-L/2, L/2), + xticks = -3:3, + yticks = -3:3, + xlabel = "x", + ylabel = "y", + title = "ϕ₃(x, y, t=" * @sprintf("%.2f", clock.t) * ")", + framestyle = :box) + +p = plot(p1, p2, ptimeϕ₁, ptimeϕ₂, ptimeϕ₃, layout=5, size = (1500, 800)) +#l = @layout Plots.grid(1, 2) +#p = plot(p1, p2, layout = l, size = (800, 360)) + + +# ## Time-stepping the `Problem` forward + +# We time-step the `Problem` forward in time. + +startwalltime = time() + +anim = @animate for j = 0:Int(nsteps/nsubs) + if j % (1000 / nsubs) == 0 + cfl = clock.dt * maximum([maximum(vars.u) / grid.dx, maximum(vars.v) / grid.dy]) + + log = @sprintf("step: %04d, t: %d, cfl: %.2f, ΔE: %.4f, ΔZ: %.4f, walltime: %.2f min", + clock.step, clock.t, cfl, E.data[E.i]/E.data[1], Z.data[Z.i]/Z.data[1], (time()-startwalltime)/60) + + println(log) + end + + p[1][1][:z] = Array(vars.ζ[:,:,1]) + p[1][:title] = "vorticity, t=" * @sprintf("%.2f", clock.t) + push!(p[2][1], E.t[E.i], E.data[E.i]/E.data[1]) + push!(p[2][2], Z.t[Z.i], Z.data[Z.i]/Z.data[1]) + p[3][1][:z] = Array(vars.ζ[:,:,2]) + p[4][1][:z] = Array(vars.ζ[:,:,3]) + p[5][1][:z] = Array(vars.ζ[:,:,4]) + + stepforward!(prob, diags, nsubs) + TwoDNavierStokesTracer.updatevars!(prob) +end + +mp4(anim, "twodturb.mp4", fps=18) + + +# Last we can save the output by calling +# ```julia +# saveoutput(out)` +# ``` + + +# ## Radial energy spectrum + +# After the simulation is done we plot the instantaneous radial energy spectrum to illustrate +# how `FourierFlows.radialspectrum` can be used, + +E = @. 0.5 * (vars.u^2 + vars.v^2) # energy density +Eh = rfft(E) # Fourier transform of energy density +kr, Ehr = FourierFlows.radialspectrum(Eh, grid, refinement=1) # compute radial specturm of `Eh` +nothing # hide + +# and we plot it. +plot(kr, abs.(Ehr), + linewidth = 2, + alpha = 0.7, + xlabel = "kᵣ", ylabel = "∫ |Ê| kᵣ dk_θ", + xlims = (5e-1, grid.nx), + xscale = :log10, yscale = :log10, + title = "Radial energy spectrum", + legend = false) From 7cb364c1f1b88f6804897f7eea6312575d6eb55d Mon Sep 17 00:00:00 2001 From: Brodie Pearson Date: Mon, 3 Jan 2022 13:12:02 -0800 Subject: [PATCH 17/29] Expands to allow stochastic forcing --- src/twodnavierstokeswithtracer.jl | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/src/twodnavierstokeswithtracer.jl b/src/twodnavierstokeswithtracer.jl index c4e813f1..3633526e 100644 --- a/src/twodnavierstokeswithtracer.jl +++ b/src/twodnavierstokeswithtracer.jl @@ -359,22 +359,22 @@ end When the problem includes forcing, calculate the forcing term ``F̂`` and add it to the nonlinear term ``N``. """ -#addforcing!(N, sol, t, clock, vars::DecayingVars, params, grid) = nothing +addforcing!(N, sol, t, clock, vars::DecayingVars, params, grid) = nothing -function addforcing!(N, sol, t, clock, vars::Vars, params, grid) - if !isnothing(vars.Fh) # Ignore if there is no forcing +function addforcing!(N, sol, t, clock, vars::ForcedVars, params, grid) + #if !isnothing(vars.Fh) # Ignore if there is no forcing params.calcF!(vars.Fh, sol, t, clock, vars, params, grid) @views @. N[:, :, 1] += vars.Fh - end + #end return nothing end function addforcing!(N, sol, t, clock, vars::StochasticForcedVars, params, grid) - #if t == clock.t # not a substep - # @. vars.prevsol = sol[:, :, 1] # sol at previous time-step is needed to compute budgets for stochastic forcing - # params.calcF!(vars.Fh, sol, t, clock, vars, params, grid) - #end - #@views @. N[:, :, 1] += vars.Fh + if t == clock.t # not a substep + @. vars.prevsol = sol[:, :, 1] # sol at previous time-step is needed to compute budgets for stochastic forcing + params.calcF!(vars.Fh, sol, t, clock, vars, params, grid) + end + @views @. N[:, :, 1] += vars.Fh return nothing end From bc0144e12845651663c2c2c0fbd8b046e9455463 Mon Sep 17 00:00:00 2001 From: Brodie Pearson Date: Mon, 3 Jan 2022 13:20:39 -0800 Subject: [PATCH 18/29] Adjusted prevsol dimension --- src/twodnavierstokeswithtracer.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/twodnavierstokeswithtracer.jl b/src/twodnavierstokeswithtracer.jl index 3633526e..5825a3df 100644 --- a/src/twodnavierstokeswithtracer.jl +++ b/src/twodnavierstokeswithtracer.jl @@ -283,8 +283,8 @@ function StochasticForcedVars(dev::Dev, grid::AbstractGrid, params) where Dev @devzeros Dev T (grid.nx, grid.ny, nlayers) ζ @devzeros Dev T (grid.nx, grid.ny) u v - @devzeros Dev Complex{T} (grid.nkr, grid.nl, nlayers) ζh prevsol - @devzeros Dev Complex{T} (grid.nkr, grid.nl) uh vh Fh + @devzeros Dev Complex{T} (grid.nkr, grid.nl, nlayers) ζh + @devzeros Dev Complex{T} (grid.nkr, grid.nl) uh vh Fh prevsol return Vars(ζ, u, v, ζh, uh, vh, Fh, prevsol) end From e92e6c01e2424222ab355581246dc982f69e169d Mon Sep 17 00:00:00 2001 From: Brodie Pearson Date: Mon, 3 Jan 2022 13:31:36 -0800 Subject: [PATCH 19/29] Update twodnavierstokeswithtracer.jl --- src/twodnavierstokeswithtracer.jl | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/src/twodnavierstokeswithtracer.jl b/src/twodnavierstokeswithtracer.jl index 5825a3df..ef691e19 100644 --- a/src/twodnavierstokeswithtracer.jl +++ b/src/twodnavierstokeswithtracer.jl @@ -369,15 +369,16 @@ function addforcing!(N, sol, t, clock, vars::ForcedVars, params, grid) return nothing end -function addforcing!(N, sol, t, clock, vars::StochasticForcedVars, params, grid) - if t == clock.t # not a substep - @. vars.prevsol = sol[:, :, 1] # sol at previous time-step is needed to compute budgets for stochastic forcing - params.calcF!(vars.Fh, sol, t, clock, vars, params, grid) - end - @views @. N[:, :, 1] += vars.Fh - return nothing -end +# function addforcing!(N, sol, t, clock, vars::StochasticForcedVars, params, grid) +# if t == clock.t # not a substep +# @. vars.prevsol = sol[:, :, 1] # sol at previous time-step is needed to compute budgets for stochastic forcing +# params.calcF!(vars.Fh, sol, t, clock, vars, params, grid) +# end +# @views @. N[:, :, 1] += vars.Fh + +# return nothing +# end # ---------------- From 82737db391dd8e5160a340bcc7bd28baee922607 Mon Sep 17 00:00:00 2001 From: Brodie Pearson Date: Mon, 3 Jan 2022 13:37:20 -0800 Subject: [PATCH 20/29] Expanded forcing AbstractArrays --- src/twodnavierstokeswithtracer.jl | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/twodnavierstokeswithtracer.jl b/src/twodnavierstokeswithtracer.jl index ef691e19..ce032eb7 100644 --- a/src/twodnavierstokeswithtracer.jl +++ b/src/twodnavierstokeswithtracer.jl @@ -232,9 +232,9 @@ struct Vars{Aphys3D, Aphys, Atrans3D, Atrans, F, P} <: AbstractVars prevsol :: P end -const DecayingVars = Vars{<:AbstractArray, <:AbstractArray, Nothing, Nothing} -const ForcedVars = Vars{<:AbstractArray, <:AbstractArray, <:AbstractArray, Nothing} -const StochasticForcedVars = Vars{<:AbstractArray, <:AbstractArray, <:AbstractArray, <:AbstractArray} +const DecayingVars = Vars{<:AbstractArray, <:AbstractArray, <:AbstractArray, <:AbstractArray, Nothing, Nothing} +const ForcedVars = Vars{<:AbstractArray, <:AbstractArray, <:AbstractArray, <:AbstractArray, <:AbstractArray, Nothing} +const StochasticForcedVars = Vars{<:AbstractArray, <:AbstractArray, <:AbstractArray, <:AbstractArray, <:AbstractArray, <:AbstractArray} """ DecayingVars(dev, grid) From 0cbcdd06822069d91c26099bc07bf6b02e0a5aa4 Mon Sep 17 00:00:00 2001 From: Brodie Pearson Date: Mon, 3 Jan 2022 13:41:33 -0800 Subject: [PATCH 21/29] Add Stochastic forcing option back in --- src/twodnavierstokeswithtracer.jl | 18 ++++++++---------- 1 file changed, 8 insertions(+), 10 deletions(-) diff --git a/src/twodnavierstokeswithtracer.jl b/src/twodnavierstokeswithtracer.jl index ce032eb7..21213e31 100644 --- a/src/twodnavierstokeswithtracer.jl +++ b/src/twodnavierstokeswithtracer.jl @@ -362,23 +362,21 @@ nonlinear term ``N``. addforcing!(N, sol, t, clock, vars::DecayingVars, params, grid) = nothing function addforcing!(N, sol, t, clock, vars::ForcedVars, params, grid) - #if !isnothing(vars.Fh) # Ignore if there is no forcing params.calcF!(vars.Fh, sol, t, clock, vars, params, grid) @views @. N[:, :, 1] += vars.Fh - #end return nothing end -# function addforcing!(N, sol, t, clock, vars::StochasticForcedVars, params, grid) -# if t == clock.t # not a substep -# @. vars.prevsol = sol[:, :, 1] # sol at previous time-step is needed to compute budgets for stochastic forcing -# params.calcF!(vars.Fh, sol, t, clock, vars, params, grid) -# end -# @views @. N[:, :, 1] += vars.Fh +function addforcing!(N, sol, t, clock, vars::StochasticForcedVars, params, grid) + if t == clock.t # not a substep + @. vars.prevsol = sol[:, :, 1] # sol at previous time-step is needed to compute budgets for stochastic forcing + params.calcF!(vars.Fh, sol, t, clock, vars, params, grid) + end + @views @. N[:, :, 1] += vars.Fh -# return nothing -# end + return nothing +end # ---------------- From 79ef3846e30a0f56c78f3ad879d51df1d60a8f48 Mon Sep 17 00:00:00 2001 From: Brodie Pearson Date: Mon, 3 Jan 2022 13:45:42 -0800 Subject: [PATCH 22/29] Copy 2D stochastic example for modification --- ...odnavierstokes_stochasticforcing_tracer.jl | 193 ++++++++++++++++++ 1 file changed, 193 insertions(+) create mode 100644 examples/twodnavierstokes_stochasticforcing_tracer.jl diff --git a/examples/twodnavierstokes_stochasticforcing_tracer.jl b/examples/twodnavierstokes_stochasticforcing_tracer.jl new file mode 100644 index 00000000..7739e80a --- /dev/null +++ b/examples/twodnavierstokes_stochasticforcing_tracer.jl @@ -0,0 +1,193 @@ +# # 2D forced-dissipative turbulence +# +#md # This example can be viewed as a Jupyter notebook via [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/literated/twodnavierstokes_stochasticforcing.ipynb). +# +# A simulation of forced-dissipative two-dimensional turbulence. We solve the +# two-dimensional vorticity equation with stochastic excitation and dissipation in +# the form of linear drag and hyperviscosity. +# +# ## Install dependencies +# +# First let's make sure we have all required packages installed. + +# ```julia +# using Pkg +# pkg"add GeophysicalFlows, Random, Printf, Plots" +# ``` + +# ## Let's begin +# Let's load `GeophysicalFlows.jl` and some other needed packages. +# +using GeophysicalFlows, Random, Printf, Plots +parsevalsum = FourierFlows.parsevalsum + +# ## Choosing a device: CPU or GPU + +dev = CPU() # Device (CPU/GPU) +nothing # hide + + +# ## Numerical, domain, and simulation parameters +# +# First, we pick some numerical and physical parameters for our model. + + n, L = 256, 2π # grid resolution and domain length + ν, nν = 2e-7, 2 # hyperviscosity coefficient and hyperviscosity order + μ, nμ = 1e-1, 0 # linear drag coefficient + dt = 0.005 # timestep +nsteps = 4000 # total number of steps + nsubs = 20 # number of steps between each plot +nothing # hide + + +# ## Forcing + +# We force the vorticity equation with stochastic excitation that is delta-correlated in time +# and while spatially homogeneously and isotropically correlated. The forcing has a spectrum +# with power in a ring in wavenumber space of radius ``k_f`` (`forcing_wavenumber`) and width +# ``δ_f`` (`forcing_bandwidth`), and it injects energy per unit area and per unit time +# equal to ``\varepsilon``. That is, the forcing covariance spectrum is proportional to +# ``\exp{[-(|\bm{k}| - k_f)^2 / (2 δ_f^2)]}``. + +forcing_wavenumber = 14.0 * 2π/L # the forcing wavenumber, `k_f`, for a spectrum that is a ring in wavenumber space +forcing_bandwidth = 1.5 * 2π/L # the width of the forcing spectrum, `δ_f` +ε = 0.1 # energy input rate by the forcing + +grid = TwoDGrid(dev, n, L) + +K = @. sqrt(grid.Krsq) # a 2D array with the total wavenumber + +forcing_spectrum = @. exp(-(K - forcing_wavenumber)^2 / (2 * forcing_bandwidth^2)) +@CUDA.allowscalar forcing_spectrum[grid.Krsq .== 0] .= 0 # ensure forcing has zero domain-average + +ε0 = parsevalsum(forcing_spectrum .* grid.invKrsq / 2, grid) / (grid.Lx * grid.Ly) +@. forcing_spectrum *= ε/ε0 # normalize forcing to inject energy at rate ε +nothing # hide + + +# We reset of the random number generator for reproducibility +if dev==CPU(); Random.seed!(1234); else; CUDA.seed!(1234); end +nothing # hide + + +# Next we construct function `calcF!` that computes a forcing realization every timestep. +# First we make sure that if `dev=GPU()`, then `CUDA.rand()` function is called for random +# numbers uniformly distributed between 0 and 1. +random_uniform = dev==CPU() ? rand : CUDA.rand + +function calcF!(Fh, sol, t, clock, vars, params, grid) + Fh .= sqrt.(forcing_spectrum) .* exp.(2π * im * random_uniform(eltype(grid), size(sol))) ./ sqrt(clock.dt) + + return nothing +end +nothing # hide + + +# ## Problem setup +# We initialize a `Problem` by providing a set of keyword arguments. The +# `stepper` keyword defines the time-stepper to be used. +prob = TwoDNavierStokes.Problem(dev; nx=n, Lx=L, ν=ν, nν=nν, μ=μ, nμ=nμ, dt=dt, stepper="ETDRK4", + calcF=calcF!, stochastic=true) +nothing # hide + +# Define some shortcuts for convenience. +sol, clock, vars, params, grid = prob.sol, prob.clock, prob.vars, prob.params, prob.grid + +x, y = grid.x, grid.y +nothing # hide + + +# First let's see how a forcing realization looks like. Function `calcF!()` computes +# the forcing in Fourier space and saves it into variable `vars.Fh`, so we first need to +# go back to physical space. + +# Note that when plotting, we decorate the variable to be plotted with `Array()` to make sure +# it is brought back on the CPU when the variable lives on the GPU. +calcF!(vars.Fh, sol, 0.0, clock, vars, params, grid) + +heatmap(x, y, Array(irfft(vars.Fh, grid.nx)'), + aspectratio = 1, + c = :balance, + clim = (-200, 200), + xlims = (-L/2, L/2), + ylims = (-L/2, L/2), + xticks = -3:3, + yticks = -3:3, + xlabel = "x", + ylabel = "y", + title = "a forcing realization", + framestyle = :box) + + +# ## Setting initial conditions + +# Our initial condition is a fluid at rest. +TwoDNavierStokes.set_ζ!(prob, ArrayType(dev)(zeros(grid.nx, grid.ny))) + + +# ## Diagnostics + +# Create Diagnostics; the diagnostics are aimed to probe the energy budget. +E = Diagnostic(TwoDNavierStokes.energy, prob, nsteps=nsteps) # energy +Z = Diagnostic(TwoDNavierStokes.enstrophy, prob, nsteps=nsteps) # enstrophy +diags = [E, Z] # a list of Diagnostics passed to `stepforward!` will be updated every timestep. +nothing # hide + + +# ## Visualizing the simulation + +# We initialize a plot with the vorticity field and the time-series of +# energy and enstrophy diagnostics. To plot energy and enstrophy on the same +# axes we scale enstrophy with ``k_f^2``. + +p1 = heatmap(x, y, Array(vars.ζ'), + aspectratio = 1, + c = :balance, + clim = (-40, 40), + xlims = (-L/2, L/2), + ylims = (-L/2, L/2), + xticks = -3:3, + yticks = -3:3, + xlabel = "x", + ylabel = "y", + title = "vorticity, t=" * @sprintf("%.2f", clock.t), + framestyle = :box) + +p2 = plot(2, # this means "a plot with two series" + label = ["energy E(t)" "enstrophy Z(t) / k_f²"], + legend = :right, + linewidth = 2, + alpha = 0.7, + xlabel = "μ t", + xlims = (0, 1.1 * μ * nsteps * dt), + ylims = (0, 0.55)) + +l = @layout Plots.grid(1, 2) +p = plot(p1, p2, layout = l, size = (900, 420)) + + +# ## Time-stepping the `Problem` forward + +# Finally, we time-step the `Problem` forward in time. + +startwalltime = time() + +anim = @animate for j = 0:round(Int, nsteps / nsubs) + if j % (1000/nsubs) == 0 + cfl = clock.dt * maximum([maximum(vars.u) / grid.dx, maximum(vars.v) / grid.dy]) + + log = @sprintf("step: %04d, t: %d, cfl: %.2f, E: %.4f, Z: %.4f, walltime: %.2f min", + clock.step, clock.t, cfl, E.data[E.i], Z.data[Z.i], (time()-startwalltime)/60) + println(log) + end + + p[1][1][:z] = Array(vars.ζ) + p[1][:title] = "vorticity, μt = " * @sprintf("%.2f", μ * clock.t) + push!(p[2][1], μ * E.t[E.i], E.data[E.i]) + push!(p[2][2], μ * Z.t[Z.i], Z.data[Z.i] / forcing_wavenumber^2) + + stepforward!(prob, diags, nsubs) + TwoDNavierStokes.updatevars!(prob) +end + +mp4(anim, "twodturb_forced.mp4", fps=18) From efa0851e575c3659aaaa8378569b897f1f436268 Mon Sep 17 00:00:00 2001 From: Brodie Pearson Date: Mon, 3 Jan 2022 14:14:54 -0800 Subject: [PATCH 23/29] Write stochastic example w/tracers --- ...odnavierstokes_stochasticforcing_tracer.jl | 89 ++++++++++++++++--- 1 file changed, 78 insertions(+), 11 deletions(-) diff --git a/examples/twodnavierstokes_stochasticforcing_tracer.jl b/examples/twodnavierstokes_stochasticforcing_tracer.jl index 7739e80a..e947067b 100644 --- a/examples/twodnavierstokes_stochasticforcing_tracer.jl +++ b/examples/twodnavierstokes_stochasticforcing_tracer.jl @@ -37,6 +37,7 @@ nothing # hide dt = 0.005 # timestep nsteps = 4000 # total number of steps nsubs = 20 # number of steps between each plot + ntracers = 3 # number of tracers nothing # hide @@ -76,7 +77,7 @@ nothing # hide random_uniform = dev==CPU() ? rand : CUDA.rand function calcF!(Fh, sol, t, clock, vars, params, grid) - Fh .= sqrt.(forcing_spectrum) .* exp.(2π * im * random_uniform(eltype(grid), size(sol))) ./ sqrt(clock.dt) + Fh .= sqrt.(forcing_spectrum) .* exp.(2π * im * random_uniform(eltype(grid), size(sol[:,:,1]))) ./ sqrt(clock.dt) return nothing end @@ -86,7 +87,7 @@ nothing # hide # ## Problem setup # We initialize a `Problem` by providing a set of keyword arguments. The # `stepper` keyword defines the time-stepper to be used. -prob = TwoDNavierStokes.Problem(dev; nx=n, Lx=L, ν=ν, nν=nν, μ=μ, nμ=nμ, dt=dt, stepper="ETDRK4", +prob = TwoDNavierStokesTracer.Problem(ntracers, dev; nx=n, Lx=L, ν=ν, nν=nν, μ=μ, nμ=nμ, dt=dt, stepper="ETDRK4", calcF=calcF!, stochastic=true) nothing # hide @@ -122,14 +123,37 @@ heatmap(x, y, Array(irfft(vars.Fh, grid.nx)'), # ## Setting initial conditions # Our initial condition is a fluid at rest. -TwoDNavierStokes.set_ζ!(prob, ArrayType(dev)(zeros(grid.nx, grid.ny))) + +ζ₀ = ArrayType(dev)(zeros(grid.nx, grid.ny)) + +# Our initial condition for tracers (``c``). +# For the first tracer we use a gaussian centered at ``(x, y) = (L_x/5, 0)``. +gaussian(x, y, σ) = exp(-(x^2 + y^2) / (2σ^2)) + +amplitude, spread = 0.5, 0.15 +for j in 1:ntracers + if j>1 + global c₀ = cat(c₀, 0.0.*ζ₀, dims=3) + else + global c₀ = 0.0.*ζ₀ + end +end +c₀[:,:,1] = [amplitude * gaussian(x[i] - 0.1 * grid.Lx, y[j], spread) for i=1:grid.nx, j=1:grid.ny] +# For the second tracer we use a wider gaussian centered at a different location +c₀[:,:,2] = [amplitude * gaussian(x[i] - 0.2 * grid.Lx, y[j] + 0.1 * grid.Ly, 4*spread) for i=1:grid.nx, j=1:grid.ny] +# For the third tracer we use a straight band in the x-direction that is an eighth as wide as the domain +width = 1/8 +[c₀[i,j,3] = amplitude for i=1:grid.nx, j=Int( (1/2) * grid.ny ):Int( (1/2 + width) * grid.ny )] + +TwoDNavierStokesTracer.set_ζ_and_tracers!(prob, ζ₀, c₀) +nothing # hide # ## Diagnostics # Create Diagnostics; the diagnostics are aimed to probe the energy budget. -E = Diagnostic(TwoDNavierStokes.energy, prob, nsteps=nsteps) # energy -Z = Diagnostic(TwoDNavierStokes.enstrophy, prob, nsteps=nsteps) # enstrophy +E = Diagnostic(TwoDNavierStokesTracer.energy, prob, nsteps=nsteps) # energy +Z = Diagnostic(TwoDNavierStokesTracer.enstrophy, prob, nsteps=nsteps) # enstrophy diags = [E, Z] # a list of Diagnostics passed to `stepforward!` will be updated every timestep. nothing # hide @@ -140,7 +164,7 @@ nothing # hide # energy and enstrophy diagnostics. To plot energy and enstrophy on the same # axes we scale enstrophy with ``k_f^2``. -p1 = heatmap(x, y, Array(vars.ζ'), +p1 = heatmap(x, y, Array(vars.ζ[:,:,1]'), aspectratio = 1, c = :balance, clim = (-40, 40), @@ -162,8 +186,48 @@ p2 = plot(2, # this means "a plot with two series" xlims = (0, 1.1 * μ * nsteps * dt), ylims = (0, 0.55)) -l = @layout Plots.grid(1, 2) -p = plot(p1, p2, layout = l, size = (900, 420)) +ptimeϕ₁ = heatmap(x, y, Array(vars.ζ[:,:,2]'), + aspectratio = 1, + c = :balance, + clim = (-0.5, 0.5), + xlims = (-L/2, L/2), + ylims = (-L/2, L/2), + xticks = -3:3, + yticks = -3:3, + xlabel = "x", + ylabel = "y", + title = "ϕ₁(x, y, t=" * @sprintf("%.2f", clock.t) * ")", + framestyle = :box) + + ptimeϕ₂ = heatmap(x, y, Array(vars.ζ[:,:,3]'), + aspectratio = 1, + c = :balance, + clim = (-0.5, 0.5), + xlims = (-L/2, L/2), + ylims = (-L/2, L/2), + xticks = -3:3, + yticks = -3:3, + xlabel = "x", + ylabel = "y", + title = "ϕ₂(x, y, t=" * @sprintf("%.2f", clock.t) * ")", + framestyle = :box) + + ptimeϕ₃ = heatmap(x, y, Array(vars.ζ[:,:,4]'), + aspectratio = 1, + c = :balance, + clim = (-0.5, 0.5), + xlims = (-L/2, L/2), + ylims = (-L/2, L/2), + xticks = -3:3, + yticks = -3:3, + xlabel = "x", + ylabel = "y", + title = "ϕ₃(x, y, t=" * @sprintf("%.2f", clock.t) * ")", + framestyle = :box) + +p = plot(p1, p2, ptimeϕ₁, ptimeϕ₂, ptimeϕ₃, layout=5, size = (1500, 800)) +#l = @layout Plots.grid(1, 2) +#p = plot(p1, p2, layout = l, size = (900, 420)) # ## Time-stepping the `Problem` forward @@ -181,13 +245,16 @@ anim = @animate for j = 0:round(Int, nsteps / nsubs) println(log) end - p[1][1][:z] = Array(vars.ζ) + p[1][1][:z] = Array(vars.ζ[:,:,1]) p[1][:title] = "vorticity, μt = " * @sprintf("%.2f", μ * clock.t) push!(p[2][1], μ * E.t[E.i], E.data[E.i]) push!(p[2][2], μ * Z.t[Z.i], Z.data[Z.i] / forcing_wavenumber^2) + p[3][1][:z] = Array(vars.ζ[:,:,2]) + p[4][1][:z] = Array(vars.ζ[:,:,3]) + p[5][1][:z] = Array(vars.ζ[:,:,4]) stepforward!(prob, diags, nsubs) - TwoDNavierStokes.updatevars!(prob) + TwoDNavierStokesTracer.updatevars!(prob) end -mp4(anim, "twodturb_forced.mp4", fps=18) +mp4(anim, "twodturb_forced_tracer.mp4", fps=18) From 9f2fa8bb8a054b7a958db6f91a653f55766049fe Mon Sep 17 00:00:00 2001 From: Brodie Pearson Date: Mon, 3 Jan 2022 15:02:24 -0800 Subject: [PATCH 24/29] Rename 2D decaying tracer video --- examples/twodnavierstokes_decaying_tracer.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/twodnavierstokes_decaying_tracer.jl b/examples/twodnavierstokes_decaying_tracer.jl index 75c3394b..44cd3e53 100644 --- a/examples/twodnavierstokes_decaying_tracer.jl +++ b/examples/twodnavierstokes_decaying_tracer.jl @@ -229,7 +229,7 @@ anim = @animate for j = 0:Int(nsteps/nsubs) TwoDNavierStokesTracer.updatevars!(prob) end -mp4(anim, "twodturb.mp4", fps=18) +mp4(anim, "twodturbtracer.mp4", fps=18) # Last we can save the output by calling From 91770fe8e86900cabeaa4ff8bf0a0f5c3d480364 Mon Sep 17 00:00:00 2001 From: Brodie Pearson Date: Mon, 3 Jan 2022 15:02:38 -0800 Subject: [PATCH 25/29] Add diffusivity to stochastic 2D tracer flow --- examples/twodnavierstokes_stochasticforcing_tracer.jl | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/examples/twodnavierstokes_stochasticforcing_tracer.jl b/examples/twodnavierstokes_stochasticforcing_tracer.jl index e947067b..4dae2afe 100644 --- a/examples/twodnavierstokes_stochasticforcing_tracer.jl +++ b/examples/twodnavierstokes_stochasticforcing_tracer.jl @@ -34,6 +34,7 @@ nothing # hide n, L = 256, 2π # grid resolution and domain length ν, nν = 2e-7, 2 # hyperviscosity coefficient and hyperviscosity order μ, nμ = 1e-1, 0 # linear drag coefficient + κ, nκ = 2e-7, 2 dt = 0.005 # timestep nsteps = 4000 # total number of steps nsubs = 20 # number of steps between each plot @@ -87,7 +88,7 @@ nothing # hide # ## Problem setup # We initialize a `Problem` by providing a set of keyword arguments. The # `stepper` keyword defines the time-stepper to be used. -prob = TwoDNavierStokesTracer.Problem(ntracers, dev; nx=n, Lx=L, ν=ν, nν=nν, μ=μ, nμ=nμ, dt=dt, stepper="ETDRK4", +prob = TwoDNavierStokesTracer.Problem(ntracers, dev; nx=n, Lx=L, ν=ν, nν=nν, μ=μ, nμ=nμ, κ=κ, nκ=nκ, dt=dt, stepper="ETDRK4", calcF=calcF!, stochastic=true) nothing # hide From 115e2501ef792413e3596ff9c6b5cb69ed887a55 Mon Sep 17 00:00:00 2001 From: Brodie Pearson Date: Mon, 3 Jan 2022 15:18:16 -0800 Subject: [PATCH 26/29] Dummy commit 1/2 --- src/surfaceqg.jl | 2 +- src/twodnavierstokes.jl | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/surfaceqg.jl b/src/surfaceqg.jl index 1b87b28c..65942fa9 100644 --- a/src/surfaceqg.jl +++ b/src/surfaceqg.jl @@ -1,7 +1,7 @@ module SurfaceQG export - Problem, + Problem set_b!, updatevars!, diff --git a/src/twodnavierstokes.jl b/src/twodnavierstokes.jl index b4826a84..e65e4aac 100644 --- a/src/twodnavierstokes.jl +++ b/src/twodnavierstokes.jl @@ -1,7 +1,7 @@ module TwoDNavierStokes export - Problem, + Problem set_ζ!, updatevars!, From 966cfa095748cc13e17875b1c31c8c2aaac18475 Mon Sep 17 00:00:00 2001 From: Brodie Pearson Date: Mon, 3 Jan 2022 15:18:43 -0800 Subject: [PATCH 27/29] Dummy commit 2/2 --- src/surfaceqg.jl | 2 +- src/twodnavierstokes.jl | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/surfaceqg.jl b/src/surfaceqg.jl index 65942fa9..1b87b28c 100644 --- a/src/surfaceqg.jl +++ b/src/surfaceqg.jl @@ -1,7 +1,7 @@ module SurfaceQG export - Problem + Problem, set_b!, updatevars!, diff --git a/src/twodnavierstokes.jl b/src/twodnavierstokes.jl index e65e4aac..b4826a84 100644 --- a/src/twodnavierstokes.jl +++ b/src/twodnavierstokes.jl @@ -1,7 +1,7 @@ module TwoDNavierStokes export - Problem + Problem, set_ζ!, updatevars!, From 49f0ae1748d9a303f707a9321ffbae915945c8c9 Mon Sep 17 00:00:00 2001 From: Brodie Pearson Date: Mon, 3 Jan 2022 15:26:47 -0800 Subject: [PATCH 28/29] Fix typos in 2D modules, revert SQG and 2D to master --- src/surfaceqg.jl | 114 ++++++++++++---------------------------- src/twodnavierstokes.jl | 80 ++++++++++++++-------------- 2 files changed, 75 insertions(+), 119 deletions(-) diff --git a/src/surfaceqg.jl b/src/surfaceqg.jl index 1b87b28c..5b7bb77e 100644 --- a/src/surfaceqg.jl +++ b/src/surfaceqg.jl @@ -7,8 +7,7 @@ export kinetic_energy, buoyancy_variance, - buoyancy_dissipation_hyperviscosity, - buoyancy_dissipation_hypoviscosity, + buoyancy_dissipation, buoyancy_work, buoyancy_advection, kinetic_energy_advection @@ -64,11 +63,9 @@ function Problem(dev::Device=CPU(); ny = nx, Lx = 2π, Ly = Lx, - # Hyper-viscosity (ν) and hypo-viscosity (μ) parameters + # Hyper-viscosity parameters ν = 0, nν = 1, - μ = 0, - nμ = 0, # Timestepper and equation options dt = 0.01, stepper = "RK4", @@ -80,7 +77,7 @@ function Problem(dev::Device=CPU(); grid = TwoDGrid(dev, nx, Lx, ny, Ly; aliased_fraction=aliased_fraction, T=T) - params = Params{T}(ν, nν, μ, nμ, calcF) + params = Params{T}(ν, nν, calcF) vars = calcF == nothingfunction ? DecayingVars(dev, grid) : (stochastic ? StochasticForcedVars(dev, grid) : ForcedVars(dev, grid)) @@ -95,7 +92,7 @@ end # ---------- """ - Params{T}(ν, nν, μ, nμ, calcF!) + Params{T}(ν, nν, calcF!) A struct containing the parameters for Surface QG dynamics. Included are: @@ -106,10 +103,6 @@ struct Params{T} <: AbstractParams ν :: T "buoyancy (hyper)-viscosity order" nν :: Int - "buoyancy (hypo)-viscosity coefficient" - μ :: T - "buoyancy (hypo)-viscosity order" - nμ :: Int "function that calculates the Fourier transform of the forcing, ``F̂``" calcF! :: Function end @@ -124,22 +117,21 @@ Params(ν, nν) = Params(ν, nν, nothingfunction) """ Equation(params, grid) -Return the `equation` for surface QG dynamics with `params` and `grid`. The linear -opeartor ``L`` includes (hyper)-viscosity of order ``n_ν`` with coefficient ``ν`` and -hypo-viscocity of order ``n_μ`` with coefficient ``μ``, +Return the `equation` for surface QG dynamics with `params` and `grid`. The linear +opeartor ``L`` includes (hyper)-viscosity of order ``n_ν`` with coefficient ``ν``, ```math -L = - ν |𝐤|^{2 n_ν} - μ |𝐤|^{2 n_μ} . +L = - ν |𝐤|^{2 n_ν} . ``` -Plain old viscocity corresponds to ``n_ν=1`` while ``n_μ=0`` corresponds to linear drag. +Plain old viscocity corresponds to ``n_ν=1``. The nonlinear term is computed via function `calcN!()`. """ function Equation(params::Params, grid::AbstractGrid) - L = @. - params.ν * grid.Krsq^params.nν - params.μ * grid.Krsq^params.nμ + L = @. - params.ν * grid.Krsq^params.nν CUDA.@allowscalar L[1, 1] = 0 - + return FourierFlows.Equation(L, calcN!, grid) end @@ -189,7 +181,7 @@ function DecayingVars(::Dev, grid::AbstractGrid) where Dev T = eltype(grid) @devzeros Dev T (grid.nx, grid.ny) b u v @devzeros Dev Complex{T} (grid.nkr, grid.nl) bh uh vh - + return Vars(b, u, v, bh, uh, vh, nothing, nothing) end @@ -202,7 +194,7 @@ function ForcedVars(dev::Dev, grid) where Dev T = eltype(grid) @devzeros Dev T (grid.nx, grid.ny) b u v @devzeros Dev Complex{T} (grid.nkr, grid.nl) bh uh vh Fh - + return Vars(b, u, v, bh, uh, vh, Fh, nothing) end @@ -213,10 +205,10 @@ Return the `vars` for stochastically forced surface QG dynamics on device `dev` """ function StochasticForcedVars(dev::Dev, grid) where Dev T = eltype(grid) - + @devzeros Dev T (grid.nx, grid.ny) b u v @devzeros Dev Complex{T} (grid.nkr, grid.nl) bh uh vh Fh prevsol - + return Vars(b, u, v, bh, uh, vh, Fh, prevsol) end @@ -228,7 +220,7 @@ end """ calcN_advection(N, sol, t, clock, vars, params, grid) -Calculate the Fourier transform of the advection term, ``- 𝖩(ψ, b)`` in conservative +Calculate the Fourier transform of the advection term, ``- 𝖩(ψ, b)`` in conservative form, i.e., ``- ∂_x[(∂_y ψ)b] - ∂_y[(∂_x ψ)b]`` and store it in `N`: ```math @@ -243,10 +235,10 @@ function calcN_advection!(N, sol, t, clock, vars, params, grid) ldiv!(vars.u, grid.rfftplan, vars.uh) ldiv!(vars.v, grid.rfftplan, vars.vh) ldiv!(vars.b, grid.rfftplan, vars.bh) - + ub, ubh = vars.u, vars.uh # use vars.u, vars.uh as scratch variables vb, vbh = vars.v, vars.vh # use vars.v, vars.vh as scratch variables - + @. ub *= vars.b # u*b @. vb *= vars.b # v*b @@ -254,7 +246,7 @@ function calcN_advection!(N, sol, t, clock, vars, params, grid) mul!(vbh, grid.rfftplan, vb) # \hat{v*b} @. N = - im * grid.kr * ubh - im * grid.l * vbh - + return nothing end @@ -273,14 +265,14 @@ 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) - + return nothing end """ addforcing!(N, sol, t, clock, vars, params, grid) -When the problem includes forcing, calculate the forcing term ``F̂`` and add it to the +When the problem includes forcing, calculate the forcing term ``F̂`` and add it to the nonlinear term ``N``. """ addforcing!(N, sol, t, clock, vars::Vars, params, grid) = nothing @@ -288,7 +280,7 @@ addforcing!(N, sol, t, clock, vars::Vars, params, grid) = nothing function addforcing!(N, sol, t, clock, vars::ForcedVars, params, grid) params.calcF!(vars.Fh, sol, t, clock, vars, params, grid) @. N += vars.Fh - + return nothing end @@ -297,9 +289,9 @@ function addforcing!(N, sol, t, clock, vars::StochasticForcedVars, params, grid) @. 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 += vars.Fh - + return nothing end @@ -321,11 +313,11 @@ function updatevars!(prob) @. vars.bh = sol @. vars.uh = im * grid.l * sqrt(grid.invKrsq) * sol @. vars.vh = - im * grid.kr * sqrt(grid.invKrsq) * sol - + ldiv!(vars.b, grid.rfftplan, deepcopy(vars.bh)) ldiv!(vars.u, grid.rfftplan, deepcopy(vars.uh)) ldiv!(vars.v, grid.rfftplan, deepcopy(vars.vh)) - + return nothing end @@ -337,9 +329,9 @@ Set the solution `sol` as the transform of `b` and update all variables. function set_b!(prob, b) mul!(prob.sol, prob.grid.rfftplan, b) CUDA.@allowscalar prob.sol[1, 1] = 0 # zero domain average - + updatevars!(prob) - + return nothing end @@ -357,10 +349,10 @@ In SQG, this is identical to half the domain-averaged surface buoyancy variance. ψh = vars.uh # use vars.uh as scratch variable kinetic_energyh = vars.bh # use vars.bh as scratch variable - + @. ψh = sqrt(grid.invKrsq) * sol @. kinetic_energyh = 1 / 2 * grid.Krsq * abs2(ψh) - + return 1 / (grid.Lx * grid.Ly) * parsevalsum(kinetic_energyh, grid) end @@ -371,17 +363,17 @@ Return the buoyancy variance, ```math \\int b² \\frac{𝖽x 𝖽y}{L_x L_y} = \\sum_{𝐤} |b̂|² . ``` -In SQG, this is identical to the velocity variance (i.e., twice the domain-averaged kinetic +In SQG, this is identical to the velocity variance (i.e., twice the domain-averaged kinetic energy). """ @inline function buoyancy_variance(prob) sol, grid = prob.sol, prob.grid - + return 1 / (grid.Lx * grid.Ly) * parsevalsum(abs2.(sol), grid) end """ - buoyancy_dissipation_hyperviscosity(prob) + buoyancy_dissipation(prob) Return the domain-averaged dissipation rate of surface buoyancy variance due to small scale (hyper)-viscosity, @@ -394,49 +386,13 @@ In SQG, this is identical to twice the rate of kinetic energy dissipation @inline function buoyancy_dissipation(prob) sol, vars, params, grid = prob.sol, prob.vars, prob.params, prob.grid buoyancy_dissipationh = vars.uh # use vars.uh as scratch variable - + @. buoyancy_dissipationh = 2 * params.ν * grid.Krsq^params.nν * abs2(sol) CUDA.@allowscalar buoyancy_dissipationh[1, 1] = 0 - - return 1 / (grid.Lx * grid.Ly) * parsevalsum(buoyancy_dissipationh, grid) -end - -""" - buoyancy_dissipation(prob, ξ, nξ) - -Return the domain-averaged dissipation rate of surface buoyancy variance due -to large scale (hypo)-viscosity, -```math -2 ξ (-1)^{n_ξ+1} \\int b ∇^{2n_ξ} b \\frac{𝖽x 𝖽y}{L_x L_y} = - 2 ν \\sum_{𝐤} |𝐤|^{2n_ξ} |b̂|² , -``` -where ``ν`` the (hyper)-viscosity coefficient ``ν`` and ``nν`` the (hyper)-viscosity order. -In SQG, this is identical to twice the rate of kinetic energy dissipation -""" -@inline function buoyancy_dissipation(prob, ξ, nξ) - sol, vars, params, grid = prob.sol, prob.vars, prob.params, prob.grid - buoyancy_dissipationh = vars.uh # use vars.uh as scratch variable - - @. buoyancy_dissipationh = 2 * ξ * grid.Krsq^nξ * abs2(sol) - CUDA.@allowscalar buoyancy_dissipationh[1, 1] = 0 - + return 1 / (grid.Lx * grid.Ly) * parsevalsum(buoyancy_dissipationh, grid) end -""" - buoyancy_dissipation_hyperviscosity(prob, ξ, νξ) - -Return the domain-averaged buoyancy dissipation rate done by the ``ν`` (hyper)-viscosity. -""" -buoyancy_dissipation_hyperviscosity(prob) = buoyancy_dissipation(prob, prob.params.ν, prob.params.nν) - -""" - buoyancy_dissipation_hypoviscosity(prob, ξ, νξ) - -Return the domain-averaged buoyancy dissipation rate done by the ``μ`` (hypo)-viscosity. -""" -buoyancy_dissipation_hypoviscosity(prob) = buoyancy_dissipation(prob, prob.params.μ, prob.params.nμ) - - """ buoyancy_work(prob) buoyancy_work(sol, vars, grid) @@ -455,7 +411,7 @@ end @inline function buoyancy_work(sol, vars::StochasticForcedVars, grid) buoyancy_workh = vars.uh # use vars.uh as scratch variable - + @. buoyancy_workh = (vars.prevsol + sol) * conj(vars.Fh) # Stratonovich return 1 / (grid.Lx * grid.Ly) * parsevalsum(buoyancy_workh, grid) end diff --git a/src/twodnavierstokes.jl b/src/twodnavierstokes.jl index b4826a84..8d96fcbb 100644 --- a/src/twodnavierstokes.jl +++ b/src/twodnavierstokes.jl @@ -130,7 +130,7 @@ Params(ν, nν) = Params(ν, nν, typeof(ν)(0), 0, nothingfunction) Equation(params, grid) Return the `equation` for two-dimensional Navier-Stokes with `params` and `grid`. The linear -operator ``L`` includes (hyper)-viscosity of order ``n_ν`` with coefficient ``ν`` and +operator ``L`` includes (hyper)-viscosity of order ``n_ν`` with coefficient ``ν`` and hypo-viscocity of order ``n_μ`` with coefficient ``μ``, ```math @@ -144,7 +144,7 @@ The nonlinear term is computed via the function `calcN!`. function Equation(params::Params, grid::AbstractGrid) L = @. - params.ν * grid.Krsq^params.nν - params.μ * grid.Krsq^params.nμ CUDA.@allowscalar L[1, 1] = 0 - + return FourierFlows.Equation(L, calcN!, grid) end @@ -188,15 +188,15 @@ const StochasticForcedVars = Vars{<:AbstractArray, <:AbstractArray, <:AbstractAr """ DecayingVars(dev, grid) -Return the `vars` for unforced two-dimensional Navier-Stokes problem on device `dev` and +Return the `vars` for unforced two-dimensional Navier-Stokes problem on device `dev` and with `grid`. """ function DecayingVars(::Dev, grid::AbstractGrid) where Dev T = eltype(grid) - + @devzeros Dev T (grid.nx, grid.ny) ζ u v @devzeros Dev Complex{T} (grid.nkr, grid.nl) ζh uh vh - + return Vars(ζ, u, v, ζh, uh, vh, nothing, nothing) end @@ -207,25 +207,25 @@ Return the `vars` for forced two-dimensional Navier-Stokes on device `dev` and w """ function ForcedVars(dev::Dev, grid::AbstractGrid) where Dev T = eltype(grid) - + @devzeros Dev T (grid.nx, grid.ny) ζ u v @devzeros Dev Complex{T} (grid.nkr, grid.nl) ζh uh vh Fh - + return Vars(ζ, u, v, ζh, uh, vh, Fh, nothing) end """ StochasticForcedVars(dev, grid) -Return the `vars` for stochastically forced two-dimensional Navier-Stokes on device `dev` and +Return the `vars` for stochastically forced two-dimensional Navier-Stokes on device `dev` and with `grid`. """ function StochasticForcedVars(dev::Dev, grid::AbstractGrid) where Dev T = eltype(grid) - + @devzeros Dev T (grid.nx, grid.ny) ζ u v @devzeros Dev Complex{T} (grid.nkr, grid.nl) ζh uh vh Fh prevsol - + return Vars(ζ, u, v, ζh, uh, vh, Fh, prevsol) end @@ -237,7 +237,7 @@ end """ calcN_advection!(N, sol, t, clock, vars, params, grid) -Calculate the Fourier transform of the advection term, ``- 𝖩(ψ, ζ)`` in conservative form, +Calculate the Fourier transform of the advection term, ``- 𝖩(ψ, ζ)`` in conservative form, i.e., ``- ∂_x[(∂_y ψ)ζ] - ∂_y[(∂_x ψ)ζ]`` and store it in `N`: ```math @@ -252,19 +252,19 @@ function calcN_advection!(N, sol, t, clock, vars, params, grid) ldiv!(vars.u, grid.rfftplan, vars.uh) ldiv!(vars.v, grid.rfftplan, vars.vh) ldiv!(vars.ζ, grid.rfftplan, vars.ζh) - + uζ = vars.u # use vars.u as scratch variable @. uζ *= vars.ζ # u*ζ vζ = vars.v # use vars.v as scratch variable @. vζ *= vars.ζ # v*ζ - + uζh = vars.uh # use vars.uh as scratch variable mul!(uζh, grid.rfftplan, uζ) # \hat{u*ζ} vζh = vars.vh # use vars.vh as scratch variable mul!(vζh, grid.rfftplan, vζ) # \hat{v*ζ} @. N = - im * grid.kr * uζh - im * grid.l * vζh - + return nothing end @@ -281,25 +281,25 @@ function calcN!(N, sol, t, clock, vars, params, grid) dealias!(sol, grid) calcN_advection!(N, sol, t, clock, vars, params, grid) - + addforcing!(N, sol, t, clock, vars, params, grid) - + return nothing end """ addforcing!(N, sol, t, clock, vars, params, grid) -When the problem includes forcing, calculate the forcing term ``F̂`` and add it to the +When the problem includes forcing, calculate the forcing term ``F̂`` and add it to the nonlinear term ``N``. """ addforcing!(N, sol, t, clock, vars::DecayingVars, params, grid) = nothing function addforcing!(N, sol, t, clock, vars::ForcedVars, params, grid) params.calcF!(vars.Fh, sol, t, clock, vars, params, grid) - + @. N += vars.Fh - + return nothing end function addforcing!(N, sol, t, clock, vars::StochasticForcedVars, params, grid) @@ -308,7 +308,7 @@ function addforcing!(N, sol, t, clock, vars::StochasticForcedVars, params, grid) params.calcF!(vars.Fh, sol, t, clock, vars, params, grid) end @. N += vars.Fh - + return nothing end @@ -330,11 +330,11 @@ function updatevars!(prob) @. vars.ζh = sol @. vars.uh = im * grid.l * grid.invKrsq * sol @. vars.vh = - im * grid.kr * grid.invKrsq * sol - + ldiv!(vars.ζ, grid.rfftplan, deepcopy(vars.ζh)) # deepcopy() since inverse real-fft destroys its input ldiv!(vars.u, grid.rfftplan, deepcopy(vars.uh)) # deepcopy() since inverse real-fft destroys its input ldiv!(vars.v, grid.rfftplan, deepcopy(vars.vh)) # deepcopy() since inverse real-fft destroys its input - + return nothing end @@ -345,18 +345,18 @@ Set the solution `sol` as the transform of `ζ` and then update variables in `va """ function set_ζ!(prob, ζ) mul!(prob.sol, prob.grid.rfftplan, ζ) - + CUDA.@allowscalar prob.sol[1, 1] = 0 # enforce zero domain average - + updatevars!(prob) - + return nothing end """ energy(prob) -Return the domain-averaged kinetic energy. Since ``u² + v² = |{\\bf ∇} ψ|²``, the domain-averaged +Return the domain-averaged kinetic energy. Since ``u² + v² = |{\\bf ∇} ψ|²``, the domain-averaged kinetic energy is ```math @@ -366,7 +366,7 @@ kinetic energy is @inline function energy(prob) sol, vars, grid = prob.sol, prob.vars, prob.grid energyh = vars.uh # use vars.uh as scratch variable - + @. energyh = 1 / 2 * grid.invKrsq * abs2(sol) return 1 / (grid.Lx * grid.Ly) * parsevalsum(energyh, grid) end @@ -393,34 +393,34 @@ Return the domain-averaged energy dissipation rate done by the viscous term, ```math - ξ (-1)^{n_ξ+1} \\int ψ ∇^{2n_ξ} ζ \\frac{𝖽x 𝖽y}{L_x L_y} = - ξ \\sum_{𝐤} |𝐤|^{2(n_ξ-1)} |ζ̂|² , ``` -where ``ξ`` and ``nξ`` could be either the (hyper)-viscosity coefficient ``ν`` and its order +where ``ξ`` and ``nξ`` could be either the (hyper)-viscosity coefficient ``ν`` and its order ``n_ν``, or the hypo-viscocity coefficient ``μ`` and its order ``n_μ``. """ @inline function energy_dissipation(prob, ξ, nξ) sol, vars, grid = prob.sol, prob.vars, prob.grid energy_dissipationh = vars.uh # use vars.uh as scratch variable - + @. energy_dissipationh = - ξ * grid.Krsq^(nξ - 1) * abs2(sol) CUDA.@allowscalar energy_dissipationh[1, 1] = 0 return 1 / (grid.Lx * grid.Ly) * parsevalsum(energy_dissipationh, grid) end """ - energy_dissipation_hyperviscosity(prob, ξ, νξ) + energy_dissipation_hyperviscosity(prob, ξ, nξ) Return the domain-averaged energy dissipation rate done by the ``ν`` (hyper)-viscosity. """ energy_dissipation_hyperviscosity(prob) = energy_dissipation(prob, prob.params.ν, prob.params.nν) """ - energy_dissipation_hypoviscosity(prob, ξ, νξ) + energy_dissipation_hypoviscosity(prob, ξ, nξ) Return the domain-averaged energy dissipation rate done by the ``μ`` (hypo)-viscosity. """ energy_dissipation_hypoviscosity(prob) = energy_dissipation(prob, prob.params.μ, prob.params.nμ) """ - enstrophy_dissipation(prob, ξ, νξ) + enstrophy_dissipation(prob, ξ, nξ) Return the domain-averaged enstrophy dissipation rate done by the viscous term, @@ -428,27 +428,27 @@ Return the domain-averaged enstrophy dissipation rate done by the viscous term, ξ (-1)^{n_ξ+1} \\int ζ ∇^{2n_ξ} ζ \\frac{𝖽x 𝖽y}{L_x L_y} = - ξ \\sum_{𝐤} |𝐤|^{2n_ξ} |ζ̂|² , ``` -where ``ξ`` and ``nξ`` could be either the (hyper)-viscosity coefficient ``ν`` and its order +where ``ξ`` and ``nξ`` could be either the (hyper)-viscosity coefficient ``ν`` and its order ``n_ν``, or the hypo-viscocity coefficient ``μ`` and its order ``n_μ``. """ @inline function enstrophy_dissipation(prob, ξ, nξ) sol, vars, grid = prob.sol, prob.vars, prob.grid enstrophy_dissipationh = vars.uh # use vars.uh as scratch variable - + @. enstrophy_dissipationh = - ξ * grid.Krsq^nξ * abs2(sol) CUDA.@allowscalar enstrophy_dissipationh[1, 1] = 0 return 1 / (grid.Lx * grid.Ly) * parsevalsum(enstrophy_dissipationh, grid) end """ - enstrophy_dissipation_hyperviscosity(prob, ξ, νξ) + enstrophy_dissipation_hyperviscosity(prob, ξ, nξ) Return the domain-averaged enstrophy dissipation rate done by the ``ν`` (hyper)-viscosity. """ enstrophy_dissipation_hyperviscosity(prob) = enstrophy_dissipation(prob, prob.params.ν, prob.params.nν) """ - enstrophy_dissipation_hypoviscosity(prob, ξ, νξ) + enstrophy_dissipation_hypoviscosity(prob, ξ, nξ) Return the domain-averaged enstrophy dissipation rate done by the ``μ`` (hypo)-viscosity. """ @@ -466,14 +466,14 @@ Return the domain-averaged rate of work of energy by the forcing ``F``, """ @inline function energy_work(sol, vars::ForcedVars, grid) energy_workh = vars.uh # use vars.uh as scratch variable - + @. energy_workh = grid.invKrsq * sol * conj(vars.Fh) return 1 / (grid.Lx * grid.Ly) * parsevalsum(energy_workh, grid) end @inline function energy_work(sol, vars::StochasticForcedVars, grid) energy_workh = vars.uh # use vars.uh as scratch variable - + @. energy_workh = grid.invKrsq * (vars.prevsol + sol) / 2 * conj(vars.Fh) return 1 / (grid.Lx * grid.Ly) * parsevalsum(energy_workh, grid) end @@ -492,14 +492,14 @@ Return the domain-averaged rate of work of enstrophy by the forcing ``F``, """ @inline function enstrophy_work(sol, vars::ForcedVars, grid) enstrophy_workh = vars.uh # use vars.uh as scratch variable - + @. enstrophy_workh = sol * conj(vars.Fh) return 1 / (grid.Lx * grid.Ly) * parsevalsum(enstrophy_workh, grid) end @inline function enstrophy_work(sol, vars::StochasticForcedVars, grid) enstrophy_workh = vars.uh # use vars.uh as scratch variable - + @. enstrophy_workh = (vars.prevsol + sol) / 2 * conj(vars.Fh) return 1 / (grid.Lx * grid.Ly) * parsevalsum(enstrophy_workh, grid) end From 8575118b71d3007ca8391eb9a99ecf3dfb0ea14b Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Mon, 4 Jul 2022 10:17:53 +1000 Subject: [PATCH 29/29] tweak some docstrings --- src/twodnavierstokeswithtracer.jl | 33 ++++++++++++++++++++----------- 1 file changed, 21 insertions(+), 12 deletions(-) diff --git a/src/twodnavierstokeswithtracer.jl b/src/twodnavierstokeswithtracer.jl index 21213e31..f46d54fb 100644 --- a/src/twodnavierstokeswithtracer.jl +++ b/src/twodnavierstokeswithtracer.jl @@ -72,8 +72,8 @@ Keyword arguments - `aliased_fraction`: the fraction of high-wavenumbers that are zero-ed out by `dealias!()`. - `T`: `Float32` or `Float64`; floating point type used for `problem` data. """ -function Problem(ntracers::Int, # number of tracers - dev::Device=CPU(); +function Problem(ntracers :: Int, # number of tracers + dev :: Device=CPU(); # Numerical parameters nx = 256, ny = nx, @@ -114,32 +114,39 @@ end """ Params{T}(ν, nν, μ, nμ, κ, nκ calcF!) -A struct containing the parameters for the two-dimensional Navier-Stokes with tracers. Included are: +A container for the parameters for the two-dimensional Navier-Stokes problem with tracers. $(TYPEDFIELDS) """ struct Params{T, Trfft} <: AbstractParams "number of tracers" - ntracers :: Int + ntracers :: Int "small-scale (hyper)-viscosity coefficient" - ν :: T + ν :: T "(hyper)-viscosity order, `nν```≥ 1``" - nν :: Int + nν :: Int "large-scale (hypo)-viscosity coefficient" - μ :: T + μ :: T "(hypo)-viscosity order, `nμ```≤ 0``" - nμ :: Int + nμ :: Int "tracer diffusivity coefficient" - κ :: T + κ :: T "tracer diffusivity order" - nκ :: Int + nκ :: Int "function that calculates the Fourier transform of the forcing, ``F̂``" - calcF! :: Function + calcF! :: Function "rfft plan for FFTs (Derived parameter)" rfftplan :: Trfft end -function Params(ntracers, ν, nν, μ, nμ, κ, nκ, grid; calcF=nothingfunction, effort=FFTW.MEASURE, dev::Device=CPU()) where TU +""" + Params(ntracers, ν, nν, μ, nμ, κ, nκ, grid; + calcF=nothingfunction, effort=FFTW.MEASURE, dev::Device=CPU()) + +Return the parameters for a the two-dimensional Navier-Stokes problem with tracers. +""" +function Params(ntracers, ν, nν, μ, nμ, κ, nκ, grid; + calcF=nothingfunction, effort=FFTW.MEASURE, dev::Device=CPU()) where TU T = eltype(grid) A = ArrayType(dev) @@ -154,12 +161,14 @@ end """ fwdtransform!(varh, var, params) + Compute the Fourier transform of `var` and store it in `varh`. """ fwdtransform!(varh, var, params::AbstractParams) = mul!(varh, params.rfftplan, var) """ invtransform!(var, varh, params) + Compute the inverse Fourier transform of `varh` and store it in `var`. """ invtransform!(var, varh, params::AbstractParams) = ldiv!(var, params.rfftplan, varh)