Mercurial > repos > public > sbplib_julia
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