Merge pull request #280 from WaterLily-jl/refactor_measure

Refactor measure

Author: weymouth

src/AutoBody.jl

@@ -28,21 +28,13 @@ The velocity is determined _solely_ from the optional `map` function.
 Skips the `n,V` calculation when `d²>fastd²`.
 """
 function measure(body::AutoBody,x,t;fastd²=Inf)
-    # eval d=f(x,t), and n̂ = ∇f
     d = sdf(body,x,t)
-    d^2>fastd² && return (d,zero(x),zero(x)) # skip n,V
-    n = ForwardDiff.gradient(x->sdf(body,x,t), x)
-    any(isnan.(n)) && return (d,zero(x),zero(x))
-
-    # correct general implicit fnc f(x₀)=0 to be a pseudo-sdf
-    #   f(x) = f(x₀)+d|∇f|+O(d²) ∴  d ≈ f(x)/|∇f|
-    m = √sum(abs2,n); d /= m; n /= m
-
-    # The velocity depends on the material change of ξ=m(x,t):
-    #   Dm/Dt=0 → ṁ + (dm/dx)ẋ = 0 ∴  ẋ =-(dm/dx)\ṁ
-    J = ForwardDiff.jacobian(x->body.map(x,t), x)
-    dot = ForwardDiff.derivative(t->body.map(x,t), t)
-    return (d,n,-J\dot)
+    d^2>fastd² && return (d,zero(x),zero(x))
+    n = ForwardDiff.gradient(ξ->body.sdf(ξ,t), body.map(x,t)) # body-frame only
+    any(isnan, n) && return (d,zero(x),zero(x))               # handle non-diff'able points
+    J = ForwardDiff.jacobian(x->body.map(x,t), x)             # for mapping n,V to x-frame
+    n = J'n; m = √sum(abs2,n); d /= m; n /= m                 # chain rule then normalise
+    return (d, n, -J\ForwardDiff.derivative(t->body.map(x,t), t))
 end
 
 using LinearAlgebra: tr

src/Body.jl

@@ -20,17 +20,19 @@ Queries the body geometry to fill the arrays:
 - `flow.μ₁`, First kernel moment scaled by the body normal
 - `flow.V`,  Body velocity
 
-at time `t` using an immersion kernel of size `ϵ`.
+at time `t` using an immersion kernel of size `ϵ`. The velocity is only filled within a narrow band
+of size `2+ϵ` around the body. This function also fills `flow.σ` with the signed distance function.
 
 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=zero(T),ϵ=1) where {N,T}
-    a.V .= zero(T); a.μ₀ .= one(T); a.μ₁ .= zero(T); d²=(2+ϵ)^2
+    a.V .= zero(T); a.μ₀ .= one(T); a.μ₁ .= zero(T); d²=T(2+ϵ)^2
+    measure_sdf!(a.σ, body, t; fastd²=d²) # measure separately to allow specialization
     @fastmath @inline function fill!(μ₀,μ₁,V,d,I)
-        d[I] = sdf(body,loc(0,I,T),t,fastd²=d²)
         if d[I]^2<d²
             for i ∈ 1:N
                 dᵢ,nᵢ,Vᵢ = measure(body,loc(i,I,T),t,fastd²=d²)
+                dᵢ = abs(dᵢ) ≤ 0.5 ? dᵢ : copysign(dᵢ,d[I]) # enforce sign consistency
                 V[I,i] = Vᵢ[i]
                 μ₀[I,i] = WaterLily.μ₀(dᵢ,ϵ)
                 for j ∈ 1:N
@@ -99,7 +101,7 @@ Base.:-(a::AbstractBody) = SetBody(-,a,NoBody())
 Base.:-(a::AbstractBody, b::AbstractBody) = a ∩ (-b)
 
 # Measurements
-function measure(body::SetBody,x,t;fastd²=Inf)
+function measure(body::SetBody,x::AbstractVector{T},t;fastd²=T(Inf)) where T 
     body.op(measure(body.a,x,t;fastd²),measure(body.b,x,t;fastd²)) # can't mapreduce within GPU kernel
 end
-measure(body::SetBody{typeof(-)},x,t;fastd²=Inf) = ((d,n,V) = measure(body.a,x,t;fastd²); (-d,-n,V))
+measure(body::SetBody{typeof(-)},x::AbstractVector{T},t;fastd²=T(Inf)) where T = ((d,n,V) = measure(body.a,x,t;fastd²); (-d,-n,V))

src/Metrics.jl

@@ -88,8 +88,8 @@ end
 
 BDIM-masked surface normal.
 """
-@inline function nds(body,x,t)
-    d,n,_ = measure(body,x,t,fastd²=1)
+@inline function nds(body,x::AbstractVector{T},t) where T
+    d,n,_ = measure(body,x,t,fastd²=one(T))
     n*WaterLily.kern(clamp(d,-1,1))
 end
 

