Logging pressure solver

Project.toml

@@ -9,6 +9,7 @@ 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"
@@ -20,12 +21,14 @@ AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 ReadVTK = "dc215faf-f008-4882-a9f7-a79a826fadc3"
 WriteVTK = "64499a7a-5c06-52f2-abe2-ccb03c286192"
+Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 
 [extensions]
 WaterLilyAMDGPUExt = "AMDGPU"
 WaterLilyCUDAExt = "CUDA"
 WaterLilyReadVTKExt = "ReadVTK"
 WaterLilyWriteVTKExt = "WriteVTK"
+WaterLilyPlotsExt = "Plots"
 
 [compat]
 DocStringExtensions = "0.9"

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'
+    @log "p"
     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¹
+    @log "c"
     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

@@ -48,7 +48,7 @@ struct MultiLevelPoisson{T,S<:AbstractArray{T},V<:AbstractArray{T}} <: AbstractP
     levels :: Vector{Poisson{T,S,V}}
     n :: Vector{Int16}
     perdir :: NTuple # direction of periodic boundary condition
-    function MultiLevelPoisson(x::AbstractArray{T},L::AbstractArray{T},z::AbstractArray{T};maxlevels=Inf,perdir=()) where T
+    function MultiLevelPoisson(x::AbstractArray{T},L::AbstractArray{T},z::AbstractArray{T};maxlevels=10,perdir=()) where T
         levels = Poisson[Poisson(x,L,z;perdir)]
         while divisible(levels[end]) && length(levels) <= maxlevels
             push!(levels,restrictML(levels[end]))
@@ -84,17 +84,16 @@ end
 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; @log ", $nᵖ, $(L∞(p)), $r₂\n"
+    while nᵖ<itmx
         Vcycle!(ml)
-        smooth!(p); r₂ = L∞(p)
-        nᵖ+=1
+        smooth!(p); r₂ = L₂(p); nᵖ+=1
+        @log ", $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/Poisson.jl

@@ -158,16 +158,14 @@ Approximate iterative solver for the Poisson matrix equation `Ax=b`.
   - `tol`: Convergence tolerance on the `L₂`-norm residual.
   - `itmx`: Maximum number of iterations.
 """
-function solver!(p::Poisson;log=false,tol=1e-4,itmx=1e3)
+function solver!(p::Poisson;tol=1e-4,itmx=1e3)
     residual!(p); r₂ = L₂(p)
-    log && (res = [r₂])
-    nᵖ=0
-    while r₂>tol && nᵖ<itmx
-        smooth!(p); r₂ = L₂(p)
-        log && push!(res,r₂)
-        nᵖ+=1
+    nᵖ=0; @log ", $nᵖ, $(L∞(p)), $r₂\n"
+    while nᵖ<itmx
+        smooth!(p); r₂ = L₂(p); nᵖ+=1
+        @log ", $nᵖ, $(L∞(p)), $r₂\n"
+        r₂<tol && break
     end
     perBC!(p.x,p.perdir)
     push!(p.n,nᵖ)
-    log && return res
 end

src/WaterLily.jl

@@ -6,7 +6,7 @@ module WaterLily
 using DocStringExtensions
 
 include("util.jl")
-export L₂,BC!,@inside,inside,δ,apply!,loc
+export L₂,BC!,@inside,inside,δ,apply!,loc,@log
 
 using Reexport
 @reexport using KernelAbstractions: @kernel,@index,get_backend
@@ -132,6 +132,15 @@ 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
+
 # Check number of threads when loading WaterLily
 """
     check_nthreads(::Val{1})
@@ -154,6 +163,7 @@ 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")
     end
     check_nthreads(Val{Threads.nthreads()}())
 end

src/util.jl

@@ -1,4 +1,27 @@
 using KernelAbstractions: get_backend, @index, @kernel
+using LoggingExtras
+
+# custom log macro
+_psolver = Logging.LogLevel(-123) # custom log level for pressure solver, needs the negative sign
+macro log(exs...)
+    quote
+        @logmsg _psolver $(map(x -> esc(x), exs)...)
+    end
+end
+"""
+    logger(fname="WaterLily")
+
+Set up a logger to write the pressure solver data to a logging file named `WaterLily.log`.
+"""
+function logger(fname::String="WaterLily")
+    ENV["JULIA_DEBUG"] = all
+    logger = FormatLogger(ifelse(fname[end-3:end]==".log",fname[1:end-4],fname)*".log"; append=false) do io, args
+        args.level == _psolver && print(io, args.message)
+    end;
+    global_logger(logger);
+    # put header in file
+    @log "p/c, iter, r∞, r₂\n"
+end
 
 @inline CI(a...) = CartesianIndex(a...)
 """