Test and benchmark pressure solver

Project.toml

@@ -9,23 +9,31 @@ EllipsisNotation = "da5c29d0-fa7d-589e-88eb-ea29b0a81949"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+LoggingExtras = "e6f89c97-d47a-5376-807f-9c37f3926c36"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
+SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 TOML = "fa267f1f-6049-4f14-aa54-33bafae1ed76"
 
 [weakdeps]
 AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+HYPRE = "b5ffcf37-a2bd-41ab-a3da-4bd9bc8ad771"
 ReadVTK = "dc215faf-f008-4882-a9f7-a79a826fadc3"
 WriteVTK = "64499a7a-5c06-52f2-abe2-ccb03c286192"
+GeometricMultigrid = "ec3bce20-3e4b-4d8a-a801-89defc82be1c"
+Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 
 [extensions]
 WaterLilyAMDGPUExt = "AMDGPU"
 WaterLilyCUDAExt = "CUDA"
+WaterLilyHYPREExt = "HYPRE"
 WaterLilyReadVTKExt = "ReadVTK"
 WaterLilyWriteVTKExt = "WriteVTK"
+WaterLilyGeometricMultigridExt = "GeometricMultigrid"
+WaterLilyPlotsExt = "Plots"
 
 [compat]
 DocStringExtensions = "0.9"
@@ -43,6 +51,7 @@ BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 GLMakie = "e9467ef8-e4e7-5192-8a1a-b1aee30e663a"
 GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
+HYPRE = "b5ffcf37-a2bd-41ab-a3da-4bd9bc8ad771"
 LaTeXStrings = "b964fa9f-0449-5b57-a5c2-d3ea65f4040f"
 Makie = "537997a7-5e4e-5d89-9595-2241ea00577e"
 PerformanceTestTools = "dc46b164-d16f-48ec-a853-60448fc869fe"
@@ -51,6 +60,8 @@ ReadVTK = "dc215faf-f008-4882-a9f7-a79a826fadc3"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 UnicodePlots = "b8865327-cd53-5732-bb35-84acbb429228"
 WriteVTK = "64499a7a-5c06-52f2-abe2-ccb03c286192"
+GeometricMultigrid = "ec3bce20-3e4b-4d8a-a801-89defc82be1c"
 
 [targets]
-test = ["Test", "BenchmarkTools", "CUDA", "AMDGPU", "GPUArrays", "WriteVTK", "ReadVTK"]
+test = ["Test", "BenchmarkTools", "CUDA", "AMDGPU", "GPUArrays", "WriteVTK",
+        "ReadVTK", "HYPRE", "GeometricMultigrid"]

ext/WaterLilyGeometricMultigridExt.jl

@@ -0,0 +1,37 @@
+module WaterLilyGeometricMultigridExt
+
+if isdefined(Base, :get_extension)
+    using GeometricMultigrid
+else
+    using ..GeometricMultigrid
+end
+
+using WaterLily
+import WaterLily: AbstractPoisson,GeomMultigridPoisson,solver!,update!,L₂,L∞
+
+struct GMGPoisson{T,Vf<:AbstractArray{T},Mf<:AbstractArray{T},St<:SolveState} <: AbstractPoisson{T,Vf,Mf}
+    x::Vf  # WaterLily approximate solution
+    L::Mf  # WaterLily lower diagonal coefficients
+    z::Vf  # WaterLily source
+    st::St
+    function GMGPoisson(x::AbstractArray{T},L::AbstractArray{T},z::AbstractArray{T};perdir=()) where T
+        st = SolveState(GeometricMultigrid.Poisson(L),FieldVector(x),FieldVector(z))
+        new{T,typeof(x),typeof(L),typeof(st)}(x,L,z,GeometricMultigrid.fill_children!(st))
+    end
+end
+GeomMultigridPoisson(x,L,z;perdir=()) = GMGPoisson(x,L,z;perdir=perdir)
+export GeomMultigridPoisson
+
+update!(p::GMGPoisson) = for I ∈ p.st.A.R
+    p.st.A.D[I] = GeometricMultigrid.calcdiag(I,p.st.A.L)
+end
+
+function solver!(p::GMGPoisson;tol=1e-4,itmx=32,kw...)
+    GeometricMultigrid.resid!(p.st.r,p.st.A,p.st.x);
+    GeometricMultigrid.iterate!(p.st,GeometricMultigrid.Vcycle!;mxiter=Int(itmx),abstol=tol,kw...)
+end
+
+L₂(p::GMGPoisson) = WaterLily.L₂(p.z)
+L∞(p::GMGPoisson) = maximum(abs,p.z)
+
+end #module

