diff --git a/.gitignore b/.gitignore
index 18027391..59fc4100 100644
--- 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
index 2541e4c5..9468c47b 100644
--- 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
index 623a870f..4bde49d2 100644
--- 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
diff --git a/src/twodnavierstokes.jl b/src/twodnavierstokes.jl
index c8b97ae9..1fe2b504 100644
--- a/src/twodnavierstokes.jl
+++ b/src/twodnavierstokes.jl
@@ -6,10 +6,13 @@ export
   updatevars!,
 
   energy,
+  energy_dissipation,
+  energy_work,
+  energy_drag,
   enstrophy,
-  dissipation,
-  work,
-  drag
+  enstrophy_dissipation,
+  enstrophy_work,
+  enstrophy_drag
 
   using
     CUDA,
@@ -46,13 +49,13 @@ function Problem(dev::Device=CPU();
            T = Float64)
 
   grid = TwoDGrid(dev, nx, Lx, ny, Ly; T=T)
-  
+
   params = Params{T}(ν, nν, μ, nμ, calcF)
-  
+
   vars = calcF == 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
 
@@ -128,7 +131,7 @@ end
 """
     ForcedVars(dev, grid)
 
-Returns the vars for forced two-dimensional turbulence on device dev and with 
+Returns the vars for forced two-dimensional turbulence on device dev and with
 `grid`.
 """
 function ForcedVars(dev::Dev, grid::AbstractGrid) where Dev
@@ -262,11 +265,11 @@ solution `sol`.
 end
 
 """
-    dissipation(prob)
+    energy_dissipation(prob)
 
 Returns the domain-averaged dissipation rate. nν must be >= 1.
 """
-@inline function dissipation(prob)
+@inline function energy_dissipation(prob)
   sol, vars, params, grid = prob.sol, prob.vars, prob.params, prob.grid
   @. vars.uh = grid.Krsq^(params.nν - 1) * abs2(sol)
   CUDA.@allowscalar vars.uh[1, 1] = 0
