Merge pull request #103 from BrodiePearson/master

Adds enstrophy budget diagnostics for 2D and barotropic QG
FourierFlows · Sep 11, 2020 · e8c39e4 · e8c39e4
2 parents 27a1374 + 670c599
commit e8c39e4
Show file tree

Hide file tree

Showing 6 changed files with 324 additions and 80 deletions.
diff --git a/.gitignore b/.gitignore
@@ -19,6 +19,7 @@ docs/Manifest.toml
 # Model output
 **/*.jld
 **/*.jld2
+**/*.mat
 
 # Plots and movies
 **/*.png

diff --git a/examples/twodnavierstokes_stochasticforcing.jl b/examples/twodnavierstokes_stochasticforcing.jl
@@ -3,19 +3,21 @@
 #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 
+# A simulation of forced-dissipative two-dimensional turbulence. We solve the
 # two-dimensional vorticity equation with stochastic excitation and dissipation in
-# the form of linear drag and hyperviscosity. As a demonstration, we compute how 
-# each of the forcing and dissipation terms contribute to the energy budget.
+# the form of linear drag and hyperviscosity. As a demonstration, we compute how
+# each of the forcing and dissipation terms contribute to the energy and the 
+# enstrophy budgets.
 
 using FourierFlows, Printf, Plots
-  
+
 using FourierFlows: parsevalsum
 using Random: seed!
 using FFTW: irfft
 
 import GeophysicalFlows.TwoDNavierStokes
-import GeophysicalFlows.TwoDNavierStokes: energy, enstrophy, dissipation, work, drag
+import GeophysicalFlows.TwoDNavierStokes: energy, energy_dissipation, energy_work, energy_drag
+import GeophysicalFlows.TwoDNavierStokes: enstrophy, enstrophy_dissipation, enstrophy_work, enstrophy_drag
 
 
 # ## Choosing a device: CPU or GPU
@@ -41,12 +43,12 @@ nothing # hide
 
 # We force the vorticity equation with stochastic excitation that is delta-correlated
 # in time and while spatially homogeneously and isotropically correlated. The forcing
-# has a spectrum with power in a ring in wavenumber space of radious $k_f$ and 
-# width $\delta k_f$, and it injects energy per unit area and per unit time equal 
+# has a spectrum with power in a ring in wavenumber space of radious $k_f$ and
+# width $\delta k_f$, and it injects energy per unit area and per unit time equal
 # to $\varepsilon$.
 
 forcing_wavenumber = 14.0    # the central forcing wavenumber for a spectrum that is a ring in wavenumber space
-forcing_bandwidth  = 1.5     # the width of the forcing spectrum 
+forcing_bandwidth  = 1.5     # the width of the forcing spectrum
 ε = 0.001                    # energy input rate by the forcing
 
 gr   = TwoDGrid(dev, n, L)
@@ -110,10 +112,14 @@ TwoDNavierStokes.set_zeta!(prob, zeros(g.nx, g.ny))
 
 # 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
-W = Diagnostic(work,        prob, nsteps=nt) # work input by forcing
-diags = [E, D, W, R] # A list of Diagnostics types passed to "stepforward!" will  be updated every timestep.
+R = Diagnostic(energy_drag,        prob, nsteps=nt) # dissipation by drag
+D = Diagnostic(energy_dissipation, prob, nsteps=nt) # dissipation by hyperviscosity
+W = Diagnostic(energy_work,        prob, nsteps=nt) # work input by forcing
+Z = Diagnostic(enstrophy,      prob, nsteps=nt) # energy
+RZ = Diagnostic(enstrophy_drag,        prob, nsteps=nt) # dissipation by drag
+DZ = Diagnostic(enstrophy_dissipation, prob, nsteps=nt) # dissipation by hyperviscosity
+WZ = Diagnostic(enstrophy_work,        prob, nsteps=nt) # work input by forcing
+diags = [E, D, W, R, Z, DZ, WZ, RZ] # A list of Diagnostics types passed to "stepforward!" will  be updated every timestep.
 nothing # hide
 
 
@@ -125,23 +131,28 @@ nothing # hide
 
 function computetendencies_and_makeplot(prob, diags)
   TwoDNavierStokes.updatevars!(prob)
-  E, D, W, R = diags
+  E, D, W, R, Z, Dᶻ, Wᶻ, Rᶻ = diags
 
   clocktime = round(μ*cl.t, digits=2)
 
   i₀ = 1
   dEdt_numerical = (E[(i₀+1):E.i] - E[i₀:E.i-1])/cl.dt #numerical first-order approximation of energy tendency
