Interface to Krylov.jl solvers

[amontoison]

Replace #3812
@glwagner @michel2323

I added all the routines needed by Krylov, but we probably want them implemented as kernels.

What I don't know is how we can overload the operators such that when we want to perform an operator-vector product with a KrylovField, we unwrap the Field in the custom vector type.
I think Gregory has an idea how we can do that.

The second question is how can we create a new KrylovField from an existing one?
Can we create a function KrylovField{T,F}(undef, n)?
I think we can't because Fields are like 3D arrays, and Krylov.jl should be modified to rely on a function like similar when a workspace is created.

[glwagner]

Most of the functions are already defined so I don't think we need to rewrite the kernels right? We just need a couple 1-liners I think.

Also let's try not to overconstrain the types since this may create issues esp wrt autodifferentiation

[amontoison]

@glwagner
Do you have an implementation of dot, norm, scal!, axpy!, axpby! and ref! with a different name in Oceananigans.jl?
I think we need a kernel for these routines.

[glwagner]

@glwagner Do you have an implementation of dot, norm, scal!, axpy!, axpby! and ref! with a different name in Oceananigans.jl? I think we need a kernel for these routines.

We have dot and norm. I think the others can be simple broadcasting and we can work on kernels once we have things working. Kernels may not be any faster than broadcasting anyways.

[amontoison]

Current version tested in Krylovable.jl.
--> glwagner/Krylovable.jl#1

[glwagner]

@amontoison here is a relatively complex example for you:

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.Simulations: NaNChecker
using Oceananigans.Utils: prettytime

N = 16
h, w = 50, 20
H, L = 100, 100
x = y = (-L/2, L/2)
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(Float32, grid.underlying_grid)
preconditioner = fft_poisson_solver(reduced_precision_grid)
pressure_solver = ConjugateGradientPoissonSolver(grid; preconditioner, maxiter=1000, regularization=1/N^3)

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

I'll post some simpler examples for you too.

[glwagner]

Here is a simple case:

using Printf
using Statistics
using Oceananigans
using Oceananigans.Grids: with_number_type
using Oceananigans.Solvers: ConjugateGradientPoissonSolver, fft_poisson_solver

N = 16
H, L = 100, 100
x = y = (-L/2, L/2)
z = (-H, 0)

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

reduced_precision_grid = with_number_type(Float32, grid.underlying_grid)
preconditioner = fft_poisson_solver(reduced_precision_grid)
pressure_solver = ConjugateGradientPoissonSolver(grid; preconditioner, maxiter=1000, regularization=1/N^3)

model = NonhydrostaticModel(; grid, pressure_solver)

ui(x, y, z) = randn()
set!(model, u=ui, v=ui, w=ui)

simulation = Simulation(model; Δt=0.1, stop_iteration=100)
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="test.jld2", schedule=IterationInterval(10), overwrite_existing=true)

run!(simulation)

I can also help with visualization the results or you can try to use a similar code as above.

[glwagner]

@xkykai you may be interested in following / helping with this work; @amontoison is replacing our in-house PCG implementation with something from Krylov, and he may be able to help with the problems we are currently having with the Poisson solver for certain problems on ImmersedBoundaryGrid

[xkykai]

Yes, I've been revisiting the PCG solver, and had encountered the solver blowing up in the middle of a simulation of flow over periodic hill. I've tried for a bit making it into a MWE but the problem was not showing up in simpler setups. This is also in a non open boundary condition setup so it is potentially separate from the issues brought up by others recently. My preliminary suspicion is that it could be a variable timestep issue and the model blows up when it is taking an almost machine precision timestep. Also @amontoison and @glwagner happy to help with whatever, just let me know!

[amontoison]

@xkykai I think you can help us to do a KrylovPoissonSolver such that we can subtitute ConjugateGradientPoissonSolver by this one.

I tested it in another repository:
glwagner/Krylovable.jl#1

But I know less Oceananigans.jl than you and you can help for sure with the integration of this new solver.

[xkykai]

I will give it a go. Where should I work on this? Krylovable.jl?

[amontoison]

I think Krylovable.jl is a nice place to prototype the new solver.

[glwagner]

Work on it here in this PR right?

[amontoison]

Ok, let's do that 👍
The prototype worked in Krylovable.jl.

[francispoulin]

Any chance we will be able to use this to find eigenvalues of a matrix?

If yes I have an idea for Kelvin Helmholtz instability example.

If not, no problem.

[amontoison]

You can do it with a Rayleigh quotient method, which requires only a few lines of code if we have the Krylov solver.
-> https://en.wikipedia.org/wiki/Rayleigh_quotient_iteration

