Add optional vertical Coriolis term to hydrostatic pressure integral

examples/linear_shear.jl

@@ -0,0 +1,98 @@
+using Oceananigans
+using Printf
+using CairoMakie
+using NCDatasets
+using Statistics
+
+grid = RectilinearGrid(CPU(),
+                      size=(32, 32), halo=(7, 7),
+                      y = (0, 1),
+                      z = (0, 1),
+                      topology=(Flat, Bounded, Bounded))
+
+
+const N² = 1.0
+const fz = 0.1
+const fy = 0.1
+const β  = 0.0
+const γ  = 0.0
+const radius = 1.0
+
+const shear_y = 1e-2
+const shear_z = 1e-2
+
+B(y, z) = -(shear_y*fz + shear_z*fy)*y + N²*z
+U(y, z) =  (shear_y*y  + shear_z*z)
+
+coriolis = NonTraditionalBetaPlane(fz=fz, fy=fy, β=β, γ=γ, radius=radius)
+model = HydrostaticFreeSurfaceModel(; grid,
+                              coriolis = coriolis,
+                              buoyancy = BuoyancyTracer(),
+                              momentum_advection = WENO(),
+                              tracer_advection = WENO(),
+                              tracers = (:b,))
+set!(model, u = U, b = B)
+
+xu, yu, zu = nodes(model.velocities.u)
+xb, yb, zb = nodes(model.tracers.b)
+
+function progress(sim)
+    umax = maximum(abs, sim.model.velocities.u)
+    bmax = maximum(sim.model.tracers.b)
+    @info @sprintf("Iter: %d, time: %.2e, max|u|: %.2e. max b: %.2e",
+       iteration(sim), time(sim), umax, bmax)
+
+    return nothing
+end
+
+simulation = Simulation(model; Δt=1e-3, stop_time=10.0)
+simulation.callbacks[:p] = Callback(progress, IterationInterval(100))
+
+u, v, w = model.velocities
+b = model.tracers.b
+outputs = (; u, b)
+
+simulation.output_writers[:fields] = NetCDFWriter(model, outputs;
+                                                        filename = "linear_shear.nc",
+                                                        schedule = TimeInterval(0.1),
+                                                        overwrite_existing = true)
+
+run!(simulation)
+
+ds = NCDataset(simulation.output_writers[:fields].filepath, "r")
+
+fig = Figure(size = (800, 600))
+
+axis_kwargs = (xlabel = "y",
+               ylabel = "z",
+               limits = ((0, grid.Lx), (0, grid.Lz)),
+               )
+
+ax_u = Axis(fig[2, 1]; title = "meridional velocity", axis_kwargs...)
+ax_b = Axis(fig[3, 1]; title = "buoyancy - N²*z", axis_kwargs...)
+
+n = Observable(1)
+
+u = @lift ds["u"][:, :, $n]
+v = @lift ds["b"][:, :, $n] - N²*zb
+
+hm_u = heatmap!(ax_u, yu, zu, u, colormap = :balance)
+Colorbar(fig[2, 2], hm_u)
+
+hm_b = heatmap!(ax_b, yb, zb, b, colormap = :balance)
+Colorbar(fig[3, 2], hm_b)
+
+times = collect(ds["time"])
+title = @lift "t = " * string(prettytime(times[$n]))
+fig[1, :] = Label(fig, title, fontsize=20, tellwidth=false)
+
+fig
+
+frames = 1:length(times)
+
+record(fig, "linear_shear_nontraditional.mp4", frames, framerate=12) do i
+    n[] = i
+end
+
+close(ds)
+

src/Models/HydrostaticFreeSurfaceModels/update_hydrostatic_free_surface_model_state.jl

@@ -75,24 +75,24 @@ function compute_auxiliaries!(model::HydrostaticFreeSurfaceModel; w_parameters =
                                                                   p_parameters = p_kernel_parameters(model.grid),
                                                                   κ_parameters = :xyz)
 
-    grid        = model.grid
-    closure     = model.closure
-    tracers     = model.tracers
-    diffusivity = model.diffusivity_fields
-    buoyancy    = model.buoyancy
-
-    P    = model.pressure.pHY′
-    arch = architecture(grid)
-
     # Update the vertical velocity to comply with the barotropic correction step
+    grid = model.grid
     update_grid_vertical_velocity!(model, grid, model.vertical_coordinate)
 
     # Advance diagnostic quantities
     compute_w_from_continuity!(model; parameters = w_parameters)
-    update_hydrostatic_pressure!(P, arch, grid, buoyancy, tracers; parameters = p_parameters)
 
-    # Update closure diffusivities
-    compute_diffusivities!(diffusivity, closure, model; parameters = κ_parameters)
+    update_hydrostatic_pressure!(model.pressure.pHY′,
+                                 grid,
+                                 model.buoyancy,
+                                 model.tracers,
+                                 model.coriolis,
+                                 model.velocities;
+                                 parameters = p_parameters)
+
+    closure = model.closure
+    diffusivity_fields = model.diffusivity_fields
+    compute_diffusivities!(diffusivity_fields, closure, model; parameters = κ_parameters)
 
     return nothing
 end

src/Models/NonhydrostaticModels/update_hydrostatic_pressure.jl

@@ -3,39 +3,52 @@ using Oceananigans.ImmersedBoundaries: PartialCellBottom, ImmersedBoundaryGrid
 using Oceananigans.Grids: topology
 using Oceananigans.Grids: XFlatGrid, YFlatGrid
 
