`ConjugateGradientPoissonSolver` blows up with `PartialCellBottom`

[xkykai]

Using PartialCellBottom causes ConjugateGradientPoissonSolver to blow up. This is independent of the preconditioner as the simulation blows up with or without a preconditioner.

Here's an MWE:

using Oceananigans
using Oceananigans.Models.NonhydrostaticModels: ConjugateGradientPoissonSolver, FFTBasedPoissonSolver, FourierTridiagonalPoissonSolver
using Oceananigans.Models.NonhydrostaticModels: nonhydrostatic_pressure_solver
using Oceananigans.Grids: with_number_type
using Printf

function initial_conditions!(model)
    h = 0.05
    x₀ = 0.5
    y₀ = 0.5
    z₀ = 0.55
    bᵢ(x, y, z) = - exp(-((x - x₀)^2 + (y - y₀)^2 + (z - z₀)^2) / 2h^2)

    set!(model, b=bᵢ)
end

function setup_grid(Nx, Ny, Nz, arch)
    grid = RectilinearGrid(arch, Float64,
                        size = (Nx, Ny, Nz), 
                        halo = (4, 4, 4),
                        x = (0, 1),
                        y = (0, 1),
                        z = (0, 1),
                        topology = (Bounded, Bounded, Bounded))

    slope(x, y) = (5 + tanh(40*(x - 1/6)) + tanh(40*(x - 2/6)) + tanh(40*(x - 3/6)) + tanh(40*(x - 4/6)) + tanh(40*(x - 5/6))) / 20 + 
                  (5 + tanh(40*(y - 1/6)) + tanh(40*(y - 2/6)) + tanh(40*(y - 3/6)) + tanh(40*(y - 4/6)) + tanh(40*(y - 5/6))) / 20

    # grid = ImmersedBoundaryGrid(grid, GridFittedBottom(slope))
    grid = ImmersedBoundaryGrid(grid, PartialCellBottom(slope))
    return grid
end

function setup_model(grid, pressure_solver)
    model = NonhydrostaticModel(; grid, pressure_solver,
                                  advection = WENO(),
                                  coriolis = FPlane(f=0.1),
                                  tracers = :b,
                                  buoyancy = BuoyancyTracer())

    initial_conditions!(model)
    return model
end

function setup_simulation(model, Δt, stop_iteration)
    return Simulation(model, Δt=Δt, stop_iteration=stop_iteration)
end

function setup_simulation(model)
    Δt = 1e-4
    stop_iteration = 10
    simulation = Simulation(model; Δt = Δt, stop_iteration = stop_iteration, minimum_relative_step = 1e-10)
    
    wall_time = Ref(time_ns())

    function progress(sim)
        pressure_solver = sim.model.pressure_solver
    
        if pressure_solver isa ConjugateGradientPoissonSolver
            pressure_iters = iteration(pressure_solver)
        else
            pressure_iters = 0
        end

        msg = @sprintf("iter: %d, time: %s, Δt: %.4f, Poisson iters: %d",
                        iteration(sim), prettytime(time(sim)), sim.Δt, pressure_iters)
    
        elapsed = 1e-9 * (time_ns() - wall_time[])
    
        u, v, w = sim.model.velocities
        d = Field(∂x(u) + ∂y(v) + ∂z(w))
        compute!(d)
    
        msg *= @sprintf(", max u: %6.3e, max w: %6.3e, max b: %6.3e, max d: %6.3e, max pressure: %6.3e, wall time: %s",
                        maximum(sim.model.velocities.u),
                        maximum(sim.model.velocities.w),
                        maximum(sim.model.tracers.b),
                        maximum(d),
                        maximum(sim.model.pressures.pNHS),
                        prettytime(elapsed))
    
        @info msg
        wall_time[] = time_ns()
    
        return nothing
    end

    simulation.callbacks[:progress] = Callback(progress, IterationInterval(1))
    
    return simulation
end

const Nx = 32
const Ny = 32
const Nz = 32

arch = CPU()
grid = setup_grid(Nx, Ny, Nz, arch)

@info "Create pressure solver"
# preconditioner = nonhydrostatic_pressure_solver(with_number_type(Float32, grid.underlying_grid))
# preconditioner = FourierTridiagonalPoissonSolver(grid.underlying_grid)
preconditioner = nothing

reltol = abstol = 1e-7
pressure_solver = ConjugateGradientPoissonSolver(
    grid, maxiter=10000, preconditioner=preconditioner; reltol, abstol)

