Interface Krylov.jl

[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.

[amontoison]

@xkykai @xkykai @michel2323
Your DiagonallyDominantPreconditioner is not SPD, so you should not use it in a symmetric Krylov solver.
A symmetric Krylov solver requires a SPD Krylov solver, otherwise everything is wrong in terms of linear algebra.
The preconditioner should definite an elliptic norm such that we always have x'Px >= 0.
In this VERY specific case, CG can eat it but don't do that 🥲
I added an absolute value in heuristic_residual and it fixed the issue.

For the FFTBasedPoissonSolver used as a preconditioner, we need some modification to ensure it is SPD by construction.

I also found a way to provide additional arguments to the operator A and the preconditioner P.
@xkykai Can probably try bicgstab + FFTBasedPoissonSolver in a problem that requires more than one iteration to see it performs well please?

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)

method = :cg
# methods = :bicgstab  # needed for `FFTBasedPoissonSolver`

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)"

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))"

[xkykai]

Your DiagonallyDominantPreconditioner is not SPD, so you should not use it in a symmetric Krylov solver.
A symmetric Krylov solver requires a SPD Krylov solver, otherwise everything is wrong in terms of linear algebra.
The preconditioner should definite an elliptic norm such that we always have x'Px >= 0.
In this VERY specific case, CG can eat it but don't do that 🥲

ah yes 😅 I suppose as a general question for you, do you know how well do non-SPD preconditioners work for CG in practice? I understand that the theoretical guarantees are lost, in your experience does it cause a lot of issues in many cases?

Also as an aside I tried solving your many gaussians test case with FFT preconditioner and :cg in the Krylov solve and it throws ERROR: DomainError with -1.844236867614985: from the line
https://github.com/JuliaSmoothOptimizers/Krylov.jl/blob/cabd85e9ef5d51b63672e06f897dece01246633a/src/cg.jl#L146

When I switched :cg out to :gmres or :fgmres it throws

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}}}

related to this line: https://github.com/JuliaSmoothOptimizers/Krylov.jl/blob/cabd85e9ef5d51b63672e06f897dece01246633a/src/gmres.jl#L218

I also found a way to provide additional arguments to the operator A and the preconditioner P.
@xkykai Can probably try bicgstab + FFTBasedPoissonSolver in a problem that requires more than one iteration to see it performs well please?

I am not 100% sure what you mean here. Did you mean using two preconditioners? How can I write that in the current KrylovSolver API?

[amontoison]

Your DiagonallyDominantPreconditioner is not SPD, so you should not use it in a symmetric Krylov solver.
A symmetric Krylov solver requires a SPD Krylov solver, otherwise everything is wrong in terms of linear algebra.
The preconditioner should definite an elliptic norm such that we always have x'Px >= 0.
In this VERY specific case, CG can eat it but don't do that 🥲

ah yes 😅 I suppose as a general question for you, do you know how well do non-SPD preconditioners work for CG in practice? I understand that the theoretical guarantees are lost, in your experience does it cause a lot of issues in many cases?

You can't use non-SPD or non-SND preconditioners in CG.
Nothing is valid in the linear algebra if you do that.
Maybe we can make something valid for indifinite diagonal matrices but it's a corner case.

Also as an aside I tried solving your many gaussians test case with FFT preconditioner and :cg in the Krylov solve and it throws ERROR: DomainError with -1.844236867614985: from the line https://github.com/JuliaSmoothOptimizers/Krylov.jl/blob/cabd85e9ef5d51b63672e06f897dece01246633a/src/cg.jl#L146

Yes, it is because the preconditioner is not SPD.

When I switched :cg out to :gmres or :fgmres it throws

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}}}

I asked for :bicgstab and not the other ones 😃
gmres and fgmres should work once I released Krylov v0.10, I worked on it last weekend for the support here.
I added a few missing routines:
JuliaSmoothOptimizers/Krylov.jl#980

related to this line: https://github.com/JuliaSmoothOptimizers/Krylov.jl/blob/cabd85e9ef5d51b63672e06f897dece01246633a/src/gmres.jl#L218

I also found a way to provide additional arguments to the operator A and the preconditioner P.

I mean that I fixed the issue with the additional args... that could be passed to the operator / preconditioner from the function solve!.

@xkykai Can probably try bicgstab + FFTBasedPoissonSolver in a problem that requires more than one iteration to see it performs well please?

I am not 100% sure what you mean here. Did you mean using two preconditioners? How can I write that in the current KrylovSolver API?

I just wanted that you try a problem with method=:bicgstab in the KrylovSolver + preconditioner based on an FFTSolver on a relevant problem where you can't apply the FFTSolver directly.