Merge pull request #198 from WaterLily-jl/turbulence_modelling

Turbulence utils: SGS model and on-the-fly temporal averaging

Author: b-fg

ext/WaterLilyJLD2Ext.jl

@@ -23,7 +23,20 @@ save!(fname, flow::Flow; dir="./") = jldsave(
 save!(fname::String, sim::AbstractSimulation; dir="./") = save!(fname, sim.flow; dir)
 
 """
-    load!(flow::Flow, fname::String; dir="./")
+    save!(fname, meanflow::MeanFlow; dir="./")
+
+Save the `meanflow::MeanFlow` time-averaged pressure, velocity, velocity-squared, and time arrays into a JLD2-formatted binary file (HDF5 compatible).
+"""
+save!(fname, meanflow::MeanFlow; dir="./") = jldsave(
+    joinpath(dir, fname);
+    P=Array(meanflow.P),
+    U=Array(meanflow.U),
+    UU=Array(meanflow.UU),
+    t=meanflow.t
+)
+
+"""
+    load!(flow::Flow; kwargs...)
 
 Load pressure, velocity, and time steps arrays from a JLD2-formatted binary file.
 Keyword arguments considered are `fname="WaterLily.jld2"` and `dir="./"`.
@@ -42,4 +55,24 @@ function load!(flow::Flow; kwargs...)
 end
 load!(sim::AbstractSimulation, ::Val{:jld2}; kwargs...) = load!(sim.flow; kwargs...)
 
+
+"""
+    load!(meanflow::MeanFlow; kwargs...)
+
+Load time-averaged pressure, velocity, velocity-square, and time arrays from a JLD2-formatted binary file.
+Keyword arguments considered are `fname="WaterLilyMean.jld2"` and `dir="./"`.
+"""
+function load!(meanflow::MeanFlow; kwargs...)
+    fname = get(Dict(kwargs), :fname, "WaterLilyMean.jld2")
+    dir = get(Dict(kwargs), :dir, "./")
+    obj = jldopen(joinpath(dir, fname))
+    @assert size(meanflow.P) == size(obj["P"]) "Simulation size does not match the size of the JLD2-stored simulation."
+    f = typeof(meanflow.P).name.wrapper
+    meanflow.P .= obj["P"] |> f
+    meanflow.U .= obj["U"] |> f
+    empty!(meanflow.t)
+    push!(meanflow.t, obj["t"]...)
+    close(obj)
+end
+
 end # module

src/Flow.jl

@@ -3,10 +3,12 @@
 @inline ϕ(a,I,f) = @inbounds (f[I]+f[I-δ(a,I)])/2
 @fastmath quick(u,c,d) = median((5c+2d-u)/6,c,median(10c-9u,c,d))
 @fastmath vanLeer(u,c,d) = (c≤min(u,d) || c≥max(u,d)) ? c : c+(d-c)*(c-u)/(d-u)
-@inline ϕu(a,I,f,u,λ=quick) = @inbounds u>0 ? u*λ(f[I-2δ(a,I)],f[I-δ(a,I)],f[I]) : u*λ(f[I+δ(a,I)],f[I],f[I-δ(a,I)])
-@inline ϕuP(a,Ip,I,f,u,λ=quick) = @inbounds u>0 ? u*λ(f[Ip],f[I-δ(a,I)],f[I]) : u*λ(f[I+δ(a,I)],f[I],f[I-δ(a,I)])
-@inline ϕuL(a,I,f,u,λ=quick) = @inbounds u>0 ? u*ϕ(a,I,f) : u*λ(f[I+δ(a,I)],f[I],f[I-δ(a,I)])
-@inline ϕuR(a,I,f,u,λ=quick) = @inbounds u<0 ? u*ϕ(a,I,f) : u*λ(f[I-2δ(a,I)],f[I-δ(a,I)],f[I])
+@fastmath cds(u,c,d) = (c+d)/2
+
+@inline ϕu(a,I,f,u,λ) = @inbounds u>0 ? u*λ(f[I-2δ(a,I)],f[I-δ(a,I)],f[I]) : u*λ(f[I+δ(a,I)],f[I],f[I-δ(a,I)])
+@inline ϕuP(a,Ip,I,f,u,λ) = @inbounds u>0 ? u*λ(f[Ip],f[I-δ(a,I)],f[I]) : u*λ(f[I+δ(a,I)],f[I],f[I-δ(a,I)])
+@inline ϕuL(a,I,f,u,λ) = @inbounds u>0 ? u*ϕ(a,I,f) : u*λ(f[I+δ(a,I)],f[I],f[I-δ(a,I)])
+@inline ϕuR(a,I,f,u,λ) = @inbounds u<0 ? u*ϕ(a,I,f) : u*λ(f[I-2δ(a,I)],f[I-δ(a,I)],f[I])
 
 @fastmath @inline function div(I::CartesianIndex{m},u) where {m}
     init=zero(eltype(u))
@@ -33,31 +35,31 @@ function median(a,b,c)
     return a
 end
 
-function conv_diff!(r,u,Φ;ν=0.1,perdir=())
+function conv_diff!(r,u,Φ,λ::F;ν=0.1,perdir=()) where {F}
     r .= zero(eltype(r))
     N,n = size_u(u)
     for i ∈ 1:n, j ∈ 1:n
         # if it is periodic direction
         tagper = (j in perdir)
         # treatment for bottom boundary with BCs
-        lowerBoundary!(r,u,Φ,ν,i,j,N,Val{tagper}())
+        lowerBoundary!(r,u,Φ,ν,i,j,N,λ,Val{tagper}())
         # inner cells
-        @loop (Φ[I] = ϕu(j,CI(I,i),u,ϕ(i,CI(I,j),u)) - ν*∂(j,CI(I,i),u);
+        @loop (Φ[I] = ϕu(j,CI(I,i),u,ϕ(i,CI(I,j),u),λ) - ν*∂(j,CI(I,i),u);
                r[I,i] += Φ[I]) over I ∈ inside_u(N,j)
         @loop r[I-δ(j,I),i] -= Φ[I] over I ∈ inside_u(N,j)
         # treatment for upper boundary with BCs
-        upperBoundary!(r,u,Φ,ν,i,j,N,Val{tagper}())
+        upperBoundary!(r,u,Φ,ν,i,j,N,λ,Val{tagper}())
     end
 end
 
 # Neumann BC Building block
-lowerBoundary!(r,u,Φ,ν,i,j,N,::Val{false}) = @loop r[I,i] += ϕuL(j,CI(I,i),u,ϕ(i,CI(I,j),u)) - ν*∂(j,CI(I,i),u) over I ∈ slice(N,2,j,2)
-upperBoundary!(r,u,Φ,ν,i,j,N,::Val{false}) = @loop r[I-δ(j,I),i] += -ϕuR(j,CI(I,i),u,ϕ(i,CI(I,j),u)) + ν*∂(j,CI(I,i),u) over I ∈ slice(N,N[j],j,2)
+lowerBoundary!(r,u,Φ,ν,i,j,N,λ,::Val{false}) = @loop r[I,i] += ϕuL(j,CI(I,i),u,ϕ(i,CI(I,j),u),λ) - ν*∂(j,CI(I,i),u) over I ∈ slice(N,2,j,2)
+upperBoundary!(r,u,Φ,ν,i,j,N,λ,::Val{false}) = @loop r[I-δ(j,I),i] += -ϕuR(j,CI(I,i),u,ϕ(i,CI(I,j),u),λ) + ν*∂(j,CI(I,i),u) over I ∈ slice(N,N[j],j,2)
 
 # Periodic BC Building block
-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)
+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)
 
 """
     accelerate!(r,t,g,U)
@@ -137,24 +139,24 @@ function project!(a::Flow{n},b::AbstractPoisson,w=1) where n
 end
 
 """
-    mom_step!(a::Flow,b::AbstractPoisson)
+    mom_step!(a::Flow,b::AbstractPoisson;λ=quick,udf=nothing,kwargs...)
 
 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;udf=nothing,kwargs...) where N
+@fastmath function mom_step!(a::Flow{N},b::AbstractPoisson;λ=quick,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)
+    conv_diff!(a.f,a.u⁰,a.σ,λ;ν=a.ν,perdir=a.perdir)
     udf!(a,udf,t₀; kwargs...)
     accelerate!(a.f,t₀,a.g,a.uBC)
     BDIM!(a); BC!(a.u,a.uBC,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.uBC,a.exitBC,a.perdir,t₁)
     # corrector u → u¹
     @log "c"
-    conv_diff!(a.f,a.u,a.σ,ν=a.ν,perdir=a.perdir)
+    conv_diff!(a.f,a.u,a.σ,λ;ν=a.ν,perdir=a.perdir)
     udf!(a,udf,t₁; kwargs...)
     accelerate!(a.f,t₁,a.g,a.uBC)
     BDIM!(a); scale_u!(a,0.5); BC!(a.u,a.uBC,a.exitBC,a.perdir,t₁)

src/Metrics.jl

@@ -9,6 +9,13 @@ Base.@propagate_inbounds @fastmath function permute(f,i)
     f(j,k)-f(k,j)
 end
 ×(a,b) = fSV(i->permute((j,k)->a[j]*b[k],i),3)
+@fastmath @inline function dot(a,b)
+    init=zero(eltype(a))
+    @inbounds for ij in eachindex(a)
+     init += a[ij] * b[ij]
+    end
+    return init
+end
 
 """
     ke(I::CartesianIndex,u,U=0)
@@ -141,4 +148,60 @@ function pressure_moment(x₀,p,df,body,t=0)
     df .= zero(Tp)
     @loop df[I,:] .= p[I]*cross(loc(0,I,Tp)-x₀,nds(body,loc(0,I,Tp),t)) over I ∈ inside(p)
     sum(To,df,dims=ntuple(i->i,ndims(p)))[:] |> Array
+end
+
+"""
+     MeanFlow{T, Sf<:AbstractArray{T}, Vf<:AbstractArray{T}, Mf<:AbstractArray{T}}
+
+Holds temporal averages of velocity, squared velocity, pressure, and Reynolds stresses.
+"""
+struct MeanFlow{T, Sf<:AbstractArray{T}, Vf<:AbstractArray{T}, Mf}
+    P :: Sf # pressure scalar field
+    U :: Vf # velocity vector field
+    UU :: Mf # squared velocity tensor, u⊗u
+    t :: Vector{T} # time steps vector
+    uu_stats :: Bool # flag to compute UU on-the-fly temporal averages
+    function MeanFlow(flow::Flow{D,T}; t_init=time(flow), uu_stats=false) where {D,T}
+        f = typeof(flow.u).name.wrapper
+        P = zeros(T, size(flow.p)) |> f
+        U = zeros(T, size(flow.u)) |> f
+        UU = uu_stats ? zeros(T, size(flow.p)...,D,D) |> f : nothing
+        new{T,typeof(P),typeof(U),typeof(UU)}(P,U,UU,T[t_init],uu_stats)
+    end
+end
+
+time(meanflow::MeanFlow) = meanflow.t[end]-meanflow.t[1]
+
+function reset!(meanflow::MeanFlow; t_init=0.0)
+    fill!(meanflow.P, 0); fill!(meanflow.U, 0)
+    !isnothing(meanflow.UU) && fill!(meanflow.UU, 0)
+    deleteat!(meanflow.t, collect(1:length(meanflow.t)))
+    push!(meanflow.t, t_init)
+end
+
+function update!(meanflow::MeanFlow, flow::Flow)
+    dt = time(flow) - meanflow.t[end]
+    ε = dt / (dt + (meanflow.t[end] - meanflow.t[1]) + eps(eltype(flow.p)))
+    @loop meanflow.P[I] = ε * flow.p[I] + (1 - ε) * meanflow.P[I] over I in CartesianIndices(flow.p)
+    @loop meanflow.U[Ii] = ε * flow.u[Ii] + (1 - ε) * meanflow.U[Ii] over Ii in CartesianIndices(flow.u)
+    if meanflow.uu_stats
+        for i in 1:ndims(flow.p), j in 1:ndims(flow.p)
+            @loop meanflow.UU[I,i,j] = ε * (flow.u[I,i] .* flow.u[I,j]) + (1 - ε) * meanflow.UU[I,i,j] over I in CartesianIndices(flow.p)
+        end
+    end
+    push!(meanflow.t, meanflow.t[end] + dt)
+end
+
+uu!(τ,a::MeanFlow) = for i in 1:ndims(a.P), j in 1:ndims(a.P)
+    @loop τ[I,i,j] = a.UU[I,i,j] - a.U[I,i,j] * a.U[I,i,j] over I in CartesianIndices(a.P)
+end
+function uu(a::MeanFlow)
+    τ = zeros(eltype(a.UU), size(a.UU)...) |> typeof(meanflow.UU).name.wrapper
+    uu!(τ,a)
+    return τ
+end
+
+function copy!(a::Flow, b::MeanFlow)
+    a.u .= b.U
+    a.p .= b.P
 end

src/WaterLily.jl

@@ -18,7 +18,7 @@ include("MultiLevelPoisson.jl")
 export MultiLevelPoisson,solver!,mult!
 
 include("Flow.jl")
-export Flow,mom_step!
+export Flow,mom_step!,quick,cds
 
 include("Body.jl")
 export AbstractBody,measure_sdf!
@@ -27,6 +27,7 @@ include("AutoBody.jl")
 export AutoBody,Bodies,measure,sdf,+,-
 
 include("Metrics.jl")
+export MeanFlow
 
 abstract type AbstractSimulation end
 """
@@ -88,24 +89,29 @@ scales.
 sim_time(sim::AbstractSimulation) = time(sim)*sim.U/sim.L
 
 """
-    sim_step!(sim::Simulation,t_end=sim(time)+Δt;max_steps=typemax(Int),remeasure=true,verbose=false)
+    sim_step!(sim::AbstractSimulation,t_end;remeasure=true,λ=quick,max_steps=typemax(Int),verbose=false,
+        udf=nothing,meanflow=nothing,kwargs...)
 
 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.
+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.
+A `::MeanFlow` can also be passed to compute on-the-fly temporal averages.
+A `λ::Function` function can be passed as a custom convective scheme, following the interface of `λ(u,c,d)` (for upstream, central,
+downstream points).
 """
-function sim_step!(sim::AbstractSimulation,t_end;remeasure=true,max_steps=typemax(Int),udf=nothing,verbose=false,kwargs...)
+function sim_step!(sim::AbstractSimulation,t_end;remeasure=true,λ=quick,max_steps=typemax(Int),verbose=false,
+        udf=nothing,meanflow=nothing,kwargs...)
     steps₀ = length(sim.flow.Δt)
     while sim_time(sim) < t_end && length(sim.flow.Δt) - steps₀ < max_steps
-        sim_step!(sim; remeasure, udf, kwargs...)
+        sim_step!(sim; remeasure, λ, udf, meanflow, kwargs...)
         verbose && sim_info(sim)
     end
 end
-function sim_step!(sim::AbstractSimulation;remeasure=true,udf=nothing,kwargs...)
+function sim_step!(sim::AbstractSimulation;remeasure=true,λ=quick,udf=nothing,meanflow=nothing,kwargs...)
     remeasure && measure!(sim)
-    mom_step!(sim.flow, sim.pois; udf, kwargs...)
+    mom_step!(sim.flow, sim.pois; λ, udf, kwargs...)
+    !isnothing(meanflow) && update!(meanflow,sim.flow)
 end
 
 """

src/util.jl

@@ -136,22 +136,23 @@ macro loop(args...)
     grab!(sym,ex)     # get arguments and replace composites in `ex`
     setdiff!(sym,[I]) # don't want to pass I as an argument
     symT = symtypes(sym) # generate a list of types for each symbol
+    symWtypes = joinsymtype(rep.(sym),symT) # symbols with types: [a::A, b::B, ...]
     @gensym(kern, kern_) # generate unique kernel function names for serial and KA execution
     @static if backend == "KernelAbstractions"
         return quote
-            @kernel function $kern_($(rep.(sym)...),@Const(I0)) # replace composite arguments
+            @kernel function $kern_($(symWtypes...),@Const(I0)) where {$(symT...)} # replace composite arguments
                 $I = @index(Global,Cartesian)
                 $I += I0
                 @fastmath @inbounds $ex
             end
-            function $kern($(joinsymtype(rep.(sym),symT)...)) where {$(symT...)}
+            function $kern($(symWtypes...)) where {$(symT...)}
                 $kern_(get_backend($(sym[1])),64)($(sym...),$R[1]-oneunit($R[1]),ndrange=size($R))
             end
             $kern($(sym...))
         end |> esc
     else # backend == "SIMD"
         return quote
-            function $kern($(joinsymtype(rep.(sym),symT)...)) where {$(symT...)}
+            function $kern($(symWtypes...)) where {$(symT...)}
                 @simd for $I ∈ $R
                     @fastmath @inbounds $ex
                 end
@@ -285,6 +286,39 @@ function interp(x::SVector{D}, varr::AbstractArray) where {D}
     return SVector{D}(interp(x+shift(i),@view(varr[..,i])) for i in 1:D)
 end
 
+"""
+    sgs!(flow, t; νₜ, S, Cs, Δ)
+
+Implements a user-defined function `udf` to model subgrid-scale LES stresses based on the Boussinesq approximation
+    τᵃᵢⱼ = τʳᵢⱼ - (1/3)τʳₖₖδᵢⱼ = -2νₜS̅ᵢⱼ
+where
+            ▁▁▁▁
+    τʳᵢⱼ =  uᵢuⱼ - u̅ᵢu̅ⱼ
+
+and we add -∂ⱼ(τᵃᵢⱼ) to the RHS as a body force (the isotropic part of the tensor is automatically modelled by the pressure gradient term).
+Users need to define the turbulent viscosity function `νₜ` and pass it as a keyword argument to this function together with rate-of-strain
+tensor array buffer `S`, Smagorinsky constant `Cs`, and filter width `Δ`.
+For example, the standard Smagorinsky–Lilly model for the sub-grid scale stresses is
+
+    νₜ = (CₛΔ)²|S̅ᵢⱼ|=(CₛΔ)²√(2S̅ᵢⱼS̅ᵢⱼ)
+
+It can be implemented as
+    `smagorinsky(I::CartesianIndex{m} where m; S, Cs, Δ) = @views (Cs*Δ)^2*sqrt(dot(S[I,:,:],S[I,:,:]))`
+and passed into `sim_step!` as a keyword argument together with the varibles than the function needs (`S`, `Cs`, and `Δ`):
+    `sim_step!(sim, ...; udf=sgs, νₜ=smagorinsky, S, Cs, Δ)`
+"""
+function sgs!(flow, t; νₜ, S, Cs, Δ)
+    N,n = size_u(flow.u)
+    @loop S[I,:,:] .= WaterLily.S(I,flow.u) over I ∈ inside(flow.σ)
+    for i ∈ 1:n, j ∈ 1:n
+        WaterLily.@loop (
+            flow.σ[I] = -νₜ(I;S,Cs,Δ)*∂(j,CI(I,i),flow.u);
+            flow.f[I,i] += flow.σ[I];
+        ) over I ∈ inside_u(N,j)
+        WaterLily.@loop flow.f[I-δ(j,I),i] -= flow.σ[I] over I ∈ WaterLily.inside_u(N,j)
+    end
+end
+
 check_fn(f,N,T,nargs) = nothing
 function check_fn(f::Function,N,T,nargs)
     @assert first(methods(f)).nargs==nargs+1 "$f signature needs $nargs arguments"
@@ -296,4 +330,4 @@ xtargs(::Val{3},N,T) = (zeros(SVector{N,T}),zero(T))
 ic_function(uBC::Function) = (i,x)->uBC(i,x,0)
 ic_function(uBC::Tuple) = (i,x)->uBC[i]
 
-squeeze(a::AbstractArray) = dropdims(a, dims = tuple(findall(size(a) .== 1)...))
+squeeze(a::AbstractArray) = dropdims(a, dims = tuple(findall(size(a) .== 1)...))

test/alloctest.jl

@@ -14,12 +14,19 @@ backend != "SIMD" && throw(ArgumentError("KernelAbstractions backend not allowed
         Simulation((20L,20L),(U,0),L,ν=U*L/Re,body=AutoBody(sdf,map),T=Float32,perdir=perdir)
     end
     sim = Sim(Float32(π/36))
-    sim_step!(sim)
+
+    sim_step!(sim) # runs with λ=quick
     b = @benchmarkable mom_step!($sim.flow, $sim.pois) 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)
 
+    sim_step!(sim; λ=cds)
+    b = @benchmarkable mom_step!($sim.flow, $sim.pois; λ=$cds) 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)
+
     sim = Sim(Float32(π/36); perdir=(2,))
     sim_step!(sim)
     b = @benchmarkable mom_step!($sim.flow, $sim.pois) samples=100; tune!(b) # check 100 times