-
Notifications
You must be signed in to change notification settings - Fork 32
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
[WIP] Add ShallowWater module #163
Draft
navidcy
wants to merge
22
commits into
main
Choose a base branch
from
ncc/shallowwater
base: main
Could not load branches
Branch not found: {{ refName }}
Loading
Could not load tags
Nothing to show
Loading
Are you sure you want to change the base?
Some commits from the old base branch may be removed from the timeline,
and old review comments may become outdated.
Draft
Changes from 3 commits
Commits
Show all changes
22 commits
Select commit
Hold shift + click to select a range
0b82b7c
adds empty ShallowWater module
navidcy 96de511
some typos and stray commas
navidcy 35d13f1
Merge branch 'master' into ncc/shallowwater
navidcy 660e2c8
fix equations in docs
navidcy 0a3b4c5
merge from master
navidcy 405af3e
first step towards defining a state
navidcy 2945ec8
add equations in conservative form
navidcy e2071ca
inviscid + dissipation in conservative form
navidcy 51b2787
it's MultiLayer now
navidcy 496a5c7
use qu, qv variables
navidcy 1f5d0db
uncomment
navidcy ba4f986
remove μ
navidcy a1c8e05
first swe tests
navidcy 3fa54b2
fix some inconsistencies
navidcy 3c2298e
calcN_advection
navidcy 0d83c07
calcN
navidcy 30140ec
adds nontrivial bathymetry
navidcy 0b68124
fix typo in bathymetry term
navidcy 33a0400
fix sign error in conservative eqs
navidcy aad7636
merge master
navidcy f905141
resolve conflicts
navidcy d7bc1ad
merge master
navidcy File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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 |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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 |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Should we write in conservative form?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Yes.
Also there are some typos, e.g., the
-g ∂η/∂y
term is missing from thev
equation.Also I'd like to include non-trivial bathymetry (which appears as a forcing term on the right-hand side in the momentum equations.)
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Here they are in conservative form:
I made an attempt to write the
calcN!
function. Bit tricky...Also, what's the form of dissipation? We want it to conserve the total
hu
,hv
andh
... Right?There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
We're having a related discussion at CliMA/Oceananigans.jl#1712