ext/WaterLilyHYPREExt.jl

@@ -0,0 +1,99 @@
+module WaterLilyHYPREExt
+
+if isdefined(Base, :get_extension)
+    using HYPRE
+else
+    using ..HYPRE
+end
+
+using SparseArrays
+using WaterLily
+import WaterLily: inside,size_u,@loop
+import WaterLily: AbstractPoisson,solver!,update!,HyprePoisson,putback!,L₂,L∞
+
+# @inline insidep(a::AbstractArray,perdir=()) = CartesianIndices(ntuple( i-> i ∈ perdir ? (1:size(a,i)) : (2:size(a,i)-1), length(axes(a))))
+
+# @fastmath @inline function filldiag(I::CartesianIndex{d},x,L,perdir) where {d}
+#     s = zero(eltype(L))
+#     for i in 1:d
+#         Ip = I+δ(i,I)
+#         (i ∈ perdir && Ip ∉ insidep(x,perdir)) && (Ip = WaterLily.CIj(i,Ip,1))
+#         s -= @inbounds(L[I,i]+L[Ip,i])
+#     end
+#     return s
+# end
+
+struct HYPREPoisson{T,V<:AbstractVector{T},M<:AbstractArray{T},
+                    Vf<:AbstractArray{T},Mf<:AbstractArray{T},
+                    SO<:HYPRE.HYPRESolver} <: AbstractPoisson{T,Vf,Mf}
+    ϵ::V   # Hypre.Vector (increment)
+    A::M   # Hypre.SparseMatrixCSC
+    r::V   # Hypre.Vector (residual)
+    x::Vf  # WaterLily approximate solution
+    L::Mf  # WaterLily lower diagonal coefficients
+    z::Vf  # WaterLily source
+    solver::SO
+end
+function HyprePoisson(x::AbstractArray{T},L::AbstractArray{T},z::AbstractArray{T};
+                      MaxIter=1000, Tol=100eps(T), PrintLevel=0, Logging=0, 
+                      Precond=HYPRE.BoomerAMG(), perdir=()) where T
+    @assert T==Float64 "Only Float64 are supported by HYPRE"
+    # create the vectors and the matrix
+    ϵ = Vector{T}(undef, prod(size(x)))
+    r = Vector{T}(undef, prod(size(z)))
+    A = SparseMatrixCSC{T,Int}(undef, prod(size(x)), prod(size(x)))
+    J = LinearIndices(x) # useful
+    N,n = WaterLily.size_u(L)
+    for I in CartesianIndices(x) # fill the vectors entirely
+        ϵ[J[I]] = x[I]; r[J[I]] = z[I]
+        A[J[I],J[I]] = one(T) # set all diagonals to one
+    end
+    for I in inside(x)
+        A[J[I],J[I]] = WaterLily.diag(I,L)
+        for i in 1:n # fill diagonal terms
+            Im = I-δ(i,I); Ip = I+δ(i,I)
+            A[J[Im],J[I]] = L[I,i]  # x[I-δ(i,I)]*L[I,i]
+            A[J[Ip],J[I]] = L[Ip,i] # x[I+δ(i,I)]*L[I+δ(i,I),i]
+        end
+    end
+    # if we have a purely Neumann problem, we must fix the pressure at one spot to satisfiability
+    # the system, we choose the first block
+    # length(perdir) == n && (A[1,1] = one(T); A[1,2:end] .= zero(T);
+    #                         A[1,1] = one(T); A[1,2:end] .= zero(T))
+    HYPRE.Init() # Init and create a conjugate gradients solver
+    solver = HYPRE.GMRES(;MaxIter,Tol,PrintLevel,Logging,Precond)
+    return HYPREPoisson(ϵ,A,r,x,L,z,solver)
+end
+export HYPREPoisson
+
+"""
+Fill back the solution `hp.x` and `hp.z` from the solution `hp.e` and `hp.r`.
+"""
+function putback!(hp::HYPREPoisson)
+    J = LinearIndices(hp.x)
+    @loop (hp.x[I] = hp.ϵ[J[I]]; hp.z[I] = hp.r[J[I]]) over I ∈ CartesianIndices(hp.x)
+end
+function Base.fill!(hp::HYPREPoisson)
+    J = LinearIndices(hp.x)
+    @loop (hp.ϵ[J[I]] = hp.x[I]; hp.r[J[I]] = hp.z[I]) over I ∈ CartesianIndices(hp.x)
+end
+update!(hp::HYPREPoisson) = nothing
+"""
+Solve the poisson problem and puts back the solution where it is expected.
+"""
+function solver!(hp::HYPREPoisson;kwargs...)
+    fill!(hp)
+    HYPRE.solve!(hp.solver, hp.ϵ,hp.A, hp.r)
+    putback!(hp); 
+    # scale the solution for zero mean
+    mean = @inbounds(sum(hp.x[inside(hp.x)]))/length(inside(hp.x))
+    @loop hp.x[I] -= mean over I ∈ inside(hp.x)
+    J = LinearIndices(hp.x)
+    @loop hp.ϵ[J[I]] -= mean over I ∈ inside(hp.x) # don't remove ean frome ghost cells
+    return nothing
+end
+
+L₂(p::HYPREPoisson) = sum(abs2,p.A*p.ϵ.-p.r)
+L∞(p::HYPREPoisson) = maximum(abs,p.A*p.ϵ.-p.r)
+
+end # module

