removes body_force and defines acceleration g as a function of g(i,x,t), hence implementing the space-varying featured previously in charge of body_force. Tests have been mdofied accordingly, an all are locally passing for CPU and GPU. Need to update examples repo.

Author: b-fg

src/Flow.jl

@@ -63,25 +63,15 @@ using EllipsisNotation
 """
     accelerate!(r,t,g,U)
 
-Add a uniform acceleration `gᵢ+dUᵢ/dt` at time `t=sum(dt)` to field `r`.
+Adds a space-time acceleration field g(i,x,t) and velocity field U(i,x,t) such that `rᵢ += gᵢ(i,x,t)+dUᵢ(i,x,t)/dt` at time `t=sum(dt)` to vector field `r`.
 """
-accelerate!(r,t,g::Nothing,U::Function) = body_force!(r,(i,x,t)->ForwardDiff.derivative(τ->U(i,x,τ),t),t)
-accelerate!(r,t,g::Function,U::Function) = body_force!(r,(i,x,t)->g(i,t)+ForwardDiff.derivative(τ->U(i,x,τ),t),t)
-accelerate!(r,t,g::Function,::Tuple) = body_force!(r,(i,x,t)->g(i,t),t)
-accelerate!(r,t,::Nothing,::Tuple) = nothing
-"""
-    body_force!(r,force,t)
-
-Adds a body force to the RHS
-
-- `force` is either a vector field or a function `f(i,x,t)` that returns the component
-  `i` of the vector field at time `t` and position `x`.
-"""
-body_force!(_,::Nothing,_) = nothing
-body_force!(r,force::AbstractArray,_) = r .+= force
-body_force!(r,force::Function,t) = for i ∈ 1:last(size(r))
-    @loop r[I,i] += force(i,loc(i,I,eltype(r)),t) over I ∈ CartesianIndices(Base.front(size(r)))
+accelerate!(_,_,::Nothing,::Union{Nothing,Tuple}) = nothing
+accelerate!(r,t,f) = for i ∈ 1:last(size(r))
+    @loop r[I,i] += f(i,loc(i,I,eltype(r)),t) over I ∈ CartesianIndices(Base.front(size(r)))
 end
+accelerate!(r,t,g::Function,::Union{Nothing,Tuple}) = accelerate!(r,t,g)
+accelerate!(r,t,::Nothing,U::Function) = accelerate!(r,t,(i,x,t)->ForwardDiff.derivative(τ->U(i,x,τ),t))
+accelerate!(r,t,g::Function,U::Function) = accelerate!(r,t,(i,x,t)->g(i,x,t)+ForwardDiff.derivative(τ->U(i,x,τ),t))
 """
     Flow{D::Int, T::Float, Sf<:AbstractArray{T,D}, Vf<:AbstractArray{T,D+1}, Tf<:AbstractArray{T,D+2}}
 
@@ -153,20 +143,20 @@ end
 Integrate the `Flow` one time step using the [Boundary Data Immersion Method](https://eprints.soton.ac.uk/369635/)
 and the `AbstractPoisson` pressure solver to project the velocity onto an incompressible flow.
 """
-@fastmath function mom_step!(a::Flow{N},b::AbstractPoisson;body_force=nothing,udf=nothing,kwargs...) where N
+@fastmath function mom_step!(a::Flow{N},b::AbstractPoisson;udf=nothing,kwargs...) where N
     a.u⁰ .= a.u; scale_u!(a,0); t₁ = sum(a.Δt); t₀ = t₁-a.Δt[end]
     # predictor u → u'
     @log "p"
     conv_diff!(a.f,a.u⁰,a.σ,ν=a.ν,perdir=a.perdir)
-    body_force!(a.f,body_force,t₀); udf!(a,udf,t₀; kwargs...)
+    udf!(a,udf,t₀; kwargs...)
     accelerate!(a.f,t₀,a.g,a.U)
     BDIM!(a); BC!(a.u,a.U,a.exitBC,a.perdir,t₁) # BC MUST be at t₁
     a.exitBC && exitBC!(a.u,a.u⁰,a.Δt[end]) # convective exit
     project!(a,b); BC!(a.u,a.U,a.exitBC,a.perdir,t₁)
     # corrector u → u¹
     @log "c"
     conv_diff!(a.f,a.u,a.σ,ν=a.ν,perdir=a.perdir)
-    body_force!(a.f,body_force,t₁); udf!(a,udf,t₁; kwargs...)
+    udf!(a,udf,t₁; kwargs...)
     accelerate!(a.f,t₁,a.g,a.U)
     BDIM!(a); scale_u!(a,0.5); BC!(a.u,a.U,a.exitBC,a.perdir,t₁)
     project!(a,b,0.5); BC!(a.u,a.U,a.exitBC,a.perdir,t₁)

