Implement an implicit free surface solver in the NonhydrostaticModel

[glwagner]

Closes #3946

[glwagner]

So here's how the algorithm changes with an implicit free surface (which is all I'd like to attempt for the time being):

1. Update velocities to obtain first predictor velocities
2. Step forward the free surface with an implicit solve
3. Update the velocities to obtain the "second" predictor velocities
4. Compute the pressure correction

I think the simplest way to implement this is therefore to add the implicit free surface solve as a precursor to the pressure correction.

[glwagner]

Okay, this script:

using Oceananigans
using Oceananigans.Models.HydrostaticFreeSurfaceModels: ImplicitFreeSurface
using GLMakie

grid = RectilinearGrid(size=(128, 32), halo=(4, 4), x=(-5, 5), z=(0, 1), topology=(Bounded, Flat, Bounded))

mountain(x) = (x - 3) / 2
grid = ImmersedBoundaryGrid(grid, GridFittedBottom(mountain))

Fu(x, z, t) = sin(t)
free_surface = ImplicitFreeSurface(gravitational_acceleration=10)
model = NonhydrostaticModel(; grid, free_surface, advection=WENO(order=5), forcing=(; u=Fu))

simulation = Simulation(model, Δt=0.1, stop_time=20*2π)
conjure_time_step_wizard!(simulation, cfl=0.7)

progress(sim) = @info string(iteration(sim), ": ", time(sim))
add_callback!(simulation, progress, IterationInterval(100))

ow = JLD2OutputWriter(model, merge(model.velocities, (; η=model.free_surface.η)),
                      filename = "nonhydrostatic_internal_tide.jld2",
                      schedule = TimeInterval(0.1),
                      overwrite_existing = true)

simulation.output_writers[:jld2] = ow

run!(simulation)

fig = Figure()

axη = Axis(fig[1, 1], xlabel="x", ylabel="Free surface \n displacement")
axw = Axis(fig[2, 1], xlabel="x", ylabel="Surface vertical velocity")
axu = Axis(fig[3, 1], xlabel="x", ylabel="z")

ut = FieldTimeSeries("nonhydrostatic_internal_tide.jld2", "u")
wt = FieldTimeSeries("nonhydrostatic_internal_tide.jld2", "w")
ηt = FieldTimeSeries("nonhydrostatic_internal_tide.jld2", "η")
Nt = length(wt)

slider = Slider(fig[4, 1], range=1:Nt, startvalue=1)
n = slider.value
Nz = size(ut.grid, 3)

u = @lift ut[$n]
η = @lift interior(ηt[$n], :, 1, 1)
w = @lift interior(wt[$n], :, 1, Nz+1)
x = xnodes(wt)

ulim = maximum(abs, ut) * 3/4

lines!(axη, x, η)
lines!(axw, x, w)
heatmap!(axu, u)

ylims!(axη, -0.1, 0.1)
ylims!(axw, -0.01, 0.01)

record(fig, "nonhydrostatic_internal_tide.mp4", 1:Nt) do nn
    @info "Drawing frame $nn of $Nt..."
    n[] = nn
end

produces this movie

nonhydrostatic_internal_tide.mp4

some weird grid artifacts in there but maybe higher resolution will help with that.

[glwagner]

Here's the script:

https://github.com/CliMA/Oceananigans.jl/blob/glw/nh-free-surface/validation/implicit_free_surface/nonhydrostatic_internal_tide.jl

[glwagner]

@shriyafruitwala let me know if this code works for the problem you are interested in.

The implementation is fairly clean, but there are a few things we could consider before merging, like tests, some validation in the constructor.

@simone-silvestri I feel this code exposes some messiness with the peformance optimization stuff regarding kernel parameters. Please check it over to make sure what I did will work and add any tests that may be missing...

[simone-silvestri]

I am unsure of the algorithm, but wouldn't this replace the w-velocity that should be prognostic?

[glwagner]

For sure. That's what is written in MITgcm docs. @jm-c can you confirm?

[simone-silvestri]

@simone-silvestri I feel this code exposes some messiness with the peformance optimization stuff regarding kernel parameters. Please check it over to make sure what I did will work and add any tests that may be missing...

it seems that you are missing required_halo_size_x and required_halo_size_y in the new file src/Models/buffer_tendency_kernel_parameters.jl. I can add it

[glwagner]

@shriyafruitwala to set η to some value try writing

η = model.free_surface.η
set!(η, value)

[glwagner]