@@ -274,34 +277,77 @@ Returns the domain-averaged dissipation rate. nν must be >= 1.
 end
 
 """
-    work(prob)
-    work(sol, v, grid)
+    enstrophy_dissipation(prob)
+
+Returns the domain-averaged dissipation rate of enstrophy. nν must be >= 1.
+"""
+@inline function enstrophy_dissipation(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
+
+"""
+    energy_work(prob)
+    energy_work(sol, v, grid)
 
 Returns the domain-averaged rate of work of energy by the forcing Fh.
 """
-@inline function work(sol, vars::ForcedVars, grid)
+@inline function energy_work(sol, vars::ForcedVars, grid)
   @. vars.uh = grid.invKrsq * sol * conj(vars.Fh)
   return 1 / (grid.Lx * grid.Ly) * parsevalsum(vars.uh, grid)
 end
 
-@inline function work(sol, vars::StochasticForcedVars, grid)
+@inline function energy_work(sol, vars::StochasticForcedVars, grid)
   @. vars.uh = grid.invKrsq * (vars.prevsol + sol) / 2 * conj(vars.Fh) # Stratonovich
   # @. vars.uh = grid.invKrsq * vars.prevsol * conj(vars.Fh)           # Ito
   return 1 / (grid.Lx * grid.Ly) * parsevalsum(vars.uh, grid)
 end
 
-@inline work(prob) = work(prob.sol, prob.vars, prob.grid)
+@inline energy_work(prob) = energy_work(prob.sol, prob.vars, prob.grid)
+
+"""
+    enstrophy_work(prob)
+    enstrophy_work(sol, v, grid)
+
+Returns the domain-averaged rate of work of enstrophy by the forcing Fh.
+"""
+@inline function enstrophy_work(sol, vars::ForcedVars, grid)
+  @. vars.uh = sol * conj(vars.Fh)
+  return 1 / (grid.Lx * grid.Ly) * parsevalsum(vars.uh, grid)
+end
+
+@inline function enstrophy_work(sol, vars::StochasticForcedVars, grid)
+  @. vars.uh = (vars.prevsol + sol) / 2 * conj(vars.Fh) # Stratonovich
+  # @. vars.uh = grid.invKrsq * vars.prevsol * conj(vars.Fh)           # Ito
+  return 1 / (grid.Lx * grid.Ly) * parsevalsum(vars.uh, grid)
+end
+
+@inline enstrophy_work(prob) = enstrophy_work(prob.sol, prob.vars, prob.grid)
 
 """
-    drag(prob)
+    energy_drag(prob)
 
 Returns the extraction of domain-averaged energy by drag/hypodrag μ.
 """
-@inline function drag(prob)
+@inline function energy_drag(prob)
   sol, vars, params, grid = prob.sol, prob.vars, prob.params, prob.grid
   @. vars.uh = grid.Krsq^(params.nμ - 1) * abs2(sol)
   CUDA.@allowscalar vars.uh[1, 1] = 0
   return params.μ / (grid.Lx * grid.Ly) * parsevalsum(vars.uh, grid)
 end
 
+"""
+    enstrophy_drag(prob)
+
+Returns the extraction of domain-averaged enstrophy by drag/hypodrag μ.
+"""
+@inline function enstrophy_drag(prob)
+  sol, vars, params, grid = prob.sol, prob.vars, prob.params, prob.grid
+  @. vars.uh = grid.Krsq^params.nμ * abs2(sol)
+  CUDA.@allowscalar vars.uh[1, 1] = 0
+  return params.μ / (grid.Lx * grid.Ly) * parsevalsum(vars.uh, grid)
+end
+
 end # module
diff --git a/test/runtests.jl b/test/runtests.jl
index 14962921..a818e11b 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -48,8 +48,10 @@ for dev in devices
 
     @test test_twodnavierstokes_advection(0.0005, "ForwardEuler", dev)
     @test test_twodnavierstokes_lambdipole(256, 1e-3, dev)
-    @test test_twodnavierstokes_stochasticforcingbudgets(dev)
-    @test test_twodnavierstokes_deterministicforcingbudgets(dev)
+    @test test_twodnavierstokes_deterministicforcing_energybudget(dev)
+    @test test_twodnavierstokes_stochasticforcing_energybudget(dev)
+    @test test_twodnavierstokes_deterministicforcing_enstrophybudget(dev)
+    @test test_twodnavierstokes_stochasticforcing_enstrophybudget(dev)
     @test test_twodnavierstokes_energyenstrophy(dev)
     @test test_twodnavierstokes_problemtype(dev, Float32)
     @test TwoDNavierStokes.nothingfunction() == nothing
diff --git a/test/test_twodnavierstokes.jl b/test/test_twodnavierstokes.jl
index 7a501f0d..52aaa4de 100644
--- a/test/test_twodnavierstokes.jl
+++ b/test/test_twodnavierstokes.jl
@@ -20,7 +20,7 @@ function test_twodnavierstokes_lambdipole(n, dt, dev::Device=CPU(); L=2π, Ue=1,
   isapprox(Ue, mean(Ue_m[2:end]), rtol=rtol_lambdipole)
 end
 
-function test_twodnavierstokes_stochasticforcingbudgets(dev::Device=CPU(); n=256, L=2π, dt=0.005, ν=1e-7, nν=2, μ=1e-1, nμ=0, tf=0.1/μ)
+function test_twodnavierstokes_stochasticforcing_energybudget(dev::Device=CPU(); n=256, L=2π, dt=0.005, ν=1e-7, nν=2, μ=1e-1, nμ=0, tf=0.1/μ)
   nt = round(Int, tf/dt)
 
   # Forcing parameters
@@ -56,20 +56,19 @@ function test_twodnavierstokes_stochasticforcingbudgets(dev::Device=CPU(); n=256
 
   TwoDNavierStokes.set_zeta!(prob, 0*x)
   E = Diagnostic(TwoDNavierStokes.energy,      prob, nsteps=nt)
-  D = Diagnostic(TwoDNavierStokes.dissipation, prob, nsteps=nt)
-  R = Diagnostic(TwoDNavierStokes.drag,        prob, nsteps=nt)
-  W = Diagnostic(TwoDNavierStokes.work,        prob, nsteps=nt)
+  D = Diagnostic(TwoDNavierStokes.energy_dissipation, prob, nsteps=nt)
+  R = Diagnostic(TwoDNavierStokes.energy_drag,        prob, nsteps=nt)
+  W = Diagnostic(TwoDNavierStokes.energy_work,        prob, nsteps=nt)
   diags = [E, D, W, R]
 
   stepforward!(prob, diags, nt)
   TwoDNavierStokes.updatevars!(prob)
 
   E, D, W, R = diags
-  t = round(μ*cl.t, digits=2)
 
   i₀ = 1
   dEdt = (E[(i₀+1):E.i] - E[i₀:E.i-1])/cl.dt
-  ii  = (i₀):E.i-1
+  ii = (i₀):E.i-1
   ii2 = (i₀+1):E.i
 
   # dEdt = W - D - R?
@@ -79,17 +78,79 @@ function test_twodnavierstokes_stochasticforcingbudgets(dev::Device=CPU(); n=256
   total = W[ii2] - D[ii] - R[ii]        # Stratonovich
 
   residual = dEdt - total
-  isapprox(mean(abs.(residual)), 0, atol=1e-4)
+
+  return isapprox(mean(abs.(residual)), 0, atol=1e-4)
 end
 
+function test_twodnavierstokes_stochasticforcing_enstrophybudget(dev::Device=CPU(); n=256, L=2π, dt=0.005, ν=1e-7, nν=2, μ=1e-1, nμ=0, tf=0.1/μ)
+  nt = round(Int, tf/dt)
+
+  # Forcing parameters
+  kf, dkf = 12.0, 2.0
+  ε = 0.1
+
+  gr  = TwoDGrid(dev, n, L)
+  x, y = gridpoints(gr)
+
+  Kr = ArrayType(dev)([CUDA.@allowscalar gr.kr[i] for i=1:gr.nkr, j=1:gr.nl])
 
-function test_twodnavierstokes_deterministicforcingbudgets(dev::Device=CPU(); n=256, dt=0.01, L=2π, ν=1e-7, nν=2, μ=1e-1, nμ=0)
+  forcingcovariancespectrum = ArrayType(dev)(zero(gr.Krsq))
+  @. forcingcovariancespectrum = exp(-(sqrt(gr.Krsq) - kf)^2 / (2 * dkf^2))
+  CUDA.@allowscalar @. forcingcovariancespectrum[gr.Krsq .< 2^2] = 0
+  CUDA.@allowscalar @. forcingcovariancespectrum[gr.Krsq .> 20^2] = 0
+  CUDA.@allowscalar @. forcingcovariancespectrum[Kr .< 2π/L] = 0
+  ε0 = parsevalsum(forcingcovariancespectrum .* gr.invKrsq / 2, gr) / (gr.Lx * gr.Ly)
+  forcingcovariancespectrum .= ε / ε0 * forcingcovariancespectrum
+
+  Random.seed!(1234)
+
+  function calcF!(Fh, sol, t, cl, v, p, g)
+    eta = ArrayType(dev)(exp.(2π * im * rand(Float64, size(sol))) / sqrt(cl.dt))
+    CUDA.@allowscalar eta[1, 1] = 0.0
+    @. Fh = eta * sqrt(forcingcovariancespectrum)
+    nothing
+  end
+
+  prob = TwoDNavierStokes.Problem(dev; nx=n, Lx=L, ν=ν, nν=nν, μ=μ, nμ=nμ, dt=dt,
+   stepper="RK4", calcF=calcF!, stochastic=true)
+
+  sol, cl, v, p, g = prob.sol, prob.clock, prob.vars, prob.params, prob.grid;
+
+  TwoDNavierStokes.set_zeta!(prob, 0*x)
+  Z = Diagnostic(TwoDNavierStokes.enstrophy,      prob, nsteps=nt)
+  D = Diagnostic(TwoDNavierStokes.enstrophy_dissipation, prob, nsteps=nt)
+  R = Diagnostic(TwoDNavierStokes.enstrophy_drag,        prob, nsteps=nt)
+  W = Diagnostic(TwoDNavierStokes.enstrophy_work,        prob, nsteps=nt)
+  diags = [Z, D, W, R]
+
+  stepforward!(prob, diags, nt)
+  TwoDNavierStokes.updatevars!(prob)
+
+  Z, D, W, R = diags
+
+  i₀ = 1
+  dZdt = (Z[(i₀+1):Z.i] - Z[i₀:Z.i-1])/cl.dt
+  ii = (i₀):Z.i-1
+  ii2 = (i₀+1):Z.i
+
+  # dZdt = W - D - R?
+  # If the Ito interpretation was used for the work
+  # then we need to add the drift term
+  # total = W[ii2]+ε - D[ii] - R[ii]      # Ito
+  total = W[ii2] - D[ii] - R[ii]        # Stratonovich
+
+  residual = dZdt - total
+
+  return isapprox(mean(abs.(residual)), 0, atol=(kf^2)*1e-4)
+end
+
+function test_twodnavierstokes_deterministicforcing_energybudget(dev::Device=CPU(); n=256, dt=0.01, L=2π, ν=1e-7, nν=2, μ=1e-1, nμ=0)
   n, L  = 256, 2π
   ν, nν = 1e-7, 2
   μ, nμ = 1e-1, 0
   dt, tf = 0.005, 0.1/μ
   nt = round(Int, tf/dt)
-  
+
   gr  = TwoDGrid(dev, n, L)
   x, y = gridpoints(gr)
 
@@ -109,9 +170,9 @@ function test_twodnavierstokes_deterministicforcingbudgets(dev::Device=CPU(); n=
   TwoDNavierStokes.set_zeta!(prob, 0*x)
 
   E = Diagnostic(TwoDNavierStokes.energy,      prob, nsteps=nt)
-  D = Diagnostic(TwoDNavierStokes.dissipation, prob, nsteps=nt)
-  R = Diagnostic(TwoDNavierStokes.drag,        prob, nsteps=nt)
-  W = Diagnostic(TwoDNavierStokes.work,        prob, nsteps=nt)
+  D = Diagnostic(TwoDNavierStokes.energy_dissipation, prob, nsteps=nt)
+  R = Diagnostic(TwoDNavierStokes.energy_drag,        prob, nsteps=nt)
+  W = Diagnostic(TwoDNavierStokes.energy_work,        prob, nsteps=nt)
   diags = [E, D, W, R]
 
   # Step forward
@@ -119,18 +180,66 @@ function test_twodnavierstokes_deterministicforcingbudgets(dev::Device=CPU(); n=
   TwoDNavierStokes.updatevars!(prob)
 
   E, D, W, R = diags
-  t = round(μ*cl.t, digits=2)
 
   i₀ = 1
   dEdt = (E[(i₀+1):E.i] - E[i₀:E.i-1])/cl.dt
-  ii  = (i₀):E.i-1
+  ii = (i₀):E.i-1
   ii2 = (i₀+1):E.i
 
   # dEdt = W - D - R?
   total = W[ii2] - D[ii] - R[ii]
   residual = dEdt - total
 
-  isapprox(mean(abs.(residual)), 0, atol=1e-8)
+  return isapprox(mean(abs.(residual)), 0, atol=1e-8)
+end
+
+function test_twodnavierstokes_deterministicforcing_enstrophybudget(dev::Device=CPU(); n=256, dt=0.01, L=2π, ν=1e-7, nν=2, μ=1e-1, nμ=0)
+  n, L  = 256, 2π
+  ν, nν = 1e-7, 2
+  μ, nμ = 1e-1, 0
+  dt, tf = 0.005, 0.1/μ
+  nt = round(Int, tf/dt)
+
+  gr  = TwoDGrid(dev, n, L)
+  x, y = gridpoints(gr)
+
+  # Forcing = 0.01cos(4x)cos(5y)cos(2t)
+  f = @. 0.01cos(4x) * cos(5y)
+  fh = rfft(f)
+  function calcF!(Fh, sol, t, cl, v, p, g::AbstractGrid{T, A}) where {T, A}
+    Fh = fh * cos(2t)
+    nothing
+  end
+
+  prob = TwoDNavierStokes.Problem(dev; nx=n, Lx=L, ν=ν, nν=nν, μ=μ, nμ=nμ, dt=dt,
+   stepper="RK4", calcF=calcF!, stochastic=false)
+
+  sol, cl, v, p, g = prob.sol, prob.clock, prob.vars, prob.params, prob.grid
+
+  TwoDNavierStokes.set_zeta!(prob, 0*x)
+
+  Z = Diagnostic(TwoDNavierStokes.enstrophy,      prob, nsteps=nt)
+  D = Diagnostic(TwoDNavierStokes.enstrophy_dissipation, prob, nsteps=nt)
+  R = Diagnostic(TwoDNavierStokes.enstrophy_drag,        prob, nsteps=nt)
+  W = Diagnostic(TwoDNavierStokes.enstrophy_work,        prob, nsteps=nt)
+  diags = [Z, D, W, R]
+
+  # Step forward
+  stepforward!(prob, diags, nt)
+  TwoDNavierStokes.updatevars!(prob)
+
+  Z, D, W, R = diags
+
+  i₀ = 1
+  dZdt = (Z[(i₀+1):Z.i] - Z[i₀:Z.i-1])/cl.dt
+  ii = (i₀):Z.i-1
+  ii2 = (i₀+1):Z.i
+
+  # dZdt = W - D - R?
+  total = W[ii2] - D[ii] - R[ii]
+  residual = dZdt - total
+
+  return isapprox(mean(abs.(residual)), 0, atol=1e-8)
 end
 
 """
@@ -194,7 +303,7 @@ function test_twodnavierstokes_energyenstrophy(dev::Device=CPU())
   enstrophy_calc = 2701/162
 
   prob = TwoDNavierStokes.Problem(dev; nx=nx, Lx=Lx, ny=ny, Ly=Ly, stepper="ForwardEuler")
-  
+
   sol, cl, v, p, g = prob.sol, prob.clock, prob.vars, prob.params, prob.grid;
 
   TwoDNavierStokes.set_zeta!(prob, ζ₀)
@@ -202,7 +311,7 @@ function test_twodnavierstokes_energyenstrophy(dev::Device=CPU())
 
   energyζ₀ = TwoDNavierStokes.energy(prob)
   enstrophyζ₀ = TwoDNavierStokes.enstrophy(prob)
-  
+
   params = TwoDNavierStokes.Params(p.ν, p.nν)
 
   (isapprox(energyζ₀, energy_calc, rtol=rtol_twodnavierstokes) &&
@@ -214,6 +323,6 @@ function test_twodnavierstokes_problemtype(dev, T)
   prob = TwoDNavierStokes.Problem(dev; T=T)
 
   A = ArrayType(dev)
-  
+
   (typeof(prob.sol)<:A{Complex{T},2} && typeof(prob.grid.Lx)==T && eltype(prob.grid.x)==T && typeof(prob.vars.u)<:A{T,2})
 end