-"""
-Update the hydrostatic pressure perturbation pHY′. This is done by integrating
-the `buoyancy_perturbationᶜᶜᶜ` downwards:
-
-    `pHY′ = ∫ buoyancy_perturbationᶜᶜᶜ dz` from `z=0` down to `z=-Lz`
-"""
-@kernel function _update_hydrostatic_pressure!(pHY′, grid, buoyancy, C)
+@kernel function _update_hydrostatic_pressure!(pHY′, grid, buoyancy, tracers, coriolis, velocities)
     i, j = @index(Global, NTuple)
 
-    @inbounds pHY′[i, j, grid.Nz] = - z_dot_g_bᶜᶜᶠ(i, j, grid.Nz+1, grid, buoyancy, C) * Δzᶜᶜᶠ(i, j, grid.Nz+1, grid)
+    dpdz⁺ = z_dot_g_bᶜᶜᶠ(i, j, grid.Nz+1, grid, buoyancy, tracers) - z_f_cross_U(i, j, grid.Nz+1, grid, coriolis, velocities)
+    @inbounds pHY′[i, j, grid.Nz] = - dpdz⁺ * Δzᶜᶜᶠ(i, j, grid.Nz+1, grid)
 
     for k in grid.Nz-1 : -1 : 1
-        @inbounds pHY′[i, j, k] = pHY′[i, j, k+1] - z_dot_g_bᶜᶜᶠ(i, j, k+1, grid, buoyancy, C) * Δzᶜᶜᶠ(i, j, k+1, grid)
+        dpdz⁺ = z_dot_g_bᶜᶜᶠ(i, j, k+1, grid, buoyancy, tracers) - z_f_cross_U(i, j, k+1, grid, coriolis, velocities)
+
+        # Using dpdz = (p⁺ - pᵏ) / Δzᶜᶜᶠ
+        @inbounds pHY′[i, j, k] = pHY′[i, j, k+1] - dpdz⁺ * Δzᶜᶜᶠ(i, j, k+1, grid)
     end
 end
 
-update_hydrostatic_pressure!(model; kwargs...) = update_hydrostatic_pressure!(model.grid, model; kwargs...)
-update_hydrostatic_pressure!(::AbstractGrid{<:Any, <:Any, <:Any, <:Flat}, model; kwargs...) = nothing
-update_hydrostatic_pressure!(grid, model; kwargs...) =
-    update_hydrostatic_pressure!(model.pressures.pHY′, model.architecture, model.grid, model.buoyancy, model.tracers; kwargs...)
+"""
+    update_hydrostatic_pressure!(pHY′, grid, buoyancy, tracers, coriolis, velocities; parameters=:xy)
+
+Update the hydrostatic pressure perturbation pHY′. This is done by integrating
+the `buoyancy_perturbationᶜᶜᶜ` downwards:
+
+```math
+pHY′ = ∫ [ b - ẑ ⋅ (f × u) ] dz′
+```
+
+from ``z′=0`` down to ``z′=z``.
+"""
+function update_hydrostatic_pressure!(pHY′, grid, buoyancy, tracers,
+                                      coriolis=nothing, velocities=nothing; parameters=:xy)
+    arch = grid.architecture
+    launch!(arch, grid, parameters,
+            _update_hydrostatic_pressure!, pHY′, grid, buoyancy, tracers, coriolis, velocities)
+    return nothing
+end
+
+# Catch some special cases
+update_hydrostatic_pressure!(pHY′, ::AbstractGrid{<:Any, <:Any, <:Any, <:Flat}, args...; kwargs...) = nothing
 
 # Partial cell "algorithm"
 const PCB = PartialCellBottom
 const PCBIBG = ImmersedBoundaryGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:PCB}
 
-update_hydrostatic_pressure!(pHY′, arch, ibg::PCBIBG, buoyancy, tracers; parameters = p_kernel_parameters(ibg.underlying_grid)) =
-    update_hydrostatic_pressure!(pHY′, arch, ibg.underlying_grid, buoyancy, tracers; parameters)
-
-update_hydrostatic_pressure!(pHY′, arch, grid, buoyancy, tracers; parameters = p_kernel_parameters(grid)) =
-    launch!(arch, grid, parameters, _update_hydrostatic_pressure!, pHY′, grid, buoyancy, tracers)
+update_hydrostatic_pressure!(pHY′, ibg::PCBIBG, args...; kwargs...) =
+    update_hydrostatic_pressure!(pHY′, ibg.underlying_grid, args...; kwargs...)
 
-update_hydrostatic_pressure!(::Nothing, arch, grid, args...; kw...) = nothing
-update_hydrostatic_pressure!(::Nothing, arch, ::PCBIBG, args...; kw...) = nothing
+update_hydrostatic_pressure!(::Nothing, grid, args...; kw...) = nothing
+update_hydrostatic_pressure!(::Nothing, ::PCBIBG, args...; kw...) = nothing
 
 # extend p kernel to compute also the boundaries
 @inline function p_kernel_parameters(grid)

src/Models/NonhydrostaticModels/update_nonhydrostatic_model_state.jl

@@ -64,7 +64,12 @@ function compute_auxiliaries!(model::NonhydrostaticModel; p_parameters = tuple(p
 
     for (ppar, κpar) in zip(p_parameters, κ_parameters)
         compute_diffusivities!(diffusivity, closure, model; parameters = κpar)
-        update_hydrostatic_pressure!(model; parameters = ppar)
+        update_hydrostatic_pressure!(model.pressures.pHY′,
+                                     model.grid,
+                                     model.buoyancy,
+                                     model.tracers,
+                                     parameters = ppar)
     end
+
     return nothing
 end