Merge pull request #129 from weymouth/AD_version

WaterLily + AutoDiff workflow

.github/workflows/ci.yml

@@ -5,21 +5,24 @@ on:
 jobs:
   test:
     if: github.event.pull_request.draft == false
-    name: Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }} - ${{ github.event_name }}
+    name: Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }} - ${{ matrix.nthreads }} - ${{ github.event_name }}
     runs-on: ${{ matrix.os }}
     strategy:
       fail-fast: false
       matrix:
         version:
           - '1.9'
-          - '1.10'
+          - '1'
         os:
           - ubuntu-latest
           - macOS-latest
           - windows-latest
         arch:
           - x64
           - x86
+        nthreads:
+          - '1'
+          - 'auto'
         exclude:
           - os: macOS-latest
             arch: x86
@@ -41,6 +44,8 @@ jobs:
             ${{ runner.os }}-
       - uses: julia-actions/julia-buildpkg@v1
       - uses: julia-actions/julia-runtest@v1
+        env:
+          JULIA_NUM_THREADS: ${{ matrix.nthreads }}
       - uses: julia-actions/julia-processcoverage@v1
       - uses: codecov/codecov-action@v1
         with:

Project.toml

@@ -55,4 +55,4 @@ UnicodePlots = "b8865327-cd53-5732-bb35-84acbb429228"
 WriteVTK = "64499a7a-5c06-52f2-abe2-ccb03c286192"
 
 [targets]
-test = ["Test", "CUDA", "AMDGPU", "GPUArrays", "WriteVTK", "ReadVTK"]
+test = ["Test", "BenchmarkTools", "CUDA", "AMDGPU", "GPUArrays", "WriteVTK", "ReadVTK"]

examples/Project.toml

@@ -1,5 +1,7 @@
 [deps]
+Optim = "429524aa-4258-5aef-a3af-852621145aeb"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 GLMakie = "e9467ef8-e4e7-5192-8a1a-b1aee30e663a"
 GR = "28b8d3ca-fb5f-59d9-8090-bfdbd6d07a71"
 GeometryBasics = "5c1252a2-5f33-56bf-86c9-59e7332b4326"

examples/TandemFoilOptim.jl

