diff --git a/docs/make.jl b/docs/make.jl index 98eaed09..75b0d26c 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -91,7 +91,8 @@ sitename = "GeophysicalFlows.jl", "modules/singlelayerqg.md", "modules/barotropicqgql.md", "modules/multilayerqg.md", - "modules/surfaceqg.md" + "modules/surfaceqg.md", + "modules/shallowwater.md" ], "Stochastic Forcing" => "stochastic_forcing.md", "Contributor's guide" => "contributing.md", diff --git a/docs/src/modules/shallowwater.md b/docs/src/modules/shallowwater.md new file mode 100644 index 00000000..6411e1f9 --- /dev/null +++ b/docs/src/modules/shallowwater.md @@ -0,0 +1,49 @@ +# ShallowWater Module + +### Basic Equations + +This module solves the 2D shallow water equations for a fluid: + +```math +\begin{aligned} +\partial_t u + \bm{u \cdot \nabla} u - f v & = - g \partial_x \eta , \\ +\partial_t v + \bm{u \cdot \nabla} v + f u & = - g \partial_y \eta , \\ +\partial_t h + \bm{\nabla \cdot} (\bm{u} h) & = 0 , +\end{aligned} +``` + +where ``\bm{u}(\bm{x}, t)`` is the horizontal fluid flow with components ``\bm{u} = (u, v)``, +``h(\bm{x}, t)`` is the total fluid depth, and ``\eta(\bm{x}, t)`` is the free surface elevation +(measured from the ``z=0`` rest-height of the free surface). The bottom bathymetry is ``b(\bm{x})`` +(also measured from the ``z=0`` rest-height of the free surface). Thus, we have that + +```math +h(\bm{x}, t) = \eta(\bm{x}, t) - b(\bm{x}) . +``` + +In terms of variables ``u``, ``v``, and ``h``: + +```math +\begin{aligned} +\partial_t u + \bm{u \cdot \nabla} u - f v & = - g \partial_x h - g \partial_x b , \\ +\partial_t v + \bm{u \cdot \nabla} v + f u & = - g \partial_y h - g \partial_y b , \\ +\partial_t h + \bm{\nabla \cdot} (\bm{u} h) & = 0 . +\end{aligned} +``` + +The bathymetric gradients ``b(\bm{x})`` enter as a forcing on the momentum equations. + +Often we write shallow water dynamics in conservative form using variables ``q_u = hu``, ``q_v = hv`` and ``h``. In these variables: + +```math +\begin{aligned} +\partial_t q_u + \partial_x \left ( \frac{q_u^2}{h} \right ) + \partial_y \left ( \frac{q_u q_v}{h} \right ) - f q_v & = - g \partial_x \left( \frac{h^2}{2} \right) - g h \partial_x b , \\ +\partial_t q_v + \partial_x \left ( \frac{q_u q_v}{h} \right ) + \partial_y \left ( \frac{q_v^2}{h} \right ) + f q_u & = - g \partial_y \left( \frac{h^2}{2} \right) - g h \partial_y b , \\ +\partial_t h + \partial_x q_u + \partial_y q_v & = 0 . +\end{aligned} +``` + + +### Implementation + +## Examples diff --git a/src/GeophysicalFlows.jl b/src/GeophysicalFlows.jl index b3401da9..4f5ac926 100644 --- a/src/GeophysicalFlows.jl +++ b/src/GeophysicalFlows.jl @@ -19,6 +19,7 @@ include("twodnavierstokes.jl") include("singlelayerqg.jl") include("multilayerqg.jl") include("surfaceqg.jl") +include("shallowwater.jl") include("barotropicqgql.jl") @reexport using GeophysicalFlows.TwoDNavierStokes diff --git a/src/shallowwater.jl b/src/shallowwater.jl new file mode 100644 index 00000000..2e299823 --- /dev/null +++ b/src/shallowwater.jl @@ -0,0 +1,305 @@ +module ShallowWater + +export + Problem, + updatevars! + +using + CUDA, + Reexport + +@reexport using FourierFlows + +using LinearAlgebra: mul!, ldiv! +using FFTW: rfft, irfft +using FourierFlows: parsevalsum + +nothingfunction(args...) = nothing + +""" + Problem(; parameters...) + +Construct a two-dimensional shallow water problem. +""" +function Problem(dev::Device=CPU(); + # Numerical parameters + nx = 256, + Lx = 2π, + ny = nx, + Ly = Lx, + dt = 0.01, + # Coriolis parameter + f = 0.0, + # Gravitational constant + g = 9.81, + # Total mean rest fluid height + H = 1.0, + # Bottom bathymetry + b = nothing, + # Drag and/or hyper-/hypo-viscosity + ν = 0, + nν = 1, + # Timestepper and equation options + stepper = "RK4", + calcF = nothingfunction, + stochastic = false, + T = Float64) + + # bottom bathymetry + b === nothing && (b = zeros(dev, T, (nx, ny))) + + grid = TwoDGrid(dev, nx, Lx, ny, Ly; T=T) + + params = Params(dev, T(f), T(g), T(H), b, T(ν), nν, calcF, grid) + + vars = calcF == nothingfunction ? Vars(dev, grid) : (stochastic ? StochasticForcedVars(dev, grid) : ForcedVars(dev, grid)) + + equation = Equation(dev, params, grid) + + return FourierFlows.Problem(equation, stepper, dt, grid, vars, params, dev) +end + + +# ---------- +# Parameters +# ---------- + +""" + Params(ν, nν, calcF!) + +Returns the params for two-dimensional shallow water. +""" +struct Params{T, A} <: AbstractParams + f :: T # Coriolis parameter + g :: T # Gravitational constant + ν :: T # Viscosity + nν :: Int # (Hyper)-viscous order + ∂b∂x :: A # x-bottom bathymetry gradient, ∂b/∂x + ∂b∂y :: A # y-bottom bathymetry gradient, ∂b/∂y + calcF! :: Function # Function that calculates the forcing `Fh` +end + +function Params(dev, f, g, H, b, ν, nν, calcF, grid) + A = ArrayType(dev) + + bh = rfft(A(b)) + ∂b∂x = irfft(im * grid.kr .* bh, grid.nx) + ∂b∂y = irfft(im * grid.l .* bh, grid.nx) + + return Params(f, g, ν, nν, ∂b∂x, ∂b∂y, calcF) +end + +# --------- +# Equations +# --------- + +""" + Equation(dev, params, grid) + +Returns the equation for two-dimensional shallow water with `params` and `grid`. +""" +function Equation(dev, params::Params, grid::AbstractGrid{T}) where T + D = @. - params.ν * grid.Krsq^params.nν + CUDA.@allowscalar D[1, 1] = 0 + + L = zeros(dev, T, (grid.nkr, grid.nl, 3)) + + L[:, :, 1] .= D # for qu equation + L[:, :, 2] .= D # for qv equation + L[:, :, 3] .= D # for h equation + + return FourierFlows.Equation(L, calcN!, grid) +end + + +# ---- +# Vars +# ---- + +abstract type ShallowWaterVars <: AbstractVars end + +struct Vars{Aphys, Atrans, S, F, P} <: ShallowWaterVars + qu :: Aphys + qv :: Aphys + u :: Aphys + v :: Aphys + h :: Aphys + quh :: Atrans + qvh :: Atrans + hh :: Atrans + state :: S + Fh :: F + prevsol :: P +end + +const ForcedVars = Vars{<:AbstractArray, <:AbstractArray, <:AbstractArray, <:AbstractArray, Nothing} +const StochasticForcedVars = Vars{<:AbstractArray, <:AbstractArray, <:AbstractArray, <:AbstractArray, <:AbstractArray} + +""" + Vars(dev, grid) + +Returns the `vars` for unforced two-dimensional shallow water on device `dev` and with `grid`. +""" +function Vars(::Dev, grid::TwoDGrid{T}) where {Dev, T} + @devzeros Dev T (grid.nx, grid.ny) qu qv u v h + @devzeros Dev Complex{T} (grid.nkr, grid.nl) quh qvh hh + @devzeros Dev Complex{T} (grid.nkr, grid.nl, 3) state + return Vars(qu, qv, u, v, h, quh, qvh, hh, state, nothing, nothing) +end + +""" + ForcedVars(dev, grid) + +Returns the vars for forced two-dimensional shallow water on device `dev` and with `grid`. +""" +function ForcedVars(dev::Dev, grid::TwoDGrid{T}) where {Dev, T} + @devzeros Dev T (grid.nx, grid.ny) qu qv u v h + @devzeros Dev Complex{T} (grid.nkr, grid.nl) quh qvh hh Fh + @devzeros Dev Complex{T} (grid.nkr, grid.nl, 3) state + return Vars(qu, qv, u, v, h, quh, qvh, hh, state, Fh, nothing) +end + +""" + StochasticForcedVars(dev, grid) + +Returns the vars for stochastically forced two-dimensional shallow water on device `dev` and with `grid`. +""" +function StochasticForcedVars(dev::Dev, grid::TwoDGrid{T}) where {Dev, T} + @devzeros Dev T (grid.nx, grid.ny) qu qv u v h + @devzeros Dev Complex{T} (grid.nkr, grid.nl) quh qvh hh Fh prevsol + @devzeros Dev Complex{T} (grid.nkr, grid.nl, 3) state + return Vars(qu, qv, u, v, h, quh, qvh, hh, state, Fh, prevsol) +end + + +# ------- +# Solvers +# ------- + +""" + calcN_advection(N, sol, t, clock, vars, params, grid) + +Calculates the advection term. +""" +function calcN!(N, sol, t, clock, vars, params, grid) + # state = vars.state + # state = (qu = view(sol, :, :, 1), qv = view(sol, :, :, 2), h = view(sol, :, :, 3)) + + vars.quh = sol[:, :, 1] + vars.qvh = sol[:, :, 2] + vars.hh = sol[:, :, 3] + + @. N[:, :, 1] = params.f * vars.qvh + @. N[:, :, 2] = - params.f * vars.quh + @. N[:, :, 3] = - im * grid.kr * vars.quh - im * grid.l * vars.qvh + + ldiv!(vars.qu, grid.rfftplan, vars.quh) + ldiv!(vars.qv, grid.rfftplan, vars.qvh) + ldiv!(vars.h, grid.rfftplan, deepcopy(vars.hh)) + + qu²_overh = vars.qu + @. qu²_overh *= vars.qu / vars.h + + qu²_overhh = vars.quh + mul!(qu²_overhh, grid.rfftplan, qu²_overh) + + quqv_overh = vars.v + @. quqv_overh = vars.qu * vars.qv / vars.h + + quqv_overhh = vars.qvh + mul!(quqv_overhh, grid.rfftplan, quqv_overh) + + @. N[:, :, 1] -= im * grid.kr * qu²_overhh + im * grid.l * quqv_overhh + @. N[:, :, 2] -= im * grid.kr * quqv_overhh + + qv²_overh = vars.qv + @. qv²_overh *= vars.qv / vars.h + + qv²_overhh = vars.qvh + mul!(qv²_overhh, grid.rfftplan, qv²_overh) + + @. N[:, :, 2] -= im * grid.l * qv²_overhh + + h² = vars.h + @. h² *= vars.h + + h²h = vars.hh + mul!(h²h, grid.rfftplan, h²) + + @. N[:, :, 1] -= im * params.g * grid.kr * h²h / 2 # - g ∂(½h²)/∂x + @. N[:, :, 2] -= im * params.g * grid.l * h²h / 2 # - g ∂(½h²)/∂y + + ldiv!(vars.h, grid.rfftplan, deepcopy(sol[:, :, 3])) + + h∂b∂x = vars.h + @. h∂b∂x *= params.∂b∂x + + h∂b∂xh = vars.hh + mul!(h∂b∂xh, grid.rfftplan, h∂b∂x) + + @. N[:, :, 1] -= params.g * h∂b∂xh + + ldiv!(vars.h, grid.rfftplan, deepcopy(sol[:, :, 3])) + + h∂b∂y = vars.h + @. h∂b∂y *= params.∂b∂y + + h∂b∂yh = vars.hh + mul!(h∂b∂yh, grid.rfftplan, h∂b∂y) + + @. N[:, :, 2] -= params.g * h∂b∂yh + + addforcing!(N, sol, t, clock, vars, params, grid) + + return nothing +end + +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 + +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 + + @. vars.quh = sol[:, :, 1] + @. vars.qvh = sol[:, :, 2] + @. vars.hh = sol[:, :, 3] + + ldiv!(vars.qu, grid.rfftplan, deepcopy(sol[:, :, 1])) # use deepcopy() because irfft destroys its input + ldiv!(vars.qv, grid.rfftplan, deepcopy(sol[:, :, 2])) # use deepcopy() because irfft destroys its input + ldiv!(vars.h, grid.rfftplan, deepcopy(sol[:, :, 3])) # use deepcopy() because irfft destroys its input + + @. vars.u = vars.qu / vars.h + @. vars.v = vars.qv / vars.h + + return nothing +end + +end # module diff --git a/test/runtests.jl b/test/runtests.jl index be762561..b30f8779 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -10,7 +10,8 @@ import # use 'import' rather than 'using' for submodules to keep namespace clean GeophysicalFlows.SingleLayerQG, GeophysicalFlows.BarotropicQGQL, GeophysicalFlows.MultiLayerQG, - GeophysicalFlows.SurfaceQG + GeophysicalFlows.SurfaceQG, + GeophysicalFlows.ShallowWater using FourierFlows: parsevalsum using GeophysicalFlows: lambdipole, peakedisotropicspectrum @@ -31,6 +32,13 @@ for dev in devices @info "testing on " * string(typeof(dev)) + @testset "ShallowWater" begin + include("test_shallowwater.jl") + + @test test_swe_problemtype(dev, Float32) + @test ShallowWater.nothingfunction() == nothing + end + @testset "Utils" begin include("test_utils.jl") diff --git a/test/test_shallowwater.jl b/test/test_shallowwater.jl new file mode 100644 index 00000000..58593d55 --- /dev/null +++ b/test/test_shallowwater.jl @@ -0,0 +1,8 @@ +function test_swe_problemtype(dev, T) + prob = ShallowWater.Problem(dev; T=T) + + A = ArrayType(dev) + + (typeof(prob.sol)<:A{Complex{T},3} && typeof(prob.grid.Lx)==T && eltype(prob.grid.x)==T && typeof(prob.vars.u)<:A{T,2}) +end +