Enhance examples + Fix some plot orientations

examples/barotropicqg_acc.jl

@@ -4,13 +4,17 @@
 #md # Also, it can be viewed as a Jupyter notebook via [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/generated/barotropicqgtopography.ipynb).
 #
 # An idealized version of the Southern Ocean. We solve the barotropic 
-# quasi-geostrophic eddy dynamics in a flud with variable depth $H-h(x,y)$. We 
+# quasi-geostrophic eddy dynamics in a flud with variable depth ``H-h(x,y)``. We 
 # also include an "Antarctic Circumpolar Current," i.e.,  a domain-average zonal 
-# velocity $U(t)$ which is forced by constant wind  stress $F$ and influenced by 
+# velocity ``U(t)`` which is forced by constant wind  stress ``F`` and influenced by 
 # bottom drag and topographic form stress. The equations solved are: 
-# $\partial_t \nabla^2\psi + \mathsf{J}(\psi-U y, \nabla^2\psi + \beta y + \eta) = -\mu\nabla^2\psi$
-# and 
-# $\partial_t U = F - \mu U -  \langle\psi\partial_x\eta\rangle$.
+# ```math
+# \partial_t \nabla^2 \psi + \mathsf{J}( \psi - U y, \nabla^2 \psi + \beta y + \eta ) = - \mu \nabla^2 \psi \, ,
+# ```
+# and
+# ```math 
+# \partial_t U = F - \mu U -  \langle \psi \partial_x \eta \rangle \, .
+# ```
 
 
 using FourierFlows, Plots, Printf
@@ -28,11 +32,11 @@ nothing # hide
 
 
 # ## Numerical parameters and time-stepping parameters
-     nx = 128            # 2D resolution = nx^2
-stepper = "FilteredRK4"  # timestepper
-    dt  = 0.1            # timestep
- nsteps = 10000          # total number of time-steps
- nsubs  = 25             # number of time-steps for intermediate logging/plotting (nsteps must be multiple of nsubs)
+      n = 128               # 2D resolution = n²
+stepper = "FilteredETDRK4"  # timestepper
+     dt = 0.1               # timestep
+ nsteps = 8000              # total number of time-steps
+ nsubs  = 25                # number of time-steps for intermediate logging/plotting (nsteps must be multiple of nsubs)
 nothing # hide
 
 
@@ -46,31 +50,31 @@ nν = 4         # viscosity order
  F = 0.0012    # normalized wind stress forcing on domain-averaged zonal flow U(t) flow
 nothing # hide
 
-# Define the topographic potential vorticity, $f_0 h(x, y)/H$
-topoPV(x, y) = 2cos(4x)*cos(4y)
+# Define the topographic potential vorticity, ``f_0 h(x, y)/H``,
+topoPV(x, y) = 2 * cos(4x) * cos(4y)
 nothing # hide
 
-# and the forcing function $F$ (here forcing is constant in time) that acts on the domain-averaged $U$ equation.
+# and the forcing function ``F`` (here forcing is constant in time) that acts on the domain-averaged ``U`` equation.
 calcFU(t) = F
 nothing # hide
 
 
 # ## Problem setup
 # We initialize a `Problem` by providing a set of keyword arguments,