src/WaterLily.jl

@@ -98,17 +98,17 @@ Can be set to `false` for static geometries to speed up simulation.
 A user-defined function `udf` can be passed to arbitrarily modify the `::Flow` during the predictor and corrector steps.
 If the `udf` user keyword arguments, these needs to be included in the `sim_step!` call as well.
 """
-function sim_step!(sim::AbstractSimulation,t_end;remeasure=true,max_steps=typemax(Int),body_force=nothing,udf=nothing,verbose=false,kwargs...)
+function sim_step!(sim::AbstractSimulation,t_end;remeasure=true,max_steps=typemax(Int),udf=nothing,verbose=false,kwargs...)
     steps₀ = length(sim.flow.Δt)
     while sim_time(sim) < t_end && length(sim.flow.Δt) - steps₀ < max_steps
-        sim_step!(sim; remeasure, body_force, udf, kwargs...)
+        sim_step!(sim; remeasure, udf, kwargs...)
         verbose && println("tU/L=",round(sim_time(sim),digits=4),
             ", Δt=",round(sim.flow.Δt[end],digits=3))
     end
 end
-function sim_step!(sim::AbstractSimulation;remeasure=true,body_force=nothing,udf=nothing,kwargs...)
+function sim_step!(sim::AbstractSimulation;remeasure=true,udf=nothing,kwargs...)
     remeasure && measure!(sim)
-    mom_step!(sim.flow, sim.pois; body_force, udf, kwargs...)
+    mom_step!(sim.flow, sim.pois; udf, kwargs...)
 end
 
 """

test/maintests.jl

@@ -170,24 +170,21 @@ end
         N = 4; a = zeros(N,N,2) |> f
         WaterLily.accelerate!(a,1,nothing,())
         @test all(a .== 0)
-        WaterLily.accelerate!(a,1,(i,t) -> i==1 ? t : 2*t,())
+        WaterLily.accelerate!(a,1,(i,x,t)->i==1 ? t : 2*t,())
         @test all(a[:,:,1] .== 1) && all(a[:,:,2] .== 2)
         WaterLily.accelerate!(a,1,nothing,(i,x,t) -> i==1 ? -t : -2*t)
         @test all(a[:,:,1] .== 0) && all(a[:,:,2] .== 0)
-        WaterLily.accelerate!(a,1,(i,t) -> i==1 ? t : 2*t,(i,x,t) -> i==1 ? -t : -2*t)
+        WaterLily.accelerate!(a,1,(i,x,t) -> i==1 ? t : 2*t,(i,x,t) -> i==1 ? -t : -2*t)
         @test all(a[:,:,1] .== 0) && all(a[:,:,2] .== 0)
-        # check applying body force
-        b = zeros(N,N,2) |> f; bf = ones(N,N,2) |> f
-        WaterLily.body_force!(b, nothing, 0)
-        @test all(b .== 0)
-        WaterLily.body_force!(b, bf, 0)
+        # check applying body force (changes in x but not t)
+        b = zeros(N,N,2) |> f
+        WaterLily.accelerate!(b,0,(i,x,t)->1,nothing)
         @test all(b .== 1)
-        apply!((i,x)->x[i], bf)
-        WaterLily.body_force!(b, bf, 0)
         a .= 0 # reset and accelerate using a non-uniform velocity field
-        WaterLily.accelerate!(a,1.,nothing,(i,x,t)->t*(x[i]+1.0))
+        WaterLily.accelerate!(a,0.,nothing,(i,x,t)->t*(x[i]+1.0))
+        WaterLily.accelerate!(b,0,(i,x,t)->x[i],nothing)
         @test all(b .== a)
-        WaterLily.body_force!(b,(i,x,t)->x[i]+1.0,1.)
+        WaterLily.accelerate!(b,1.,(i,x,t)->x[i]+1.0,nothing)
         WaterLily.accelerate!(a,1.,nothing,(i,x,t)->t*(x[i]+1.0))
         @test all(b .== a)
     end
@@ -304,7 +301,7 @@ function acceleratingFlow(N;use_g=false,T=Float64,perdir=(1,),jerk=4,mem=Array)
     # assuming gravitational scale is 1 and Fr is 1, U scale is Fr*√gL
     UScale = √N  # this is also initial U
     # constant jerk in x, zero acceleration in y
-    g(i,t) = i==1 ? t*jerk : 0
+    g(i,x,t) = i==1 ? t*jerk : 0
     !use_g && (g = nothing)
     return WaterLily.Simulation(
         (N,N), (UScale,0.), N; ν=0.001,g,Δt=0.001,perdir,T,mem