# Linear rotating shallow water dynamics

This example can be run online via . Also, it can be viewed as a Jupyter notebook via .

This example solves the linear 1D rotating shallow water equations for the $u(x, t)$, $v(x, t)$ and the surface surface elevation $\eta(x, t)$, for a fluid with constant rest-depth $H$. That is, the total fluid's depth is $H + \eta(x, t)$ with $|\eta| \ll H$.

The linearized equations for the evolution of $u$, $v$, $\eta$ are:

\begin{aligned} \partial_t u - f v & = - g \partial_x \eta - \mathrm{D} u, \\ \partial_t v + f u & = - \mathrm{D} v, \\ \partial_t \eta + H \partial_x u & = - \mathrm{D} \eta. \end{aligned}

Above, $g$ is the gravitational acceleration, $f$ is the Coriolis parameter, and $\mathrm{D}$ indicates a hyperviscous linear operator of the form $(-1)^{n_ν} ν \nabla^{2 n_ν}$, with $ν$ the viscosity coefficient and $n_ν$ the order of the operator.

Rotation introduces the deformation length scale, $L_d = \sqrt{g H} / f$. Disturbances with length scales much smaller than $L_d$ don't "feel" the rotation and propagate as inertia-gravity waves. Disturbances with length scales comparable or larger than $L_d$ should be approximately in geostrophic balance, i.e., the Coriolis acceleration $f \widehat{\bm{z}} \times \bm{u}$ should be in approximate balance with the pressure gradient $-g \bm{\nabla} \eta$.

using FourierFlows, CairoMakie, Printf, Random, JLD2
using LinearAlgebra: mul!, ldiv!

## Coding up the equations

### A demonstration of FourierFlows.jl framework

What follows is a step-by-step tutorial demonstrating how you can create your own solver for an equation of your liking.

The basic building blocks for a FourierFlows.Problem are:

• Grid struct containining the physical and wavenumber grid for the problem,
• Params struct containining all the parameters of the problem,
• Vars struct containining arrays with the variables used in the problem,
• Equation struct containining the coefficients of the linear operator $L$ and the function that computes the nonlinear terms, usually named calcN!().

The Grid structure is provided by FourierFlows.jl. We simply have to call one of either OneDGrid(), TwoDGrid(), or ThreeDGrid() constructors, depending on the dimensionality of the problem. All other structs mentioned above are problem-specific and need to be constructed for every set of equations we want to solve.

First let's construct the Params struct that contains all parameters of the problem.

struct Params{T} <: AbstractParams
ν :: T         # Hyperviscosity coefficient
nν :: Int       # Order of the hyperviscous operator
g :: T         # Gravitational acceleration
H :: T         # Fluid depth
f :: T         # Coriolis parameter
end

Now the Vars struct that contains all variables used in this problem. For this problem Vars includes the representations of the flow fields in physical space u, v and η and their Fourier transforms uh, vh, and ηh.

struct Vars{Aphys, Atrans} <: AbstractVars
u :: Aphys
v :: Aphys
η :: Aphys
uh :: Atrans
vh :: Atrans
ηh :: Atrans
end

A constructor populates empty arrays based on the dimension of the grid and then creates Vars struct.

"""
Vars(grid)

Construct the Vars for 1D linear shallow water dynamics based on the dimensions of the grid arrays.
"""
function Vars(grid)
Dev = typeof(grid.device)
T = eltype(grid)

@devzeros Dev T grid.nx u v η
@devzeros Dev Complex{T} grid.nkr uh vh ηh

return Vars(u, v, η, uh, vh, ηh)
end

In Fourier space, the 1D linear shallow water dynamics read:

\begin{aligned} \frac{\partial \hat{u}}{\partial t} & = \underbrace{ f \hat{v} - i k g \hat{\eta} }_{N_u} \; \underbrace{- \nu k^2 }_{L_u} \hat{u} , \\ \frac{\partial \hat{v}}{\partial t} & = \underbrace{ - f \hat{u} }_{N_v} \; \underbrace{- \nu k^2 }_{L_v} \hat{v} , \\ \frac{\partial \hat{\eta}}{\partial t} & = \underbrace{ - i k H \hat{u} }_{N_{\eta}} \; \underbrace{- \nu k^2 }_{L_{\eta}} \hat{\eta} . \end{aligned}

