diff src/SbpOperators/stencil.jl @ 867:313648b01504 feature/variable_derivatives

Start implementing nested stencils
author Jonatan Werpers <jonatan@werpers.com>
date Wed, 19 Jan 2022 21:49:34 +0100
parents 4433be383840
children e37ee63bf9ac
line wrap: on
line diff
--- a/src/SbpOperators/stencil.jl	Wed Jan 19 11:08:43 2022 +0100
+++ b/src/SbpOperators/stencil.jl	Wed Jan 19 21:49:34 2022 +0100
@@ -1,4 +1,5 @@
 export CenteredStencil
+export CenteredNestedStencil
 
 struct Stencil{T,N}
     range::Tuple{Int,Int}
@@ -78,3 +79,50 @@
     end
     return w
 end
+
+
+struct NestedStencil{T,N}
+    s::Stencil{Stencil{T,N},N}
+end
+
+NestedStencil(s::Vararg{Stencil}; center) = NestedStencil(Stencil(s... ; center))
+
+function NestedStencil(weights::Vararg{Tuple}; center)
+    inner_stencils = map((w,c) -> Stencil(w...,center=c), weights, 1:length(weights))
+    return NestedStencil(Stencil(inner_stencils... ; center))
+end
+
+CenteredNestedStencil(s...) = NestedStencil(CenteredStencil(s...))
+
+Base.eltype(::NestedStencil{T}) where T = T
+
+function flip(ns::NestedStencil)
+    s_flip = flip(ns.s)
+    return NestedStencil(Stencil(s_flip.range, flip.(s_flip.weights)))
+end
+
+Base.getindex(ns::NestedStencil, i::Int) = ns.s[i]
+
+"Apply inner stencils to `c` and get a concrete stencil"
+Base.@propagate_inbounds function apply_stencil(ns::NestedStencil, c::AbstractVector, i::Int)
+    weights = apply_stencil.(ns.s.weights, Ref(c), i)
+    return Stencil(ns.s.range, weights)
+end
+
+"Apply the whole nested stencil"
+Base.@propagate_inbounds function apply_stencil(ns::NestedStencil, c::AbstractVector, v::AbstractVector, i::Int)
+    s = apply_stencil(ns,c,i)
+    return apply_stencil(s, v, i)
+end
+
+"Apply inner stencils backwards to `c` and get a concrete stencil"
+Base.@propagate_inbounds @inline function apply_stencil_backwards(ns::NestedStencil, c::AbstractVector, v::AbstractVector, i::Int)
+    weights = apply_stencil_backwards.(ns.s.weights, Ref(c), i)
+    return Stencil(ns.s.range, weights)
+end
+
+"Apply the whole nested stencil backwards"
+Base.@propagate_inbounds @inline function apply_stencil_backwards(ns::NestedStencil, c::AbstractVector, v::AbstractVector, i::Int)
+    s = apply_stencil_backwards(ns,c,i)
+    return apply_stencil_backwards(s, v, i)
+end