[michel2323]

Is there a showstopper for this PR @glwagner ?

[glwagner]

No, I think this is great! It'd be nice to have a few tests before merging.

I also think it would be nice to see some benchmarks that demonstrate the superiority of Krylov here; we want to do this work anyways, and so its well-motivated to report that. It might stir up some excitement. Moreover (but it doesn't have to be done immediately) we should try to use a Krylov solver with FFT-based preconditioner for a nonhydrostatic problem in a complex domain. There are a few issues around with some tricky test problems that I think have stymied the current PCG solver. @xkykai can probably advise on this.

A final consideration is whether Krylov integration would be better implemented in an extension. I want to raise this because it is clearly feasible, but also say that my opinion is probably that it should not be an extension. The argument for not doing an extension is: if Krylov is superior, I think we can motivate eliminating the other iterative solvers completely. In that case then, we want krylov solvers in src.

[xkykai]

Forgive me if I am using this incorrectly but when I try to use methods other than :cg it doesn't work? MWE is like

using Oceananigans
using Oceananigans.BoundaryConditions: fill_halo_regions!
using Oceananigans.Solvers: FFTBasedPoissonSolver, solve!
using Oceananigans.Solvers: ConjugateGradientPoissonSolver, DiagonallyDominantPreconditioner, compute_laplacian!
using Oceananigans.Solvers: KrylovSolver
using LinearAlgebra
using Statistics
using Random
Random.seed!(123)

arch = CPU()
Nx = Ny = Nz = 64
topology = (Periodic, Periodic, Periodic)

x = y = z = (0, 2π)

grid = RectilinearGrid(arch; x, y, z, topology, size=(Nx, Ny, Nz))

δ = 0.1 # Gaussian width
gaussian(ξ, η, ζ) = exp(-(ξ^2 + η^2 + ζ^2) / 2δ^2)
Ngaussians = 17
Ξs = [2π * rand(3) for _ = 1:Ngaussians]

function many_gaussians(ξ, η, ζ)
    val = zero(ξ)
    for Ξ₀ in Ξs
        ξ₀, η₀, ζ₀ = Ξ₀
        val += gaussian(ξ - ξ₀, η - η₀, ζ - ζ₀)
    end
    return val
end

# Note: we use in-place transforms, so the RHS has to be AbstractArray{Complex{T}}.
# So, we first fill up "b" and then copy it into "bc = fft_solver.storage",
# which has the correct type.
b = CenterField(grid)
set!(b, many_gaussians)
parent(b) .-= mean(interior(b))

xpcg = CenterField(grid)

reltol = abstol = 1e-7

# preconditioner = nothing
# preconditioner = DiagonallyDominantPreconditioner()
preconditioner = FFTBasedPoissonSolver(grid)

preconditioner_str = preconditioner === nothing ? "nothing" : preconditioner isa FFTBasedPoissonSolver ? "FFT" : "DiagonallyDominant"

pcg_solver = ConjugateGradientPoissonSolver(grid, maxiter=1000; reltol, abstol, preconditioner)

# run it once for JIT
solve!(xpcg, pcg_solver.conjugate_gradient_solver, b)
xpcg .= 0
t_pcg = @timed solve!(xpcg, pcg_solver.conjugate_gradient_solver, b)
@info "PCG iteration $(pcg_solver.conjugate_gradient_solver.iteration), time $(t_pcg.time), preconditioner $(preconditioner_str)"

method = :fgmres
krylov_solver = KrylovSolver(compute_laplacian!; template_field=b, reltol, abstol, preconditioner, method)
xpcg_krylov = CenterField(grid)

# run it once for JIT
solve!(xpcg_krylov, krylov_solver, b)
xpcg_krylov .= 0
t_krylov = @timed solve!(xpcg_krylov, krylov_solver, b)
@info "Krylov.jl iteration $(krylov_solver.workspace.stats.niter), time $(t_krylov.time), preconditioner $(preconditioner_str)"

@info "Solution difference: $(norm(xpcg .- xpcg_krylov))"

eg. here I use :fgmres and get error of

ERROR: CanonicalIndexError: setindex! not defined for Oceananigans.Solvers.KrylovField{Float64, Field{Center, Center, Center, Nothing, RectilinearGrid{Float64, Periodic, Periodic, Periodic, Oceananigans.Grids.StaticVerticalDiscretization{OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, Float64, Float64}, Float64, Float64, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, 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}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Flux, Nothing}}, Nothing, Oceananigans.Fields.FieldBoundaryBuffers{Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}}}
Stacktrace:
  [1] error_if_canonical_setindex(::IndexCartesian, A::Oceananigans.Solvers.KrylovField{Float64, Field{…}}, ::Int64)
    @ Base .\abstractarray.jl:1423
  [2] setindex!
    @ .\abstractarray.jl:1412 [inlined]
  [3] _setindex!
    @ .\abstractarray.jl:1448 [inlined]
  [4] setindex!
    @ .\abstractarray.jl:1413 [inlined]
  [5] macro expansion
    @ .\broadcast.jl:973 [inlined]
  [6] macro expansion
    @ .\simdloop.jl:77 [inlined]
  [7] copyto!
    @ .\broadcast.jl:972 [inlined]
  [8] copyto!
    @ .\broadcast.jl:925 [inlined]
  [9] materialize!
    @ .\broadcast.jl:883 [inlined]
 [10] materialize!
    @ .\broadcast.jl:880 [inlined]
 [11] kdivcopy!(n::Int64, y::Oceananigans.Solvers.KrylovField{…}, x::Oceananigans.Solvers.KrylovField{…}, s::Float64)
    @ Krylov C:\Users\xinle\.julia\packages\Krylov\3JTOh\src\krylov_utils.jl:334
 [12] fgmres!(solver::Krylov.FgmresSolver{…}, A::Oceananigans.Solvers.KrylovOperator{…}, b::Oceananigans.Solvers.KrylovField{…}; M::Oceananigans.Solvers.KrylovPreconditioner{…}, N::UniformScaling{…}, ldiv::Bool, restart::Bool, reorthogonalization::Bool, atol::Float64, rtol::Float64, itmax::Int64, timemax::Float64, verbose::Int64, history::Bool, callback::Krylov.var"#1032#1040", iostream::Core.CoreSTDOUT)
    @ Krylov C:\Users\xinle\.julia\packages\Krylov\3JTOh\src\fgmres.jl:224
 [13] fgmres!
    @ C:\Users\xinle\.julia\packages\Krylov\3JTOh\src\fgmres.jl:115 [inlined]
 [14] solve!
    @ C:\Users\xinle\.julia\packages\Krylov\3JTOh\src\krylov_solve.jl:159 [inlined]
 [15] #solve!#54
    @ c:\Users\xinle\MIT\Krylov_solver\Oceananigans.jl\src\Solvers\krylov_solver.jl:145 [inlined]
 [16] solve!(::Field{…}, ::KrylovSolver{…}, ::Field{…})
    @ Oceananigans.Solvers c:\Users\xinle\MIT\Krylov_solver\Oceananigans.jl\src\Solvers\krylov_solver.jl:142
 [17] top-level scope
    @ c:\Users\xinle\MIT\Krylov_solver\Oceananigans.jl\krylov_preconditioner_mwe.jl:63