It's been discussed elsewhere (and there are notes lying around -- hopefully they can be shared here) that the implementation in this PR is incorrect. In particular, the pressure satisfies a Robin boundary condition (not a Neumann condition) at the surface; implementing a Robin boundary condition requires changing the pressure solver. It's possible that the current implementation represents some kind of approximation to the true problem which may be revealed by comparing to an exact solution.

In thinking about a different problem, I realized that --- I think --- a Robin boundary condition can be implemented using BatchedTridiagonalSolver fairly easily. This requires modifying FourierTridiagonalPoissonSolver, and also the pressure correction algorithm (I believe there will also be a change in the contribution to the Poisson equation RHS at the top).

If true then it should be relatively straightforward to develop a direct solver for the flat bottom case. Hopefully, this would help in developing a CG solver for the case with bathymetry.

I'd also like to point out there seem to be a few issues with the CG solver currently, as documented in #4007 and #3848.

cc @joernc @shriyafruitwala

[shriyafruitwala]

Hey everyone! I have set up a test of the nonhydrostatic free surface code of a deep water surface gravity wave initialized with a Gaussian bump. Here, G = 0, so per the MITgcm docs, the free surface equation (2.52) essentially becomes a diffusion equation, with K = gH \Delta t. We see exactly this in the test case, where the velocities are essentially zero, and the evolution of the free surface is diffusive.
We would also expect that the time scale of evolution, L^2/K, matches the evolution of the wave in the animation. \Delta t is 0.1, so K should be about 50 m^2/s. L is the width of the bump (10m), so the diffusivity should be about 2 s. The time scale of the animation is also 0.1, so with a frame rate of 24 fps, 1 second of time in the animation would correspond to 2.4 seconds in the simulation, which is about what we see. I have also included the hydrostatic free surface version to compare.

nonhydrostatic_deepwater_test.jld2.mp4

hydrostatic_deepwater_test.jld2.mp4

[glwagner]

Great work. Can you include the code that you used to run the simulation? (The code that generates the animation may be helpful too if anyone would like to reproduce your work.)

The hydrostatic model can also be used with free_surface = ExplicitFreeSurface() which may be useful for comparing results.

[shriyafruitwala]

Sure! Here is the code:

using Oceananigans
using Oceananigans.Models.HydrostaticFreeSurfaceModels: ImplicitFreeSurface
using GLMakie
using Oceananigans.Units

Nx, Nz = 50, 50
const H = 50meters
const L = 50meters
const g = 10
k = 2
ω² = g*abs(k)
ω = sqrt(ω²)
coriolis = FPlane(latitude=28)
f=coriolis.f

grid = RectilinearGrid(size = (Nx, Nz),
                            x = (-25meters, 25meters),
                            z = (-H, 0),
                            halo = (4,4),
                            topology = (Periodic, Flat, Bounded))

free_surface = ImplicitFreeSurface(gravitational_acceleration=10)
model = NonhydrostaticModel(; grid, free_surface, coriolis, advection=WENO(order=5))
#model = HydrostaticFreeSurfaceModel(; grid, free_surface, coriolis, momentum_advection=WENO(order=5))

#initialize with gaussian bump surface
η = model.free_surface.η
bump(x) = exp(-(x^2)/(2 * (5meters)^2))
x_vals = LinRange(-L/2, L/2, Nx)
η₀_vals = bump.(x_vals)
η₀ = reshape(η₀_vals, (Nx,1,1))
set!(η, η₀)

#sets up simulation run
simulation = Simulation(model, Δt=0.1, stop_time=50)

progress(sim) = @info string(iteration(sim), ": ", time(sim))
add_callback!(simulation, progress, IterationInterval(100))
output_file = "hydrostatic_deepwater_test_k2_implicit.jld2"
ow = JLD2OutputWriter(model, merge(model.velocities, (; η=model.free_surface.η)),
                      filename = output_file,  
                      schedule = TimeInterval(0.1),
                      overwrite_existing = true)

simulation.output_writers[:jld2] = ow

run!(simulation)

#produces animation
fig = Figure()

axη = Axis(fig[1, 1], xlabel="x", ylabel="η", width=400, height=75, title="sea surface height")
axw = Axis(fig[2, 1], xlabel="x", ylabel="z", width=400, height=75, title = "w velocity")
axu = Axis(fig[3, 1], xlabel="x", ylabel="z", width=400, height=75, title = "u velocity")
                      
