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]

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

ace06eb attempts to address 1.

I only wrote two methods for calcN_advection! that dispatch on the type of the param.Ul I haven't checked whether the terms written there for non-zero U are correct. But doing a test as point 2 suggests will ensure that the terms are correct!

[navidcy]

Btw, we also need to update docs/src/modules/singlelayerqg.md.

[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

https://github.com/user-attachments/assets/391a0e5e-54bb-4f1f-9a9a-4c2cc1c00430

[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.