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