@@ -0,0 +1,50 @@
+using WaterLily,StaticArrays
+
+function make_foils(φ;two=true,L=32,Re=1e3,St=0.3,αₘ=-π/18,U=1,n=8,m=4)
+    # Map from simulation coordinate x to surface coordinate ξ
+    nose,pivot = SA[2L,m*L//2],SA[L//4,0]
+    θ₀ = αₘ+atan(π*St); h₀=L; ω=π*St*U/h₀
+    function map(x,t)
+        back = two && x[1]>nose[1]+2L # back body?
+        ϕ = back ? φ : zero(φ)        # phase shift
+        S = back ? 3L : zero(L)       # horizontal shift
+        s,c = sincos(θ₀*cos(ω*t+ϕ))   # sin & cos of angle
+        h = SA[S,h₀*sin(ω*t+ϕ)]       # position
+        # move to origin and align with x-axis
+        ξ = SA[c -s; s c]*(x-nose-h-pivot)+pivot 
+        return SA[ξ[1],abs(ξ[2])]    # reflect to positive y
+    end
+
+    # Line segment SDF
+    function sdf(ξ,t)
+        p = ξ-SA[clamp(ξ[1],0,L),0] # vector from closest point on [0,L] segment to ξ 
+        √(p'*p)-2                   # distance (with thickness offset)
+    end
+
+    Simulation((n*L,m*L),(U,0),L;ν=U*L/Re,body=AutoBody(sdf,map),T=typeof(φ))
+end
+
+drag(flow,body,t) = sum(inside(flow.p)) do I
+    d,n,_ = measure(body,WaterLily.loc(0,I),t)
+    flow.p[I]*n[1]*WaterLily.kern(clamp(d,-1,1))
+end
+
+function Δimpulse!(sim)
+    Δt = sim.flow.Δt[end]*sim.U/sim.L
+    sim_step!(sim)
+    Δt*drag(sim.flow,sim.body,WaterLily.time(sim))
+end
+
+function mean_drag(φ,two=true,St=0.3,N=3,period=2N/St)
+    sim = make_foils(φ;two,St)
+    sim_step!(sim,period) # warm-in transient period
+    impulse = 0           # integrate impulse
+    while sim_time(sim)<2period
+        impulse += Δimpulse!(sim)
+    end
+    impulse/period        # return mean drag
+end
+
+using Optim
+θ = Optim.minimizer(optimize(x->-mean_drag(first(x)), [0f0], Newton(),
+    Optim.Options(show_trace=true,f_tol=1e-2); autodiff = :forward))

src/Body.jl

@@ -27,13 +27,13 @@ at time `t` using an immersion kernel of size `ϵ`.
 
 See Maertens & Weymouth, doi:[10.1016/j.cma.2014.09.007](https://doi.org/10.1016/j.cma.2014.09.007).
 """
-function measure!(a::Flow{N,T},body::AbstractBody;t=T(0),ϵ=1) where {N,T}
+function measure!(a::Flow{N,T},body::AbstractBody;t=zero(T),ϵ=1) where {N,T}
     a.V .= 0; a.μ₀ .= 1; a.μ₁ .= 0
     @fastmath @inline function fill!(μ₀,μ₁,V,d,I)
-        d[I] = sdf(body,loc(0,I,T),t)
+        d[I] = sdf(body,loc(0,I),t)
         if abs(d[I])<2+ϵ
             for i ∈ 1:N
-                dᵢ,nᵢ,Vᵢ = measure(body,WaterLily.loc(i,I,T),t)
+                dᵢ,nᵢ,Vᵢ = measure(body,loc(i,I),t)
                 V[I,i] = Vᵢ[i]
                 μ₀[I,i] = WaterLily.μ₀(dᵢ,ϵ)
                 μ₁[I,i,:] .= WaterLily.μ₁(dᵢ,ϵ)*nᵢ
@@ -43,7 +43,7 @@ function measure!(a::Flow{N,T},body::AbstractBody;t=T(0),ϵ=1) where {N,T}
         end
     end
     @loop fill!(a.μ₀,a.μ₁,a.V,a.σ,I) over I ∈ inside(a.p)
-    BC!(a.μ₀,zeros(SVector{N,T}),a.exitBC,a.perdir) # BC on μ₀, don't fill normal component yet
+    BC!(a.μ₀,zeros(SVector{N,T}),false,a.perdir) # BC on μ₀, don't fill normal component yet
     BC!(a.V ,zeros(SVector{N,T}),a.exitBC,a.perdir)
 end
 

src/Flow.jl

@@ -33,7 +33,7 @@ function median(a,b,c)
     return a
 end
 
-function conv_diff!(r,u,Φ;ν=0.1,perdir=(0,))
+function conv_diff!(r,u,Φ;ν=0.1,perdir=())
     r .= 0.
     N,n = size_u(u)
     for i ∈ 1:n, j ∈ 1:n
@@ -61,21 +61,23 @@ upperBoundary!(r,u,Φ,ν,i,j,N,::Val{true}) = @loop r[I-δ(j,I),i] -= Φ[CIj(j,I
 
 using EllipsisNotation
 """
-    accelerate!(r,t,g)
+    accelerate!(r,dt,g)
 
-This function adds a uniform acceleration field `g` at time `t` to `r`.
-If `g ≠ nothing`, then `g(i,t)=dUᵢ/dt`.
+Add a uniform acceleration `gᵢ+dUᵢ/dt` at time `t=sum(dt)` to field `r`.
 """
-accelerate!(r,t,g::Function,::Tuple) = for i ∈ 1:last(size(r))
+accelerate!(r,dt,g::Function,::Tuple,t=sum(dt)) = for i ∈ 1:last(size(r))
     r[..,i] .+= g(i,t)
 end
-accelerate!(r,t,g::Nothing,U::Function) = for i ∈ 1:last(size(r))
-    r[..,i] .+= ForwardDiff.derivative(τ->U(i,τ),t)
-end
-accelerate!(r,t,g::Function,U::Function) = for i ∈ 1:last(size(r))
-    r[..,i] .+= g(i,t) + ForwardDiff.derivative(τ->U(i,τ),t)
-end
-accelerate!(r,t,::Nothing,::Tuple) = nothing
+accelerate!(r,dt,g::Nothing,U::Function) = accelerate!(r,dt,(i,t)->ForwardDiff.derivative(τ->U(i,τ),t),())
+accelerate!(r,dt,g::Function,U::Function) = accelerate!(r,dt,(i,t)->g(i,t)+ForwardDiff.derivative(τ->U(i,τ),t),())
+accelerate!(r,dt,::Nothing,::Tuple) = nothing
+"""
+    BCTuple(U,dt,N)
+
+Return BC tuple `U(i∈1:N, t=sum(dt))`.
+"""
+BCTuple(f::Function,dt,N,t=sum(dt))=ntuple(i->f(i,t),N)
+BCTuple(f::Tuple,dt,N)=f
 
 """
     Flow{D::Int, T::Float, Sf<:AbstractArray{T,D}, Vf<:AbstractArray{T,D+1}, Tf<:AbstractArray{T,D+2}}
@@ -106,21 +108,19 @@ struct Flow{D, T, Sf<:AbstractArray{T}, Vf<:AbstractArray{T}, Tf<:AbstractArray{
     exitBC :: Bool # Convection exit
     perdir :: NTuple # tuple of periodic direction
     function Flow(N::NTuple{D}, U; f=Array, Δt=0.25, ν=0., g=nothing,
-                  uλ::Function=(i, x) -> 0., perdir=(0,), exitBC=false, T=Float64) where D
+                  uλ::Function=(i, x) -> 0., perdir=(), exitBC=false, T=Float64) where D
         Ng = N .+ 2
         Nd = (Ng..., D)
         u = Array{T}(undef, Nd...) |> f; apply!(uλ, u);
         BC!(u,BCTuple(U,0.,D),exitBC,perdir); exitBC!(u,u,BCTuple(U,0.,D),0.)
         u⁰ = copy(u);
         fv, p, σ = zeros(T, Nd) |> f, zeros(T, Ng) |> f, zeros(T, Ng) |> f
         V, μ₀, μ₁ = zeros(T, Nd) |> f, ones(T, Nd) |> f, zeros(T, Ng..., D, D) |> f
+        BC!(μ₀,ntuple(zero, D),false,perdir)
         new{D,T,typeof(p),typeof(u),typeof(μ₁)}(u,u⁰,fv,p,σ,V,μ₀,μ₁,U,T[Δt],ν,g,exitBC,perdir)
     end
 end
 
-time(flow::Flow) = sum(flow.Δt[1:end-1])
-timeNext(flow::Flow) = sum(flow.Δt)
-
 function BDIM!(a::Flow)
     dt = a.Δt[end]
     @loop a.f[Ii] = a.u⁰[Ii]+dt*a.f[Ii]-a.V[Ii] over Ii in CartesianIndices(a.f)
@@ -146,16 +146,16 @@ 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'
-    U = BCTuple(a.U,time(a),N)
+    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,time(a),a.g,a.U)
+    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¹
-    U = BCTuple(a.U,timeNext(a),N)
+    U = BCTuple(a.U,a.Δt,N)
     conv_diff!(a.f,a.u,a.σ,ν=a.ν,perdir=a.perdir)
-    accelerate!(a.f,timeNext(a),a.g,a.U)
+    accelerate!(a.f,a.Δt,a.g,a.U)
     BDIM!(a); scale_u!(a,0.5); BC!(a.u,U,a.exitBC,a.perdir)
     project!(a,b,0.5); BC!(a.u,U,a.exitBC,a.perdir)
     push!(a.Δt,CFL(a))

src/Metrics.jl

@@ -81,7 +81,7 @@ end
 Surface normal integral of field `p` over the `body`.
 """
 function ∮nds(p::AbstractArray{T,N},df::AbstractArray{T},body::AbstractBody,t=0) where {T,N}
-    @loop df[I,:] .= p[I]*nds(body,loc(0,I,T),t) over I ∈ inside(p)
+    @loop df[I,:] .= p[I]*nds(body,loc(0,I),t) over I ∈ inside(p)
     [sum(@inbounds(df[inside(p),i])) for i ∈ 1:N] |> Array
 end
 @inline function nds(body::AbstractBody,x,t)

src/MultiLevelPoisson.jl

@@ -23,7 +23,7 @@ function restrictML(b::Poisson)
     restrictL!(aL,b.L,perdir=b.perdir)
     Poisson(ax,aL,copy(ax);b.perdir)
 end
-function restrictL!(a::AbstractArray{T},b;perdir=(0,)) where T
+function restrictL!(a::AbstractArray{T},b;perdir=()) where T
     Na,n = size_u(a)
     for i ∈ 1:n
         @loop a[I,i] = restrictL(I,i,b) over I ∈ CartesianIndices(map(n->2:n-1,Na))
@@ -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=(0,)) where T
+    function MultiLevelPoisson(x::AbstractArray{T},L::AbstractArray{T},z::AbstractArray{T};maxlevels=Inf,perdir=()) where T
         levels = Poisson[Poisson(x,L,z;perdir)]
         while divisible(levels[end]) && length(levels) <= maxlevels
             push!(levels,restrictML(levels[end]))
@@ -78,7 +78,6 @@ function Vcycle!(ml::MultiLevelPoisson;l=1)
     smooth!(coarse)
     # correct fine
     prolongate!(fine.ϵ,coarse.x)
-    BC!(fine.ϵ;perdir=fine.perdir)
     increment!(fine)
 end
 
@@ -87,16 +86,15 @@ residual!(ml::MultiLevelPoisson,x) = residual!(ml.levels[1],x)
 
 function solver!(ml::MultiLevelPoisson;tol=2e-4,itmx=32)
     p = ml.levels[1]
-    BC!(p.x;perdir=p.perdir)
     residual!(p); r₀ = r₂ = L∞(p); r₂₀ = L₂(p)
     nᵖ=0
     while r₂>tol && nᵖ<itmx
         Vcycle!(ml)
         smooth!(p); r₂ = L∞(p)
         nᵖ+=1
     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
-    BC!(p.x;perdir=p.perdir)
     push!(ml.n,nᵖ);
 end

src/Poisson.jl

@@ -28,7 +28,7 @@ struct Poisson{T,S<:AbstractArray{T},V<:AbstractArray{T}} <: AbstractPoisson{T,S
     z :: S # source
     n :: Vector{Int16} # pressure solver iterations
     perdir :: NTuple # direction of periodic boundary condition
-    function Poisson(x::AbstractArray{T},L::AbstractArray{T},z::AbstractArray{T};perdir=(0,)) where T
+    function Poisson(x::AbstractArray{T},L::AbstractArray{T},z::AbstractArray{T};perdir=()) where T
         @assert axes(x) == axes(z) && axes(x) == Base.front(axes(L)) && last(axes(L)) == eachindex(axes(x))
         r = similar(x); fill!(r,0)
         ϵ,D,iD = copy(r),copy(r),copy(r)
@@ -37,6 +37,8 @@ struct Poisson{T,S<:AbstractArray{T},V<:AbstractArray{T}} <: AbstractPoisson{T,S
     end
 end
 
+using ForwardDiff: Dual,Tag
+Base.eps(::Type{D}) where D<:Dual{Tag{G,T}} where {G,T} = eps(T)
 function set_diag!(D,iD,L)
     @inside D[I] = diag(I,L)
     @inside iD[I] = abs2(D[I])<2eps(eltype(D)) ? 0. : inv(D[I])
@@ -59,6 +61,7 @@ Fills `p.z = p.A x` with 0 in the ghost cells.
 """
 function mult!(p::Poisson,x)
     @assert axes(p.z)==axes(x)
+    perBC!(x,p.perdir)
     fill!(p.z,0)
     @inside p.z[I] = mult(I,p.L,p.D,x)
     return p.z
@@ -86,14 +89,18 @@ Note: These corrections mean `x` is not strictly solving `Ax=z`, but
 without the corrections, no solution exists.
 """
 function residual!(p::Poisson) 
+    perBC!(p.x,p.perdir)
     @inside p.r[I] = ifelse(p.iD[I]==0,0,p.z[I]-mult(I,p.L,p.D,p.x))
-    s = sum(p.r)/length(p.r[inside(p.r)])
+    s = sum(p.r)/length(inside(p.r))
     abs(s) <= 2eps(eltype(s)) && return
     @inside p.r[I] = p.r[I]-s
 end
 
-increment!(p::Poisson) = @loop (p.r[I] = p.r[I]-mult(I,p.L,p.D,p.ϵ);
-                                p.x[I] = p.x[I]+p.ϵ[I]) over I ∈ inside(p.x)
+function increment!(p::Poisson) 
+    perBC!(p.ϵ,p.perdir)
+    @loop (p.r[I] = p.r[I]-mult(I,p.L,p.D,p.ϵ);
+           p.x[I] = p.x[I]+p.ϵ[I]) over I ∈ inside(p.x)
+end
 """
     Jacobi!(p::Poisson; it=1)
 
@@ -102,7 +109,6 @@ Note: This runs for general backends, but is _very_ slow to converge.
 """
 @fastmath Jacobi!(p;it=1) = for _ ∈ 1:it
     @inside p.ϵ[I] = p.r[I]*p.iD[I]
-    BC!(p.ϵ;perdir=p.perdir)
     increment!(p)
 end
 
@@ -117,18 +123,17 @@ Note: This runs for general backends and is the default smoother.
 function pcg!(p::Poisson{T};it=6) where T
     x,r,ϵ,z = p.x,p.r,p.ϵ,p.z
     @inside z[I] = ϵ[I] = r[I]*p.iD[I]
-    insideI = inside(x) # [insideI]
-    rho = T(r⋅z)
+    rho = r⋅z
     abs(rho)<10eps(T) && return
     for i in 1:it
-        BC!(ϵ;perdir=p.perdir)
+        perBC!(ϵ,p.perdir)
         @inside z[I] = mult(I,p.L,p.D,ϵ)
-        alpha = rho/T(z[insideI]⋅ϵ[insideI])
+        alpha = rho/(z⋅ϵ)
         @loop (x[I] += alpha*ϵ[I];
                r[I] -= alpha*z[I]) over I ∈ inside(x)
         (i==it || abs(alpha)<1e-2) && return
         @inside z[I] = r[I]*p.iD[I]
-        rho2 = T(r⋅z)
+        rho2 = r⋅z
         abs(rho2)<10eps(T) && return
         beta = rho2/rho
         @inside ϵ[I] = beta*ϵ[I]+z[I]
@@ -138,7 +143,7 @@ end
 smooth!(p) = pcg!(p)
 
 L₂(p::Poisson) = p.r ⋅ p.r # special method since outside(p.r)≡0
-L∞(p::Poisson) = maximum(abs.(p.r))
+L∞(p::Poisson) = maximum(abs,p.r)
 
 """
     solver!(A::Poisson;log,tol,itmx)
@@ -154,7 +159,6 @@ Approximate iterative solver for the Poisson matrix equation `Ax=b`.
   - `itmx`: Maximum number of iterations.
 """
 function solver!(p::Poisson;log=false,tol=1e-4,itmx=1e3)
-    BC!(p.x;perdir=p.perdir)
     residual!(p); r₂ = L₂(p)
     log && (res = [r₂])
     nᵖ=0
@@ -163,7 +167,7 @@ function solver!(p::Poisson;log=false,tol=1e-4,itmx=1e3)
         log && push!(res,r₂)
         nᵖ+=1
     end
-    BC!(p.x;perdir=p.perdir)
+    perBC!(p.x,p.perdir)
     push!(p.n,nᵖ)
     log && return res
 end