-prob = BarotropicQG.Problem(dev; nx=nx, Lx=Lx, β=β, eta=topoPV,
+prob = BarotropicQG.Problem(dev; nx=n, Lx=Lx, β=β, eta=topoPV,
                   calcFU=calcFU, ν=ν, nν=nν, μ=μ, dt=dt, stepper=stepper)
 nothing # hide
 
 # and define some shortcuts.
-sol, cl, vs, pr, gr = prob.sol, prob.clock, prob.vars, prob.params, prob.grid
-x, y = gr.x, gr.y
+sol, clock, vars, params, grid = prob.sol, prob.clock, prob.vars, prob.params, prob.grid
+x, y = grid.x, grid.y
 nothing # hide
 
 
 # ## Setting initial conditions
 
 # Our initial condition is simply fluid at rest.
-BarotropicQG.set_zeta!(prob, zeros(gr.nx, gr.ny))
+BarotropicQG.set_zeta!(prob, zeros(grid.nx, grid.ny))
 
 
 # ## Diagnostics
@@ -100,7 +104,7 @@ nothing # hide
 
 # and then create Output.
 get_sol(prob) = sol # extracts the Fourier-transformed solution
-get_u(prob) = irfft(im*gr.lr.*gr.invKrsq.*sol, gr.nx)
+get_u(prob) = irfft(im * grid.lr .* grid.invKrsq .* sol, grid.nx)
 out = Output(prob, filename, (:sol, get_sol), (:u, get_u))
 nothing # hide
 
@@ -111,13 +115,15 @@ nothing # hide
 # of energy and enstrophy.
 
 function plot_output(prob)
-
-  pq = heatmap(x, y, vs.q,
+
+  grid = prob.grid
+
+  pq = heatmap(grid.x, grid.y, vars.q',
                  c = :balance,
               clim = (-2, 2),
        aspectratio = 1,
-             xlims = (-gr.Lx/2, gr.Lx/2),
-             ylims = (-gr.Ly/2, gr.Ly/2),
+             xlims = (-grid.Lx/2, grid.Lx/2),
+             ylims = (-grid.Ly/2, grid.Ly/2),
             xticks = -3:3,
             yticks = -3:3,
             xlabel = "x",
@@ -129,19 +135,19 @@ function plot_output(prob)
              label = ["eddy energy" "mean energy"],
          linewidth = 2,
              alpha = 0.7,
-             xlims = (-0.1, 10.1),
+             xlims = (-0.1, 1.01 * μ * nsteps * dt),
              ylims = (0, 0.0008),
             xlabel = "μt")
 
   pQ = plot(2,
              label = ["eddy enstrophy" "mean enstrophy"],
          linewidth = 2,
              alpha = 0.7,
-             xlims = (-0.1, 10.1),
+             xlims = (-0.1, 1.01 * μ * nsteps * dt),
              ylims = (-0.02, 0.12),
             xlabel = "μt")
 
-  l = @layout [ a{0.5w} grid(2, 1) ]
+  l = @layout [a{0.5w} Plots.grid(2, 1)]
   p = plot(pq, pE, pQ, layout=l, size = (900, 600))
 
   return p
@@ -157,34 +163,35 @@ p = plot_output(prob)
 
 startwalltime = time()
 
-anim = @animate for j=0:Int(nsteps/nsubs)
+anim = @animate for j = 0:round(Int, nsteps / nsubs)
 
-  cfl = cl.dt*maximum([maximum(vs.U.+vs.u)/gr.dx, maximum(vs.v)/gr.dy])
+  if j % (2000 / nsubs) == 0
+    cfl = clock.dt * maximum([maximum(vars.U .+ vars.u) / grid.dx, maximum(vars.v) / grid.dy])
+
+    log = @sprintf("step: %04d, t: %d, cfl: %.2f, E: %.4f, Q: %.4f, walltime: %.2f min",
+      clock.step, clock.t, cfl, E.data[E.i], Q.data[Q.i], (time()-startwalltime)/60)
+
+    println(log)
+  end
 
-  log = @sprintf("step: %04d, t: %d, cfl: %.2f, E: %.4f, Q: %.4f, walltime: %.2f min",
-    cl.step, cl.t, cfl, E.data[E.i], Q.data[Q.i], (time()-startwalltime)/60)
-
-  if j%(2000/nsubs)==0; println(log) end
-
-  p[1][1][:z] = Array(vs.q)
-  p[1][:title] = "∇²ψ + η, μt="*@sprintf("%.2f", μ*cl.t)
-  push!(p[2][1], μ*E.t[E.i], E.data[E.i])
-  push!(p[2][2], μ*Emean.t[Emean.i], Emean.data[Emean.i])
-  push!(p[3][1], μ*Q.t[Q.i], Q.data[Q.i])
-  push!(p[3][2], μ*Qmean.t[Qmean.i], Qmean.data[Qmean.i])
+  p[1][1][:z] = vars.q
+  p[1][:title] = "∇²ψ + η, μt = "*@sprintf("%.2f", μ * clock.t)
+  push!(p[2][1], μ * E.t[E.i], E.data[E.i])
+  push!(p[2][2], μ * Emean.t[Emean.i], Emean.data[Emean.i])
+  push!(p[3][1], μ * Q.t[Q.i], Q.data[Q.i])
+  push!(p[3][2], μ * Qmean.t[Qmean.i], Qmean.data[Qmean.i])
 
   stepforward!(prob, diags, nsubs)
   BarotropicQG.updatevars!(prob)
-
 end
 
 mp4(anim, "barotropicqg_acc.mp4", fps=18)
 
-# Note that since mean flow enstrophy is $Q_U = \beta U$ it can attain negative values. 
+# Note that since mean flow enstrophy is ``Q_U = \beta U`` it can attain negative values. 
 
 
 # ## Save
 
 # Finally save the last snapshot.
-savename = @sprintf("%s_%09d.png", joinpath(plotpath, plotname), cl.step)
+savename = @sprintf("%s_%09d.png", joinpath(plotpath, plotname), clock.step)
 savefig(savename)

examples/barotropicqg_betadecay.jl

@@ -22,7 +22,7 @@ nothing # hide
 
 # ## Numerical parameters and time-stepping parameters
 
-      n = 128            # 2D resolution = n^2
+      n = 128            # 2D resolution = n²
 stepper = "FilteredRK4"  # timestepper
      dt = 0.05           # timestep
  nsteps = 2000           # total number of time-steps
@@ -39,70 +39,70 @@ nothing # hide
 
 # ## Problem setup
 # We initialize a `Problem` by providing a set of keyword arguments. Not providing
-# a viscosity coefficient ν leads to the module's default value: ν=0. In this
-# example numerical instability due to accumulation of enstrophy in high wavenumbers
+# a viscosity coefficient `ν` leads to the module's default value: `ν=0`. In this
+# example numerical instability due to accumulation of enstrophy at high wavenumbers
 # is taken care with the `FilteredTimestepper` we picked. 
 prob = BarotropicQG.Problem(dev; nx=n, Lx=L, β=β, μ=μ, dt=dt, stepper=stepper)
 nothing # hide
 
 # and define some shortcuts
-sol, cl, vs, pr, gr = prob.sol, prob.clock, prob.vars, prob.params, prob.grid
-x, y = gr.x, gr.y
+sol, clock, vars, params, grid = prob.sol, prob.clock, prob.vars, prob.params, prob.grid
+x, y = grid.x, grid.y
 nothing # hide
 
 
 # ## Setting initial conditions
 
 # Our initial condition consist of a flow that has power only at wavenumbers with
-# $8 < \frac{L}{2\pi} \sqrt{k_x^2+k_y^2} < 10$ and initial energy $E_0$:
+# ``8 < \frac{L}{2\pi} \sqrt{k_x^2 + k_y^2} < 10`` and initial energy ``E_0``:
 
-E0 = 0.1 # energy of initial condition
+E₀ = 0.1 # energy of initial condition
 
-K = @. sqrt(gr.Krsq)                      # a 2D array with the total wavenumber
-k = [gr.kr[i] for i=1:gr.nkr, j=1:gr.nl]  # a 2D array with the zonal wavenumber
+K = @. sqrt(grid.Krsq)                          # a 2D array with the total wavenumber
+k = [grid.kr[i] for i=1:grid.nkr, j=1:grid.nl]  # a 2D array with the zonal wavenumber
 
 Random.seed!(1234)
-qih = randn(Complex{eltype(gr)}, size(sol))
-qih[K .<  8 * 2π/L] .= 0
-qih[K .> 10 * 2π/L] .= 0
-qih[k .== 0] .= 0         # no power at zonal wavenumber k=0 component
-Ein = energy(qih, gr)     # compute energy of qi
-qih *= sqrt(E0 / Ein)     # normalize qi to have energy E0
-qi  = irfft(qih, gr.nx)
+qih = randn(Complex{eltype(grid)}, size(sol))
+@. qih = ifelse(K < 2, 0, qih)
+@. qih = ifelse(K > 10, 0, qih)
+@. qih = ifelse(k == 0, 0, qih)   # no power at zonal wavenumber k=0 component
+Ein = energy(qih, grid)           # compute energy of qi
+qih *= sqrt(E₀ / Ein)             # normalize qi to have energy E₀
+qi = irfft(qih, grid.nx)
 
 BarotropicQG.set_zeta!(prob, qi)
 nothing #hide
 
 # Let's plot the initial vorticity field:
 
-p1 = heatmap(x, y, vs.q,
+p1 = heatmap(x, y, vars.q',
          aspectratio = 1,
               c = :balance,
            clim = (-12, 12),
-          xlims = (-gr.Lx/2, gr.Lx/2),
-          ylims = (-gr.Ly/2, gr.Ly/2),
+          xlims = (-grid.Lx/2, grid.Lx/2),
+          ylims = (-grid.Ly/2, grid.Ly/2),
          xticks = -3:3,
          yticks = -3:3,
          xlabel = "x",
          ylabel = "y",
           title = "initial vorticity ζ=∂v/∂x-∂u/∂y",
      framestyle = :box)
 
-p2 = contourf(x, y, vs.psi,
+p2 = contourf(x, y, vars.psi',
         aspectratio = 1,
              c = :viridis,
         levels = range(-0.65, stop=0.65, length=10), 
           clim = (-0.65, 0.65),
-         xlims = (-gr.Lx/2, gr.Lx/2),
-         ylims = (-gr.Ly/2, gr.Ly/2),
+         xlims = (-grid.Lx/2, grid.Lx/2),
+         ylims = (-grid.Ly/2, grid.Ly/2),
         xticks = -3:3,
         yticks = -3:3,
         xlabel = "x",
         ylabel = "y",
          title = "initial streamfunction ψ",
     framestyle = :box)
 
-l = @layout grid(1, 2)
+l = @layout Plots.grid(1, 2)
 p = plot(p1, p2, layout=l, size=(900, 400))
 
 
@@ -131,14 +131,14 @@ nothing # hide
 
 # and then create Output.
 get_sol(prob) = sol # extracts the Fourier-transformed solution
-get_u(prob) = irfft(im*gr.l.*gr.invKrsq.*sol, gr.nx)
+get_u(prob) = irfft(im * grid.l .* grid.invKrsq .* sol, grid.nx)
 out = Output(prob, filename, (:sol, get_sol), (:u, get_u))
 nothing # hide
 
 
 # ## Visualizing the simulation
 
-# We define a function that plots the vorticity and streamfunction fields and 
+# We define a function that plots the vorticity and streamfunction and 
 # their corresponding zonal mean structure.
 
 function plot_output(prob)
@@ -147,28 +147,28 @@ function plot_output(prob)
   ζ̄ = mean(ζ, dims=1)'
   ū = mean(prob.vars.u, dims=1)'
 
-  pζ = heatmap(x, y, ζ,
+  pζ = heatmap(x, y, ζ',
        aspectratio = 1,
             legend = false,
                  c = :balance,
               clim = (-12, 12),
-             xlims = (-gr.Lx/2, gr.Lx/2),
-             ylims = (-gr.Ly/2, gr.Ly/2),
+             xlims = (-grid.Lx/2, grid.Lx/2),
+             ylims = (-grid.Ly/2, grid.Ly/2),
             xticks = -3:3,
             yticks = -3:3,
             xlabel = "x",
             ylabel = "y",
              title = "vorticity ζ=∂v/∂x-∂u/∂y",
         framestyle = :box)
 
-  pψ = contourf(x, y, ψ,
+  pψ = contourf(x, y, ψ',
        aspectratio = 1,
             legend = false,
                  c = :viridis,
             levels = range(-0.65, stop=0.65, length=10), 
               clim = (-0.65, 0.65),
-             xlims = (-gr.Lx/2, gr.Lx/2),
-             ylims = (-gr.Ly/2, gr.Ly/2),
+             xlims = (-grid.Lx/2, grid.Lx/2),
+             ylims = (-grid.Ly/2, grid.Ly/2),
             xticks = -3:3,
             yticks = -3:3,
             xlabel = "x",
@@ -196,7 +196,7 @@ function plot_output(prob)
             ylabel = "y")
   plot!(pum, 0*y, y, linestyle=:dash, linecolor=:black)
 
-  l = @layout grid(2, 2)
+  l = @layout Plots.grid(2, 2)
   p = plot(pζ, pζm, pψ, pum, layout = l, size = (900, 800))
 
   return p
@@ -212,29 +212,32 @@ startwalltime = time()
 
 p = plot_output(prob)
 
-anim = @animate for j=0:Int(nsteps/nsubs)
+anim = @animate for j = 0:round(Int, nsteps/nsubs)
 
-  log = @sprintf("step: %04d, t: %d, E: %.4f, Q: %.4f, walltime: %.2f min",
-    cl.step, cl.t, E.data[E.i], Z.data[Z.i], (time()-startwalltime)/60)
+  if j % (1000 / nsubs) == 0
+    cfl = clock.dt * maximum([maximum(vars.u) / grid.dx, maximum(vars.v) / grid.dy])
 
-  if j%(1000/nsubs)==0; println(log) end  
+    log = @sprintf("step: %04d, t: %d, cfl: %.2f, E: %.4f, Q: %.4f, walltime: %.2f min",
+      clock.step, clock.t, cfl, E.data[E.i], Z.data[Z.i], (time()-startwalltime)/60)
 
+    println(log)
+  end  
 
-  p[1][1][:z] = Array(vs.zeta)
-  p[1][:title] = "vorticity, t="*@sprintf("%.2f", cl.t)
-  p[3][1][:z] = Array(vs.psi)
-  p[2][1][:x] = mean(vs.zeta, dims=1)'
-  p[4][1][:x] = mean(vs.u, dims=1)'
+  p[1][1][:z] = vars.zeta
+  p[1][:title] = "vorticity, t="*@sprintf("%.2f", clock.t)
+  p[3][1][:z] = vars.psi
+  p[2][1][:x] = mean(vars.zeta, dims=1)'
+  p[4][1][:x] = mean(vars.u, dims=1)'
 
   stepforward!(prob, diags, nsubs)
   BarotropicQG.updatevars!(prob)
 
 end
 
-mp4(anim, "barotropicqg_betadecay.mp4", fps=14)
+mp4(anim, "barotropicqg_betadecay.mp4", fps=12)
 
 # ## Save
 
 # Finally save the last snapshot.
-savename = @sprintf("%s_%09d.png", joinpath(plotpath, plotname), cl.step)
+savename = @sprintf("%s_%09d.png", joinpath(plotpath, plotname), clock.step)
 savefig(savename)