ext/WaterLilyPlotsExt.jl

@@ -0,0 +1,102 @@
+module WaterLilyPlotsExt
+
+if isdefined(Base, :get_extension)
+    using Plots; gr()
+else
+    using ..Plots; gr()
+end
+
+using WaterLily
+import WaterLily: flood,addbody,body_plot!,sim_gif!,plot_logger
+
+"""
+    flood(f)
+
+Plot a filled contour plot of the 2D array `f`. The keyword arguments are passed to `Plots.contourf`.
+"""
+function flood(f::Array;shift=(0.,0.),cfill=:RdBu_11,clims=(),levels=10,kv...)
+    if length(clims)==2
+        @assert clims[1]<clims[2]
+        @. f=min(clims[2],max(clims[1],f))
+    else
+        clims = (minimum(f),maximum(f))
+    end
+    Plots.contourf(axes(f,1).+shift[1],axes(f,2).+shift[2],f'|>Array,
+                   linewidth=0, levels=levels, color=cfill, clims = clims, 
+                   aspect_ratio=:equal; kv...)
+end
+
+addbody(x,y;c=:black) = Plots.plot!(Shape(x,y), c=c, legend=false)
+function body_plot!(sim;levels=[0],lines=:black,R=inside(sim.flow.p))
+    WaterLily.measure_sdf!(sim.flow.σ,sim.body,WaterLily.time(sim))
+    contour!(sim.flow.σ[R]'|>Array;levels,lines)
+end
+
+"""
+    sim_gif!(sim;duration=1,step=0.1,verbose=true,R=inside(sim.flow.p),
+                    remeasure=false,plotbody=false,kv...)
+
+Make a gif of the simulation `sim` for `duration` seconds with `step` time steps. The keyword arguments are passed to `flood` and `body_plot!`.
+"""
+function sim_gif!(sim;duration=1,step=0.1,verbose=true,R=inside(sim.flow.p),
+                    remeasure=false,plotbody=false,kv...)
+    t₀ = round(WaterLily.sim_time(sim))
+    @time @gif for tᵢ in range(t₀,t₀+duration;step)
+        WaterLily.sim_step!(sim,tᵢ;remeasure)
+        @WaterLily.inside sim.flow.σ[I] = WaterLily.curl(3,I,sim.flow.u)*sim.L/sim.U
+        flood(sim.flow.σ[R]; kv...)
+        plotbody && body_plot!(sim)
+        verbose && println("tU/L=",round(tᵢ,digits=4),
+                           ", Δt=",round(sim.flow.Δt[end],digits=3))
+    end
+end
+
+
+"""
+    plot_logger(fname="WaterLily.log")
+
+Plot the residuals and MG iterations from the log file `fname`.
+"""
+function plot_logger(fname="WaterLily.log")
+    predictor = []; corrector = []
+    open(ifelse(fname[end-3:end]==".log",fname[1:end-4],fname)*".log","r") do f
+        readline(f) # read first line and dump it
+        which = "p"
+        while ! eof(f)  
+            s = split(readline(f) , ",")          
+            which = s[1] != "" ? s[1] : which
+            push!(which == "p" ? predictor : corrector,parse.(Float64,s[2:end]))
+        end
+    end
+    predictor = reduce(hcat,predictor)
+    corrector = reduce(hcat,corrector)
+
+    # get index of all time steps
+    idx = findall(==(0.0),@views(predictor[1,:]))
+    # plot inital L∞ and L₂ norms of residuals for the predictor step
+    p1=plot(1:length(idx),predictor[2,idx],color=:1,ls=:dash,alpha=0.8,label="predictor initial r∞",yaxis=:log,size=(800,400),dpi=600,
+            xlabel="Time step",ylabel="L∞-norm",title="Residuals",ylims=(1e-8,1e0),xlims=(0,length(idx)))
+    p2=plot(1:length(idx),predictor[2,idx],color=:1,ls=:dash,alpha=0.8,label="predictor initial r₂",yaxis=:log,size=(800,400),dpi=600,
+            xlabel="Time step",ylabel="L₂-norm",title="Residuals",ylims=(1e-8,1e0),xlims=(0,length(idx)))
+    # plot final L∞ and L₂norms of residuals for the predictor
+    plot!(p1,1:length(idx),vcat(predictor[2,idx[2:end].-1],predictor[2,end]),color=:1,lw=2,alpha=0.8,label="predictor r∞")
+    plot!(p2,1:length(idx),vcat(predictor[3,idx[2:end].-1],predictor[3,end]),color=:1,lw=2,alpha=0.8,label="predictor r₂")
+    # plot the MG iterations for the predictor
+    p3=plot(1:length(idx),clamp.(vcat(predictor[1,idx[2:end].-1],predictor[1,end]),√1/2,32),lw=2,alpha=0.8,label="predictor",size=(800,400),dpi=600,
+            xlabel="Time step",ylabel="Iterations",title="MG Iterations",ylims=(√1/2,32),xlims=(0,length(idx)),yaxis=:log2)
+    yticks!([√1/2,1,2,4,8,16,32],["0","1","2","4","8","16","32"])
+    # get index of all time steps
+    idx = findall(==(0.0),@views(corrector[1,:]))
+    # plot inital L∞ and L₂ norms of residuals for the corrector step
+    plot!(p1,1:length(idx),corrector[2,idx],color=:2,ls=:dash,alpha=0.8,label="corrector initial r∞",yaxis=:log)
+    plot!(p2,1:length(idx),corrector[3,idx],color=:2,ls=:dash,alpha=0.8,label="corrector initial r₂",yaxis=:log)
+    # plot final L∞ and L₂ norms of residuals for the corrector step
+    plot!(p1,1:length(idx),vcat(corrector[2,idx[2:end].-1],corrector[2,end]),color=:2,lw=2,alpha=0.8,label="corrector r∞")
+    plot!(p2,1:length(idx),vcat(corrector[3,idx[2:end].-1],corrector[3,end]),color=:2,lw=2,alpha=0.8,label="corrector r₂")
+    # plot MG iterations of the corrector
+    plot!(p3,1:length(idx),clamp.(vcat(corrector[1,idx[2:end].-1],corrector[1,end]),√1/2,32),lw=2,alpha=0.8,label="corrector")
+    # plot all together
+    plot(p1,p2,p3,layout=@layout [a b c])
+end
+
+end # module

src/Flow.jl

@@ -153,13 +153,15 @@ and the `AbstractPoisson` pressure solver to project the velocity onto an incomp
 @fastmath function mom_step!(a::Flow{N},b::AbstractPoisson) where N
     a.u⁰ .= a.u; scale_u!(a,0)
     # predictor u → u'
+    @debug "mlp"
     U = BCTuple(a.U,@view(a.Δt[1:end-1]),N)
     conv_diff!(a.f,a.u⁰,a.σ,ν=a.ν,perdir=a.perdir)
     accelerate!(a.f,@view(a.Δt[1:end-1]),a.g,a.U)
     BDIM!(a); BC!(a.u,U,a.exitBC,a.perdir)
     a.exitBC && exitBC!(a.u,a.u⁰,U,a.Δt[end]) # convective exit
     project!(a,b); BC!(a.u,U,a.exitBC,a.perdir)
     # corrector u → u¹
+    @debug "mlc"
     U = BCTuple(a.U,a.Δt,N)
     conv_diff!(a.f,a.u,a.σ,ν=a.ν,perdir=a.perdir)
     accelerate!(a.f,a.Δt,a.g,a.U)

src/MultiLevelPoisson.jl

@@ -81,20 +81,23 @@ function Vcycle!(ml::MultiLevelPoisson;l=1)
     increment!(fine)
 end
 
+L₂(p::MultiLevelPoisson) = L₂(p.levels[1])
+L∞(p::MultiLevelPoisson) = L∞(p.levels[1])
+
 mult!(ml::MultiLevelPoisson,x) = mult!(ml.levels[1],x)
 residual!(ml::MultiLevelPoisson,x) = residual!(ml.levels[1],x)
 
-function solver!(ml::MultiLevelPoisson;tol=2e-4,itmx=32)
+function solver!(ml::MultiLevelPoisson;tol=1e-4,itmx=32)
     p = ml.levels[1]
-    residual!(p); r₀ = r₂ = L∞(p); r₂₀ = L₂(p)
-    nᵖ=0
-    while r₂>tol && nᵖ<itmx
+    residual!(p); r₂ = L₂(p)
+    nᵖ=0; @debug "ml, $nᵖ, $(L∞(p)), $r₂\n"
+    while nᵖ<itmx
         Vcycle!(ml)
-        smooth!(p); r₂ = L∞(p)
+        smooth!(p); r₂ = L₂(p)
         nᵖ+=1
+        @debug "ml, $nᵖ, $(L∞(p)), $r₂\n"
+        r₂<tol && break
     end
     perBC!(p.x,p.perdir)
-    (nᵖ<2 && length(ml.levels)>5) && pop!(ml.levels); # remove coarsest level if this was easy
-    (nᵖ>4 && divisible(ml.levels[end])) && push!(ml.levels,restrictML(ml.levels[end])) # add a level if this was hard
     push!(ml.n,nᵖ);
-end
+end

src/WaterLily.jl

@@ -65,15 +65,15 @@ mutable struct Simulation
     function Simulation(dims::NTuple{N}, u_BC, L::Number;
                         Δt=0.25, ν=0., g=nothing, U=nothing, ϵ=1, perdir=(),
                         uλ=nothing, exitBC=false, body::AbstractBody=NoBody(),
-                        T=Float32, mem=Array) where N
+                        T=Float32, mem=Array, psolver=:MultiLevelPoisson) where N
         @assert !(isa(u_BC,Function) && isa(uλ,Function)) "`u_BC` and `uλ` cannot be both specified as Function"
         @assert !(isnothing(U) && isa(u_BC,Function)) "`U` must be specified if `u_BC` is a Function"
         isa(u_BC,Function) && @assert all(typeof.(ntuple(i->u_BC(i,zero(T)),N)).==T) "`u_BC` is not type stable"
         uλ = isnothing(uλ) ? ifelse(isa(u_BC,Function),(i,x)->u_BC(i,0.),(i,x)->u_BC[i]) : uλ
         U = isnothing(U) ? √sum(abs2,u_BC) : U # default if not specified
         flow = Flow(dims,u_BC;uλ,Δt,ν,g,T,f=mem,perdir,exitBC)
         measure!(flow,body;ϵ)