src/MultiLevelPoisson.jl

@@ -23,12 +23,12 @@ 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=()) where T
+function restrictL!(a::AbstractArray{T,M},b;perdir=()) where {T,M}
     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))
     end
-    BC!(a,zeros(SVector{n,T}),false,perdir)  # correct μ₀ @ boundaries
+    BC!(a,zero(SVector{M-1,T}),false,perdir)  # correct μ₀ @ boundaries
 end
 restrict!(a,b) = @inside a[I] = restrict(I,b)
 prolongate!(a,b) = @inside a[I] = b[down(I)]

src/RigidMap.jl

@@ -25,18 +25,21 @@ for n in 1:10
 end
 ```
 """
-struct RigidMap{A<:AbstractVector,R} <: Function
+struct RigidMap{A<:AbstractVector,R,M} <: Function
     x₀ :: A   # center of translation
     θ  :: R   # rotation (angle in 2D, euler angles in 3D)
     xₚ :: A   # rotation offset
     V  :: A   # linear velocity of the center
     ω  :: R   # angular velocity (scalar in 2D, vector in 3D)
+    R̂  :: M   # rotation matrix (precomputed for efficiency)
 end
-RigidMap(x₀::SVector,θ;xₚ=zero(x₀),V=zero(x₀),ω=zero(θ)) = RigidMap(x₀, θ, xₚ, V, ω)
+RigidMap(x₀::SVector,θ;xₚ=zero(x₀),V=zero(x₀),ω=zero(θ)) = RigidMap(x₀, θ, xₚ, V, ω, rotation(θ))
 
 # this is the function map(x,t) AND derivative(t->map(x,t),t)
-(map::RigidMap)(x::SVector,t=0) = rotation(map.θ)*(x-map.x₀-map.xₚ)+map.xₚ
-(map::RigidMap)(x::SVector,::ForwardDiff.Dual{Tag}) where Tag = Dual{Tag}.(map(x),-rotation(map.θ)*(map.V+map.ω×(x-map.x₀-map.xₚ)))
+(m::RigidMap)(x::SVector,t=0) = m.R̂*(x-m.x₀-m.xₚ)+m.xₚ
+(m::RigidMap)(x::SVector,t::ForwardDiff.Dual{Tag}) where Tag = Dual{Tag}.(m(x),map_velocity(m, x, t))
+map_jacobian(m::RigidMap, x, t) = m.R̂
+map_velocity(m::RigidMap, x, t) = -m.R̂*(m.V + m.ω×(x-m.x₀-m.xₚ))
 
 # cross product in 2D and rotation matrix in 2D and 3D
 import WaterLily: ×
@@ -45,7 +48,9 @@ rotation(θ::T) where T = SA{T}[cos(θ) sin(θ); -sin(θ) cos(θ)]
 rotation(θ::SVector{3,T}) where T = SA{T}[cos(θ[3])*cos(θ[2]) cos(θ[3])*sin(θ[2])*sin(θ[1])+sin(θ[3])*cos(θ[1]) -cos(θ[3])*sin(θ[2])*cos(θ[1])+sin(θ[3])*sin(θ[1]);
                                          -sin(θ[3])*cos(θ[2]) -sin(θ[3])*sin(θ[2])*sin(θ[1])+cos(θ[3])*cos(θ[1]) sin(θ[3])*sin(θ[2])*cos(θ[1])+cos(θ[3])*sin(θ[1]);
                                                 sin(θ[2])                         -cos(θ[2])*sin(θ[1])                               cos(θ[2])*cos(θ[1])]
-import ConstructionBase: setproperties
+
+import ConstructionBase: setproperties, constructorof
+constructorof(::Type{<:RigidMap}) = (x₀,θ,xₚ,V,ω,_) -> RigidMap(x₀,θ;xₚ,V,ω) # force precomputation of R̂
 setmap(body::AbstractBody; kwargs...) = setproperties(body,map=setproperties(body.map; kwargs...))
 setmap(body::SetBody; kwargs...) = SetBody(body.op,setmap(body.a; kwargs...),setmap(body.b; kwargs...))
 setmap(body::NoBody; kwargs...) = NoBody()

test/alloctest.jl

@@ -33,4 +33,9 @@ backend != "SIMD" && throw(ArgumentError("KernelAbstractions backend not allowed
     r = run(b)
     println("▶ Allocated "*@sprintf("%.0f", r.memory/1e3)*" KiB")
     @test r.memory < 50000 # less than 50 KiB allocated on the best mom_step! run (commit f721343 ≈ 8 KiB)
+
+    b = @benchmarkable measure!($sim) samples=100; tune!(b) # check 100 times
+    r = run(b)
+    println("▶ Allocated "*@sprintf("%.0f", r.memory/1e3)*" KiB")
+    @test r.memory < 50000 # less than 50 KiB allocated on the best mom_step! run (commit f721343 ≈ 8 KiB)
 end