Some type information was truncated. Use `show(err)` to see complete types.

I've tried adding this Krylov branch JuliaSmoothOptimizers/Krylov.jl@cabd85e but it doesn't seem to solve the issue.

Ignore me if this is not the case, I might have gotten confused about the right version of packages to use due to the merge conflict on Project.toml.

[amontoison]

@xkykai You will need the release 0.10 of Krylov.jl for some solvers.
I will release it and quite soon.
:cg and :bicgstab work for sure now.

[Yixiao-Zhang]

It might be off-topic, but I want to show how I treat the non-SPD matrices.

I think that shifting the mean value of r (instead of p) to zero at each iteration enables CG work for non-SPD matrices. I include the following patch in all my simulations that use the immersed boundary condition, and so far they been running smoothly without any numerical issues.

Without this patch, due to round-off errors, the mean value of r will drift away from zero at each iteration, and then Ax = r will not have any solution.

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

[amontoison]

@xkykai Do you just solve (A + a I) x = b instead of Ax = b where a = mean_r ?
I can add an option an Krylov.jl to directly provide this shift. It is quite common is symmetric Krylov solvers.
But you should always ensure positive definiteness of your linear operator A.
Otherwise you minimize the quadratic x'Ax + b'x that is unbounded below with CG...

[tomchor]

I'm a bit lost on the overall purpose of this PR. Are Krylov solvers going to be used for iterative methods when solving for pressure?

[amontoison]

More comments or ready to go?

[glwagner]

Good to go for me

[amontoison]

I can't merge this PR because of the failing build, but anyone that have the rights to do, please go ahead.

[francispoulin]

Hello Alexis (@amontoison), now that this has merged, any chance we could chat about how to use this library? My email is [email protected]. Thanks in advance!

[amontoison]

Hello @francispoulin! I will traveling for conferences the next four weeks but send me an email on my Canadian address [email protected].
I will find a moment to discuss!