-        new(U,L,ϵ,flow,body,MultiLevelPoisson(flow.p,flow.μ₀,flow.σ;perdir))
+        new(U,L,ϵ,flow,body,eval(psolver)(flow.p,flow.μ₀,flow.σ;perdir))
     end
 end
 
@@ -132,6 +132,24 @@ function restart_sim! end
 # export
 export restart_sim!
 
+#default Plots functions
+function flood end
+function addbody end
+function body_plot! end
+function sim_gif! end
+function plot_logger end
+# export
+export flood,addbody,body_plot!,sim_gif!,plot_logger
+
+#default HYPRE functions
+function HyprePoisson end
+function putback! end
+export HyprePoisson,putback!,HYPREPoisson
+
+# default GeometricMultigrid functions
+function GeomMultigridPoisson end
+export GeomMultigridPoisson
+
 # Check number of threads when loading WaterLily
 """
     check_nthreads(::Val{1})
@@ -154,6 +172,9 @@ function __init__()
         @require CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" include("../ext/WaterLilyCUDAExt.jl")
         @require WriteVTK = "64499a7a-5c06-52f2-abe2-ccb03c286192" include("../ext/WaterLilyWriteVTKExt.jl")
         @require ReadVTK = "dc215faf-f008-4882-a9f7-a79a826fadc3" include("../ext/WaterLilyReadVTKExt.jl")
+        @require Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80" include("../ext/WaterLilyPlotsExt.jl")
+        @require HYPRE = "b5ffcf37-a2bd-41ab-a3da-4bd9bc8ad771" include("../ext/WaterLilyHYPREExt.jl")
+        @require GeometricMultigrid = "ec3bce20-3e4b-4d8a-a801-89defc82be1c" include("../ext/WaterLilyGeometricMultigridExt.jl")
     end
     check_nthreads(Val{Threads.nthreads()}())
 end