Merge ea26575 into 6e70e88

.gitignore

@@ -12,3 +12,4 @@
 docs/build/
 docs/site/
 coverage/
+docs/Manifest.toml

Manifest.toml

@@ -65,7 +65,7 @@ version = "0.7.0"
 deps = ["InteractiveUtils", "OrderedCollections", "Random", "Serialization", "Test"]
 git-tree-sha1 = "8fc6e166e24fda04b2b648d4260cdad241788c54"
 uuid = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
-version = "0.14.0"
+version = "0.15.0"
 
 [[Dates]]
 deps = ["Printf"]
@@ -203,9 +203,9 @@ version = "1.18.5"
 
 [[PyPlot]]
 deps = ["Colors", "Compat", "LaTeXStrings", "PyCall", "VersionParsing"]
-git-tree-sha1 = "777c48006c57c8bd322d2032267085983a5fd50f"
+git-tree-sha1 = "dd9f4a0bba3933f907640a567e1ce2a8ab44c58c"
 uuid = "d330b81b-6aea-500a-939a-2ce795aea3ee"
-version = "2.6.3"
+version = "2.7.0"
 
 [[REPL]]
 deps = ["InteractiveUtils", "Markdown", "Sockets"]

README.md

@@ -21,9 +21,6 @@
 
 </p>
 
-
-
-
 This package leverages the [FourierFlows.jl]() framework to provide modules for solving problems in
 Geophysical Fluid Dynamics on periodic domains using Fourier-based pseudospectral methods.
 
@@ -43,6 +40,7 @@ All modules provide solvers on two-dimensional domains. We currently provide
 * `TwoDTurb`: the two-dimensional vorticity equation.
 * `BarotropicQG`: the barotropic quasi-geostrophic equation, which generalizes `TwoDTurb` to cases with topography and Coriolis parameters of the form `f = f₀ + βy`.
 * `BarotropicQGQL`: the quasi-linear barotropic quasi-geostrophic equation.
+* `MultilayerQG`: a multi-layer quasi-geostrophic model over topography and with ability to impose a zonal flow `U_n(y)` in each layer.
 
 
 [FourierFlows.jl]: https://github.com/FourierFlows/FourierFlows.jl

docs/make.jl

