WaterLily-jl · weymouth · Apr 19, 2025 · Sep 6, 2024 · Sep 6, 2024 · Sep 9, 2024
diff --git a/src/Flow.jl b/src/Flow.jl
@@ -59,26 +59,18 @@ lowerBoundary!(r,u,Φ,ν,i,j,N,::Val{true}) = @loop (
     Φ[I] = ϕuP(j,CIj(j,CI(I,i),N[j]-2),CI(I,i),u,ϕ(i,CI(I,j),u)) -ν*∂(j,CI(I,i),u); r[I,i] += Φ[I]) over I ∈ slice(N,2,j,2)
 upperBoundary!(r,u,Φ,ν,i,j,N,::Val{true}) = @loop r[I-δ(j,I),i] -= Φ[CIj(j,I,2)] over I ∈ slice(N,N[j],j,2)
 
-using EllipsisNotation
 """
-    accelerate!(r,dt,g)
+    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,dt,g::Function,::Tuple,t=sum(dt)) = for i ∈ 1:last(size(r))
-    r[..,i] .+= g(i,t)
+accelerate!(r,t,::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,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
-
+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}}
 
@@ -108,11 +100,11 @@ 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=(), exitBC=false, T=Float64) where D
+                  uλ::Function=(i, x) -> 0., perdir=(), exitBC=false, T=Float32) 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.)
+        BC!(u,U,exitBC,perdir); exitBC!(u,u,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
@@ -150,21 +142,23 @@ 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) where N
-    a.u⁰ .= a.u; scale_u!(a,0); U = BCTuple(a.U,a.Δt,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)
-    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)
+    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)
-    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)
+    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₁)
     push!(a.Δt,CFL(a))
 end
 scale_u!(a,scale) = @loop a.u[Ii] *= scale over Ii ∈ inside_u(size(a.p))
@@ -180,3 +174,12 @@ end
     end
     return s
 end
+
+"""
+    udf!(flow::Flow,udf::Function,t)
+
+User defined function using `udf::Function` to operate on `flow::Flow` during the predictor and corrector step, in sync with time `t`.
+Keyword arguments must be passed to `sim_step!` for them to be carried over the actual function call.
+"""
+udf!(flow,::Nothing,t; kwargs...) = nothing
+udf!(flow,force!::Function,t; kwargs...) = force!(flow,t; kwargs...)
diff --git a/src/Metrics.jl b/src/Metrics.jl
@@ -100,12 +100,12 @@ function pressure_force(p,df,body,t=0,T=promote_type(Float64,eltype(p)))
 end
 
 """
-    ∇²u(I::CartesianIndex,u)
+    S(I::CartesianIndex,u)
 
 Rate-of-strain tensor.
 """
-∇²u(I::CartesianIndex{2},u) = @SMatrix [∂(i,j,I,u)+∂(j,i,I,u) for i ∈ 1:2, j ∈ 1:2]
-∇²u(I::CartesianIndex{3},u) = @SMatrix [∂(i,j,I,u)+∂(j,i,I,u) for i ∈ 1:3, j ∈ 1:3]
+S(I::CartesianIndex{2},u) = @SMatrix [0.5*(∂(i,j,I,u)+∂(j,i,I,u)) for i ∈ 1:2, j ∈ 1:2]
+S(I::CartesianIndex{3},u) = @SMatrix [0.5*(∂(i,j,I,u)+∂(j,i,I,u)) for i ∈ 1:3, j ∈ 1:3]
 """
    viscous_force(sim::Simulation)
 
@@ -115,7 +115,7 @@ 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,T=promote_type(Float64,eltype(u)))
     df .= zero(eltype(u))
-    @loop df[I,:] .= -ν*∇²u(I,u)*nds(body,loc(0,I,T),t) over I ∈ inside_u(u)
+    @loop df[I,:] .= -2ν*S(I,u)*nds(body,loc(0,I,T),t) over I ∈ inside_u(u)
     sum(T,df,dims=ntuple(i->i,ndims(u)-1))[:] |> Array
 end
 

diff --git a/src/WaterLily.jl b/src/WaterLily.jl
@@ -40,7 +40,7 @@ Constructor for a WaterLily.jl simulation:
 
   - `dims`: Simulation domain dimensions.
   - `u_BC`: Simulation domain velocity boundary conditions, either a
-            tuple `u_BC[i]=uᵢ, i=eachindex(dims)`, or a time-varying function `f(i,t)`
+            tuple `u_BC[i]=uᵢ, i=eachindex(dims)`, or a time and space-varying function `u_BC(i,x,t)`
   - `L`: Simulation length scale.
   - `U`: Simulation velocity scale.
   - `Δt`: Initial time step.