Although, e.g., terms involving the Coriolis accelaration are, in principle, linear we include them in the nonlinear term $N$ because they render the linear operator $L$ non-diagonal.

With these in mind, we construct function calcN! that computes the nonlinear terms.

"""
calcN!(N, sol, t, clock, vars, params, grid)

Compute the nonlinear terms for 1D linear shallow water dynamics.
"""
function calcN!(N, sol, t, clock, vars, params, grid)
@. vars.uh = sol[:, 1]
@. vars.vh = sol[:, 2]
@. vars.ηh = sol[:, 3]

@. N[:, 1] =   params.f * vars.vh - im * grid.kr * params.g * vars.ηh    #  + f v - g ∂η/∂x
@. N[:, 2] = - params.f * vars.uh                                        #  - f u
@. N[:, 3] = - im * grid.kr * params.H * vars.uh                         #  - H ∂u/∂x

dealias!(N, grid)

return nothing
end

Next we construct the Equation struct:

"""
Equation(params, grid)

Construct the equation: the linear part, in this case the hyperviscous dissipation,
and the nonlinear part, which is computed by calcN! function.
"""
function Equation(params, grid)
T = eltype(grid)
dev = grid.device

L = zeros(dev, T, (grid.nkr, 3))
D = @. - params.ν * grid.kr^(2*params.nν)

L[:, 1] .= D # for u equation
L[:, 2] .= D # for v equation
L[:, 3] .= D # for η equation

return FourierFlows.Equation(L, calcN!, grid)
end

We now have all necessary building blocks to construct a FourierFlows.Problem. It would be useful, however, to define some more "helper functions". For example, a function that updates all variables given the solution sol which comprises $\hat{u}$, $\hat{v}$ and $\hat{\eta}$:

"""
updatevars!(prob)

Update the variables in prob.vars using the solution in prob.sol.
"""
function updatevars!(prob)
vars, grid, sol = prob.vars, prob.grid, prob.sol

@. vars.uh = sol[:, 1]
@. vars.vh = sol[:, 2]
@. vars.ηh = sol[:, 3]

ldiv!(vars.u, grid.rfftplan, deepcopy(sol[:, 1])) # use deepcopy() because irfft destroys its input
ldiv!(vars.v, grid.rfftplan, deepcopy(sol[:, 2])) # use deepcopy() because irfft destroys its input
ldiv!(vars.η, grid.rfftplan, deepcopy(sol[:, 3])) # use deepcopy() because irfft destroys its input

return nothing
end

Another useful function is one that prescribes an initial condition to the state variable sol.

"""
set_uvη!(prob, u0, v0, η0)

Set the state variable prob.sol as the Fourier transforms of u0, v0, and η0
and update all variables in prob.vars.
"""
function set_uvη!(prob, u0, v0, η0)
vars, grid, sol = prob.vars, prob.grid, prob.sol

A = typeof(vars.u) # determine the type of vars.u

# below, e.g., A(u0) converts u0 to the same type as vars expects
# (useful when u0 is a CPU array but grid.device is GPU)
mul!(vars.uh, grid.rfftplan, A(u0))
mul!(vars.vh, grid.rfftplan, A(v0))
mul!(vars.ηh, grid.rfftplan, A(η0))

@. sol[:, 1] = vars.uh
@. sol[:, 2] = vars.vh
@. sol[:, 3] = vars.ηh

updatevars!(prob)

return nothing
end

## Let's prescibe parameter values and solve the PDE

We are now ready to write up a program that sets up parameter values, constructs the problem prob, # time steps the solutions prob.sol and plots it.

## Choosing a device: CPU or GPU

dev = CPU()    # Device (CPU/GPU)

## Numerical parameters and time-stepping parameters

     nx = 512            # grid resolution
stepper = "FilteredRK4"  # timestepper
dt = 20.0           # timestep (s)
nsteps = 320            # total number of time-steps

## Physical parameters

Lx = 500e3      # Domain length (m)
g  = 9.8        # Gravitational acceleration (m s⁻²)
H  = 200.0      # Fluid depth (m)
f  = 1e-2       # Coriolis parameter (s⁻¹)
ν  = 100.0      # Viscosity (m² s⁻¹)
nν = 1          # Viscosity order (nν = 1 means Laplacian ∇²)

## Construct the structs and you are ready to go!

