Error related to conjugate gradient Poisson solver when setting initial condition

[liuchihl]

I've received an error possibly related to the preconditioner, and it appears when I set the initial condition, e.g., set!(model, u=uᵢ):

ERROR: MethodError: no method matching cpu_fourier_tridiagonal_preconditioner_rhs!(::KernelAbstractions.CompilerMetadata{…}, ::Array{…}, ::Oceananigans.Grids.ZDirection, ::Field{…})

Closest candidates are:
  cpu_fourier_tridiagonal_preconditioner_rhs!(::Any, ::Any, ::Oceananigans.Grids.ZDirection, ::Any, ::Any)
   @ Oceananigans none:0
  cpu_fourier_tridiagonal_preconditioner_rhs!(::Any, ::Any, ::Oceananigans.Grids.YDirection, ::Any, ::Any)
   @ Oceananigans none:0
  cpu_fourier_tridiagonal_preconditioner_rhs!(::Any, ::Any, ::Oceananigans.Grids.XDirection, ::Any, ::Any)
   @ Oceananigans none:0

Stacktrace:
  [1] __thread_run(tid::Int64, len::Int64, rem::Int64, obj::KernelAbstractions.Kernel{…}, ndrange::Nothing, iterspace::KernelAbstractions.NDIteration.NDRange{…}, args::Tuple{…}, dynamic::KernelAbstractions.NDIteration.DynamicCheck)
    @ KernelAbstractions ~/.julia/packages/KernelAbstractions/491pi/src/cpu.jl:144
  [2] __run(obj::KernelAbstractions.Kernel{…}, ndrange::Nothing, iterspace::KernelAbstractions.NDIteration.NDRange{…}, args::Tuple{…}, dynamic::KernelAbstractions.NDIteration.DynamicCheck, static_threads::Bool)
    @ KernelAbstractions ~/.julia/packages/KernelAbstractions/491pi/src/cpu.jl:111
  [3] (::KernelAbstractions.Kernel{…})(::Array{…}, ::Vararg{…}; ndrange::Nothing, workgroupsize::Nothing)
    @ KernelAbstractions ~/.julia/packages/KernelAbstractions/491pi/src/cpu.jl:46
  [4] (::KernelAbstractions.Kernel{…})(::Array{…}, ::Vararg{…})
    @ KernelAbstractions ~/.julia/packages/KernelAbstractions/491pi/src/cpu.jl:39
  [5] #_launch!#12
    @ ~/.julia/packages/Oceananigans/HPOLD/src/Utils/kernel_launching.jl:286 [inlined]
  [6] _launch!
    @ ~/.julia/packages/Oceananigans/HPOLD/src/Utils/kernel_launching.jl:268 [inlined]
  [7] launch!
    @ ~/.julia/packages/Oceananigans/HPOLD/src/Utils/kernel_launching.jl:251 [inlined]
  [8] compute_preconditioner_rhs!(solver::Oceananigans.Solvers.FourierTridiagonalPoissonSolver{…}, rhs::Field{…})
    @ Oceananigans.Solvers ~/.julia/packages/Oceananigans/HPOLD/src/Solvers/conjugate_gradient_poisson_solver.jl:109
  [9] precondition!(p::Field{…}, preconditioner::Oceananigans.Solvers.FourierTridiagonalPoissonSolver{…}, r::Field{…}, args::Field{…})
    @ Oceananigans.Solvers ~/.julia/packages/Oceananigans/HPOLD/src/Solvers/conjugate_gradient_poisson_solver.jl:117
 [10] #construct_regionally#62
    @ ~/.julia/packages/Oceananigans/HPOLD/src/Utils/multi_region_transformation.jl:148 [inlined]
 [11] construct_regionally
    @ ~/.julia/packages/Oceananigans/HPOLD/src/Utils/multi_region_transformation.jl:143 [inlined]
 [12] macro expansion
    @ ~/.julia/packages/Oceananigans/HPOLD/src/Utils/multi_region_transformation.jl:221 [inlined]
 [13] iterate!(::Field{…}, ::Oceananigans.Solvers.ConjugateGradientSolver{…}, ::Field{…})
    @ Oceananigans.Solvers ~/.julia/packages/Oceananigans/HPOLD/src/Solvers/conjugate_gradient_solver.jl:192
 [14] solve!(::Field{…}, ::Oceananigans.Solvers.ConjugateGradientSolver{…}, ::Field{…})
    @ Oceananigans.Solvers ~/.julia/packages/Oceananigans/HPOLD/src/Solvers/conjugate_gradient_solver.jl:177
 [15] solve_for_pressure!(pressure::Field{…}, solver::ConjugateGradientPoissonSolver{…}, Δt::Float64, Ũ::@NamedTuple{…})
    @ Oceananigans.Models.NonhydrostaticModels ~/.julia/packages/Oceananigans/HPOLD/src/Models/NonhydrostaticModels/solve_for_pressure.jl:89
 [16] calculate_pressure_correction!(model::NonhydrostaticModel{…}, Δt::Float64)
    @ Oceananigans.Models.NonhydrostaticModels ~/.julia/packages/Oceananigans/HPOLD/src/Models/NonhydrostaticModels/pressure_correction.jl:15
 [17] set!(model::NonhydrostaticModel{…}; enforce_incompressibility::Bool, kwargs::@Kwargs{…})
    @ Oceananigans.Models.NonhydrostaticModels ~/.julia/packages/Oceananigans/HPOLD/src/Models/NonhydrostaticModels/set_nonhydrostatic_model.jl:54
 [18] top-level scope
    @ ~/code/internal-tide-mixing/test_MWE_conjugate_gradient.jl:28
