Error in ConjugateGradientPoissonSolver for nonuniform grids

[liuchihl]

I am attempting to useConjugateGradientPoissonSolver in my simulation, but the error specifically occurs when the grids are stretched. The error happens when running either fft_poisson_solver(grid.underlying_grid) or ConjugateGradientPoissonSolver(grid; preconditioner, maxiter=20).

The error message:

ERROR: type RectilinearGrid has no field underlying_grid
Stacktrace:
 [1] getproperty
   @ ./Base.jl:37 [inlined]
 [2] fft_poisson_solver(grid::RectilinearGrid{…})
   @ Oceananigans.Solvers ~/code/Oceananigans.jl/src/Solvers/Solvers.jl:55
 [3] top-level scope
   @ REPL[7]:1

I don't understand why there is no underlying_grid, which clearly exists.
Here is the MWE that I slightly modified from #3831 (comment), thanks to @ali-ramadhan, @glwagner

using Printf
using Statistics
using Oceananigans
using Oceananigans.Grids: with_number_type
using Oceananigans.BoundaryConditions: FlatExtrapolationOpenBoundaryCondition
using Oceananigans.Solvers: ConjugateGradientPoissonSolver, fft_poisson_solver
using Oceananigans.Utils: prettytime

N = 16
h, w = 50, 20
H, L = 100, 100
x = y = (-L/2, L/2)

# Create stretched vertical grid
kwarp(k, N) = (N + 1 - k) / N
# Linear near-surface generator
ζ(k, N, refinement) = 1 + (kwarp(k, N) - 1) / refinement
# Bottom-intensified stretching function
Σ(k, N, stretching) = (1 - exp(-stretching * kwarp(k, N))) / (1 - exp(-stretching))
# Generating function
z_faces(k) = - H * (ζ(k, N, 1.2) * Σ(k, N, 15) - 1)
z = z_faces

# uniform vertical grid
# z = (-H, 0)

grid = RectilinearGrid(size=(N, N, N); x, y, z, halo=(2, 2, 2), topology=(Bounded, Periodic, Bounded))

mount(x, y=0) = h * exp(-(x^2 + y^2) / 2w^2)
bottom(x, y=0) = -H + mount(x, y)
grid = ImmersedBoundaryGrid(grid, GridFittedBottom(bottom))

prescribed_flow = OpenBoundaryCondition(0.01)
extrapolation_bc = FlatExtrapolationOpenBoundaryCondition()
u_bcs = FieldBoundaryConditions(west = prescribed_flow,
                                east = extrapolation_bc)
                                #east = prescribed_flow)

boundary_conditions = (; u=u_bcs)
# reduced_precision_grid = with_number_type(grid.underlying_grid)
preconditioner = fft_poisson_solver(grid.underlying_grid)
pressure_solver = ConjugateGradientPoissonSolver(grid; preconditioner, maxiter=20)

model = NonhydrostaticModel(; grid, boundary_conditions, pressure_solver)
simulation = Simulation(model; Δt=0.1, stop_iteration=1000)
conjure_time_step_wizard!(simulation, cfl=0.5)

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

function progress(sim)
    model = sim.model
    u, v, w = model.velocities
    @printf("Iter: %d, time: %.1f, Δt: %.2e, max|δ|: %.2e",
            iteration(sim), time(sim), sim.Δt, maximum(abs, δ))

    r = model.pressure_solver.conjugate_gradient_solver.residual
    @printf(", solver iterations: %d, max|r|: %.2e\n",
            iteration(model.pressure_solver), maximum(abs, r))
end

add_callback!(simulation, progress)

simulation.output_writers[:fields] =
    JLD2OutputWriter(model, model.velocities; filename="3831.jld2", schedule=IterationInterval(10), overwrite_existing=true)

run!(simulation)

using GLMakie

ds = FieldDataset("3831.jld2")
fig = Figure(size=(1000, 500))

n = Observable(1)
times = ds["u"].times
title = @lift @sprintf("time = %s", prettytime(times[$n]))

Nx, Ny, Nz = size(grid)
j = round(Int, Ny/2)
k = round(Int, Nz/2)
u_surface = @lift view(ds["u"][$n], :, :, k)
u_slice = @lift view(ds["u"][$n], :, j, :)

ax1 = Axis(fig[1, 1]; title = "u (xy)", xlabel="x", ylabel="y")
hm1 = heatmap!(ax1, u_surface, colorrange=(-0.01, 0.01), colormap=:balance)
Colorbar(fig[1, 2], hm1, label="m/s")

ax2 = Axis(fig[1, 3]; title = "u (xz)", xlabel="x", ylabel="z")
hm2 = heatmap!(ax2, u_slice, colorrange=(-0.01, 0.01), colormap=:balance)
Colorbar(fig[1, 4], hm2, label="m/s")

fig[0, :] = Label(fig, title)

record(fig, "3831.mp4", 1:length(times), framerate=10) do i
    n[] = i
end

[glwagner]

Nice find!

@liuchihl please please please try to come up with a short snippet / MWE that isolates the issue. It makes the job of fixing so much easier! I know it is a bit of extra work, but it helps a lot.

I created one for this:

using Oceananigans
using Oceananigans.Solvers: fft_poisson_solver

N = 2
x = y = (0, 1)
z = [0, 0.2, 1]
grid = RectilinearGrid(size=(N, N, N); x, y, z, halo=(2, 2, 2), topology=(Bounded, Periodic, Bounded))
fft_solver = fft_poisson_solver(grid)

[liuchihl]

Thanks for creating this!