Create a grid and also params, vars, and the equation structs. Then give them all as input to the FourierFlows.Problem() constructor to get a problem struct, prob, that contains all of the above.

    grid = OneDGrid(dev; nx, Lx)
params = Params(ν, nν, g, H, f)
vars = Vars(grid)
equation = Equation(params, grid)

prob = FourierFlows.Problem(equation, stepper, dt, grid, vars, params)

## Setting initial conditions

For initial condition we take the fluid at rest ($u = v = 0$). The free surface elevation is perturbed from its rest position ($\eta=0$); the disturbance we impose a Gaussian bump with half-width greater than the deformation radius and on top of that we superimpose some random noise with scales smaller than the deformation radius. We mask the small-scale perturbations so that it only applies in the central part of the domain by applying

The system develops geostrophically-balanced jets around the Gaussian bump, while the smaller-scale noise propagates away as inertia-gravity waves.

First let's construct the Gaussian bump.

gaussian_width = 6e3
gaussian_amplitude = 3.0
gaussian_bump = @. gaussian_amplitude * exp( - grid.x^2 / (2*gaussian_width^2) )

fig = Figure(size = (600, 260))
ax =  Axis(fig[1, 1];
xlabel = "x [km]",
ylabel = "η [m]",
title = "A gaussian bump with half-width ≈ " * string(gaussian_width/1e3) * " km",
limits = ((-Lx/2e3, Lx/2e3), nothing))

lines!(ax, grid.x/1e3, gaussian_bump;    # divide with 1e3 to convert m -> km
color = (:black, 0.7),
linewidth = 2)

Next the noisy perturbation. The mask is simply a product of hyperbolic tangent functions.

mask = @. 1/4 * (1 + tanh( -(grid.x - 100e3) / 10e3)) * (1 + tanh( (grid.x + 100e3) / 10e3))

noise_amplitude = 0.1 # the amplitude of the noise for η(x, t=0) (m)
η_noise = noise_amplitude * Random.randn(size(grid.x))

fig = Figure(size = (600, 520))

kwargs = (xlabel = "x [km]", limits = ((-Lx/2e3, Lx/2e3), nothing))

ax1 =  Axis(fig[1, 1]; ylabel = "η [m]", title = "small-scale noise", kwargs...)

ax2 =  Axis(fig[2, 1]; ylabel = "mask", kwargs...)

lines!(ax1, grid.x/1e3, η_noise;      # divide with 1e3 to convert m -> km
color = (:black, 0.7),
linewidth = 3)

lines!(ax2, grid.x/1e3, mask;         # divide with 1e3 to convert m -> km
color = (:gray, 0.7),
linewidth = 2)

Sum the Gaussian bump and the noise and then call set_uvη!() to set the initial condition to the problem prob.

η0 = @. gaussian_bump + η_noise
u0 = zeros(grid.nx)
v0 = zeros(grid.nx)

set_uvη!(prob, u0, v0, η0)

fig = Figure(size = (600, 260))

ax =  Axis(fig[1, 1];
xlabel = "x [km]",
ylabel = "η [m]",
title = "initial surface elevation, η(x, t=0)",
limits = ((-Lx/2e3, Lx/2e3), nothing))

lines!(ax, grid.x/1e3, η0;    # divide with 1e3 to convert m -> km
color = (:black, 0.7),
linewidth = 2)

## Saving output

filepath = "."
filename = joinpath(filepath, "linear_swe.jld2")

get_sol(prob) = prob.sol

out = Output(prob, filename, (:sol, get_sol))
Output
├──── prob: FourierFlows.Problem{DataType, Matrix{ComplexF64}, Float64, Matrix{Float64}}
├──── path: ./linear_swe_1.jld2
└── fields: Dict{Symbol, Function}(:sol => Main.get_sol)

We call saveproblem to we write the problem's configuration parameters to the .jld2 file.

saveproblem(out)

## Time-stepping the Problem forward

for j = 0:nsteps
updatevars!(prob)
stepforward!(prob)
saveoutput(out)
end

## Visualizing the simulation

First we load the saved output files.

using JLD2

file = jldopen(out.path)
iterations = parse.(Int, keys(file["snapshots/t"]))

nx = file["grid/nx"]
x = file["grid/x"]
512-element Vector{Float64}:
-250000.0
-249023.4375
-248046.875
-247070.3125
-246093.75
-245117.1875
-244140.625
-243164.0625
-242187.5
-241210.9375
⋮
241210.9375
242187.5
243164.0625
244140.625
245117.1875
246093.75
247070.3125
248046.875
249023.4375