@@ -19,7 +19,8 @@ sitename = "GeophysicalFlows.jl",
             "Home" => "index.md",
             "Modules" => Any[
               "modules/twodturb.md",
-              "modules/barotropicqg.md"
+              "modules/barotropicqg.md",
+              "modules/multilayerqg.md"
             ],
             "DocStrings" => Any[
             "man/types.md",

docs/src/modules/multilayerqg.md

@@ -0,0 +1,103 @@
+# MultilayerQG Module
+
+```math
+\newcommand{\J}{\mathsf{J}}
+```
+
+### Basic Equations
+
+This module solves the layered quasi-geostrophic equations on a beta-plane of variable fluid depth $H-h(x,y)$. The flow in each layer is obtained through a streamfunction $\psi_j$ as $(u_j, \upsilon_j) = (-\partial_y\psi_j, \partial_x\psi_j)$, $j=1,...,n$, where $n$ is the number of fluid layers.
+
+The QGPV in each layer is
+
+```math
+\mathrm{QGPV}_j = q_j  + \underbrace{f_0+\beta y}_{\textrm{planetary PV}} + \delta_{j,n}\underbrace{\frac{f_0 h}{H_n}}_{\textrm{topographic PV}},\quad j=1,...,n.
+```
+
+where
+
+```math
+q_1 = \nabla^2\psi_1 + F_{3/2, 1} (\psi_2-\psi_1),\\
+q_j = \nabla^2\psi_j + F_{j-1/2, j} (\psi_{j-1}-\psi_j) + F_{j+1/2, j} (\psi_{j+1}-\psi_j),\quad j=2,\dots,n-1,\\
+q_n = \nabla^2\psi_n + F_{n-1/2, n} (\psi_{n-1}-\psi_n).
+```
+
+with
+
+```math
+F_{j+1/2, k} = \frac{f_0^2}{g'_{j+1/2} H_k}\quad\text{and}\quad
+g'_{j+1/2} = g\frac{\rho_{j+1}-\rho_j}{\rho_{j+1}} .
+```
+
+Therefore, in Fourier space the $q$'s and $\psi$'s are related through
+
+```math
+\begin{pmatrix} \widehat{q}_{\boldsymbol{k},1}\\\vdots\\\widehat{q}_{\boldsymbol{k},n} \end{pmatrix} =
+\underbrace{\left(-|\boldsymbol{k}|^2\mathbb{1} + \mathbb{F} \right)}_{\equiv \mathbb{S}_{\boldsymbol{k}}}
+\begin{pmatrix} \widehat{\psi}_{\boldsymbol{k},1}\\\vdots\\\widehat{\psi}_{\boldsymbol{k},n} \end{pmatrix}
+```
+
+where
+
+```math
+\mathbb{F} \equiv \begin{pmatrix}
+ -F_{3/2, 1} &              F_{3/2, 1}  &   0   &  \cdots    & 0\\
+  F_{3/2, 2} & -(F_{3/2, 2}+F_{5/2, 2}) & F_{5/2, 2} &       & \vdots\\
+ 0           &                  \ddots  &   \ddots   & \ddots & \\
+ \vdots      &                          &            &        &  0 \\
+ 0           &       \cdots             &   0   & F_{n-1/2, n} & -F_{n-1/2, n}
+\end{pmatrix}
+```
+
+The kinetic energy in each layer is:
+
+```math
+\textrm{KE}_j = \dfrac{H_j}{H} \int \dfrac1{2} |\boldsymbol{\nabla}\psi_j|^2 \frac{\mathrm{d}^2\boldsymbol{x}}{L_x L_y},\quad j=1,\dots,n,
+```
+
+while the potential energy related to each of fluid interface is
+
+```math
+\textrm{PE}_{j+1/2} = \int \dfrac1{2} \dfrac{f_0^2}{g'_{j+1/2}} (\psi_j-\psi_{j+1})^2 \frac{\mathrm{d}^2\boldsymbol{x}}{L_x L_y},\quad j=1,\dots,n-1.
+```
+
+
+Including an imposed zonal flow $U_j(y)$ in each layer the equations of motion are:
+
+```math
+\partial_t q_j + \J(\psi_j, q_j ) + (U_j - \partial_y\psi_j) \partial_x Q_j +  U_j \partial_x q_j  + (\partial_y Q_j)(\partial_x\psi_j) = -\delta_{j,n}\mu\nabla^2\psi_n - \nu(-1)^{n_\nu} \nabla^{2n_\nu} q_j,
+```
+
+with
+
+```math
+\partial_y Q_j \equiv \beta - \partial_y^2 U_j - (1-\delta_{j,1})F_{j-1/2, j} (U_{j-1}-U_j) - (1-\delta_{j,n})F_{j+1/2, j} (U_{j+1}-U_j) + \delta_{j,n}\partial_y\eta, \\
+\partial_x Q_j \equiv \delta_{j,n}\partial_x\eta.
+```
+
+
+### Implementation
+
+Matrices $\mathbb{S}_{\boldsymbol{k}}$ as well as $\mathbb{S}^{-1}_{\boldsymbol{k}}$ are included in `params` as `params.S` and `params.invS` respectively.
+
+You can get $\widehat{\psi}_j$ from $\widehat{q}_j$ with `streamfunctionfrompv!(psih, qh, invS, grid)`, while to go from $\widehat{\psi}_j$ back to $\widehat{q}_j$ `pvfromstreamfunction!(qh, psih, S, grid)`.
+
+
+
+
+The equations are time-stepped forward in Fourier space:
+
+```math
+\partial_t \widehat{q}_j = - \widehat{\J(\psi_j, q_j)}  - \widehat{U_j \partial_x Q_j} - \widehat{U_j \partial_x q_j}
++ \widehat{(\partial_y\psi_j) \partial_x Q_j}  - \widehat{(\partial_x\psi_j)(\partial_y Q_j)} + \delta_{j,n}\mu k^{2} \widehat{\psi}_n - \nu k^{2n_\nu} \widehat{q}_j
+```
+
+In doing so the Jacobian is computed in the conservative form: $\J(f,g) =
+\partial_y [ (\partial_x f) g] -\partial_x[ (\partial_y f) g]$.
+
+
+Thus:
+
+$$\mathcal{L} = - \nu k^{2n_\nu}\ ,$$
+$$\mathcal{N}(\widehat{q}_j) = - \widehat{\J(\psi_j, q_j)} - \widehat{U_j \partial_x Q_j} - \widehat{U_j \partial_x q_j}
+ + \widehat{(\partial_y\psi_j)(\partial_x Q_j)} - \widehat{(\partial_x\psi_j)(\partial_y Q_j)} + \delta_{j,n}\mu k^{2} \widehat{\psi}_n\ .$$

examples/barotropicqg_acc.jl

@@ -133,5 +133,5 @@ println((time()-startwalltime))
 
 plot_output(prob, fig, axs; drawcolorbar=false)
 
-savename = @sprintf("%s_%09d.png", joinpath(plotpath, plotname), prob.step)
+savename = @sprintf("%s_%09d.png", joinpath(plotpath, plotname), cl.step)
 savefig(savename, dpi=240)

examples/barotropicqg_betadecay.jl

@@ -152,5 +152,5 @@ println((time()-startwalltime))
 plot_output(prob, fig, axs; drawcolorbar=false)
 
 # save the figure as png
-savename = @sprintf("%s_%09d.png", joinpath(plotpath, plotname), prob.step)
+savename = @sprintf("%s_%09d.png", joinpath(plotpath, plotname), cl.step)
 savefig(savename)

examples/barotropicqg_betaforced.jl

@@ -180,5 +180,5 @@ println((time()-startwalltime))
 plot_output(prob, fig, axs; drawcolorbar=false)
 
 # save the figure as png
-savename = @sprintf("%s_%09d.png", joinpath(plotpath, plotname), prob.step)
+savename = @sprintf("%s_%09d.png", joinpath(plotpath, plotname), cl.step)
 savefig(savename)

examples/barotropicqgql_betaforced.jl

@@ -187,5 +187,5 @@ figure(2); pcolormesh(zMean.t, y[1, :], UM)
 
 
 # save the figure as png
-# savename = @sprintf("%s_%09d.png", joinpath(plotpath, plotname), prob.step)
+# savename = @sprintf("%s_%09d.png", joinpath(plotpath, plotname), cl.step)
 # savefig(savename)

examples/multilayerqg_2layer.jl

@@ -0,0 +1,145 @@
+using
+  PyPlot,
+  JLD2,
+  Printf,
+  FourierFlows
+
+using FFTW: ifft
+
+import GeophysicalFlows.MultilayerQG
+import GeophysicalFlows.MultilayerQG: energies, fluxes
+
+
+# Numerical parameters and time-stepping parameters
+nx = 64          # 2D resolution = nx^2
+ny = nx
+
+stepper = "FilteredRK4"   # timestepper
+dt  = 1e-2     # timestep
+nsteps = 20000 # total number of time-steps
+nsubs  = 200   # number of time-steps for plotting
+               # (nsteps must be multiple of nsubs)
+
+# Physical parameters
+Lx  = 2π      # domain size
+mu  = 1e-2     # bottom drag
+beta = 5       # the y-gradient of planetary PV
+
+nlayers = 2       # these choice of parameters give the
+f0, g = 1, 1      # desired PV-streamfunction relations
+  H = [0.2, 0.8]  # q1 = Δψ1 + 25*(ψ2-ψ1), and
+rho = [4.0, 5.0]  # q2 = Δψ2 + 25/4*(ψ1-ψ2).
+  U = zeros(nlayers)
+  U[1] = 0.5
+  U[2] = 0.0
+
+gr = TwoDGrid(nx, Lx)
+x, y = gridpoints(gr)
+k0, l0 = gr.k[2], gr.l[2] # fundamental wavenumbers
+
+# Initialize problem
+prob = MultilayerQG.InitialValueProblem(nlayers=nlayers, nx=nx, Lx=Lx, f0=f0, g=g, H=H, rho=rho, U=U, dt=dt, stepper=stepper, mu=mu, beta=beta)
+sol, cl, pr, vs, gr = prob.sol, prob.clock, prob.params, prob.vars, prob.grid
+
+# Files
+filepath = "."
+plotpath = "./plots_2layer"
+plotname = "snapshots"
+filename = joinpath(filepath, "2layer.jl.jld2")
+
+# File management
+if isfile(filename); rm(filename); end
+if !isdir(plotpath); mkdir(plotpath); end
+
+# Initialize with zeros
+MultilayerQG.set_q!(prob, 1e-2randn(Float64, (nx, ny, nlayers)))
+
+
+# Create Diagnostics
+E = Diagnostic(energies, prob; nsteps=nsteps)
+diags = [E]
+
+# Create Output
+get_sol(prob) = sol # extracts the Fourier-transformed solution
+function get_u(prob)
+  @. v.qh = sol
+  streamfunctionfrompv!(v.psih, v.qh, p.invS, g)
+  @. v.uh = -im*g.l *v.psih
+  invtransform!(v.u, v.uh, p)
+  v.u
+end
+out = Output(prob, filename, (:sol, get_sol), (:u, get_u))
+
+
+function plot_output(prob, fig, axs; drawcolorbar=false)
+
+  # Plot the PV field and the evolution of energy and enstrophy.
+
+  sol, v, p, g = prob.sol, prob.vars, prob.params, prob.grid
+  MultilayerQG.updatevars!(prob)
+
+  for j in 1:nlayers
+    sca(axs[j])
+    pcolormesh(x, y, v.q[:, :, j])
+    axis("square")
+    xlim(-Lx/2, Lx/2)
+    ylim(-Lx/2, Lx/2)
+    title(L"$\nabla^2\psi + \eta$ (part of the domain)")
+    if drawcolorbar==true
+      colorbar()
+    end
+
+    sca(axs[j+2])
+    cla()
+    contourf(x, y, v.psi[:, :, j])
+    contour(x, y, v.psi[:, :, j], colors="k")
+    axis("square")
+    xlim(-Lx/2, Lx/2)
+    ylim(-Lx/2, Lx/2)
+    title(L"$\nabla^2\psi + \eta$ (part of the domain)")
+    if drawcolorbar==true
+      colorbar()
+    end
+  end
+
+  sca(axs[5])
+  plot(mu*E.t[E.i], E.data[E.i][1][1], ".", color="b", label=L"$KE_1$")
+  plot(mu*E.t[E.i], E.data[E.i][1][2], ".", color="r", label=L"$KE_2$")
+  xlabel(L"\mu t")
+  ylabel(L"KE")
+  # legend()
+
+  sca(axs[6])
+  plot(mu*E.t[E.i], E.data[E.i][2][1], ".", color="k", label=L"$PE_1$")
+  xlabel(L"\mu t")
+  ylabel(L"PE")
+  # legend()
+
+  pause(0.001)
+end
+
+
+fig, axs = subplots(ncols=3, nrows=2, figsize=(15, 8))
+plot_output(prob, fig, axs; drawcolorbar=true)
+
+
+# Step forward
+startwalltime = time()
+
+while cl.step < nsteps
+  stepforward!(prob, diags, nsubs)
+
+  # Message
+  log = @sprintf("step: %04d, t: %d, KE1: %.4f, KE2: %.4f, PE: %.4f, τ: %.2f min", cl.step, cl.t, E.data[E.i][1][1], E.data[E.i][1][2], E.data[E.i][2][1], (time()-startwalltime)/60)
+
+  println(log)
+
+  plot_output(prob, fig, axs; drawcolorbar=false)
+
+end
+println((time()-startwalltime))
+
+plot_output(prob, fig, axs; drawcolorbar=false)
+
+savename = @sprintf("%s_%09d.png", joinpath(plotpath, plotname), cl.step)
+savefig(savename, dpi=240)

src/GeophysicalFlows.jl

@@ -11,5 +11,6 @@ include("utils.jl")
 include("twodturb.jl")
 include("barotropicqg.jl")
 include("barotropicqgql.jl")
+include("multilayerqg.jl")
 
 end # module

src/barotropicqg.jl

@@ -263,7 +263,7 @@ end
 """
     updatevars!(v, s, g)
 
-Update the vars in v on the grid g with the solution in s.sol.
+Update the vars in v on the grid g with the solution in sol.
 """
 function updatevars!(sol, v, p, g)
   v.U[] = sol[1, 1].re
@@ -289,7 +289,7 @@ updatevars!(prob) = updatevars!(prob.sol, prob.vars, prob.params, prob.grid)
     set_zeta!(prob, zeta)
     set_zeta!(s, v, g, zeta)
 
-Set the solution s.sol as the transform of zeta and update variables v
+Set the solution sol as the transform of zeta and update variables v
 on the grid g.
 """
 function set_zeta!(sol, v::Vars, p, g, zeta)
@@ -316,9 +316,9 @@ set_zeta!(prob, zeta) = set_zeta!(prob.sol, prob.vars, prob.params, prob.grid, z
 
 """
     set_U!(prob, U)
-    set_U!(s, v, g, U)
+    set_U!(sol, v, g, U)
 
-Set the (kx, ky)=(0, 0) part of solution s.sol as the domain-average zonal flow U.
+Set the (kx, ky)=(0, 0) part of solution sol as the domain-average zonal flow U.
 """
 function set_U!(sol, v, p, g, U::Float64)
   sol[1, 1] = U

src/barotropicqgql.jl

@@ -7,8 +7,6 @@ export
 
   energy,
   enstrophy,
-  energy00,
-  enstrophy00,
   dissipation,
   work,
   drag