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.

[shriyafruitwala]

I am trying to test the Robin BC in the FourierTridiagonalPoissonSolver (using the same deep water wave that I posted earlier), but I am having a bit of trouble with toggling the right solver. When setting the ImplicitFreeSurface, it seems to accept solver_method=:fourier_tridiagonal_poisson_solver as an argument with no issue, but when I then try and build the nonhydrostatic model (model = NonhydrostaticModel(; grid, free_surface, coriolis, advection=WENO(order=5))), I get the following error message:

MethodError: no method matching build_implicit_step_solver(::Val{:fourier_tridiagonal_poisson_solver}, ::RectilinearGrid{Float64, Periodic, Flat, Bounded, Float64, Float64, Float64, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, Nothing, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, CPU}, ::Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}, ::Float64)
The function build_implicit_step_solver exists, but no method is defined for this combination of argument types.

[glwagner]

I am trying to test the Robin BC in the FourierTridiagonalPoissonSolver (using the same deep water wave that I posted earlier), but I am having a bit of trouble with toggling the right solver. When setting the ImplicitFreeSurface, it seems to accept solver_method=:fourier_tridiagonal_poisson_solver as an argument with no issue, but when I then try and build the nonhydrostatic model (model = NonhydrostaticModel(; grid, free_surface, coriolis, advection=WENO(order=5))), I get the following error message:

MethodError: no method matching build_implicit_step_solver(::Val{:fourier_tridiagonal_poisson_solver}, ::RectilinearGrid{Float64, Periodic, Flat, Bounded, Float64, Float64, Float64, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, Nothing, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, CPU}, ::Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}, ::Float64) The function build_implicit_step_solver exists, but no method is defined for this combination of argument types.

The first thing is that some of the changes that I made need to be walked back. For example, these four lines need to be commented out or deleted:

Oceananigans.jl/src/Models/NonhydrostaticModels/pressure_correction.jl

Lines 12 to 16 in 42b425c

	if !isnothing(model.free_surface)
	step_free_surface!(model.free_surface, model, model.timestepper, Δt)
	# "First" barotropic pressure correction
	pressure_correct_velocities!(model, model.free_surface, Δt)
	end

The next step is to change this line:

Oceananigans.jl/src/Models/NonhydrostaticModels/NonhydrostaticModels.jl

Line 35 in dac06ae

nonhydrostatic_pressure_solver(arch, grid::XYZRegularRG) = FFTBasedPoissonSolver(grid)

to

nonhydrostatic_pressure_solver(arch, grid::XYZRegularRG) =
    FourierTridiagonalPoissonSolver(grid)

This way, no matter what grid you use with NonhydrostaticModel, you will get a FourierTridiagonalPoissonSolver.

Next, hard-code your boundary condition changes into the existing FourierTridiagonalPoissonSolver:

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

You can verify that the modifications are working as expected by directly calling solve! using the new pressure solver (eg model.pressure_solver). You can also try taking a single time step.

I think you will also need to implement a free surface update, which could go after these lines:

Oceananigans.jl/src/Models/NonhydrostaticModels/pressure_correction.jl

Line 23 in 42b425c

fill_halo_regions!(model.pressures.pNHS)

