Merge pull request #71 from FourierFlows/EnhanceForcedBetaplane

Simplifies examples + adds binder/nbviewer intro

docs/Project.toml

@@ -9,4 +9,4 @@ Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 PyPlot = "d330b81b-6aea-500a-939a-2ce795aea3ee"
 
 [compat]
-FFTW = "≥ 0.4.1"
+FourierFlows = "≥ 0.4.1"

docs/make.jl

@@ -40,7 +40,7 @@ end
 Timer(t -> println(" "), 0, interval=240)
 
 format = Documenter.HTML(
-  collapselevel = 1,
+  collapselevel = 2,
      prettyurls = get(ENV, "CI", nothing) == "true",
       canonical = "https://fourierflows.github.io/GeophysicalFlows.jl/dev/"
 )

examples/README.md

@@ -1,17 +1,20 @@
 # GeophysicalFlows.jl/examples
 
 
-These are some basic examples of the various modules included in GeophysicalFlows.jl.
+These are some basic examples of the various modules included in GeophysicalFlows.jl. The best way to go through the examples is via the <a href="https://fourierflows.github.io/GeophysicalFlows.jl/dev/">documentation <img src="https://img.shields.io/badge/docs-dev-blue.svg"></a>.
 
-## Usage
 
-Run
+## Run the examples
+
+You can run the examples in an executable environment via [Binder](https://mybinder.org) by clicking on the [![badge](https://img.shields.io/badge/binder-badge-579ACA.svg?logo=data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAFkAAABZCAMAAABi1XidAAAB8lBMVEX///9XmsrmZYH1olJXmsr1olJXmsrmZYH1olJXmsr1olJXmsrmZYH1olL1olJXmsr1olJXmsrmZYH1olL1olJXmsrmZYH1olJXmsr1olL1olJXmsrmZYH1olL1olJXmsrmZYH1olL1olL0nFf1olJXmsrmZYH1olJXmsq8dZb1olJXmsrmZYH1olJXmspXmspXmsr1olL1olJXmsrmZYH1olJXmsr1olL1olJXmsrmZYH1olL1olLeaIVXmsrmZYH1olL1olL1olJXmsrmZYH1olLna31Xmsr1olJXmsr1olJXmsrmZYH1olLqoVr1olJXmsr1olJXmsrmZYH1olL1olKkfaPobXvviGabgadXmsqThKuofKHmZ4Dobnr1olJXmsr1olJXmspXmsr1olJXmsrfZ4TuhWn1olL1olJXmsqBi7X1olJXmspZmslbmMhbmsdemsVfl8ZgmsNim8Jpk8F0m7R4m7F5nLB6jbh7jbiDirOEibOGnKaMhq+PnaCVg6qWg6qegKaff6WhnpKofKGtnomxeZy3noG6dZi+n3vCcpPDcpPGn3bLb4/Mb47UbIrVa4rYoGjdaIbeaIXhoWHmZYHobXvpcHjqdHXreHLroVrsfG/uhGnuh2bwj2Hxk17yl1vzmljzm1j0nlX1olL3AJXWAAAAbXRSTlMAEBAQHx8gICAuLjAwMDw9PUBAQEpQUFBXV1hgYGBkcHBwcXl8gICAgoiIkJCQlJicnJ2goKCmqK+wsLC4usDAwMjP0NDQ1NbW3Nzg4ODi5+3v8PDw8/T09PX29vb39/f5+fr7+/z8/Pz9/v7+zczCxgAABC5JREFUeAHN1ul3k0UUBvCb1CTVpmpaitAGSLSpSuKCLWpbTKNJFGlcSMAFF63iUmRccNG6gLbuxkXU66JAUef/9LSpmXnyLr3T5AO/rzl5zj137p136BISy44fKJXuGN/d19PUfYeO67Znqtf2KH33Id1psXoFdW30sPZ1sMvs2D060AHqws4FHeJojLZqnw53cmfvg+XR8mC0OEjuxrXEkX5ydeVJLVIlV0e10PXk5k7dYeHu7Cj1j+49uKg7uLU61tGLw1lq27ugQYlclHC4bgv7VQ+TAyj5Zc/UjsPvs1sd5cWryWObtvWT2EPa4rtnWW3JkpjggEpbOsPr7F7EyNewtpBIslA7p43HCsnwooXTEc3UmPmCNn5lrqTJxy6nRmcavGZVt/3Da2pD5NHvsOHJCrdc1G2r3DITpU7yic7w/7Rxnjc0kt5GC4djiv2Sz3Fb2iEZg41/ddsFDoyuYrIkmFehz0HR2thPgQqMyQYb2OtB0WxsZ3BeG3+wpRb1vzl2UYBog8FfGhttFKjtAclnZYrRo9ryG9uG/FZQU4AEg8ZE9LjGMzTmqKXPLnlWVnIlQQTvxJf8ip7VgjZjyVPrjw1te5otM7RmP7xm+sK2Gv9I8Gi++BRbEkR9EBw8zRUcKxwp73xkaLiqQb+kGduJTNHG72zcW9LoJgqQxpP3/Tj//c3yB0tqzaml05/+orHLksVO+95kX7/7qgJvnjlrfr2Ggsyx0eoy9uPzN5SPd86aXggOsEKW2Prz7du3VID3/tzs/sSRs2w7ovVHKtjrX2pd7ZMlTxAYfBAL9jiDwfLkq55Tm7ifhMlTGPyCAs7RFRhn47JnlcB9RM5T97ASuZXIcVNuUDIndpDbdsfrqsOppeXl5Y+XVKdjFCTh+zGaVuj0d9zy05PPK3QzBamxdwtTCrzyg/2Rvf2EstUjordGwa/kx9mSJLr8mLLtCW8HHGJc2R5hS219IiF6PnTusOqcMl57gm0Z8kanKMAQg0qSyuZfn7zItsbGyO9QlnxY0eCuD1XL2ys/MsrQhltE7Ug0uFOzufJFE2PxBo/YAx8XPPdDwWN0MrDRYIZF0mSMKCNHgaIVFoBbNoLJ7tEQDKxGF0kcLQimojCZopv0OkNOyWCCg9XMVAi7ARJzQdM2QUh0gmBozjc3Skg6dSBRqDGYSUOu66Zg+I2fNZs/M3/f/Grl/XnyF1Gw3VKCez0PN5IUfFLqvgUN4C0qNqYs5YhPL+aVZYDE4IpUk57oSFnJm4FyCqqOE0jhY2SMyLFoo56zyo6becOS5UVDdj7Vih0zp+tcMhwRpBeLyqtIjlJKAIZSbI8SGSF3k0pA3mR5tHuwPFoa7N7reoq2bqCsAk1HqCu5uvI1n6JuRXI+S1Mco54YmYTwcn6Aeic+kssXi8XpXC4V3t7/ADuTNKaQJdScAAAAAElFTkSuQmCC)](https://mybinder.org)
+ at the top of each example page in the documentation.
+
+Alternatively, you can run these scripts directly. Run
 ```julia
 ] instantiate
 ```
-to install dependencies. Then you can run any of the examples by
+to install dependencies. Then run any of the examples by
 ```julia
 include("example_script.jl")
 ```
 
-

examples/barotropicqg_acc.jl

@@ -1,8 +1,8 @@
-#md # [![](https://mybinder.org/badge_logo.svg)](@__BINDER_ROOT_URL__/generated/barotropicqgtopography.ipynb)
-#md # [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/generated/barotropicqgtopography.ipynb)
-
 # # Barotropic QG beta-plane turbulence over topography
 #
+#md # This example can be run online via [![](https://mybinder.org/badge_logo.svg)](@__BINDER_ROOT_URL__/generated/barotropicqgtopography.ipynb).
+#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 
 # also include an "Antarctic Circumpolar Current," i.e.,  a domain-average zonal 
@@ -76,11 +76,11 @@ BarotropicQG.set_zeta!(prob, 0*x)
 
 # ## Diagnostics
 
-# Create Diagnostics -- `energy` and `enstrophy` functions are imported at the top.
+# Create Diagnostics -- `energy`, `meanenergy`, `enstrophy`, and `meanenstrophy` functions are imported at the top.
 E = Diagnostic(energy, prob; nsteps=nsteps)
 Q = Diagnostic(enstrophy, prob; nsteps=nsteps)
 Emean = Diagnostic(meanenergy, prob; nsteps=nsteps)
-Qmean = Diagnostic(meanenergy, prob; nsteps=nsteps)
+Qmean = Diagnostic(meanenstrophy, prob; nsteps=nsteps)
 diags = [E, Emean, Q, Qmean]
 nothing # hide
 
@@ -173,6 +173,7 @@ fig, axs = subplots(ncols=3, nrows=1, figsize=(15, 4), dpi=200)
 plot_output(prob, fig, axs; drawcolorbar=true)
 gcf() # hide
 
-# and finally save the figure
+# Note that since mean flow enstrophy is $Q_U = \beta U$ it can attain negative values. 
+# Finally we save the figure.
 savename = @sprintf("%s_%09d.png", joinpath(plotpath, plotname), cl.step)
 savefig(savename)

examples/barotropicqg_betadecay.jl

@@ -1,8 +1,8 @@
-#md # [![](https://mybinder.org/badge_logo.svg)](@__BINDER_ROOT_URL__/generated/barotropicqg_betadecay.ipynb)
-#md # [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/generated/barotropicqg_betadecay.ipynb)
-
 # # Decaying barotropic QG beta-plane turbulence
 #
+#md # This example can be run online via [![](https://mybinder.org/badge_logo.svg)](@__BINDER_ROOT_URL__/generated/barotropicqg_betadecay.ipynb). 
+#md # Also, it can be viewed as a Jupyter notebook via [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/generated/barotropicqg_betadecay.ipynb).
+# 
 # An example of decaying barotropic quasi-geostrophic turbulence on a beta plane.
 
 using FourierFlows, PyPlot, Printf, Random
@@ -37,7 +37,7 @@ Lx = 2π        # domain size
 nν = 1         # viscosity order
  β = 15.0      # planetary PV gradient
  μ = 0.0       # bottom drag
-
+nothing # hide
 
 # ## Problem setup
 # We initialize a `Problem` by providing a set of keyword arguments,
@@ -57,16 +57,12 @@ nothing # hide
 
 Random.seed!(1234)
 E0 = 0.1
-modk = ones(g.nkr, g.nl)
-modk[real.(g.Krsq).<(8*2*pi/g.Lx)^2] .= 0
-modk[real.(g.Krsq).>(10*2*pi/g.Lx)^2] .= 0
-modk[1, :] .= 0
-psih = (randn(Float64, size(sol)) .+ im*randn(Float64, size(sol))).*modk
-psih = @. psih*prob.timestepper.filter
-Ein = real(sum(g.Krsq.*abs2.(psih)/(g.nx*g.ny)^2))
-psih = psih*sqrt(E0/Ein)
-qi = -irfft(g.Krsq.*psih, g.nx)
-E0 = FourierFlows.parsevalsum(g.Krsq.*abs2.(psih), g)
+qih = randn(Complex{Float64}, size(sol))
+qih[ g.Krsq .< (8*2π /g.Lx)^2] .= 0
+qih[ g.Krsq .> (10*2π/g.Lx)^2] .= 0
+Ein = energy(qih, g) # compute energy of qi
+qih = qih*sqrt(E0/Ein) # normalize qi to have energy E0
+qi  = irfft(qih, g.nx) 
 
 BarotropicQG.set_zeta!(prob, qi)
 nothing #hide

examples/barotropicqg_betaforced.jl

@@ -1,8 +1,8 @@
-#md # [![](https://mybinder.org/badge_logo.svg)](@__BINDER_ROOT_URL__/generated/barotropicqg_betaforced.ipynb)
-#md # [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/generated/barotropicqg_betaforced.ipynb)
-
 # # Forced-dissipative barotropic QG beta-plane turbulence
 #
+#md # This example can be run online via [![](https://mybinder.org/badge_logo.svg)](@__BINDER_ROOT_URL__/generated/barotropicqg_betaforced.ipynb). 
+#md # Also, it can be viewed as a Jupyter notebook via [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/generated/barotropicqg_betaforced.ipynb).
+#
 # A simulation of forced-dissipative barotropic quasi-geostrophic turbulence on 
 # a beta plane. The dynamics include linear drag and stochastic excitation.
 
@@ -90,6 +90,20 @@ sol, cl, v, p, g = prob.sol, prob.clock, prob.vars, prob.params, prob.grid
 nothing # hide
 
 
+# First let's see how a forcing realization looks like.
+calcFq!(v.Fqh, sol, 0.0, cl, v, p, g)
+
+fig = figure(figsize=(3, 2), dpi=150)
+pcolormesh(x, y, irfft(v.Fqh, g.nx))
+axis("square")
+xticks(-3:1:3)
+yticks(-3:1:3)
+title("a forcing realization")
+colorbar()
+clim(-8, 8)
+gcf() # hide
+
+
 # ## Setting initial conditions
 
 # Our initial condition is simply fluid at rest.
@@ -98,7 +112,7 @@ BarotropicQG.set_zeta!(prob, 0*x)
 
 # ## Diagnostics
 
-# Create Diagnostic -- "energy" and "enstrophy" are functions imported at the top.
+# Create Diagnostic -- `energy` and `enstrophy` are functions imported at the top.
 E = Diagnostic(energy, prob; nsteps=nsteps)
 Z = Diagnostic(enstrophy, prob; nsteps=nsteps)
 diags = [E, Z] # A list of Diagnostics types passed to "stepforward!" will  be updated every timestep.
@@ -139,8 +153,8 @@ function plot_output(prob, fig, axs; drawcolorbar=false)
   cla()
   pcolormesh(x, y, v.q)
   axis("square")
-  xticks(-2:2:2)
-  yticks(-2:2:2)
+  xticks(-3:1:3)
+  yticks(-3:1:3)
   title(L"vorticity $\zeta = \partial_x v - \partial_y u$")
   if drawcolorbar==true
     colorbar()
@@ -153,8 +167,8 @@ function plot_output(prob, fig, axs; drawcolorbar=false)
     contour(x, y, v.psi, colors="k")
   end
   axis("square")
-  xticks(-2:2:2)
-  yticks(-2:2:2)
+  xticks(-3:1:3)
+  yticks(-3:1:3)
   title(L"streamfunction $\psi$")
   if drawcolorbar==true
     colorbar()

examples/barotropicqgql_betaforced.jl

@@ -1,8 +1,8 @@
-#md # [![](https://mybinder.org/badge_logo.svg)](@__BINDER_ROOT_URL__/generated/barotropicqgql_betaforced.ipynb)
-#md # [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/generated/barotropicqgql_betaforced.ipynb)
-
 # # Quasi-Linear forced-dissipative barotropic QG beta-plane turbulence
 #
+#md # This example can be run online via [![](https://mybinder.org/badge_logo.svg)](@__BINDER_ROOT_URL__/generated/barotropicqgql_betaforced.ipynb). 
+#md # Also, it can be viewed as a Jupyter notebook via [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/generated/barotropicqgql_betaforced.ipynb).
+#
 # A simulation of forced-dissipative barotropic quasi-geostrophic turbulence on 
 # a beta plane under the *quasi-linear approximation*. The dynamics include 
 # linear drag and stochastic excitation.
@@ -84,6 +84,20 @@ sol, cl, v, p, g = prob.sol, prob.clock, prob.vars, prob.params, prob.grid
 nothing # hide
 
 
+# First let's see how a forcing realization looks like.
+calcF!(v.Fh, sol, 0.0, cl, v, p, g)
+
+fig = figure(figsize=(3, 2), dpi=150)
+pcolormesh(x, y, irfft(v.Fh, g.nx))
+axis("square")
+xticks(-3:1:3)
+yticks(-3:1:3)
+title("a forcing realization")
+colorbar()
+clim(-8, 8)
+gcf() # hide
+
+
 # ## Setting initial conditions
 
 # Our initial condition is simply fluid at rest.

examples/multilayerqg_2layer.jl

@@ -1,16 +1,16 @@
-#md # [![](https://mybinder.org/badge_logo.svg)](@__BINDER_ROOT_URL__/generated/multilayerqg_2layer.ipynb)
-#md # [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/generated/multilayerqg_2layer.ipynb)
-
 # # Phillips model of Baroclinic Instability
 #
+#md # This example can be run online via [![](https://mybinder.org/badge_logo.svg)](@__BINDER_ROOT_URL__/generated/multilayerqg_2layer.ipynb).
+#md # Also, it can be viewed as a Jupyter notebook via [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/generated/multilayerqg_2layer.ipynb).
+#
 # A simulation of the growth of barolinic instability in the Phillips 2-layer model
 # when we impose a vertical mean flow shear as a difference $\Delta U$ in the
 # imposed, domain-averaged, zonal flow at each layer.
 
 using FourierFlows, PyPlot, Printf
 
 import GeophysicalFlows.MultilayerQG
-import GeophysicalFlows.MultilayerQG: energies, fluxes
+import GeophysicalFlows.MultilayerQG: energies
 
 
 # ## Numerical parameters and time-stepping parameters
@@ -61,7 +61,7 @@ nothing # hide
 
 # ## Diagnostics
 
-# Create Diagnostics -- `energy` function is imported at the top.
+# Create Diagnostics -- `energies` function is imported at the top.
 E = Diagnostic(energies, prob; nsteps=nsteps)
 diags = [E] # A list of Diagnostics types passed to "stepforward!" will  be updated every timestep.
 nothing # hide

examples/twodnavierstokes_decaying.jl

@@ -1,8 +1,8 @@
-#md # [![](https://mybinder.org/badge_logo.svg)](@__BINDER_ROOT_URL__/generated/twodnavierstokes_decaying.ipynb)
-#md # [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/generated/twodnavierstokes_decaying.ipynb)
-
 # # 2D decaying turbulence
 #
+#md # This example can be run online via [![](https://mybinder.org/badge_logo.svg)](@__BINDER_ROOT_URL__/generated/twodnavierstokes_decaying.ipynb).
+#md # Also, it can be viewed as a Jupyter notebook via [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/generated/twodnavierstokes_decaying.ipynb).
+#
 # A simulations of decaying two-dimensional turbulence.
 
 using FourierFlows, PyPlot, Printf, Random

examples/twodnavierstokes_stochasticforcing.jl

@@ -1,14 +1,15 @@
-#md # [![](https://mybinder.org/badge_logo.svg)](@__BINDER_ROOT_URL__/generated/twodnavierstokes_stochasticforcing.ipynb)
-#md # [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/generated/twodnavierstokes_stochasticforcing.ipynb)
-
 # # 2D forced-dissipative turbulence
 #
+#md # This example can be run online via [![](https://mybinder.org/badge_logo.svg)](@__BINDER_ROOT_URL__/generated/twodnavierstokes_stochasticforcing.ipynb).
+#md # Also, it can be viewed as a Jupyter notebook via [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/generated/twodnavierstokes_stochasticforcing.ipynb).
+#
 # A simulation of forced-dissipative two-dimensional turbulence. We solve the 
 # two-dimensional vorticity equation with linear drag and stochastic excitation.
 
 using PyPlot, FourierFlows, Printf
 
 using Random: seed!
+using FFTW: irfft
 
 import GeophysicalFlows.TwoDNavierStokes
 import GeophysicalFlows.TwoDNavierStokes: energy, enstrophy, dissipation, work, drag
@@ -81,6 +82,20 @@ sol, cl, v, p, g = prob.sol, prob.clock, prob.vars, prob.params, prob.grid
 nothing # hide
 
 
+# First let's see how a forcing realization looks like.
+calcF!(v.Fh, sol, 0.0, cl, v, p, g)
+
+fig = figure(figsize=(3, 2), dpi=150)
+pcolormesh(x, y, irfft(v.Fh, g.nx))
+axis("square")
+xticks(-3:1:3)
+yticks(-3:1:3)
+title("a forcing realization")
+colorbar()
+clim(-200, 200)
+gcf() # hide
+
+
 # ## Setting initial conditions
 
 # Our initial condition is simply fluid at rest.
@@ -89,7 +104,7 @@ TwoDNavierStokes.set_zeta!(prob, 0*x)
 
 # ## Diagnostics
 
-# Create Diagnostics; the diagnostics here will probe the energy budget.
+# Create Diagnostics; the diagnostics are aimed to probe the energy budget.
 E = Diagnostic(energy,      prob, nsteps=nt) # energy
 R = Diagnostic(drag,        prob, nsteps=nt) # dissipation by drag
 D = Diagnostic(dissipation, prob, nsteps=nt) # dissipation by hyperviscosity
@@ -111,10 +126,12 @@ function makeplot(prob, diags)
   t = round(μ*cl.t, digits=2)
   sca(axs[1]); cla()
   pcolormesh(x, y, v.zeta)
+  axis("square")
+  xticks(-3:1:3)
+  yticks(-3:1:3)
   xlabel(L"$x$")
   ylabel(L"$y$")
-  title("\$\\nabla^2\\psi(x,y,\\mu t= $t )\$")
-  axis("square")
+  title("\$\\nabla^2\\psi(x,y,\\mu t= $t)\$")
 
   sca(axs[3]); cla()
 

src/barotropicqg.jl

@@ -316,31 +316,40 @@ set_U!(prob, U::Float64) = set_U!(prob.sol, prob.vars, prob.params, prob.grid, U
 
 
 """
-Calculate the domain-averaged kinetic energy.
+    energy(sol, g)
+    energy(prob)
+
+Returns the domain-averaged kinetic energy of sol.
 """
-function energy(prob)
-  sol, g = prob.sol, prob.grid
+function energy(sol, g::AbstractGrid)
   0.5*(parsevalsum2(g.kr.*g.invKrsq.*sol, g)
         + parsevalsum2(g.l.*g.invKrsq.*sol, g))/(g.Lx*g.Ly)
 end
-
+energy(prob) = energy(prob.sol, prob.grid)
 
 """
-Returns the domain-averaged enstrophy.
+    enstrophy(sol, g, v)
+    enstrophy(prob)
+
+Returns the domain-averaged enstrophy of sol.
 """
-function enstrophy(prob)
-  sol, v, g = prob.sol, prob.vars, prob.grid
+function enstrophy(sol, g::AbstractGrid, v::AbstractVars)
   @. v.uh = sol
   v.uh[1, 1] = 0
   0.5*parsevalsum2(v.uh, g)/(g.Lx*g.Ly)
 end
+enstrophy(prob) = enstrophy(prob.sol, prob.grid, prob.vars)
 
 """
+    meanenergy(prob)
+    
 Returns the energy of the domain-averaged U.
 """
 meanenergy(prob) = real(0.5*prob.sol[1, 1].^2)
 
 """
+    meanenstrophy(prob)
+    
 Returns the enstrophy of the domain-averaged U.
 """
 meanenstrophy(prob) = real(prob.params.β*prob.sol[1, 1])