Add background zonal flow U or U(y) in the SingleLayerQG module

[mncrowe]

[edited] Add a constant background zonal flow in the $U$ which can also be varying in $y$, i.e., $U(y)$.
When the background advection is constant (not varying with $y$) it's included in the linear term L of the equation; when it's varying with $y$ it's included in the nonlinear term N alongside with contributions from background PV gradient $\partial_y^2 U$.

[mncrowe]

I've added a draft of the full functionality with U = U(y) (periodic in y). One question: when evaluating CalcN! there's an assignment Uy = real.(ifft(im * grid.l .* fft(params.U))). This seems to be required as there's no 1D fft plan in the grid that could be used with mul!. My questions:

1. Is this the best way to do this derivative?
2. An identical value is calculated and assigned each timestep. As it's only in the y direction it shouldn't take too much memory, however, is this step (in particular, assigning GPU memory) likely to slow down the timestepping?
3. We could have Uy saved in params but this seems unnecessary. Do you agree?

[navidcy]

@mncrowe is it OK if I push few edits?

[navidcy]

@mncrowe this looks very good! Thanks!

I changed the tests for when U is not varying with y to utilize the Rossby wave test.

There are two things left to do:

1. remove the if params.U isa Number ... else ... statement from calc_advection!
   https://github.com/mncrowe/GeophysicalFlows.jl/blob/dd3c1ad4b6b3a3b2866525ad3ae993f3f3329270/src/singlelayerqg.jl#L352-L368

which reduces performance; and

2. add a test for a background flow U(y). This can be done via the test_1layerqg_nonlinearadvection test,

https://github.com/mncrowe/GeophysicalFlows.jl/blob/dd3c1ad4b6b3a3b2866525ad3ae993f3f3329270/test/test_singlelayerqg.jl#L257

similarly to how this was done for the MultilayerQG module; see

https://github.com/mncrowe/GeophysicalFlows.jl/blob/dd3c1ad4b6b3a3b2866525ad3ae993f3f3329270/test/test_multilayerqg.jl#L215-L305

[navidcy]

I think I see a minor problem. When U is just a number, don't we need a term - U ∂η/∂x? At the moment this term is not included.

I think I'll try to refactor the code to mimic the MultiLayerQG where the background PV gradients include the ∂η/∂x and ∂η/∂y terms and the ∂²U/∂y² (if U is not a constant).

[navidcy]

@mncrowe I think I did everything! Have a look and let me know if you find any issues running the code or if the code is not behaving as you expect? Otherwise I'm happy to merge this!

[navidcy]

I would just appreciate your look on the code. In principle I can approve the PR but it'd be nice if somebody else has a look because I made few changes.

[mncrowe]

Thanks for the very thorough reviewing and rewrite. Is there anything outstanding for me to do?

[navidcy]

Does it work as expected?

[navidcy]

A flow at rest advected by a background imposed U=0.5 over some topography.

using GeophysicalFlows, GLMakie, Printf

    dev = CPU()
      n = 256
stepper = "FilteredRK4"
     dt = 0.02
 nsteps = 4000
 nsubs  = 5
      L = 2π

σx, σy = 0.4, 0.8
topographicPV(x, y) = 3exp(-(x - 1)^2 / 2σx^2 - (y - 1)^2 / 2σy^2) - 2exp(- (x + 1)^2 / 2σx^2 - (y + 1)^2 / 2σy^2)

prob = SingleLayerQG.Problem(dev; nx=n, Lx=L, eta=topographicPV, U=0.5, dt, stepper, aliased_fraction=0)
sol, clock, vars, params, grid = prob.sol, prob.clock, prob.vars, prob.params, prob.grid

x,  y  = grid.x,  grid.y
Lx, Ly = grid.Lx, grid.Ly

q = Observable(Array(vars.q))

fig = Figure(size=(380, 300))

axis_kwargs = (xlabel = "x", ylabel = "y", aspect = 1,
               limits = ((-Lx/2, Lx/2), (-Ly/2, Ly/2)))

title_q = Observable("initial vorticity ∂v/∂x-∂u/∂y")
axq = Axis(fig[1, 1]; title = title_q, axis_kwargs...)

hm = heatmap!(axq, x, y, q, colormap = :balance, colorrange = (-5, 5))

Colorbar(fig[1, 2], hm)

contour!(axq, x, y, params.eta;
         levels = collect(0.5:1:2.5), linewidth = 2, color = (:black, 0.5))
contour!(axq, x, y, params.eta;
         levels = collect(-2.5:1:-0.5), linewidth = 2, color = (:grey, 0.7), linestyle = :dash)

title_q[] = "vorticity, t=" * @sprintf("%.2f", clock.t)

record(fig, "singlelayerqg_topography_U0.mp4", 0:round(Int, nsteps/nsubs), framerate = 40) do j
  q[] = irfft(sol, grid.nx)
  title_q[] = "vorticity, t=" * @sprintf("%.2f", clock.t)
  stepforward!(prob, nsubs)
end

singlelayerqg_topography_U0.mp4

[mncrowe]

Does it work as expected?

All my scripts are producing the same results as before, and match my earlier results using Dedalus.

[navidcy]

Nice! I’m merging then.