fix mean pressure in FFT preconditioner for ConjugateGradientPoissonSolver for immersed boundaries

[xkykai]

Closes #4553.

This fixes the immersed-aware mean pressure to be zero at every preconditioning step, effectively avoiding numerical blowup due to floating point error accumulation.

However this also means that without using the preconditioner, the error in mean pressure can and will accumulate and grow over time. However this is perhaps an unlikely use case as it garners no benefits vs. using the FFT preconditioner.

Here I show the results with FFT preconditioner with the test case in #4553

[ Info: Initializing simulation...
[ Info: Iter: 0, time: 0.000e+00, Δt: 1.600e-04, Poisson iters: 11, max u: 6.646e+01, max v: 0.000e+00, max w: 3.451e+01, max d: 1.653e-05, max pressure: 1.516e+01, mean pressure: -9.107e-18
[ Info:     ... simulation initialization complete (2.146 seconds)
[ Info: Executing initial time step...
[ Info: Iter: 1, time: 1.600e-04, Δt: 1.600e-04, Poisson iters: 11, max u: 6.743e+01, max v: 0.000e+00, max w: 3.522e+01, max d: 2.001e-07, max pressure: 2.320e+03, mean pressure: -7.460e-16
[ Info:     ... initial time step complete (1.005 seconds).
[ Info: Iter: 2, time: 3.200e-04, Δt: 1.600e-04, Poisson iters: 10, max u: 6.774e+01, max v: 0.000e+00, max w: 3.553e+01, max d: 1.412e-06, max pressure: 2.295e+03, mean pressure: -7.460e-16
[ Info: Iter: 3, time: 4.800e-04, Δt: 1.600e-04, Poisson iters: 10, max u: 6.743e+01, max v: 0.000e+00, max w: 3.541e+01, max d: 9.215e-07, max pressure: 2.254e+03, mean pressure: -1.492e-15
[ Info: Iter: 4, time: 6.400e-04, Δt: 1.600e-04, Poisson iters: 10, max u: 6.688e+01, max v: 0.000e+00, max w: 3.506e+01, max d: 5.800e-07, max pressure: 2.240e+03, mean pressure: 0.000e+00
[ Info: Iter: 5, time: 8.000e-04, Δt: 1.600e-04, Poisson iters: 10, max u: 6.618e+01, max v: 0.000e+00, max w: 3.459e+01, max d: 4.424e-07, max pressure: 2.214e+03, mean pressure: -3.730e-16
[ Info: Iter: 6, time: 9.600e-04, Δt: 1.600e-04, Poisson iters: 11, max u: 6.541e+01, max v: 0.000e+00, max w: 3.406e+01, max d: 8.807e-08, max pressure: 2.160e+03, mean pressure: 0.000e+00
[ Info: Iter: 7, time: 1.120e-03, Δt: 1.600e-04, Poisson iters: 11, max u: 6.473e+01, max v: 0.000e+00, max w: 3.351e+01, max d: 6.340e-08, max pressure: 2.095e+03, mean pressure: 3.730e-16
[ Info: Iter: 8, time: 1.280e-03, Δt: 1.600e-04, Poisson iters: 11, max u: 6.423e+01, max v: 0.000e+00, max w: 3.298e+01, max d: 7.278e-08, max pressure: 2.026e+03, mean pressure: -1.492e-15
[ Info: Iter: 9, time: 1.440e-03, Δt: 1.600e-04, Poisson iters: 11, max u: 6.382e+01, max v: 0.000e+00, max w: 3.252e+01, max d: 7.359e-08, max pressure: 1.955e+03, mean pressure: 0.000e+00
[ Info: Simulation is stopping after running for 3.875 seconds.
[ Info: Model iteration 10 equals or exceeds stop iteration 10.
[ Info: Iter: 10, time: 1.600e-03, Δt: 1.600e-04, Poisson iters: 11, max u: 6.348e+01, max v: 0.000e+00, max w: 3.219e+01, max d: 8.471e-08, max pressure: 1.888e+03, mean pressure: -7.460e-16

Here we see that the mean pressure is effectively maintained at zero.

Alternatively, to solve the issue in cases where no preconditioner is used as well, we can use @Yixiao-Zhang 's proposal which is to modify https://github.com/CliMA/Oceananigans.jl/blob/f6463bd4a8280dfa2d7182ccdc2d8e438e49d384/src/Solvers/conjugate_gradient_solver.jl#L242C1-L242C55:

function update_solution_and_residuals!(x, r, q, p, α)
    xp = parent(x)
    rp = parent(r)
    qp = parent(q)
    pp = parent(p)

    xp .+= α .* pp
    rp .-= α .* qp

    mean_r = mean(r)
    grid = r.grid
    arch = architecture(grid)
    launch!(arch, grid, :xyz, subtract_and_mask!, r, grid, mean_r)

    return nothing
end

This touches the internals of the ConjugateGradientSolver, and this will fail other arbitrary problems solvable by conjugate gradient where there is no gauge invariance. We can perhaps argue that solving the pressure Poisson equation with Neumann boundary conditions is a special use case of the ConjugateGradientSolver. If we decide to address the issue within the conjugate gradient iteration we should perhaps think of modifying the API slightly to not do it for all arbitrary problems. Also if we take this approach we should also do the subtract_and_mask! on x as well, no reason not to.

Also for reference, here is what the MWE script does:

vortex_sheet.mp4

[glwagner]

This touches the internals of the ConjugateGradientSolver, and this will fail other arbitrary problems solvable by conjugate gradient where there is no gauge invariance. We can perhaps argue that solving the pressure Poisson equation with Neumann boundary conditions is a special use case of the ConjugateGradientSolver.

I only see a modification to the preconditioner for the Poisson solver. So I don't understand how this PR changes other problems. Can you explain?

In any case, rather than hard coding something into a general solver, add a property to ConjugateGradientSolver called, say, enforce_gauge_condition!. Then future solver implementations with other / no gauge condition can use the interface. This is vastly preferred to hard coding something not only because it keeps the generality of the solver, but also because it is "self-documenting" because it explicitly declares the new step as enforcing a gauge condition that is specific to the problem being solved.

[xkykai]

I only see a modification to the preconditioner for the Poisson solver. So I don't understand how this PR changes other problems. Can you explain?

I'm referring to doing this alternative approach which I have not implemented:

Alternatively, to solve the issue in cases where no preconditioner is used as well, we can use @Yixiao-Zhang 's proposal which is to modify f6463bd/src/Solvers/conjugate_gradient_solver.jl#L242C1-L242C55:

function update_solution_and_residuals!(x, r, q, p, α)
    xp = parent(x)
    rp = parent(r)
    qp = parent(q)
    pp = parent(p)

    xp .+= α .* pp
    rp .-= α .* qp

    mean_r = mean(r)
    grid = r.grid
    arch = architecture(grid)
    launch!(arch, grid, :xyz, subtract_and_mask!, r, grid, mean_r)

    return nothing
end

but yes I agree doing something like enforce_gauge_condition! will be a good idea

[xkykai]

Is this a good interface to apply this gauge condition? @glwagner happy to hear your recommendations

[glwagner]

No, I would implement a generic interface, because we don't know what the gauge condition will be for an arbitrary linear_operation!. Just like we can't have "apply preconditioner" as a boolean for all linear operations, we also need to have some kind of generic interface for enforcing gauge conditions that will work with any linear operation / preconditioner.

[glwagner]

For example, the implementation of linear_operation! is generic --- this can be a function or object which computes the result of applying a "linear operation" to an object.

similarly you can have enforce_gauge_condition! which is a function operating on something. Or you can use an object-function interface, eg a property gauge_enforcer, which is then involved when calling

enforce_gauge_condition!(modified_args..., gauge_enforcer::UserDefinedThing, other_args...)

[xkykai]

I have changed enforce_gauge_condition! to take a function: I think this provides the greatest ease for the user to define a gauge condition they wish. The API for it would be

enforce_gauge_condition!(x, r)

Would there be a case where one might need to access more internals than x and r? Else this is ready to merge methinks.

[xkykai]

Is this ready to go? @glwagner @tomchor @simone-silvestri

[tomchor]

Given the importance of the pressure solver, I'm thinking we should add a test for this. @xkykai what do you think?

[xkykai]

Given the importance of the pressure solver, I'm thinking we should add a test for this. @xkykai what do you think?

Sounds great! Should we adopt it from what you've written here https://github.com/CliMA/Oceananigans.jl/blob/tc/perturb-adv-obc-tweaks/test/test_conjugate_gradient_poisson_solver.jl? What else do you wish to see in the test?

[tomchor]

Given the importance of the pressure solver, I'm thinking we should add a test for this. @xkykai what do you think?

Sounds great! Should we adopt it from what you've written here https://github.com/CliMA/Oceananigans.jl/blob/tc/perturb-adv-obc-tweaks/test/test_conjugate_gradient_poisson_solver.jl? What else do you wish to see in the test?

I just tested and those tests will fail here. I'd say let's keep it intuitive and test that the immersed-aware mean pressure is zero at every preconditioning step, which is what this PR is designed to do.

I'll add those other tests after the issue of open boundaries is fixed (probably after #4582?)

[tomchor]

For the record, running

using Oceananigans
using Oceananigans.BoundaryConditions: PerturbationAdvectionOpenBoundaryCondition
using Oceananigans.Solvers: ConjugateGradientPoissonSolver

Lx, Lz = 10, 3
Nx = Nz = 8

grid_base = RectilinearGrid(topology = (Bounded, Flat, Bounded), size = (Nx, Nz), x = (0, Lx), z = (0, Lz))
flat_bottom(x) = 1
grid = ImmersedBoundaryGrid(grid_base, PartialCellBottom(flat_bottom))
U = 1
inflow_timescale = 0.0
outflow_timescale = Inf
u_boundaries = FieldBoundaryConditions(
    west   = PerturbationAdvectionOpenBoundaryCondition(U; inflow_timescale, outflow_timescale),
    east   = PerturbationAdvectionOpenBoundaryCondition(U; inflow_timescale, outflow_timescale),)
    
model = NonhydrostaticModel(; grid,
                              boundary_conditions = (; u = u_boundaries),
                              pressure_solver = ConjugateGradientPoissonSolver(grid, maxiter=10),
                              advection = WENO(; grid, order=5)
                             )
u₀(x, z) = U + 1e-2*rand()
set!(model, u=u₀, enforce_incompressibility=true)

in this branch produces peaks in velocity approaching infinity, indicating that there's still an issue with open boundaries:

julia> model.velocities.u
9×1×8 Field{Face, Center, Center} on ImmersedBoundaryGrid on CPU
├── grid: 8×1×8 ImmersedBoundaryGrid{Float64, Bounded, Flat, Bounded} on CPU with 4×0×4 halo
├── boundary conditions: FieldBoundaryConditions
│   └── west: Open{Oceananigans.BoundaryConditions.PerturbationAdvection{Float64}}, east: Open{Oceananigans.BoundaryConditions.PerturbationAdvection{Float64}}, south: Nothing, north: Nothing, bottom: ZeroFlux, top: ZeroFlux, immersed: ZeroFlux
└── data: 17×1×16 OffsetArray(::Array{Float64, 3}, -3:13, 1:1, -3:12) with eltype Float64 with indices -3:13×1:1×-3:12
    └── max=2.11106e13, min=-2.11106e13, mean=1.00442

[xkykai]

For the record, running

using Oceananigans
using Oceananigans.BoundaryConditions: PerturbationAdvectionOpenBoundaryCondition
using Oceananigans.Solvers: ConjugateGradientPoissonSolver

Lx, Lz = 10, 3
Nx = Nz = 8

grid_base = RectilinearGrid(topology = (Bounded, Flat, Bounded), size = (Nx, Nz), x = (0, Lx), z = (0, Lz))
flat_bottom(x) = 1
grid = ImmersedBoundaryGrid(grid_base, PartialCellBottom(flat_bottom))
U = 1
inflow_timescale = 0.0
outflow_timescale = Inf
u_boundaries = FieldBoundaryConditions(
    west   = PerturbationAdvectionOpenBoundaryCondition(U; inflow_timescale, outflow_timescale),
    east   = PerturbationAdvectionOpenBoundaryCondition(U; inflow_timescale, outflow_timescale),)
    
model = NonhydrostaticModel(; grid,
                              boundary_conditions = (; u = u_boundaries),
                              pressure_solver = ConjugateGradientPoissonSolver(grid, maxiter=10),
                              advection = WENO(; grid, order=5)
                             )
u₀(x, z) = U + 1e-2*rand()
set!(model, u=u₀, enforce_incompressibility=true)

in this branch produces peaks in velocity approaching infinity, indicating that there's still an issue with open boundaries:

julia> model.velocities.u
9×1×8 Field{Face, Center, Center} on ImmersedBoundaryGrid on CPU
├── grid: 8×1×8 ImmersedBoundaryGrid{Float64, Bounded, Flat, Bounded} on CPU with 4×0×4 halo
├── boundary conditions: FieldBoundaryConditions
│   └── west: Open{Oceananigans.BoundaryConditions.PerturbationAdvection{Float64}}, east: Open{Oceananigans.BoundaryConditions.PerturbationAdvection{Float64}}, south: Nothing, north: Nothing, bottom: ZeroFlux, top: ZeroFlux, immersed: ZeroFlux
└── data: 17×1×16 OffsetArray(::Array{Float64, 3}, -3:13, 1:1, -3:12) with eltype Float64 with indices -3:13×1:1×-3:12
    └── max=2.11106e13, min=-2.11106e13, mean=1.00442

Sorry! I just tested this on #4563 (which should be the supposed PR to rule them all) the MWE works fine, with no explosion of velocities:

julia> model.velocities.u
9×1×8 Field{Face, Center, Center} on ImmersedBoundaryGrid on CPU
├── grid: 8×1×8 ImmersedBoundaryGrid{Float64, Bounded, Flat, Bounded} on CPU with 4×0×4 halo
├── boundary conditions: FieldBoundaryConditions
│   └── west: Open{Oceananigans.BoundaryConditions.PerturbationAdvection{Float64}}, east: Open{Oceananigans.BoundaryConditions.PerturbationAdvection{Float64}}, south: Nothing, north: Nothing, bottom: ZeroFlux, top: ZeroFlux, immersed: ZeroFlux
└── data: 17×1×16 OffsetArray(::Array{Float64, 3}, -3:13, 1:1, -3:12) with eltype Float64 with indices -3:13×1:1×-3:12
    └── max=1.00987, min=0.996446, mean=1.00329

[xkykai]

I've added simple tests and I think it is ready to go!