From 0f700672bcd1db014375bdcb19e564ccc49ccc3b Mon Sep 17 00:00:00 2001
From: Gabriel Weymouth <gabriel.weymouth@gmail.com>
Date: Fri, 21 Jun 2024 17:12:49 +0200
Subject: [PATCH] faster and type general

---
 src/AutoBody.jl | 19 ++++++++-----------
 src/Metrics.jl  | 14 +++++++-------
 2 files changed, 15 insertions(+), 18 deletions(-)

diff --git a/src/AutoBody.jl b/src/AutoBody.jl
index 1c09b77..e0fa284 100644
--- a/src/AutoBody.jl
+++ b/src/AutoBody.jl
@@ -100,26 +100,23 @@ sdf(a::Bodies,x,t) = sdf_map_d(a.bodies,a.ops,x,t)[end]
 
 using ForwardDiff
 """
-    d,n,V = measure(body::AutoBody,x,t)
-    d,n,V = measure(body::Bodies,x,t)
+    d,n,V = measure(body::AutoBody,x,t;fast=false)
+    d,n,V = measure(body::Bodies,x,t;fast=false)
 
 Determine the implicit geometric properties from the `sdf` and `map`.
 The gradient of `d=sdf(map(x,t))` is used to improve `d` for pseudo-sdfs.
 The velocity is determined _solely_ from the optional `map` function.
+Using `fast=true` skips the `n,V` calculation when |d|>1.
 """
-measure(body::AutoBody,x,t) = measure(body.sdf,body.map,x,t)
-function measure(a::Bodies,x,t)
+measure(body::AutoBody,x,t;fast=false) = measure(body.sdf,body.map,x,t,fast)
+function measure(a::Bodies,x,t;fast=false)
     sdf, map, _ = sdf_map_d(a.bodies,a.ops,x,t)
-    measure(sdf,map,x,t)
+    measure(sdf,map,x,t,fast)
 end
-# measures the distance, normal, and velocity accuratly for sdf≤1
-function measure_fast(body::AutoBody,x,t)
-    abs(body.sdf(x,t))>1 && return zero(eltype(x)),zero(x),zero(x)
-    measure(body.sdf,body.map,x,t)
-end
-function measure(sdf,map,x,t)
+function measure(sdf,map,x,t,fast)
     # eval d=f(x,t), and n̂ = ∇f
     d = sdf(x,t)
+    fast && abs(d)>1 && return (d,zero(x),zero(x)) # skip n,V
     n = ForwardDiff.gradient(x->sdf(x,t), x)
     any(isnan.(n)) && return (d,zero(x),zero(x))
 
diff --git a/src/Metrics.jl b/src/Metrics.jl
index c02b866..d1f96d6 100644
--- a/src/Metrics.jl
+++ b/src/Metrics.jl
@@ -82,7 +82,7 @@ end
 BDIM-masked surface normal.
 """
 @inline function nds(body,x,t)
-    d,n,_ = measure_fast(body,x,t)
+    d,n,_ = measure(body,x,t,fast=true)
     n*WaterLily.kern(clamp(d,-1,1))
 end
 
@@ -93,10 +93,10 @@ Compute the pressure force on an immersed body.
 """
 pressure_force(sim) = pressure_force(sim.flow,sim.body)
 pressure_force(flow,body) = pressure_force(flow.p,flow.f,body,time(flow))
-function pressure_force(p,df,body,t=0)
+function pressure_force(p,df,body,t=0,T=promote_type(Float64,eltype(p)))
     df .= zero(eltype(p))
     @loop df[I,:] .= p[I]*nds(body,loc(0,I),t) over I ∈ inside(p)
-    sum(Float64,df,dims=ntuple(i->i,ndims(p)))[:] |> Array
+    sum(T,df,dims=ntuple(i->i,ndims(p)))[:] |> Array
 end
 
 """
@@ -113,10 +113,10 @@ Compute the viscous force on an immersed body.
 """
 viscous_force(sim) = viscous_force(sim.flow,sim.body)
 viscous_force(flow,body) = viscous_force(flow.u,flow.ν,flow.f,body,time(flow))
-function viscous_force(u,ν,df,body,t=0)
+function viscous_force(u,ν,df,body,t=0,T=promote_type(Float64,eltype(u)))
     df .= zero(eltype(u))
     @loop df[I,:] .= -ν*∇²u(I,u)*nds(body,loc(0,I),t) over I ∈ inside_u(u)
-    sum(Float64,df,dims=ntuple(i->i,ndims(u)-1))[:] |> Array
+    sum(T,df,dims=ntuple(i->i,ndims(u)-1))[:] |> Array
 end
 
 """
@@ -134,8 +134,8 @@ Computes the pressure moment on an immersed body relative to point x₀.
 """
 pressure_moment(x₀,sim) = pressure_moment(x₀,sim.flow,sim.body)
 pressure_moment(x₀,flow,body) = pressure_moment(x₀,flow.p,flow.f,body,time(flow))
-function pressure_moment(x₀,p,df,body,t=0)
+function pressure_moment(x₀,p,df,body,t=0,T=promote_type(Float64,eltype(p)))
     df .= zero(eltype(p))
     @loop df[I,:] .= p[I]*cross(loc(0,I)-x₀,nds(body,loc(0,I),t)) over I ∈ inside(p)
-    sum(Float64,df,dims=ntuple(i->i,ndims(p)))[:] |> Array
+    sum(T,df,dims=ntuple(i->i,ndims(p)))[:] |> Array
 end
\ No newline at end of file