[WIP] Add ShallowWater module

docs/make.jl

@@ -90,7 +90,8 @@ sitename = "GeophysicalFlows.jl",
               "modules/barotropicqg.md",
               "modules/barotropicqgql.md",
               "modules/multilayerqg.md",
-              "modules/surfaceqg.md"
+              "modules/surfaceqg.md",
+              "modules/shallowwater.md"
             ],
             "Forcing" => "forcing.md",
             "DocStrings" => Any[

docs/src/modules/shallowwater.md

@@ -0,0 +1,17 @@
+# ShallowWater Module
+
+### Basic Equations
+
+This module solves the 2D shallow water equations:
+
+```math
+\begin{aligned}
+\frac{\partial u}{\partial t} + \boldsymbol{u \cdot \nabla} u - f v & = - g \frac{\partial \eta}{\partial x} - \mathrm{D} u, \\
+\frac{\partial v}{\partial t} + \boldsymbol{u \cdot \nabla} v + f u & = - \mathrm{D} v, \\
+\frac{\partial \eta}{\partial t} + \boldsymbol{\nabla \cdot} ( \boldsymbol{u} h ) & = - \mathrm{D} \eta.
+\end{aligned}
+```
+
+### Implementation
+
+## Examples

src/GeophysicalFlows.jl

@@ -14,5 +14,6 @@ include("barotropicqg.jl")
 include("barotropicqgql.jl")
 include("multilayerqg.jl")
 include("surfaceqg.jl")
+include("shallowwater.jl")
 
 end # module

src/shallowwater.jl

@@ -0,0 +1,203 @@
+module ShallowWater
+
+export
+  Problem,
+  updatevars!,
+
+  energy
+
+using
+  CUDA,
+  Reexport
+
+@reexport using FourierFlows
+
+using LinearAlgebra: mul!, ldiv!
+using FourierFlows: parsevalsum
+
+nothingfunction(args...) = nothing
+
+"""
+    Problem(; parameters...)
+
+Construct a 2D shallow water problem.
+"""
+function Problem(dev::Device=CPU();
+  # Numerical parameters
+          nx = 256,
+          Lx = 2π,
+          ny = nx,
+          Ly = Lx,
+          dt = 0.01,
+  # Drag and/or hyper-/hypo-viscosity
+           ν = 0,
+          nν = 1,
+           μ = 0,
+          nμ = 0,
+  # Timestepper and equation options
+     stepper = "RK4",
+       calcF = nothingfunction,
+  stochastic = false,
+           T = Float64)
+
+  grid = TwoDGrid(dev, nx, Lx, ny, Ly; T=T)
+
+  params = Params{T}(ν, nν, μ, nμ, calcF)
+
+  vars = calcF == nothingfunction ? Vars(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(ν, nν, μ, nμ, calcF!)
+
+Returns the params for two-dimensional turbulence.
+"""
+struct Params{T} <: AbstractParams
+       ν :: T         # Vorticity viscosity
+      nν :: Int       # Vorticity hyperviscous order
+       μ :: T         # Bottom drag or hypoviscosity
+      nμ :: Int       # Order of hypodrag
+  calcF! :: Function  # Function that calculates the forcing F
+end
+
+
+# ---------
+# Equations
+# ---------
+
+"""
+    Equation(params, grid)
+
+Returns the equation for two-dimensional turbulence with `params` and `grid`.
+"""
+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
+
+struct Vars{Aphys, Atrans, F, P} <: TwoDNavierStokesVars
+        u :: Aphys
+        v :: Aphys
+        η :: Aphys
+       uh :: Atrans
+       vh :: Atrans
+       ηh :: Atrans
+       Fh :: F
+  prevsol :: P
+end
+
+const ForcedVars = Vars{<:AbstractArray, <:AbstractArray, <:AbstractArray, Nothing}
+const StochasticForcedVars = Vars{<:AbstractArray, <:AbstractArray, <:AbstractArray, <:AbstractArray}
+
+"""
+    Vars(dev, grid)
+
+Returns the `vars` for unforced two-dimensional turbulence on device `dev` and with `grid`.
+"""
+function Vars(::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) uh vh ηh
+  return Vars(u, v, η, uh, vh, ηh, nothing, nothing)
+end
+
+"""
+    ForcedVars(dev, grid)
+
+Returns the vars for forced two-dimensional turbulence 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) uh vh ηh Fh
+  return Vars(u, v, η, uh, vh, ηh, Fh, nothing)
+end
+
+"""
+    StochasticForcedVars(dev, grid)
+
+Returns the vars for stochastically forced two-dimensional turbulence 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) uh vh ηh Fh prevsol
+  return Vars(u, v, η, uh, vh, ηh, Fh, prevsol)
+end
+
+
+# -------
+# Solvers
+# -------
+
+"""
+    calcN_advection(N, sol, t, clock, vars, params, grid)
+
+Calculates the advection term.
+"""
+function calcN_advection!(N, sol, t, clock, vars, params, grid)
+  @. N = 0 * sol
+
+  return nothing
+end
+
+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::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
+
+  return nothing
+end
+
+end # module

test/runtests.jl

@@ -11,7 +11,8 @@ import # use 'import' rather than 'using' for submodules to keep namespace clean
   GeophysicalFlows.BarotropicQG,
   GeophysicalFlows.BarotropicQGQL,
   GeophysicalFlows.MultilayerQG,
-  GeophysicalFlows.SurfaceQG
+  GeophysicalFlows.SurfaceQG,
+  GeophysicalFlows.ShallowWater
 
 using FourierFlows: parsevalsum
 using GeophysicalFlows: lambdipole, peakedisotropicspectrum