@@ -69,9 +69,12 @@ mutable struct Simulation <: AbstractSimulation
                         T=Float32, mem=Array) where N
         @assert !(isa(u_BC,Function) && isa(uλ,Function)) "`u_BC` and `uλ` cannot be both specified as Function"
         @assert !(isnothing(U) && isa(u_BC,Function)) "`U` must be specified if `u_BC` is a Function"
-        isa(u_BC,Function) && @assert all(typeof.(ntuple(i->u_BC(i,zero(T)),N)).==T) "`u_BC` is not type stable"
-        uλ = isnothing(uλ) ? ifelse(isa(u_BC,Function),(i,x)->u_BC(i,0.),(i,x)->u_BC[i]) : uλ
-        U = isnothing(U) ? √sum(abs2,u_BC) : U # default if not specified
+        isa(g,Function) && @assert first(methods(g)).nargs==4 "g::Function needs to be defined as g(i,x,t)"
+        isa(g,Function) && @assert all(typeof.(ntuple(i->g(i,zeros(SVector{N}),zero(T)),N)).==T) "`g` is not type stable"
+        isa(u_BC,Function) && @assert first(methods(u_BC)).nargs==4 "u_BC::Function needs to be defined as u_BC(i,x,t)"
+        isa(u_BC,Function) && @assert all(typeof.(ntuple(i->u_BC(i,zeros(SVector{N}),zero(T)),N)).==T) "`u_BC` is not type stable"
+        isnothing(uλ) && (uλ = isa(u_BC, Function) ? uλ = (i,x)->u_BC(i,x,0) : uλ = (i,_)->u_BC[i])
+        isnothing(U) && (U = √sum(abs2,u_BC))
         flow = Flow(dims,u_BC;uλ,Δt,ν,g,T,f=mem,perdir,exitBC)
         measure!(flow,body;ϵ)
         new(U,L,ϵ,flow,body,MultiLevelPoisson(flow.p,flow.μ₀,flow.σ;perdir))
@@ -94,18 +97,20 @@ sim_time(sim::AbstractSimulation) = time(sim)*sim.U/sim.L
 Integrate the simulation `sim` up to dimensionless time `t_end`.
 If `remeasure=true`, the body is remeasured at every time step.
 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),verbose=false)
+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)
+        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)
+function sim_step!(sim::AbstractSimulation;remeasure=true,udf=nothing,kwargs...)
     remeasure && measure!(sim)
-    mom_step!(sim.flow,sim.pois)
+    mom_step!(sim.flow, sim.pois; udf, kwargs...)
 end
 
 """

diff --git a/src/util.jl b/src/util.jl
@@ -189,7 +189,8 @@
 boundary. For example `aₓ(x)=Aₓ ∀ x ∈ minmax(X)`. A zero Neumann condition
 is applied to the tangential components.
 """
-function BC!(a,A,saveexit=false,perdir=())
+BC!(a,U,saveexit=false,perdir=(),t=0) = BC!(a,(i,x,t)->U[i],saveexit,perdir,t)
+function BC!(a,u_BC::Function,saveexit=false,perdir=(),t=0)
     N,n = size_u(a)
     for i ∈ 1:n, j ∈ 1:n
         if j in perdir
@@ -198,26 +199,28 @@
         else
             if i==j # Normal direction, Dirichlet
                 for s ∈ (1,2)
-                    @loop a[I,i] = A[i] over I ∈ slice(N,s,j)
+                    @loop a[I,i] = u_BC(i,loc(i,I),t) over I ∈ slice(N,s,j)
                 end
-                (!saveexit || i>1) && (@loop a[I,i] = A[i] over I ∈ slice(N,N[j],j)) # overwrite exit
+                (!saveexit || i>1) && (@loop a[I,i] = u_BC(i,loc(i,I),t) over I ∈ slice(N,N[j],j)) # overwrite exit
             else    # Tangential directions, Neumann
                 @loop a[I,i] = a[I+δ(j,I),i] over I ∈ slice(N,1,j)
                 @loop a[I,i] = a[I-δ(j,I),i] over I ∈ slice(N,N[j],j)
             end
         end
     end
 end
+
 """
     exitBC!(u,u⁰,U,Δt)
 
 Apply a 1D convection scheme to fill the ghost cell on the exit of the domain.
 """
-function exitBC!(u,u⁰,U,Δt)
+function exitBC!(u,u⁰,Δt)
     N,_ = size_u(u)
     exitR = slice(N.-1,N[1],1,2)              # exit slice excluding ghosts
