comparison src/SbpOperators/stencil.jl @ 1052:ca718fd4e816 feature/nested_stencils

Start implementing nested stencils (grafted from 313648b015042773f7f1b6375a700e357231e984)
author Jonatan Werpers <jonatan@werpers.com>
date Wed, 19 Jan 2022 21:49:34 +0100
parents 4433be383840
children e37ee63bf9ac
comparison
equal deleted inserted replaced
1050:396278072f18 1052:ca718fd4e816
1 export CenteredStencil 1 export CenteredStencil
2 export CenteredNestedStencil
2 3
3 struct Stencil{T,N} 4 struct Stencil{T,N}
4 range::Tuple{Int,Int} 5 range::Tuple{Int,Int}
5 weights::NTuple{N,T} 6 weights::NTuple{N,T}
6 7
76 @simd for k ∈ N-1:-1:1 77 @simd for k ∈ N-1:-1:1
77 w += s.weights[k]*v[i - s.range[1] - k + 1] 78 w += s.weights[k]*v[i - s.range[1] - k + 1]
78 end 79 end
79 return w 80 return w
80 end 81 end
82
83
84 struct NestedStencil{T,N}
85 s::Stencil{Stencil{T,N},N}
86 end
87
88 NestedStencil(s::Vararg{Stencil}; center) = NestedStencil(Stencil(s... ; center))
89
90 function NestedStencil(weights::Vararg{Tuple}; center)
91 inner_stencils = map((w,c) -> Stencil(w...,center=c), weights, 1:length(weights))
92 return NestedStencil(Stencil(inner_stencils... ; center))
93 end
94
95 CenteredNestedStencil(s...) = NestedStencil(CenteredStencil(s...))
96
97 Base.eltype(::NestedStencil{T}) where T = T
98
99 function flip(ns::NestedStencil)
100 s_flip = flip(ns.s)
101 return NestedStencil(Stencil(s_flip.range, flip.(s_flip.weights)))
102 end
103
104 Base.getindex(ns::NestedStencil, i::Int) = ns.s[i]
105
106 "Apply inner stencils to `c` and get a concrete stencil"
107 Base.@propagate_inbounds function apply_stencil(ns::NestedStencil, c::AbstractVector, i::Int)
108 weights = apply_stencil.(ns.s.weights, Ref(c), i)
109 return Stencil(ns.s.range, weights)
110 end
111
112 "Apply the whole nested stencil"
113 Base.@propagate_inbounds function apply_stencil(ns::NestedStencil, c::AbstractVector, v::AbstractVector, i::Int)
114 s = apply_stencil(ns,c,i)
115 return apply_stencil(s, v, i)
116 end
117
118 "Apply inner stencils backwards to `c` and get a concrete stencil"
119 Base.@propagate_inbounds @inline function apply_stencil_backwards(ns::NestedStencil, c::AbstractVector, v::AbstractVector, i::Int)
120 weights = apply_stencil_backwards.(ns.s.weights, Ref(c), i)
121 return Stencil(ns.s.range, weights)
122 end
123
124 "Apply the whole nested stencil backwards"
125 Base.@propagate_inbounds @inline function apply_stencil_backwards(ns::NestedStencil, c::AbstractVector, v::AbstractVector, i::Int)
126 s = apply_stencil_backwards(ns,c,i)
127 return apply_stencil_backwards(s, v, i)
128 end