@info "Create model"
model = setup_model(grid, pressure_solver)

simulation = setup_simulation(model)

run!(simulation)

which gives the output of

[ Info: Initializing simulation...
[ Info: iter: 0, time: 0 seconds, Δt: 0.0001, Poisson iters: 0, max u: 0.000e+00, max w: 0.000e+00, max b: -2.750e-66, max d: 0.000e+00, max pressure: 0.000e+00, wall time: 444.359 ms
[ Info:     ... simulation initialization complete (576.466 ms)
[ Info: Executing initial time step...
[ Info: iter: 1, time: 100 μs, Δt: 0.0001, Poisson iters: 125, max u: 3.783e-05, max w: 2.261e-05, max b: -2.750e-66, max d: 3.888e-03, max pressure: 3.194e-02, wall time: 511.381 ms
[ Info:     ... initial time step complete (803.167 ms).
[ Info: iter: 2, time: 200.000 μs, Δt: 0.0001, Poisson iters: 117, max u: 4.387e-05, max w: 1.381e-05, max b: -2.750e-66, max d: 2.169e-03, max pressure: 2.018e-02, wall time: 550.837 ms
[ Info: iter: 3, time: 300.000 μs, Δt: 0.0001, Poisson iters: 198, max u: 2.443e-04, max w: 1.401e-04, max b: -2.750e-66, max d: 4.469e-02, max pressure: 2.528e-01, wall time: 627.441 ms
[ Info: iter: 4, time: 400.000 μs, Δt: 0.0001, Poisson iters: 866, max u: 1.114e-03, max w: 3.694e-04, max b: -2.750e-66, max d: 1.158e-01, max pressure: 1.645e+00, wall time: 2.619 seconds
[ Info: iter: 5, time: 500.000 μs, Δt: 0.0001, Poisson iters: 420, max u: 4.251e-02, max w: 5.466e-03, max b: -2.750e-66, max d: 1.696e+00, max pressure: 9.528e+00, wall time: 2.286 seconds
[ Info: iter: 6, time: 600 μs, Δt: 0.0001, Poisson iters: 247, max u: 4.238e-02, max w: 1.437e-02, max b: 7.223e-52, max d: 5.950e+01, max pressure: 6.749e+03, wall time: 988.661 ms
[ Info: iter: 7, time: 700 μs, Δt: 0.0001, Poisson iters: 234, max u: 1.741e+02, max w: 2.057e-01, max b: 9.315e-45, max d: 5.466e+03, max pressure: 5.876e+02, wall time: 893.501 ms
[ Info: iter: 8, time: 800.000 μs, Δt: 0.0001, Poisson iters: 205, max u: 2.560e+04, max w: 6.915e+01, max b: 1.092e-39, max d: 2.371e+06, max pressure: 4.060e+07, wall time: 949.581 ms
[ Info: iter: 9, time: 900.000 μs, Δt: 0.0001, Poisson iters: 177, max u: 1.623e+20, max w: 1.931e+18, max b: 1.698e+02, max d: 2.387e+22, max pressure: 2.208e+23, wall time: 750.430 ms
[ Info: Simulation is stopping after running for 13.094 seconds.
[ Info: Model iteration 10 equals or exceeds stop iteration 10.
[ Info: iter: 10, time: 1000.000 μs, Δt: 0.0001, Poisson iters: 158, max u: 4.397e+147, max w: 1.564e+145, max b: 5.441e+108, max d: 5.818e+149, max pressure: 4.489e+150, wall time: 629.644 ms

When PartialCellBottom is replaced with GridFittedBottom then the simulation does not blow up.

[xkykai]

Also tested with #4563 and the same problem occurs. My current suspicion is that there is an issue somewhere within PartialCellBottom since I've also tested nonuniform grids in the z direction with GridFittedBottom with #4563 and main and it doesn't seem to cause trouble.

[glwagner]

Perhaps the laplacian operator is not correct?

[xkykai]

I tried reducing the problem into 2D and even further down so that it is just a flat immersed boundary, where the boundary is in between the two grid points, and I still find the same issue. Here's an MWE:

using Oceananigans
using Oceananigans.Models.NonhydrostaticModels: ConjugateGradientPoissonSolver, FFTBasedPoissonSolver, FourierTridiagonalPoissonSolver
using Oceananigans.Models.NonhydrostaticModels: nonhydrostatic_pressure_solver
using Oceananigans.Grids: with_number_type
using Printf
using CairoMakie
using Oceananigans.Utils: launch!
using KernelAbstractions
using Oceananigans.Architectures: architecture
using Oceananigans.Operators: Vᶜᶜᶜ, Azᶜᶜᶠ, Axᶠᶜᶜ
using Oceananigans.Grids: inactive_cell

function initial_conditions!(model)
    h = 0.05
    x₀ = 0.5
    z₀ = 0.55
    bᵢ(x, z) = - exp(-((x - x₀)^2 + (z - z₀)^2) / 2h^2)

    set!(model, b=bᵢ)
end

function setup_grid(Nx, Nz, arch)
    grid = RectilinearGrid(arch, Float64,
                        size = (Nx, Nz), 
                        halo = (4, 4),
                        x = (0, 1),
                        z = (0, 1),
                        topology = (Bounded, Flat, Bounded))

    slope(x) = 1 / Nz * 10.5

    grid = ImmersedBoundaryGrid(grid, PartialCellBottom(slope))
    return grid
end

function setup_model(grid, pressure_solver)
    model = NonhydrostaticModel(; grid, pressure_solver,
                                  advection = WENO(),
                                  coriolis = FPlane(f=0.1),
                                  tracers = :b,
                                  buoyancy = BuoyancyTracer())

    initial_conditions!(model)
    return model
end

function setup_simulation(model, Δt, stop_iteration)
    return Simulation(model, Δt=Δt, stop_iteration=stop_iteration)
end

function setup_simulation(model)
    Δt = 1e-4
    stop_iteration = 10
    simulation = Simulation(model; Δt = Δt, stop_iteration = stop_iteration, minimum_relative_step = 1e-10)
    
    wall_time = Ref(time_ns())

    function progress(sim)
        pressure_solver = sim.model.pressure_solver
    
        if pressure_solver isa ConjugateGradientPoissonSolver
            pressure_iters = iteration(pressure_solver)
        else
            pressure_iters = 0
        end

        msg = @sprintf("iter: %d, time: %s, Δt: %.4f, Poisson iters: %d",
                        iteration(sim), prettytime(time(sim)), sim.Δt, pressure_iters)
    
        elapsed = 1e-9 * (time_ns() - wall_time[])
    
        u, v, w = sim.model.velocities
        d = Field(∂x(u) + ∂y(v) + ∂z(w))
        compute!(d)
    
        msg *= @sprintf(", max u: %6.3e, max w: %6.3e, max b: %6.3e, max d: %6.3e, max pressure: %6.3e, wall time: %s",
                        maximum(sim.model.velocities.u),
                        maximum(sim.model.velocities.w),
                        maximum(sim.model.tracers.b),
                        maximum(d),
                        maximum(sim.model.pressures.pNHS),
                        prettytime(elapsed))
    
        @info msg
        wall_time[] = time_ns()
    
        return nothing
    end

    simulation.callbacks[:progress] = Callback(progress, IterationInterval(1))
    
    return simulation
end

const Nx = 32
const Nz = 32

arch = CPU()
grid = setup_grid(Nx, Nz, arch)

@info "Create pressure solver"
# preconditioner = nonhydrostatic_pressure_solver(with_number_type(Float32, grid.underlying_grid))
# preconditioner = FourierTridiagonalPoissonSolver(grid.underlying_grid)
preconditioner = nothing

reltol = abstol = 1e-7 / (max(grid.Nx, grid.Ny, grid.Nz))^3
pressure_solver = ConjugateGradientPoissonSolver(
    grid, maxiter=10000, preconditioner=preconditioner; reltol, abstol)

@info "Create model"
model = setup_model(grid, pressure_solver)

simulation = setup_simulation(model)

run!(simulation)

The flow is still divergent in this case and the divergence grows over time. I noted that the divergence is at the immersed boundary:

[ Info: Initializing simulation...
[ Info: iter: 0, time: 0 seconds, Δt: 0.0001, Poisson iters: 0, max u: 0.000e+00, max w: 0.000e+00, max b: -1.708e-37, max d: 0.000e+00, max pressure: 0.000e+00, wall time: 339.482 ms
[ Info:     ... simulation initialization complete (199.696 ms)
[ Info: Executing initial time step...
[ Info: iter: 1, time: 100 μs, Δt: 0.0001, Poisson iters: 67, max u: 2.274e-05, max w: 1.104e-05, max b: -1.708e-37, max d: 5.338e-04, max pressure: 2.260e-02, wall time: 88.841 ms
[ Info:     ... initial time step complete (106.962 ms).
[ Info: iter: 2, time: 200.000 μs, Δt: 0.0001, Poisson iters: 68, max u: 3.315e-05, max w: 2.219e-05, max b: -1.708e-37, max d: 5.844e-04, max pressure: 2.354e-02, wall time: 98.658 ms
[ Info: iter: 3, time: 300.000 μs, Δt: 0.0001, Poisson iters: 68, max u: 4.731e-05, max w: 3.333e-05, max b: -1.708e-37, max d: 5.510e-04, max pressure: 2.371e-02, wall time: 91.828 ms
[ Info: iter: 4, time: 400.000 μs, Δt: 0.0001, Poisson iters: 68, max u: 6.190e-05, max w: 4.448e-05, max b: -1.708e-37, max d: 5.163e-04, max pressure: 2.370e-02, wall time: 93.552 ms
[ Info: iter: 5, time: 500.000 μs, Δt: 0.0001, Poisson iters: 68, max u: 7.650e-05, max w: 5.562e-05, max b: -1.708e-37, max d: 4.817e-04, max pressure: 2.370e-02, wall time: 100.688 ms
[ Info: iter: 6, time: 600 μs, Δt: 0.0001, Poisson iters: 68, max u: 9.109e-05, max w: 6.677e-05, max b: -1.708e-37, max d: 4.471e-04, max pressure: 2.371e-02, wall time: 95.165 ms
[ Info: iter: 7, time: 700 μs, Δt: 0.0001, Poisson iters: 68, max u: 1.057e-04, max w: 7.791e-05, max b: -1.708e-37, max d: 4.125e-04, max pressure: 2.370e-02, wall time: 96.781 ms
[ Info: iter: 8, time: 800.000 μs, Δt: 0.0001, Poisson iters: 68, max u: 1.203e-04, max w: 8.906e-05, max b: -1.708e-37, max d: 3.876e-04, max pressure: 2.371e-02, wall time: 94.556 ms
[ Info: iter: 9, time: 900.000 μs, Δt: 0.0001, Poisson iters: 68, max u: 1.349e-04, max w: 1.002e-04, max b: -1.708e-37, max d: 3.763e-04, max pressure: 2.371e-02, wall time: 93.568 ms
[ Info: Simulation is stopping after running for 1.135 seconds.
[ Info: Model iteration 10 equals or exceeds stop iteration 10.
[ Info: iter: 10, time: 1000.000 μs, Δt: 0.0001, Poisson iters: 68, max u: 1.495e-04, max w: 1.113e-04, max b: -1.708e-37, max d: 3.650e-04, max pressure: 2.371e-02, wall time: 95.135 ms

I believe that the problem is likely coming from some form of spacing/area/volume operators, since if I set the immersed boundary to be exactly at a face cell location by doing

    slope(x) = 1 / Nz * 10.5

one ends up with a well behaved simulation:

[ Info: Initializing simulation...
[ Info: iter: 0, time: 0 seconds, Δt: 0.0001, Poisson iters: 0, max u: 0.000e+00, max w: 0.000e+00, max b: -1.708e-37, max d: 0.000e+00, max pressure: 0.000e+00, wall time: 383.381 ms
[ Info:     ... simulation initialization complete (211.374 ms)
[ Info: Executing initial time step...
[ Info: iter: 1, time: 100 μs, Δt: 0.0001, Poisson iters: 63, max u: 2.092e-05, max w: 1.659e-05, max b: -1.708e-37, max d: 3.091e-10, max pressure: 1.868e-02, wall time: 105.294 ms
[ Info:     ... initial time step complete (113.884 ms).
[ Info: iter: 2, time: 200.000 μs, Δt: 0.0001, Poisson iters: 59, max u: 3.417e-05, max w: 2.218e-05, max b: -1.708e-37, max d: 2.825e-10, max pressure: 1.747e-02, wall time: 90.051 ms
[ Info: iter: 3, time: 300.000 μs, Δt: 0.0001, Poisson iters: 13, max u: 4.790e-05, max w: 3.332e-05, max b: -1.708e-37, max d: 4.226e-10, max pressure: 1.736e-02, wall time: 45.930 ms
[ Info: iter: 4, time: 400.000 μs, Δt: 0.0001, Poisson iters: 4, max u: 6.249e-05, max w: 4.447e-05, max b: -1.708e-37, max d: 2.855e-10, max pressure: 1.736e-02, wall time: 9.603 ms
[ Info: iter: 5, time: 500.000 μs, Δt: 0.0001, Poisson iters: 2, max u: 7.708e-05, max w: 5.561e-05, max b: -1.708e-37, max d: 2.310e-10, max pressure: 1.736e-02, wall time: 6.356 ms
[ Info: iter: 6, time: 600 μs, Δt: 0.0001, Poisson iters: 8, max u: 9.167e-05, max w: 6.676e-05, max b: -1.708e-37, max d: 2.573e-10, max pressure: 1.736e-02, wall time: 11.071 ms
[ Info: iter: 7, time: 700 μs, Δt: 0.0001, Poisson iters: 2, max u: 1.063e-04, max w: 7.790e-05, max b: -1.708e-37, max d: 3.987e-10, max pressure: 1.736e-02, wall time: 9.973 ms
[ Info: iter: 8, time: 800.000 μs, Δt: 0.0001, Poisson iters: 1, max u: 1.209e-04, max w: 8.905e-05, max b: -1.708e-37, max d: 3.270e-10, max pressure: 1.736e-02, wall time: 5.798 ms
[ Info: iter: 9, time: 900.000 μs, Δt: 0.0001, Poisson iters: 2, max u: 1.354e-04, max w: 1.002e-04, max b: -1.708e-37, max d: 2.606e-10, max pressure: 1.736e-02, wall time: 12.333 ms
[ Info: Simulation is stopping after running for 577.120 ms.
[ Info: Model iteration 10 equals or exceeds stop iteration 10.
[ Info: iter: 10, time: 1000.000 μs, Δt: 0.0001, Poisson iters: 1, max u: 1.500e-04, max w: 1.113e-04, max b: -1.708e-37, max d: 3.964e-10, max pressure: 1.736e-02, wall time: 6.917 ms

My suspects were the geometric operators used in the Laplacian computation i.e. Vᶜᶜᶜ(i, j, k, grid), inactive_cell(i, j, k, grid), Azᶜᶜᶠ(i, j, k, grid), Ayᶜᶠᶜ(i, j, k, grid), and Axᶠᶜᶜ(i, j, k, grid). I tested out the kernel functions on my original grid where the immersed boundary is in at a cell center:

V = CenterField(grid)
@kernel function _geometric_operator!(V, grid)
    i, j, k = @index(Global, NTuple)
    active = !inactive_cell(i, j, k, grid)
    @inbounds V[i, j, k] = Vᶜᶜᶜ(i, j, k, grid) * active
    # @inbounds V[i, j, k] = Azᶜᶜᶠ(i, j, k, grid) * active
    # @inbounds V[i, j, k] = Axᶠᶜᶜ(i, j, k, grid) * active
end

function compute_geometric_property!(V, grid)
    arch = architecture(grid)
    launch!(arch, grid, :xyz, _geometric_operator!, V, grid)
    return nothing
end

compute_geometric_property!(V, grid)

Plotting them out, for Vᶜᶜᶜ(i, j, k, grid) I get

For Azᶜᶜᶠ(i, j, k, grid):

For Ayᶜᶠᶜ(i, j, k, grid):

For Axᶠᶜᶜ(i, j, k, grid):

They all appear to be as expected. @glwagner @simone-silvestri do you have suggestions on where else I should look perhaps? If you have time for a quick debugging session that would be wonderful as well! Thanks a lot

[glwagner]

I think there are more operators involved, since the definition of the Laplacian is

Oceananigans.jl/src/Operators/laplacian_operators.jl

Lines 36 to 40 in baad8a3

	@inline function ∇²ᶜᶜᶜ(i, j, k, grid, c)
	return V⁻¹ᶜᶜᶜ(i, j, k, grid) * (δxᶜᵃᵃ(i, j, k, grid, Ax_∂xᶠᶜᶜ, c) +
	δyᵃᶜᵃ(i, j, k, grid, Ay_∂yᶜᶠᶜ, c) +
	δzᵃᵃᶜ(i, j, k, grid, Az_∂zᶜᶜᶠ, c))
	end

Note that its the reciprocal volume rather than the volume which appears directly... also the reciprocal spacings appear in the derivative operators, eg Δy⁻¹ᶜᶠᶜ, etc.