-    @loop u[I,1] = u⁰[I,1]-U[1]*Δt*(u⁰[I,1]-u⁰[I-δ(1,I),1]) over I ∈ exitR
-    ∮u = sum(u[exitR,1])/length(exitR)-U[1]   # mass flux imbalance
+    U = sum(@view(u[slice(N.-1,2,1,2),1]))/length(exitR) # inflow mass flux
+    @loop u[I,1] = u⁰[I,1]-U*Δt*(u⁰[I,1]-u⁰[I-δ(1,I),1]) over I ∈ exitR
+    ∮u = sum(@view(u[exitR,1]))/length(exitR)-U   # mass flux imbalance
     @loop u[I,1] -= ∮u over I ∈ exitR         # correct flux
 end
 """
@@ -250,6 +253,7 @@
     end
     return s
 end
+using EllipsisNotation
 function interp(x::SVector{D}, varr::AbstractArray) where {D}
     # Shift to align with each staggered grid component and interpolate
     @inline shift(i) = SVector{D}(ifelse(i==j,0.5,0.0) for j in 1:D)

diff --git a/test/maintests.jl b/test/maintests.jl
@@ -42,9 +42,20 @@ using ReadVTK, WriteVTK
         BC!(u,U,true) # save exit values
         @test GPUArrays.@allowscalar all(u[end, :, 1] .== 3)
 
-        WaterLily.exitBC!(u,u,U,0) # conservative exit check
+        WaterLily.exitBC!(u,u,0) # conservative exit check
         @test GPUArrays.@allowscalar all(u[end,2:end-1, 1] .== U[1])
 
+        # test BC with function
+        Ubc(i,x,t) = i==1 ? 1.0 : 0.5
+        v = rand(Ng..., D) |> f # vector
+        BC!(v,Ubc,false); BC!(u,U,false) # make sure we apply the same
+        @test all(v[1, :, 1] .≈ u[1, :, 1]) && all(v[2, :, 1] .≈ u[2, :, 1]) && all(v[end, :, 1] .≈ u[end, :, 1])
+        @test all(v[:, 1, 2] .≈ u[:, 1, 2]) && all(v[:, 2, 2] .≈ u[:, 2, 2]) && all(v[:, end, 2] .≈ u[:, end, 2])
+        # test exit bc
+        GPUArrays.@allowscalar v[end,:,1] .= 3
+        BC!(v,Ubc,true) # save exit values
+        @test GPUArrays.@allowscalar all(v[end, :, 1] .== 3)
+
         BC!(u,U,true,(2,)) # periodic in y and save exit values
         @test GPUArrays.@allowscalar all(u[:, 1:2, 1] .== u[:, end-1:end, 1]) && all(u[:, 1:2, 1] .== u[:,end-1:end,1])
         WaterLily.perBC!(σ,(1,2)) # periodic in two directions
@@ -154,20 +165,30 @@ end
     Ip = WaterLily.CIj(1,I,length(f)-2); # make periodic
     @test ϕuP(1,Ip,I,f,1)==λ(f[Ip],f[I-δ(1,I)],f[I])
 
-    @test all(WaterLily.BCTuple((1,2,3),[0],3).==WaterLily.BCTuple((i,t)->i,0,3))
-    @test all(WaterLily.BCTuple((i,t)->t,[1.234],3).==ntuple(i->1.234,3))
-
     # check applying acceleration
     for f ∈ arrays
         N = 4; a = zeros(N,N,2) |> f
-        WaterLily.accelerate!(a,[1],nothing,())
+        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,t) -> i==1 ? -t : -2*t)
+        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,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 (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)
+        WaterLily.accelerate!(b,1.,(i,x,t)->0,(i,x,t)->t)
+        @test all(b .== 2)
+        a .= 0; b .= 1 # reset and accelerate using a non-uniform velocity field
+        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.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
     # Impulsive flow in a box
     U = (2/3, -1/3)
@@ -277,34 +298,47 @@ end
     @test derivative(lift,2.0) ≈ (lift(2+h)-lift(2-h))/2h rtol=√h
 end
 
-function acceleratingFlow(N;T=Float64,perdir=(1,),jerk=4,mem=Array)
+function acceleratingFlow(N;use_g=false,T=Float64,perdir=(1,),jerk=4,mem=Array)
     # periodic in x, Neumann in y
     # 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
     ),jerk
 end
+gravity!(flow::Flow,t; jerk=4) = for i ∈ 1:last(size(flow.f))
+    WaterLily.@loop flow.f[I,i] += i==1 ? t*jerk : 0 over I ∈ CartesianIndices(Base.front(size(flow.f)))
+end
 @testset "Flow.jl with increasing body force" begin
     for f ∈ arrays
         N = 8
-        sim,jerk = acceleratingFlow(N;mem=f)
+        sim,jerk = acceleratingFlow(N;use_g=true,mem=f)
         sim_step!(sim,1.0); u = sim.flow.u |> Array
         # Exact uₓ = uₓ₀ + ∫ a dt = uₓ₀ + ∫ jerk*t dt = uₓ₀ + 0.5*jerk*t^2
         uFinal = sim.flow.U[1] + 0.5*jerk*WaterLily.time(sim)^2
         @test (
             WaterLily.L₂(u[:,:,1].-uFinal) < 1e-4 &&
             WaterLily.L₂(u[:,:,2].-0) < 1e-4
         )
+
+        # Test with user defined function instead of acceleration
+        sim_udf,_ = acceleratingFlow(N;mem=f)
+        sim_step!(sim_udf,1.0; udf=gravity!, jerk=jerk); u_udf = sim_udf.flow.u |> Array
+        uFinal = sim_udf.flow.U[1] + 0.5*jerk*WaterLily.time(sim_udf)^2
+        @test (
+            WaterLily.L₂(u_udf[:,:,1].-uFinal) < 1e-4 &&
+            WaterLily.L₂(u_udf[:,:,2].-0) < 1e-4
+        )
     end
 end
 
 @testset "Circle in accelerating flow" begin
     for f ∈ arrays
         make_accel_circle(radius=32,H=16) = Simulation(radius.*(2H,2H),
-            (i,t)-> i==1 ? t : zero(t), radius; U=1, mem=f,
+            (i,x,t)-> i==1 ? t : zero(t), radius; U=1, mem=f,
             body=AutoBody((x,t)->√sum(abs2,x .-H*radius)-radius))
         sim = make_accel_circle(); sim_step!(sim)
         @test isapprox(WaterLily.pressure_force(sim)/(π*sim.L^2),[-1,0],atol=0.04)
@@ -347,16 +381,16 @@ import WaterLily: ×
         # stress tensor
         u₂ = zeros(N,N,2) |> f
         u₃ = zeros(N,N,N,3) |> f
-        @test GPUArrays.@allowscalar all(WaterLily.∇²u(CartesianIndex(N÷2,N÷2),u₂) .≈ 0)
-        @test GPUArrays.@allowscalar all(WaterLily.∇²u(CartesianIndex(N÷2,N÷2,N÷2),u₃) .≈ 0)
+        @test GPUArrays.@allowscalar all(2WaterLily.S(CartesianIndex(N÷2,N÷2),u₂) .≈ 0)
+        @test GPUArrays.@allowscalar all(2WaterLily.S(CartesianIndex(N÷2,N÷2,N÷2),u₃) .≈ 0)
         apply!((i,x)->x[i],u₂) # uniform gradient
         apply!((i,x)->x[i],u₃)
-        @test GPUArrays.@allowscalar all(WaterLily.∇²u(CartesianIndex(N÷2,N÷2),u₂) .≈ SA[2 0; 0 2])
-        @test GPUArrays.@allowscalar all(WaterLily.∇²u(CartesianIndex(N÷2,N÷2,N÷2),u₃) .≈ SA[2 0 0; 0 2 0; 0 0 2])
+        @test GPUArrays.@allowscalar all(2WaterLily.S(CartesianIndex(N÷2,N÷2),u₂) .≈ SA[2 0; 0 2])
+        @test GPUArrays.@allowscalar all(2WaterLily.S(CartesianIndex(N÷2,N÷2,N÷2),u₃) .≈ SA[2 0 0; 0 2 0; 0 0 2])
         apply!((i,x)->x[i%2+1],u₂) # shear
         apply!((i,x)->x[i%3+1],u₃)
-        @test GPUArrays.@allowscalar all(WaterLily.∇²u(CartesianIndex(N÷2,N÷2),u₂) .≈ SA[0 2; 2 0])
-        @test GPUArrays.@allowscalar all(WaterLily.∇²u(CartesianIndex(N÷2,N÷2,N÷2),u₃) .≈ SA[0 1 1; 1 0 1; 1 1 0])
+        @test GPUArrays.@allowscalar all(2WaterLily.S(CartesianIndex(N÷2,N÷2),u₂) .≈ SA[0 2; 2 0])
+        @test GPUArrays.@allowscalar all(2WaterLily.S(CartesianIndex(N÷2,N÷2,N÷2),u₃) .≈ SA[0 1 1; 1 0 1; 1 1 0])
         # viscous force
         u₂ .= 0; u₃ .= 0
         @test all(WaterLily.viscous_force(u₂,1.0,df₂,body) .≈ 0)