Then we animate the output. We use Makie's Observable to animate the data. To dive into how Observables work we refer to Makie.jl's Documentation.

 n = Observable(1)

u = @lift irfft(file[string("snapshots/sol/", iterations[$n])][:, 1], nx) v = @lift irfft(file[string("snapshots/sol/", iterations[$n])][:, 2], nx)
η = @lift irfft(file[string("snapshots/sol/", iterations[$n])][:, 3], nx) toptitle = @lift "t = " * @sprintf("%.1f", file[string("snapshots/t/", iterations[$n])]/60) * " min"

fig = Figure(size = (600, 800))

kwargs_η = (xlabel = "x [km]", limits = ((-Lx/2e3, Lx/2e3), nothing))
kwargs_uv = (xlabel = "x [km]", limits = ((-Lx/2e3, Lx/2e3), (-0.3, 0.3)))

ax_η =  Axis(fig[2, 1]; ylabel = "η [m]", title = toptitle, kwargs_η...)

ax_u =  Axis(fig[3, 1]; ylabel = "u [m s⁻¹]", kwargs_uv...)

ax_v =  Axis(fig[4, 1]; ylabel = "v [m s⁻¹]", kwargs_uv...)

Ld = @sprintf "%.2f" sqrt(g * H) / f /1e3     # divide with 1e3 to convert m -> km
title = "Deformation radius √(gh) / f = "*string(Ld)*" km"

fig[1, 1] = Label(fig, title, fontsize=24, tellwidth=false)

lines!(ax_η, grid.x/1e3, η; # divide with 1e3 to convert m -> km
color = (:blue, 0.7))

lines!(ax_u, grid.x/1e3, u; # divide with 1e3 to convert m -> km
color = (:red, 0.7))

lines!(ax_v, grid.x/1e3, v; # divide with 1e3 to convert m -> km
color = (:green, 0.7))

frames = 1:length(iterations)

record(fig, "onedshallowwater.mp4", frames, framerate=18) do i
n[] = i
end

## Geostrophic balance

It is instructive to compare the solution for $\bm{u}$ with its geostrophically balanced approximation, $f \widehat{\bm{z}} \times \bm{u}_{\rm geostrophic} = - g \bm{\nabla} \eta$, i.e.,

\begin{aligned} v_{\rm geostrophic} & = \frac{g}{f} \frac{\partial \eta}{\partial x} \ , \\ u_{\rm geostrophic} & = - \frac{g}{f} \frac{\partial \eta}{\partial y} = 0 \ . \end{aligned}
u_geostrophic = zeros(grid.nx)  # -g/f ∂η/∂y = 0
v_geostrophic = params.g / params.f * irfft(im * grid.kr .* vars.ηh, grid.nx)  #g/f ∂η/∂x

The geostrophic solution should capture well the the behavior of the flow in the center of the domain, after small-scale disturbances propagate away. Let's plot and see!

fig = Figure(size = (600, 600))

kwargs = (xlabel = "x [km]", limits = ((-Lx/2e3, Lx/2e3), (-0.3, 0.3)))

ax_u =  Axis(fig[2, 1]; ylabel = "u [m s⁻¹]", kwargs...)

ax_v =  Axis(fig[3, 1]; ylabel = "v [m s⁻¹]", kwargs...)

fig[1, 1] = Label(fig, "Geostrophic balance", fontsize=24, tellwidth=false)

lines!(ax_u, grid.x/1e3, vars.u; # divide with 1e3 to convert m -> km
label = "u",
linewidth = 3,
color = (:red, 0.7))

lines!(ax_u, grid.x/1e3, u_geostrophic; # divide with 1e3 to convert m -> km
label = "- g/f ∂η/∂y",
linewidth = 3,
color = (:purple, 0.7))

axislegend(ax_u)

lines!(ax_v, grid.x/1e3, vars.v; # divide with 1e3 to convert m -> km
label = "v",
linewidth = 3,
color = (:green, 0.7))

lines!(ax_v, grid.x/1e3, v_geostrophic; # divide with 1e3 to convert m -> km
label = "g/f ∂η/∂x",
linewidth = 3,
color = (:purple, 0.7))

axislegend(ax_v)