+  dZdt_numerical = (Z[(i₀+1):Z.i] - Z[i₀:Z.i-1])/cl.dt #numerical first-order approximation of enstrophy tendency
   ii = (i₀):E.i-1
   ii2 = (i₀+1):E.i
 
   t = E.t[ii]
   dEdt_computed = W[ii2] - D[ii] - R[ii]        # Stratonovich interpretation
-
-  residual = dEdt_computed - dEdt_numerical
+  dZdt_computed = Wᶻ[ii2] - Dᶻ[ii] - Rᶻ[ii]
+
+  residual_E = dEdt_computed - dEdt_numerical
+  residual_Z = dZdt_computed - dZdt_numerical
+
+  εᶻ = parsevalsum(forcing_spectrum / 2, g) / (g.Lx * g.Ly)
 
-  l = @layout grid(2, 2)
+  l = @layout grid(2,3)
 
-  p1 = heatmap(x, y, v.zeta,
+  pzeta = heatmap(x, y, v.zeta,
             aspectratio = 1,
             legend = false,
                  c = :viridis,
@@ -155,30 +166,57 @@ function computetendencies_and_makeplot(prob, diags)
              title = "∇²ψ(x, y, t="*@sprintf("%.2f", cl.t)*")",
         framestyle = :box)
 
-  p2 = plot(μ*t, [W[ii2] ε.+0*t -D[ii] -R[ii]],
-             label = ["work, W" "ensemble mean work, <W>" "dissipation, D" "drag, D=-2μE"],
+  pζ = plot(pzeta, size = (400, 400))
+
+  p1 = plot(μ*t, [W[ii2] ε.+0*t -D[ii] -R[ii]],
+             label = ["work, W" "ensemble mean work, <W>" "dissipation, D" "drag, R=-2μE"],
          linestyle = [:solid :dash :solid :solid],
          linewidth = 2,
              alpha = 0.8,
             xlabel = "μt",
             ylabel = "energy sources and sinks")
 
-  p3 = plot(μ*t, [dEdt_computed[ii], dEdt_numerical],
-             label = ["computed W-D" "numerical dE/dt"],
-         linestyle = [:solid :dashdotdot],
-         linewidth = 2,
-             alpha = 0.8,
-            xlabel = "μt",
-            ylabel = "dE/dt")
-
-  p4 = plot(μ*t, residual,
-             label = "residual dE/dt = computed - numerical",
-         linewidth = 2,
-             alpha = 0.7,
-            xlabel = "μt")
-
-  p = plot(p1, p2, p3, p4, layout=l, size = (900, 800))
-  return p
+  p2 = plot(μ*t, [dEdt_computed[ii], dEdt_numerical],
+           label = ["computed W-D" "numerical dE/dt"],
+       linestyle = [:solid :dashdotdot],
+       linewidth = 2,
+           alpha = 0.8,
+          xlabel = "μt",
+          ylabel = "dE/dt")
+
+  p3 = plot(μ*t, residual_E,
+           label = "residual dE/dt = computed - numerical",
+       linewidth = 2,
+           alpha = 0.7,
+          xlabel = "μt")
+
+  p4 = plot(μ*t, [Wᶻ[ii2] εᶻ.+0*t -Dᶻ[ii] -Rᶻ[ii]],
+           label = ["Enstrophy work, Wᶻ" "mean enstrophy work, <Wᶻ>" "enstrophy dissipation, Dᶻ" "enstrophy drag, Rᶻ=-2μZ"],
+       linestyle = [:solid :dash :solid :solid],
+       linewidth = 2,
+           alpha = 0.8,
+          xlabel = "μt",
+          ylabel = "enstrophy sources and sinks")
+
+
+  p5 = plot(μ*t, [dZdt_computed[ii], dZdt_numerical],
+         label = ["computed Wᶻ-Dᶻ" "numerical dZ/dt"],
+     linestyle = [:solid :dashdotdot],
+     linewidth = 2,
+         alpha = 0.8,
+        xlabel = "μt",
+        ylabel = "dZ/dt")
+
+  p6 = plot(μ*t, residual_Z,
+         label = "residual dZ/dt = computed - numerical",
+     linewidth = 2,
+         alpha = 0.7,
+        xlabel = "μt")
+
+
+  pbudgets = plot(p1, p2, p3, p4, p5, p6, layout=l, size = (1300, 900))
+
+  return pζ, pbudgets
 end
 nothing # hide
 
@@ -198,11 +236,13 @@ for i = 1:ns
   log = @sprintf("step: %04d, t: %.1f, cfl: %.3f, walltime: %.2f min", cl.step, cl.t,
         cfl, (time()-startwalltime)/60)
 
-  println(log)        
+  println(log)
 end
 
 
 # ## Plot
 # And now let's see what we got. We plot the output.
 
-p = computetendencies_and_makeplot(prob, diags)
+pζ, pbudgets = computetendencies_and_makeplot(prob, diags)
+
+pbudgets
diff --git a/src/barotropicqg.jl b/src/barotropicqg.jl
@@ -11,7 +11,10 @@ export
   meanenstrophy,
   dissipation,
   work,
-  drag
+  drag,
+  drag_ens,
+  work_ens,
+  dissipation_ens
 
 using
   CUDA,
@@ -62,9 +65,9 @@ function Problem(dev::Device=CPU();
   params = !(typeof(eta)<:ArrayType(dev)) ? Params(grid, β, eta, μ, ν, nν, calcFU, calcFq) : Params(β, eta, rfft(eta), μ, ν, nν, calcFU, calcFq)
 
   vars = (calcFq == nothingfunction && calcFU == 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
 
@@ -206,7 +209,7 @@ function calcN_advection!(N, sol, t, clock, vars, params, grid)
   @. vars.v = vars.v * vars.q # v*q
 
   mul!(vars.uh, grid.rfftplan, vars.u)      # \hat{(u+U)*q}
-  
+
   # Nonlinear advection term for q (part 1)
   @. N = -im * grid.kr * vars.uh            # -∂[(U+u)q]/∂x
   mul!(vars.uh, grid.rfftplan, vars.v)      # \hat{v*q}
@@ -341,18 +344,30 @@ enstrophy(prob) = enstrophy(prob.sol, prob.grid, prob.vars)
 
 """
     meanenergy(prob)
-    
+
 Returns the energy of the domain-averaged U.
 """
 meanenergy(prob) = CUDA.@allowscalar real(0.5 * prob.sol[1, 1].^2)
 
 """
     meanenstrophy(prob)
-    
+
 Returns the enstrophy of the domain-averaged U.
 """
 meanenstrophy(prob) = CUDA.@allowscalar real(prob.params.β * prob.sol[1, 1])
 
+"""
+    dissipation_ens(prob)
+
+Returns the domain-averaged dissipation rate of enstrophy. nν must be >= 1.
+"""
+@inline function dissipation_ens(prob)
+  sol, vars, params, grid = prob.sol, prob.vars, prob.params, prob.grid
+  @. vars.uh = grid.Krsq^params.nν * abs2(sol)
+  vars.uh[1, 1] = 0
+  return params.ν / (grid.Lx * grid.Ly) * parsevalsum(vars.uh, grid)
+end
+
 """
     dissipation(prob)
     dissipation(sol, vars, params, grid)
@@ -386,6 +401,25 @@ end
 
 @inline work(prob) = work(prob.sol, prob.vars, prob.grid)
 
+"""
+    work_ens(prob)
+    work_ens(sol, v, grid)
+
+Returns the domain-averaged rate of work of enstrophy by the forcing Fh.
+"""
+@inline function work_ens(sol, vars::ForcedVars, grid)
+  @. vars.uh = sol * conj(vars.Fqh)
+  return 1 / (grid.Lx * grid.Ly) * parsevalsum(vars.uh, grid)
+end
+
+@inline function work_ens(sol, vars::StochasticForcedVars, grid)
+  @. vars.uh = (vars.prevsol + sol) / 2 * conj(vars.Fqh) # Stratonovich
+  # @. vars.uh = grid.invKrsq * vars.prevsol * conj(vars.Fh)           # Ito
+  return 1 / (grid.Lx * grid.Ly) * parsevalsum(vars.uh, grid)
+end
+
+@inline work_ens(prob) = work_ens(prob.sol, prob.vars, prob.grid)
+
 """
     drag(prob)
     drag(sol, vars, params, grid)
@@ -399,4 +433,16 @@ Returns the extraction of domain-averaged energy by drag μ.
   return params.μ / (grid.Lx * grid.Ly) * parsevalsum(vars.uh, grid)
 end
 
+"""
+    drag_ens(prob)
+
+Returns the extraction of domain-averaged enstrophy by drag/hypodrag μ.
+"""
+@inline function drag_ens(prob)
+  sol, vars, params, grid = prob.sol, prob.vars, prob.params, prob.grid
+  @. vars.uh = grid.Krsq^0.0 * abs2(sol)
+  CUDA.@allowscalar vars.uh[1, 1] = 0
+  return params.μ / (grid.Lx * grid.Ly) * parsevalsum(vars.uh, grid)
+end
+
 end # module