using the vertical component of the predictor velocity field (there are other choices but the crucial thing right now is just to validate the algorithm; afterwards we can shuffle code around.

If you make a pull request that is pointed at this branch, we will be able to see the changes and I can see your code and comment on specific lines which may help speed things up.

[shriyafruitwala]

Great, thanks! I still seem to be getting an error when I try to build the model that has to do with the grid (stretched_dimensions - I've pasted the error message below). Is there something that needs to be altered in rectilinear_grid.jl?

MethodError: stretched_dimensions(::RectilinearGrid{Float64, Periodic, Flat, Bounded, Float64, Float64, Float64, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, Nothing, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, CPU}) is ambiguous.

[glwagner]

Great, thanks! I still seem to be getting an error when I try to build the model that has to do with the grid (stretched_dimensions - I've pasted the error message below). Is there something that needs to be altered in rectilinear_grid.jl?

MethodError: stretched_dimensions(::RectilinearGrid{Float64, Periodic, Flat, Bounded, Float64, Float64, Float64, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, Nothing, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, CPU}) is ambiguous.

Can you paste the whole error so I can find the file?

[glwagner]

Great, thanks! I still seem to be getting an error when I try to build the model that has to do with the grid (stretched_dimensions - I've pasted the error message below). Is there something that needs to be altered in rectilinear_grid.jl?

MethodError: stretched_dimensions(::RectilinearGrid{Float64, Periodic, Flat, Bounded, Float64, Float64, Float64, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, Nothing, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, CPU}) is ambiguous.

As a stop gap, change the problem you are working on so that the z-direction is an array (which will be interpreted as "stretched"), for example if you are using

z = (-Lz, 0)

when you build the grid then change this to

dz = Lz / Nz
z = -Lz:dz:0

[shriyafruitwala]

Great, thanks! I still seem to be getting an error when I try to build the model that has to do with the grid (stretched_dimensions - I've pasted the error message below). Is there something that needs to be altered in rectilinear_grid.jl?
MethodError: stretched_dimensions(::RectilinearGrid{Float64, Periodic, Flat, Bounded, Float64, Float64, Float64, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, Nothing, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, CPU}) is ambiguous.

Can you paste the whole error so I can find the file?

yes, here it is!

MethodError: stretched_dimensions(::RectilinearGrid{Float64, Periodic, Flat, Bounded, Float64, Float64, Float64, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, Nothing, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, CPU}) is ambiguous.

Candidates:

stretched_dimensions(::RectilinearGrid{<:Any, <:Any, <:Any, <:Any, <:Number, <:Number})

@ Oceananigans.Grids c:\Users\shriy\Documents\Research\Oceananigans.jl\src\Grids\rectilinear_grid.jl:67

stretched_dimensions(::RectilinearGrid{<:Any, <:Any, <:Any, <:Any, <:Number, <:Any, <:Number})

@ Oceananigans.Grids c:\Users\shriy\Documents\Research\Oceananigans.jl\src\Grids\rectilinear_grid.jl:66

stretched_dimensions(::RectilinearGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:Number, <:Number})

@ Oceananigans.Grids c:\Users\shriy\Documents\Research\Oceananigans.jl\src\Grids\rectilinear_grid.jl:65

Possible fix, define

stretched_dimensions(::RectilinearGrid{<:Any, <:Any, <:Any, <:Any, <:Number, <:Number, <:Number})

Stacktrace:

[1] Oceananigans.Solvers.FourierTridiagonalPoissonSolver(grid::RectilinearGrid{Float64, Periodic, Flat, Bounded, Float64, Float64, Float64, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, Nothing, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, CPU}, planner_flag::UInt32)

@ Oceananigans.Solvers c:\Users\shriy\Documents\Research\Oceananigans.jl\src\Solvers\fourier_tridiagonal_poisson_solver.jl:64

[2] nonhydrostatic_pressure_solver(arch::CPU, grid::RectilinearGrid{Float64, Periodic, Flat, Bounded, Float64, Float64, Float64, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, Nothing, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, CPU})

@ Oceananigans.Models.NonhydrostaticModels c:\Users\shriy\Documents\Research\Oceananigans.jl\src\Models\NonhydrostaticModels\NonhydrostaticModels.jl:38

[3] nonhydrostatic_pressure_solver(grid::RectilinearGrid{Float64, Periodic, Flat, Bounded, Float64, Float64, Float64, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, Nothing, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, CPU})

@ Oceananigans.Models.NonhydrostaticModels c:\Users\shriy\Documents\Research\Oceananigans.jl\src\Models\NonhydrostaticModels\NonhydrostaticModels.jl:66

[4] NonhydrostaticModel(; grid::RectilinearGrid{Float64, Periodic, Flat, Bounded, Float64, Float64, Float64, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, Nothing, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, CPU}, clock::Clock{Float64, Float64}, advection::WENO{3, Float64, Nothing, Nothing, Nothing, Nothing, WENO{2, Float64, Nothing, Nothing, Nothing, Nothing, UpwindBiased{1, Float64, Nothing, Nothing, Nothing, Nothing, Centered{1, Float64, Nothing, Nothing, Nothing, Nothing}}, Centered{1, Float64, Nothing, Nothing, Nothing, Nothing}}, Centered{2, Float64, Nothing, Nothing, Nothing, Centered{1, Float64, Nothing, Nothing, Nothing, Nothing}}}, buoyancy::Nothing, coriolis::FPlane{Float64}, stokes_drift::Nothing, forcing::@NamedTuple{}, closure::Nothing, free_surface::ImplicitFreeSurface{Nothing, Int64, Nothing, Nothing, Symbol, @kwargs{}}, boundary_conditions::@NamedTuple{}, tracers::Tuple{}, timestepper::Symbol, background_fields::@NamedTuple{}, particles::Nothing, biogeochemistry::Nothing, velocities::Nothing, hydrostatic_pressure_anomaly::Oceananigans.Models.NonhydrostaticModels.DefaultHydrostaticPressureAnomaly, nonhydrostatic_pressure::Field{Center, Center, Center, Nothing, RectilinearGrid{Float64, Periodic, Flat, Bounded, Float64, Float64, Float64, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, Nothing, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, CPU}, Tuple{Colon, Colon, Colon}, OffsetArrays.OffsetArray{Float64, 3, Array{Float64, 3}}, Float64, FieldBoundaryConditions{BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, Nothing, Nothing, BoundaryCondition{Oceananigans.BoundaryConditions.Flux, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Flux, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Flux, Nothing}}, Nothing, Oceananigans.Fields.FieldBoundaryBuffers{Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}}, diffusivity_fields::Nothing, pressure_solver::Nothing, auxiliary_fields::@NamedTuple{})

@ Oceananigans.Models.NonhydrostaticModels c:\Users\shriy\Documents\Research\Oceananigans.jl\src\Models\NonhydrostaticModels\nonhydrostatic_model.jl:228

[5] top-level scope

@ c:\Users\shriy\Documents\Research\Oceananigans.jl\jl_notebook_cell_df34fa98e69747e1a8f8a730347b8e2f_W1sZmlsZQ==.jl:20

[glwagner]

Thank you!

PS when you paste code in a comment, surround it in triple backticks (```) to make it easy to read.

[glwagner]

I think we need to bring this branch up to date with main. Then I can put a PR to main to fix the issue.

Let me know if the "stop gap" works or not in the meantime.

[glwagner]

I can also try to give some more substantial help the week after next. But next week I will be on vacation. We also need to solve the merge conflicts on this PR, and we should merge all the changes into a single PR.

[glwagner]

@shriyafruitwala can you check if the issue is solved by using a stretched vertical grid?

[glwagner]

#4583 may resolve issues here. However we can't use that until we resolve the conflicts here, so the workaround in the mean time is to use a stretched vertical grid.

[shriyafruitwala]

Hi Greg, I reverted back to the rigid-lid condition to make sure things looked reasonable, and I had a few questions. The Poisson equation does not hold at the top cell, but it does at the bottom. Are the top and bottom boundary being handled differently in the rigid-lid case? Additionally, similar to before, it looks like a switch happens between some time steps, where the residual (∇²ϕ - ∇ ⋅ u* / Δt) is very large (~1e10) at the top cell and nonzero for a couple grid points below. This does not happen at the bottom boundary. Do you have a guess for why this happens?

[glwagner]

Are the top and bottom boundary being handled differently in the rigid-lid case? Additionally, similar to before, it looks like a switch happens between some time steps, where the residual (∇²ϕ - ∇ ⋅ u* / Δt) is very large (~1e10) at the top cell and nonzero for a couple grid points below.

Can you post the script you are working with? I'd like to reproduce your results to get a handle on what you're observing. For a rigid lid case with homogeneous boundary conditions on ϕ, I think we expect the residual ∇²ϕ - ∇ ⋅ u* / Δt to vanish --- provided that the halo regions of ϕ and u* are filled prior to evaluating the operators. Another check is to compute the divergence, which should be close to machine eps (this is the kind of check that I've performed more often, since its easier to compute offline).

[shriyafruitwala]

Sure! The divergence looks right (0 to machine precision). Here is the script - I put in an initial buoyancy anomaly to make internal waves.

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

coriolis = FPlane(latitude=28)
f=coriolis.f
dz = H/Nz
z = -H:dz:0
grid = RectilinearGrid(size = (Nx, Nz),
                            x = (-25meters, 25meters), 
                            z = z,
                            halo = (4,4),
                            topology = (Periodic, Flat, Bounded))
poisson = Field{Center, Center, Face}(grid)
divergence = Field{Center, Center, Face}(grid)
model = NonhydrostaticModel(; grid, coriolis, advection=WENO(order=5), auxiliary_fields = (; divergence, poisson), buoyancy = BuoyancyTracer(), tracers = :b)

N² = 1e-4
b₀(z) = N² * z

x₀, z₀ = 0.0, -H/2
σ = 5
b_anomaly(x,z) = exp(-((x-x₀)^2 / (2σ^2) + (z-z₀)^2 / (2σ^2)))   
bᵢ(x,z) = b₀(z) + b_anomaly(x,z)
set!(model.tracers[:b], bᵢ)

using Statistics
using Printf

simulation = Simulation(model, Δt=0.01, stop_time=12)
function print_stats(sim)
    w_max = maximum(model.velocities.w)
    w_min = minimum(model.velocities.w)
    w_mean = mean(model.velocities.w)

    u_max = maximum(model.velocities.u)
    u_min = minimum(model.velocities.u)
    u_mean = mean(model.velocities.u)

    @info @sprintf("Time: %.2f, max(w): %.3e, min(w): %.3e, mean(w): %.3e",
                   time(sim), w_max, w_min, w_mean)
    @info @sprintf("Time: %.2f, max(u): %.3e, min(u): %.3e, mean(u): %.3e",
                   time(sim), u_max, u_min, u_mean)

    return nothing
end

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

output_file = "nonhydro_rigidlid_.01s.jld2"
ow = JLD2Writer(model, merge(model.velocities, (; b = model.tracers[:b], 
                      divergence = model.auxiliary_fields[:divergence], poisson = model.auxiliary_fields[:poisson])),
                      filename = output_file,  
                      schedule = TimeInterval(0.01),
                      overwrite_existing = true)

simulation.output_writers[:jld2] = ow

run!(simulation)

ut = FieldTimeSeries(output_file, "u") 
wt = FieldTimeSeries(output_file, "w") 
dt = FieldTimeSeries(output_file, "divergence") 
pt = FieldTimeSeries(output_file, "poisson") 

Nt = length(wt)

To calculate the Poisson residual, I put the following before this line in pressure_correction.jl (not sure if that's the easiest way to do it).

    u = model.velocities.u
    w = model.velocities.w
    pNHS = model.pressures.pNHS

    # check if poisson eqn is satisfied - pre velocity corrector step:
    #     ∇²ϕ = ∇ ⋅ u* / Δt   -->    ∇²ϕ - ∇ ⋅ u* / Δt = 0
    if hasfield(typeof(model.auxiliary_fields), :poisson)
        poisson = model.auxiliary_fields.poisson

        for k in 1:model.grid.Nz 
            for i in 1:model.grid.Nx
                if k == 1
                    p = (pNHS[i-1,1,k] - 2*pNHS[i,1,k] + pNHS[i+1,1,k])/Δxᶜᶜᶜ(i, 1, k, model.grid)^2 + (pNHS[i,1,k+1] - pNHS[i,1,k]) / Δzᶜᶜᶜ(i, 1, k, model.grid)^2
                    rhs = (u[i+1,1,k] - u[i,1,k])/(Δt*Δxᶜᶜᶜ(i, 1, k, model.grid)) + (w[i,1,k+1] - w[i,1,k]) / (Δt * Δzᶜᶜᶜ(i, 1, k, model.grid))
                elseif k == model.grid.Nz
                    p = (pNHS[i-1,1,k] - 2*pNHS[i,1,k] + pNHS[i+1,1,k])/Δxᶜᶜᶜ(i, 1, k, model.grid)^2 + (pNHS[i,1,k-1] - pNHS[i,1,k]) / Δzᶜᶜᶜ(i, 1, k, model.grid)^2
                    rhs = (u[i+1,1,k] - u[i,1,k])/(Δt*Δxᶜᶜᶜ(i, 1, k, model.grid)) + (w[i,1,k] - w[i,1,k-1]) / (Δt * Δzᶜᶜᶜ(i, 1, k, model.grid))
                else
                    p = (pNHS[i-1,1,k] - 2*pNHS[i,1,k] + pNHS[i+1,1,k])/Δxᶜᶜᶜ(i, 1, k, model.grid)^2 + (pNHS[i,1,k-1] - 2*pNHS[i,1,k] + pNHS[i,1,k+1])/Δzᶜᶜᶜ(i, 1, k, model.grid)^2
                    rhs = (u[i+1,1,k] - u[i,1,k])/(Δt*Δxᶜᶜᶜ(i, 1, k, model.grid)) + (w[i,1,k+1] - w[i,1,k])/(Δt*Δzᶜᶜᶜ(i, 1, k, model.grid))
                end
                poisson.data[i,1,k] = p - rhs
            end
        end

    end

[glwagner]

You should be able to compute the residual using abstract operations, something like

∇²p = ∂x(∂x(p)) + ∂y(∂y(p)) + ∂z(∂z(p))
δ = ∂x(u) + ∂y(v) + ∂z(w)
r = Field(∇²p - δ / Δt)
@show r

or to store in a precomputed field (I'm not sure this will be more efficient but it could be)

r = ∇²p - δ / Δt
poisson .= r

you may need fill_halo_regions!(p) beforehand if it wasn't already called (I don't think it is called after solve! by default)

[glwagner]

poisson should be CenterField right? I don't know if this matters. The same with divergence, both are defined at cell centers on the C-grid.

[glwagner]

I modified the mwe slightly here:

using Oceananigans
using Oceananigans.Units
using Statistics
using Printf

H = 50
Nx, Nz = 50, 50
z = -H:(H/Nz):0

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

residual = CenterField(grid)

model = NonhydrostaticModel(; grid,
                            advection = WENO(),
                            auxiliary_fields = (; residual),
                            buoyancy = BuoyancyTracer(),
                            tracers = :b)

N² = 1e-4
x₀, z₀, σ = 0.0, -H/2, 5
bᵢ(x, z) = N² * z + exp(-((x - x₀)^2 / 2σ^2 + (z - z₀)^2 / 2σ^2))
set!(model, b=bᵢ)

simulation = Simulation(model, Δt=0.01, stop_time=12)

u, v, w = model.velocities
δ = Field(∂x(u) + ∂y(v) + ∂z(w))

function progress(sim)
    w_max = maximum(model.velocities.w)
    w_min = minimum(model.velocities.w)
    w_mean = mean(model.velocities.w)

    u_max = maximum(model.velocities.u)
    u_min = minimum(model.velocities.u)
    u_mean = mean(model.velocities.u)

    compute!(δ)

    msg = @sprintf("Time: %.2f, max|δ|: %.2e, max(w): %.3e, min(w): %.3e, mean(w): %.3e",
                   time(sim), maximum(abs, δ), w_max, w_min, w_mean)

    msg *= @sprintf(", max(u): %.3e, min(u): %.3e, mean(u): %.3e",
                    u_max, u_min, u_mean)

    @info msg

    return nothing
end

add_callback!(simulation, progress, IterationInterval(1))

run!(simulation)

I then ran this from this branch: https://github.com/CliMA/Oceananigans.jl/tree/glw/test-pressure-correction

which is main + 383ffaa

This produces the output:

julia> include("mwe.jl")
Precompiling Oceananigans finished.
  1 dependency successfully precompiled in 11 seconds. 140 already precompiled.
[ Info: Oceananigans will use 16 threads
[ Info: Precompiling OceananigansMakieExt [8b7e02c2-18e1-5ade-af7b-cfb5875075c8]
r = 50×1×50 Field{Center, Center, Center} on RectilinearGrid on CPU
├── grid: 50×1×50 RectilinearGrid{Float64, Periodic, Flat, Bounded} on CPU with 4×0×4 halo
├── boundary conditions: FieldBoundaryConditions
│   └── west: Periodic, east: Periodic, south: Nothing, north: Nothing, bottom: ZeroFlux, top: ZeroFlux, immersed: ZeroFlux
└── data: 58×1×58 OffsetArray(::Array{Float64, 3}, -3:54, 1:1, -3:54) with eltype Float64 with indices -3:54×1:1×-3:54
    └── max=0.0, min=0.0, mean=0.0
[ Info: Initializing simulation...
[ Info: Time: 0.00, max|δ|: 0.00e+00, max(w): 0.000e+00, min(w): 0.000e+00, mean(w): 0.000e+00, max(u): 0.000e+00, min(u): 0.000e+00, mean(u): 0.000e+00
[ Info:     ... simulation initialization complete (4.924 seconds)
[ Info: Executing initial time step...
r = 50×1×50 Field{Center, Center, Center} on RectilinearGrid on CPU
├── grid: 50×1×50 RectilinearGrid{Float64, Periodic, Flat, Bounded} on CPU with 4×0×4 halo
├── boundary conditions: FieldBoundaryConditions
│   └── west: Periodic, east: Periodic, south: Nothing, north: Nothing, bottom: ZeroFlux, top: ZeroFlux, immersed: ZeroFlux
└── data: 58×1×58 OffsetArray(::Array{Float64, 3}, -3:54, 1:1, -3:54) with eltype Float64 with indices -3:54×1:1×-3:54
    └── max=5.42101e-17, min=-4.65737e-17, mean=2.07703e-23
r = 50×1×50 Field{Center, Center, Center} on RectilinearGrid on CPU
├── grid: 50×1×50 RectilinearGrid{Float64, Periodic, Flat, Bounded} on CPU with 4×0×4 halo
├── boundary conditions: FieldBoundaryConditions
│   └── west: Periodic, east: Periodic, south: Nothing, north: Nothing, bottom: ZeroFlux, top: ZeroFlux, immersed: ZeroFlux
└── data: 58×1×58 OffsetArray(::Array{Float64, 3}, -3:54, 1:1, -3:54) with eltype Float64 with indices -3:54×1:1×-3:54
    └── max=1.32273e-17, min=-1.69136e-17, mean=-6.1664e-23
r = 50×1×50 Field{Center, Center, Center} on RectilinearGrid on CPU
├── grid: 50×1×50 RectilinearGrid{Float64, Periodic, Flat, Bounded} on CPU with 4×0×4 halo
├── boundary conditions: FieldBoundaryConditions
│   └── west: Periodic, east: Periodic, south: Nothing, north: Nothing, bottom: ZeroFlux, top: ZeroFlux, immersed: ZeroFlux
└── data: 58×1×58 OffsetArray(::Array{Float64, 3}, -3:54, 1:1, -3:54) with eltype Float64 with indices -3:54×1:1×-3:54
    └── max=2.92735e-17, min=-2.45436e-17, mean=-8.63974e-23
[ Info: Time: 0.01, max|δ|: 2.94e-17, max(w): 4.606e-03, min(w): -1.460e-03, mean(w): -1.911e-20, max(u): 1.468e-03, min(u): -1.468e-03, mean(u): 1.187e-21
[ Info:     ... initial time step complete (12.814 seconds).
r = 50×1×50 Field{Center, Center, Center} on RectilinearGrid on CPU
├── grid: 50×1×50 RectilinearGrid{Float64, Periodic, Flat, Bounded} on CPU with 4×0×4 halo
├── boundary conditions: FieldBoundaryConditions
│   └── west: Periodic, east: Periodic, south: Nothing, north: Nothing, bottom: ZeroFlux, top: ZeroFlux, immersed: ZeroFlux
└── data: 58×1×58 OffsetArray(::Array{Float64, 3}, -3:54, 1:1, -3:54) with eltype Float64 with indices -3:54×1:1×-3:54
    └── max=6.5269e-17, min=-4.79217e-17, mean=-5.6243e-23
r = 50×1×50 Field{Center, Center, Center} on RectilinearGrid on CPU
├── grid: 50×1×50 RectilinearGrid{Float64, Periodic, Flat, Bounded} on CPU with 4×0×4 halo
├── boundary conditions: FieldBoundaryConditions
│   └── west: Periodic, east: Periodic, south: Nothing, north: Nothing, bottom: ZeroFlux, top: ZeroFlux, immersed: ZeroFlux
└── data: 58×1×58 OffsetArray(::Array{Float64, 3}, -3:54, 1:1, -3:54) with eltype Float64 with indices -3:54×1:1×-3:54
    └── max=1.30646e-17, min=-1.09504e-17, mean=1.47723e-22
r = 50×1×50 Field{Center, Center, Center} on RectilinearGrid on CPU
├── grid: 50×1×50 RectilinearGrid{Float64, Periodic, Flat, Bounded} on CPU with 4×0×4 halo
├── boundary conditions: FieldBoundaryConditions
│   └── west: Periodic, east: Periodic, south: Nothing, north: Nothing, bottom: ZeroFlux, top: ZeroFlux, immersed: ZeroFlux
└── data: 58×1×58 OffsetArray(::Array{Float64, 3}, -3:54, 1:1, -3:54) with eltype Float64 with indices -3:54×1:1×-3:54
    └── max=3.92481e-17, min=-2.73219e-17, mean=-2.11419e-22
[ Info: Time: 0.02, max|δ|: 3.92e-17, max(w): 9.213e-03, min(w): -2.920e-03, mean(w): -3.457e-20, max(u): 2.936e-03, min(u): -2.936e-03, mean(u): -1.225e-21
r = 50×1×50 Field{Center, Center, Center} on RectilinearGrid on CPU
├── grid: 50×1×50 RectilinearGrid{Float64, Periodic, Flat, Bounded} on CPU with 4×0×4 halo
├── boundary conditions: FieldBoundaryConditions
│   └── west: Periodic, east: Periodic, south: Nothing, north: Nothing, bottom: ZeroFlux, top: ZeroFlux, immersed: ZeroFlux
└── data: 58×1×58 OffsetArray(::Array{Float64, 3}, -3:54, 1:1, -3:54) with eltype Float64 with indices -3:54×1:1×-3:54
    └── max=5.11743e-17, min=-4.85723e-17, mean=-6.05798e-22
r = 50×1×50 Field{Center, Center, Center} on RectilinearGrid on CPU
├── grid: 50×1×50 RectilinearGrid{Float64, Periodic, Flat, Bounded} on CPU with 4×0×4 halo
├── boundary conditions: FieldBoundaryConditions
│   └── west: Periodic, east: Periodic, south: Nothing, north: Nothing, bottom: ZeroFlux, top: ZeroFlux, immersed: ZeroFlux
└── data: 58×1×58 OffsetArray(::Array{Float64, 3}, -3:54, 1:1, -3:54) with eltype Float64 with indices -3:54×1:1×-3:54
    └── max=1.64257e-17, min=-1.14925e-17, mean=-1.76183e-23
r = 50×1×50 Field{Center, Center, Center} on RectilinearGrid on CPU
├── grid: 50×1×50 RectilinearGrid{Float64, Periodic, Flat, Bounded} on CPU with 4×0×4 halo
├── boundary conditions: FieldBoundaryConditions
│   └── west: Periodic, east: Periodic, south: Nothing, north: Nothing, bottom: ZeroFlux, top: ZeroFlux, immersed: ZeroFlux
└── data: 58×1×58 OffsetArray(::Array{Float64, 3}, -3:54, 1:1, -3:54) with eltype Float64 with indices -3:54×1:1×-3:54
    └── max=3.01408e-17, min=-2.68882e-17, mean=-6.24772e-22
[ Info: Time: 0.03, max|δ|: 3.02e-17, max(w): 1.382e-02, min(w): -4.380e-03, mean(w): -3.598e-20, max(u): 4.404e-03, min(u): -4.404e-03, mean(u): 8.159e-21

so at least on that commit, the residual is machine eps. (Note that I had to adjust the residual computation, because on main we solve for p * dt rather than just p).

[shriyafruitwala]

Yes, I see. It looks like solving for p*dt instead of p (in solve_for_pressure! and make_pressure_correction!) makes a difference in the residual. If I don’t do it this way and instead divide the rhs by dt, some time steps still have a nonzero residual in the top cell.

[glwagner]

Yes, I see. It looks like solving for p*dt instead of p (in solve_for_pressure! and make_pressure_correction!) makes a difference in the residual. If I don’t do it this way and instead divide the rhs by dt, some time steps still have a nonzero residual in the top cell.

True but if you don't account for the division by dt, then you aren't computing the residual correctly, right?

[shriyafruitwala]

Yes, I see. It looks like solving for p*dt instead of p (in solve_for_pressure! and make_pressure_correction!) makes a difference in the residual. If I don’t do it this way and instead divide the rhs by dt, some time steps still have a nonzero residual in the top cell.

True but if you don't account for the division by dt, then you aren't computing the residual correctly, right?

Right, but before I was accounting for the division by dt in the calculation of the rhs (e.g. rhs[i, j, k] = active * Δzᶜᶜᶜ(i, j, k, grid) * δ / Δt), so then I would need to put in a factor dt in the residual calculation. Doing it this way seems to yield a nonzero residual in the top cell for certain timesteps, for some reason.

[glwagner]

Yes, I see. It looks like solving for p*dt instead of p (in solve_for_pressure! and make_pressure_correction!) makes a difference in the residual. If I don’t do it this way and instead divide the rhs by dt, some time steps still have a nonzero residual in the top cell.

True but if you don't account for the division by dt, then you aren't computing the residual correctly, right?

Right, but before I was accounting for the division by dt in the calculation of the rhs (e.g. rhs[i, j, k] = active * Δzᶜᶜᶜ(i, j, k, grid) * δ / Δt), so then I would need to put in a factor dt in the residual calculation. Doing it this way seems to yield a nonzero residual in the top cell for certain timesteps, for some reason.

Are you using adaptive time stepping or do you have a TimeInterval schedule that would cause the time step to change? Also I am not sure what code / branch you are using. Does your branch solve for p*dt, or p?

[shriyafruitwala]

Yes, I see. It looks like solving for p*dt instead of p (in solve_for_pressure! and make_pressure_correction!) makes a difference in the residual. If I don’t do it this way and instead divide the rhs by dt, some time steps still have a nonzero residual in the top cell.

True but if you don't account for the division by dt, then you aren't computing the residual correctly, right?

Right, but before I was accounting for the division by dt in the calculation of the rhs (e.g. rhs[i, j, k] = active * Δzᶜᶜᶜ(i, j, k, grid) * δ / Δt), so then I would need to put in a factor dt in the residual calculation. Doing it this way seems to yield a nonzero residual in the top cell for certain timesteps, for some reason.

Are you using adaptive time stepping or do you have a TimeInterval schedule that would cause the time step to change? Also I am not sure what code / branch you are using. Does your branch solve for p*dt, or p?

My dt should be constant. Sorry, I wasn't being very clear - my branch now solves for p*dt, mirroring what you did above, so the residual looks good. Solving for p seemed to mess up the top cell.

[shriyafruitwala]

Reverting back to the free surface, there seems to be a problem with computing the diagonal in the fourier_tridiagonal_poisson_solver. The Robin boundary condition depends on Δt, which was not being passed into the function, so I hardcoded it. This creates a mismatch in RK3 - where Δt is different for each substep. As I understand, the solver builds the matrix once, now using the hardcoded Δt, I’m not sure how to best resolve this - maybe by precomputing 3 versions of the matrix (which would only work for constant Δt) or by rebuilding the matrix every substep?

[glwagner]

Reverting back to the free surface, there seems to be a problem with computing the diagonal in the fourier_tridiagonal_poisson_solver. The Robin boundary condition depends on Δt, which was not being passed into the function, so I hardcoded it. This creates a mismatch in RK3 - where Δt is different for each substep. As I understand, the solver builds the matrix once, now using the hardcoded Δt, I’m not sure how to best resolve this - maybe by precomputing 3 versions of the matrix (which would only work for constant Δt) or by rebuilding the matrix every substep?

Ok, this sounds like a core issue. First, I would try using AB2 instead to avoid the substep issue. Is that possible? It's achieved by passing timestepper=:QuasiAdamsBashforth2 to the model constructor.

To allow variable time steps, we may have to use a different formulation of the tridiagonal coefficients that avoids fully specifying arrays. The BatchedTridiagonalSolver can be extended for this using the get_coefficient interface:

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

for example:

Oceananigans.jl/src/Solvers/batched_tridiagonal_solver.jl

Lines 219 to 245 in e1f1080

	@inline function solve_batched_tridiagonal_system_z!(i, j, Nz, ϕ, a, b, c, f, t, grid, p, args, tridiagonal_direction)
	@inbounds begin
	β = get_coefficient(i, j, 1, grid, b, p, tridiagonal_direction, args...)
	f₁ = get_coefficient(i, j, 1, grid, f, p, tridiagonal_direction, args...)
	ϕ[i, j, 1] = f₁ / β
	for k = 2:Nz
	cᵏ⁻¹ = get_coefficient(i, j, k-1, grid, c, p, tridiagonal_direction, args...)
	bᵏ = get_coefficient(i, j, k, grid, b, p, tridiagonal_direction, args...)
	aᵏ⁻¹ = get_coefficient(i, j, k-1, grid, a, p, tridiagonal_direction, args...)
	t[i, j, k] = cᵏ⁻¹ / β
	β = bᵏ - aᵏ⁻¹ * t[i, j, k]
	fᵏ = get_coefficient(i, j, k, grid, f, p, tridiagonal_direction, args...)
	# If the problem is not diagonally-dominant such that `β ≈ 0`,
	# the algorithm is unstable and we elide the forward pass update of `ϕ`.
	definitely_diagonally_dominant = abs(β) > 10 * eps(float_eltype(ϕ))
	ϕ★ = (fᵏ - aᵏ⁻¹ * ϕ[i, j, k-1]) / β
	ϕ[i, j, k] = ifelse(definitely_diagonally_dominant, ϕ★, ϕ[i, j, k])
	end
	for k = Nz-1:-1:1
	ϕ[i, j, k] -= t[i, j, k+1] * ϕ[i, j, k+1]
	end
	end
	end

This interface is invoked by the vertical diffusion solver, so an example of using it is provided here:

Oceananigans.jl/src/TurbulenceClosures/vertically_implicit_diffusion_solver.jl

Lines 160 to 162 in e1f1080

	@inline get_coefficient(i, j, k, grid, ::VerticallyImplicitDiffusionLowerDiagonal, p, ::ZDirection, args...) = _ivd_lower_diagonal(i, j, k, grid, args...)
	@inline get_coefficient(i, j, k, grid, ::VerticallyImplicitDiffusionUpperDiagonal, p, ::ZDirection, args...) = _ivd_upper_diagonal(i, j, k, grid, args...)
	@inline get_coefficient(i, j, k, grid, ::VerticallyImplicitDiffusionDiagonal, p, ::ZDirection, args...) = ivd_diagonal(i, j, k, grid, args...)

here:

Oceananigans.jl/src/TurbulenceClosures/vertically_implicit_diffusion_solver.jl

Lines 150 to 153 in e1f1080

	z_solver = BatchedTridiagonalSolver(grid;
	lower_diagonal = VerticallyImplicitDiffusionLowerDiagonal(),
	diagonal = VerticallyImplicitDiffusionDiagonal(),
	upper_diagonal = VerticallyImplicitDiffusionUpperDiagonal())

and here:

Oceananigans.jl/src/TurbulenceClosures/vertically_implicit_diffusion_solver.jl

Lines 202 to 204 in e1f1080

	return solve!(field, implicit_solver, field,
	# ivd_*_diagonal gets called with these args after (i, j, k, grid):
	vi_closure, vi_diffusivity_fields, tracer_index, LX(), LY(), LZ(), Δt, clock; kwargs...)

For the FourierTridiagonalPoissonSolver, I think we want to use something similar, except we also want to store the static/precomputed part of the diagonal, and then add the part that changes depending on dt (which also needs to be passed into the pressure solver, and then passed on to the tridiagonal solver).

I am wondering if we need to start from scratch here, given the large number of conflicts on this PR and the need for fairly involved extension of FourierTridiagonalPoissonSolver I just described (which I can help with). It might make sense just to test AB2 for now to see what that gives us? Let me know what you think @shriyafruitwala , and also if you see a path forward.

[shriyafruitwala]

I think it makes sense to try and get this working with AB2 before moving to variable dt. I tested the free surface case with AB2, and it blows up almost immediately. As a sanity check, I went back to the rigid lid case, and that seems to yield the same results with AB2 as RK3, so there is probably another issue with the free surface implementation besides the time stepper.

One question that might indicate what is going wrong is the difference between calculating p and p*dt in the rigid lid case. Calculating p*dt seems to give a reasonable answer (as discussed above), but solving for p yields a nonzero Poisson residual in the top cell for certain time steps. Dividing the equation by dt shouldn’t yield a different answer right? Also it seems odd that this only happens at the top boundary and not also at the bottom.

I also tested the free surface case with AB2 and calculating p instead of p*dt. This runs without blowing up (the flickering in w that we were seeing before is no longer there). The divergence is pretty consistently nonzero, which is a problem, but w and the free surface evolution look reasonable. I think that understanding this p vs p*dt discrepancy would give us a clue for what’s going wrong.

[glwagner]

Dividing the equation by dt shouldn’t yield a different answer right? Also it seems odd that this only happens at the top boundary and not also at the bottom.

Correct. There's no mathematical difference between solving for p*dt or p for fixed non-small dt. There are only differences for dt near machine epsilon.

[shriyafruitwala]

Dividing the equation by dt shouldn’t yield a different answer right? Also it seems odd that this only happens at the top boundary and not also at the bottom.

Correct. There's no mathematical difference between solving for p*dt or p for fixed non-small dt. There are only differences for dt near machine epsilon.

Do you know why solving for p vs p*dt could result in different answers?

[glwagner]

Dividing the equation by dt shouldn’t yield a different answer right? Also it seems odd that this only happens at the top boundary and not also at the bottom.

Correct. There's no mathematical difference between solving for p*dt or p for fixed non-small dt. There are only differences for dt near machine epsilon.

Do you know why solving for p vs p*dt could result in different answers?

It should not result in different answers. When we made the change on main some time ago, our regression tests all passed indicating that the solution is identical to within round-off error for a variety of cases.

[shriyafruitwala]

Dividing the equation by dt shouldn’t yield a different answer right? Also it seems odd that this only happens at the top boundary and not also at the bottom.

Correct. There's no mathematical difference between solving for p*dt or p for fixed non-small dt. There are only differences for dt near machine epsilon.

Do you know why solving for p vs p*dt could result in different answers?

It should not result in different answers. When we made the change on main some time ago, our regression tests all passed indicating that the solution is identical to within round-off error for a variety of cases.

That’s interesting - in the rigid lid case, I get a nonzero Poisson residual in the top grid cell for certain time steps when calculating p. This does not happen for p*dt, and I don’t understand why.

[glwagner]

That’s interesting - in the rigid lid case, I get a nonzero Poisson residual in the top grid cell for certain time steps when calculating p. This does not happen for p*dt, and I don’t understand why.

Are you assuming a constant dt in the residual calculation? In the residual for p*dt, the time-step dt does not appear, but dt does appear in the formula for p. dt can change if there is a TimeInterval schedule somewhere in your Simulation (or if you use adaptive time stepping).

[shriyafruitwala]

That’s interesting - in the rigid lid case, I get a nonzero Poisson residual in the top grid cell for certain time steps when calculating p. This does not happen for p*dt, and I don’t understand why.

Are you assuming a constant dt in the residual calculation? In the residual for p*dt, the time-step dt does not appear, but dt does appear in the formula for p. dt can change if there is a TimeInterval schedule somewhere in your Simulation (or if you use adaptive time stepping).

I believe dt is constant in my simulation. I set it up with simulation = Simulation(model, Δt=0.01, stop_time=10) and TimeInterval is only called in the output writer setup, which should only control how often simulation data is saved right?

ow = JLD2Writer(model, merge(model.velocities, (; b = model.tracers[:b], η=model.free_surface.η, 
                      divergence = model.auxiliary_fields[:divergence], poisson = model.auxiliary_fields[:poisson])),
                      filename = output_file,  
                      schedule = TimeInterval(0.01),
                      overwrite_existing = true)

simulation.output_writers[:jld2] = ow

[glwagner]

That’s interesting - in the rigid lid case, I get a nonzero Poisson residual in the top grid cell for certain time steps when calculating p. This does not happen for p*dt, and I don’t understand why.

Are you assuming a constant dt in the residual calculation? In the residual for p*dt, the time-step dt does not appear, but dt does appear in the formula for p. dt can change if there is a TimeInterval schedule somewhere in your Simulation (or if you use adaptive time stepping).

I believe dt is constant in my simulation. I set it up with simulation = Simulation(model, Δt=0.01, stop_time=10) and TimeInterval is only called in the output writer setup, which should only control how often simulation data is saved right?

ow = JLD2Writer(model, merge(model.velocities, (; b = model.tracers[:b], η=model.free_surface.η, 
                      divergence = model.auxiliary_fields[:divergence], poisson = model.auxiliary_fields[:poisson])),
                      filename = output_file,  
                      schedule = TimeInterval(0.01),
                      overwrite_existing = true)

simulation.output_writers[:jld2] = ow

No, TimeInterval can incur a change in the time-step so that output occurs exactly at the scheduled time. Because of round off error this can occasionally result in very small time-steps. You should use IterationInterval until we have a scheme that works for any time-step.

[shriyafruitwala]

No, TimeInterval can incur a change in the time-step so that output occurs exactly at the scheduled time. Because of round off error this can occasionally result in very small time-steps. You should use IterationInterval until we have a scheme that works for any time-step.

Oh I see, that makes sense.

[shriyafruitwala]

Looks like it’s working now with ab2! I have plotted the Oceananigans free surface compared to the Fourier analytical solution and provided the code below.

nonhydro_robin_ab2_pdt.jld2_compare_eta.mp4

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

#grid setup: flat bottom, free surface
Nx, Nz = 50, 50
H = 50meters
L = 50meters
g = 10.0

coriolis = FPlane(latitude=28)
f=coriolis.f
dz = H/Nz
z = -H:dz:0
grid = RectilinearGrid(size = (Nx, Nz),
                            x = (-25meters, 25meters), 
                            z = z,
                            halo = (4,4),
                            topology = (Periodic, Flat, Bounded))

divergence = CenterField(grid)
free_surface = ImplicitFreeSurface(gravitational_acceleration=10)
model = NonhydrostaticModel(; grid, free_surface, coriolis, advection=WENO(order=5), auxiliary_fields = (; divergence), buoyancy = BuoyancyTracer(), tracers = :b, timestepper=:QuasiAdamsBashforth2) 

#initialize with gaussian bump surface
η = model.free_surface.η
bump(x) = 1*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
using Statistics
using Printf

simulation = Simulation(model, Δt=0.01, stop_time=10)

function print_stats(sim)
    w_max = maximum(model.velocities.w)
    w_min = minimum(model.velocities.w)
    w_mean = mean(model.velocities.w)

    u_max = maximum(model.velocities.u)
    u_min = minimum(model.velocities.u)
    u_mean = mean(model.velocities.u)    

    @info @sprintf("Time: %.2f, max(w): %.3e, min(w): %.3e, mean(w): %.3e",
                   time(sim), w_max, w_min, w_mean)
    @info @sprintf("Time: %.2f, max(u): %.3e, min(u): %.3e, mean(u): %.3e",
                   time(sim), u_max, u_min, u_mean)
    return nothing
end

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

output_file = "nonhydro_robin_ab2_pdt.jld2"
ow = JLD2Writer(model, merge(model.velocities, (; b = model.tracers[:b], η=model.free_surface.η, 
                      divergence = model.auxiliary_fields[:divergence])),
                      filename = output_file,  
                      schedule = IterationInterval(1),
                      overwrite_existing = true)

simulation.output_writers[:jld2] = ow
run!(simulation)

# produces animation comparing simulation with analytical solution
using FFTW
using Printf

# load field
ηt = FieldTimeSeries(output_file, "η") 
Nt = length(ηt)

# parameters
x_sim = xnodes(ηt)
N = length(x_sim)
L = 50.0
a = 1.0
δ = 5.0

times = ηt.times
k = [8π * i / L for i in 0:(N÷2)]
η0 = a .* exp.(-(x_sim .^ 2) ./ (2 * δ^2))
η0_hat = rfft(η0) ./ 2

function η_fourier(t)
    ω = sqrt.(g .* abs.(k))
    η_hat_t = η0_hat .* exp.(-1im .* ω .* t) .+ η0_hat .* exp.(1im .* ω .* t)
    return irfft(η_hat_t, N)
end

#plot
fig = Figure(; size=(1000, 400))

ax_sim = Axis(fig[1, 1], xlabel="x (m)", ylabel="η", title="Oceananigans Free Surface")
ax_fft = Axis(fig[1, 2], xlabel="x (m)", ylabel="η", title="Fourier-Analytical Free Surface")
n = Observable(1)
η_sim = @lift interior(ηt[$n], :, 1, 1)
η_fft = Observable(η_fourier(0.0))

lines!(ax_sim, x_sim, η_sim)
lines!(ax_fft, x_sim, η_fft)

ylims!(ax_sim, -1.5, 1.5)
ylims!(ax_fft, -1.5, 1.5)

record(fig, output_file * "_compare_eta.mp4", 1:Nt; framerate=24) do frame
    @info "Drawing frame $frame of $Nt..."
    n[] = frame
    η_fft[] = η_fourier(times[frame])
end

In terms of next steps, I think we need to clean up my code a bit (e.g. not have gravity hard coded, use the right dz, etc.). For my problem, I use an ImmersedBoundary, so I need to apply the Robin boundary condition to the CG solver. Where in the solver does the boundary condition get set?

Also, there remains the problem with RK3…

[glwagner]

Looks like it’s working now with ab2! I have plotted the Oceananigans free surface compared to the Fourier analytical solution and provided the code below.

nonhydro_robin_ab2_pdt.jld2_compare_eta.mp4

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

#grid setup: flat bottom, free surface
Nx, Nz = 50, 50
H = 50meters
L = 50meters
g = 10.0

coriolis = FPlane(latitude=28)
f=coriolis.f
dz = H/Nz
z = -H:dz:0
grid = RectilinearGrid(size = (Nx, Nz),
                            x = (-25meters, 25meters), 
                            z = z,
                            halo = (4,4),
                            topology = (Periodic, Flat, Bounded))

divergence = CenterField(grid)
free_surface = ImplicitFreeSurface(gravitational_acceleration=10)
model = NonhydrostaticModel(; grid, free_surface, coriolis, advection=WENO(order=5), auxiliary_fields = (; divergence), buoyancy = BuoyancyTracer(), tracers = :b, timestepper=:QuasiAdamsBashforth2) 

#initialize with gaussian bump surface
η = model.free_surface.η
bump(x) = 1*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
using Statistics
using Printf

simulation = Simulation(model, Δt=0.01, stop_time=10)

function print_stats(sim)
    w_max = maximum(model.velocities.w)
    w_min = minimum(model.velocities.w)
    w_mean = mean(model.velocities.w)

    u_max = maximum(model.velocities.u)
    u_min = minimum(model.velocities.u)
    u_mean = mean(model.velocities.u)    

    @info @sprintf("Time: %.2f, max(w): %.3e, min(w): %.3e, mean(w): %.3e",
                   time(sim), w_max, w_min, w_mean)
    @info @sprintf("Time: %.2f, max(u): %.3e, min(u): %.3e, mean(u): %.3e",
                   time(sim), u_max, u_min, u_mean)
    return nothing
end

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

output_file = "nonhydro_robin_ab2_pdt.jld2"
ow = JLD2Writer(model, merge(model.velocities, (; b = model.tracers[:b], η=model.free_surface.η, 
                      divergence = model.auxiliary_fields[:divergence])),
                      filename = output_file,  
                      schedule = IterationInterval(1),
                      overwrite_existing = true)

simulation.output_writers[:jld2] = ow
run!(simulation)

# produces animation comparing simulation with analytical solution
using FFTW
using Printf

# load field
ηt = FieldTimeSeries(output_file, "η") 
Nt = length(ηt)

# parameters
x_sim = xnodes(ηt)
N = length(x_sim)
L = 50.0
a = 1.0
δ = 5.0

times = ηt.times
k = [8π * i / L for i in 0:(N÷2)]
η0 = a .* exp.(-(x_sim .^ 2) ./ (2 * δ^2))
η0_hat = rfft(η0) ./ 2

function η_fourier(t)
    ω = sqrt.(g .* abs.(k))
    η_hat_t = η0_hat .* exp.(-1im .* ω .* t) .+ η0_hat .* exp.(1im .* ω .* t)
    return irfft(η_hat_t, N)
end

#plot
fig = Figure(; size=(1000, 400))

ax_sim = Axis(fig[1, 1], xlabel="x (m)", ylabel="η", title="Oceananigans Free Surface")
ax_fft = Axis(fig[1, 2], xlabel="x (m)", ylabel="η", title="Fourier-Analytical Free Surface")
n = Observable(1)
η_sim = @lift interior(ηt[$n], :, 1, 1)
η_fft = Observable(η_fourier(0.0))

lines!(ax_sim, x_sim, η_sim)
lines!(ax_fft, x_sim, η_fft)

ylims!(ax_sim, -1.5, 1.5)
ylims!(ax_fft, -1.5, 1.5)

record(fig, output_file * "_compare_eta.mp4", 1:Nt; framerate=24) do frame
    @info "Drawing frame $frame of $Nt..."
    n[] = frame
    η_fft[] = η_fourier(times[frame])
end

In terms of next steps, I think we need to clean up my code a bit (e.g. not have gravity hard coded, use the right dz, etc.). For my problem, I use an ImmersedBoundary, so I need to apply the Robin boundary condition to the CG solver. Where in the solver does the boundary condition get set?

Also, there remains the problem with RK3…

Really good news!

Here is my thought: there are a lot of conflicts with main already on this PR, such that we should consider starting the implementation here from scratch. We can use what we've learned in this effort to design the new implementation to meet all the necessary requirements. I believe what we need is the ability to compute some part of the diagonals used in FourierTridiagonalPoissonSolver on the fly. When we have that we can pass in both g and dt, and support both RK3 and variable time stepping.

As for problems in complex domains: I suggest setting up a test problem and running simulations with the "naive" solver first. We will want to have that as a baseline before developing a CG solver anyways. We need to develop a new boundary condition type to express Robin boundary conditions, as well as the associated halo filling algorithm.

[shriyafruitwala]

Really good news!

Here is my thought: there are a lot of conflicts with main already on this PR, such that we should consider starting the implementation here from scratch. We can use what we've learned in this effort to design the new implementation to meet all the necessary requirements. I believe what we need is the ability to compute some part of the diagonals used in FourierTridiagonalPoissonSolver on the fly. When we have that we can pass in both g and dt, and support both RK3 and variable time stepping.

As for problems in complex domains: I suggest setting up a test problem and running simulations with the "naive" solver first. We will want to have that as a baseline before developing a CG solver anyways. We need to develop a new boundary condition type to express Robin boundary conditions, as well as the associated halo filling algorithm.

Okay, great! I had actually pulled in changes from main a little while ago and merged them into my branch. I pushed updates this morning, so the branch might not be as far behind as it seems, but might be worth double-checking before starting from scratch.

Also, it might be useful to use this same test case with the deep water waves to test the CG solver, since we have a ground truth for what that should look like. Where do boundary conditions get implemented in the CG solver?

[glwagner]

There are conflicts in about 10 files now:

sometimes these can be a little tricky / time consuming to resolve. The original implementation from scratch took about 1-2 hour of programming time. We can double that to account for the new features needed, but I am not sure after all it will end up being more work than will take to resolve so many conflicts.

[shriyafruitwala]

There are conflicts in about 10 files now:
sometimes these can be a little tricky / time consuming to resolve. The original implementation from scratch took about 1-2 hour of programming time. We can double that to account for the new features needed, but I am not sure after all it will end up being more work than will take to resolve so many conflicts.

Okay, that makes sense! What are the next steps?