Some type information was truncated. Use `show(err)` to see complete types.

Here's a minimal working example (I tried to make it as minimal as possible, apologies if it's not minimal enough). I set the bottom topography to zero everywhere except for one grid cell, which is set to 0.2. Notably, when I change this value from 0.2 to 0.01, the error no longer appears. It is not clear to me what the issue might be.

using Oceananigans
using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid, GridFittedBoundary
using Oceananigans.Solvers: ConjugateGradientPoissonSolver, fft_poisson_solver

underlying_grid = RectilinearGrid(size=(4, 4, 4), 
       x = (0, 1),
       y = (0, 1), 
       z = [0, 0.2,0.4,0.6, 2],
       halo = (2,2,2),
       topology = (Oceananigans.Periodic, Oceananigans.Periodic, Oceananigans.Bounded)
)
    bottom =   [0.0  0.0  0.0   0.0;
                0.0  0.0  0.0   0.0;
                0.0  0.0  0.2   0.0;
                0.0  0.0  0.0   0.0]

    grid = ImmersedBoundaryGrid(underlying_grid, GridFittedBottom(bottom))

    uᵢ(x, y, z) = 0.1

    model = NonhydrostaticModel(;
        grid=grid,
        pressure_solver = ConjugateGradientPoissonSolver(
                        grid; preconditioner = fft_poisson_solver(underlying_grid)),
        advection = WENO(),
        timestepper = :RungeKutta3,
    )
   set!(model, u=uᵢ)

[glwagner]

That's a decent minimal example! Are you sure that the error requires advection=WENO() and timestepper=:RungeKutta3? The latter cannot be necessary since its the default (so omitting it has the same effect as including it).

I find I can reproduce the error without ImmersedBoundaryGrid at all.

About the error. The top of the message says

ERROR: MethodError: no method matching cpu_fourier_tridiagonal_preconditioner_rhs!

This means that the kernel function fourier_tridiagonal_preconditioner_rhs! is being called with the wrong arguments. For example:

julia> f(x, y) = x + y
f (generic function with 1 method)

julia> f(1)
ERROR: MethodError: no method matching f(::Int64)

Closest candidates are:
  f(::Any, ::Any)
   @ Main REPL[1]:1

Stacktrace:
 [1] top-level scope
   @ REPL[2]:1

The stacktrace shows

  [8] compute_preconditioner_rhs!(solver::Oceananigans.Solvers.FourierTridiagonalPoissonSolver{…}, rhs::Field{…})
    @ Oceananigans.Solvers ~/.julia/packages/Oceananigans/HPOLD/src/Solvers/conjugate_gradient_poisson_solver.jl:109

let's look at that line:

Oceananigans.jl/src/Solvers/conjugate_gradient_poisson_solver.jl

Lines 109 to 110 in 6c40d7e

	launch!(arch, grid, :xyz, fourier_tridiagonal_preconditioner_rhs!,
	solver.storage, tridiagonal_dir, rhs)

This uses the Oceananigans utility launch! which launches the kernel fourier_tridiagonal_preconditioner_rhs! with the arguments solver.storage, tridiagonal_dir, rhs. However, looking at the function fourier_tridiagonal_preconditioner_rhs a few lines above

Oceananigans.jl/src/Solvers/conjugate_gradient_poisson_solver.jl

Line 93 in 6c40d7e

@kernel function fourier_tridiagonal_preconditioner_rhs!(preconditioner_rhs, ::ZDirection, grid, rhs)

we see that the function has 4 arguments, not 3. Hence the error.

To summarize the analysis method, the key is to find the function that causes the error in the source code (fourier_tridiagonal_preconditioner_rhs) and then identify where it is called, and how it should be called.

Here's an updated MWE from your nice one @liuchihl :

using Oceananigans
using Oceananigans.Solvers: ConjugateGradientPoissonSolver, 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)
pressure_solver = ConjugateGradientPoissonSolver(grid, preconditioner=fft_poisson_solver(grid))
model = NonhydrostaticModel(; grid, pressure_solver)
set!(model, u=1)

[glwagner]

Using that MWE I fixed a number of bugs in this commit: aa0a43e

and the MWE now works.

[liuchihl]

That's a decent minimal example! Are you sure that the error requires advection=WENO() and timestepper=:RungeKutta3? The latter cannot be necessary since its the default (so omitting it has the same effect as including it).

I find I can reproduce the error without ImmersedBoundaryGrid at all.

You are absolutely right! Those are not needed.

Thank you so much for the detailed explanation, I get it now :)