ut = FieldTimeSeries(output_file, "u") 
wt = FieldTimeSeries(output_file, "w") 
ηt = FieldTimeSeries(output_file, "η") 
Nt = length(wt)

n = Observable(1)

u = @lift ut[$n]
η = @lift interior(ηt[$n], :, 1, 1)
w = @lift wt[$n]
x = xnodes(wt)

ulim = maximum(abs, ut)
wlim = maximum(abs, wt)

lines!(axη, x, η)
hm_w = heatmap!(axw, w, colorrange=(-wlim,wlim))
Colorbar(fig[2, 2], hm_w, label = "m s⁻¹")
hm_u = heatmap!(axu, u, colorrange=(-ulim,ulim))
Colorbar(fig[3, 2], hm_u, label = "m s⁻¹")

ylims!(axη, -1, 1)

record(fig, output_file * "title.mp4", 1:Nt) do nn
    @info "Drawing frame $nn of $Nt..."
    n[] = nn
end

Using free_surface = ExplicitFreeSurface(gravitational_acceleration=10) yields the following (we would expect something less dissipative, which is indeed what we see):

hydrostatic_deepwater_test_k2_explicit.mp4

[shriyafruitwala]

I also wanted to point out that for deep water waves, the evolution of the free surface should look something like this (in contrast to the pure diffusion we see in the Oceananigans nonhydrostatic test I posted earlier):

deepwater_amplitude.mp4

[glwagner]

That's great to have a ground truth. Do we have a plan to fix the nonhydrostatic solver?

[shriyafruitwala]

That's great to have a ground truth. Do we have a plan to fix the nonhydrostatic solver?

I think that as we discussed, we want the pressure to satisfy a Robin boundary condition, which from your earlier post seemed to be doable using the BatchedTridiagonalSolver. Besides this, there is no plan at the moment as I am not familiar with the model codebase and would likely be inefficient at going about it myself. Would you be open to implementing this? Happy to discuss more and really appreciate your help!

[glwagner]

I am happy to offer guidance but cannot take the lead on implementing it. I agree I would be faster, but this argument breaks down quickly --- with that argument, I should implement everything! That said, I do not think you will need to change very much. The solvers are in place and I believe the implementation is a matter of rearranging things and perhaps a few critical lines here and there. Honestly I am not exactly sure what needs to be changed, and figuring out precisely what code needs to change is one of the major pieces you can take the lead on.

I think we can start by rehashing the algorithm that we would like to implement here. I can't remember the specifics, and we need our plan to be documented in this PR.

[joernc]

I think a good place to start would be to implement the Robin boundary condition in the Fourier pressure solver, so we can simulate deep-water waves over a flat bottom. I'm attaching notes on the algorithm, which show two versions. The first version solves for the pressure field in one go, whereas the second splits the pressure field into hydrostatic and nonhydrostatic parts, following the MITgcm practice. The first version is much simpler and may be a good place to start. The main argument in favor of the second version is that the 3D solve might converge more quickly in nearly hydrostatic conditions, though that should be tested, I suppose.

So, I would propose we do the Fourier solver first with the first version of the algorithm. That requires a Robin BC but not much else, as far as I can tell. Once we have this working, we can work on the CG solver to allow for topography, which is what Shriya is after in the end. @glwagner, does that sounds reasonable?

[glwagner]

I think that will work. I will just clarify, I am not aware of a pure Fourier algorithm that will work for the Robin BC. However, I believe that the Robin BC can be implemented in the FourierTridiagonalPoissonSolver:

https://github.com/CliMA/Oceananigans.jl/blob/main/src/Solvers/fourier_tridiagonal_poisson_solver.jl

which uses Fourier transforms in the horizontal and a tridiagonal solve in the vertical. I believe the modification requires deducing the changes needed here:

Oceananigans.jl/src/Solvers/fourier_tridiagonal_poisson_solver.jl

Lines 45 to 46 in 859c36f

	# Using a homogeneous Neumann (zero Gradient) boundary condition:
	@inbounds D[i, j, 1] = -1 / Δzᵃᵃᶠ(i, j, 2, grid) - Δzᵃᵃᶜ(i, j, 1, grid) * (λx[i] + λy[j])

The batched tridiagonal solver is implemented here for reference:

https://github.com/CliMA/Oceananigans.jl/blob/main/src/Solvers/batched_tridiagonal_solver.jl

We'll also have to design an interface for specifying which boundary condition we'd like to use when building FourierTridiagonalPoissonSolver (so that we can continue to use it for the rigid lid case as well), but we can worry about that part once something is working.