Mercurial > repos > public > sbplib_julia
diff src/SbpOperators/stencil.jl @ 2057:8a2a0d678d6f feature/lazy_tensors/pretty_printing
Merge default
| author | Jonatan Werpers <jonatan@werpers.com> |
|---|---|
| date | Tue, 10 Feb 2026 22:41:19 +0100 |
| parents | 244311761969 |
| children |
line wrap: on
line diff
--- a/src/SbpOperators/stencil.jl Mon May 23 07:20:27 2022 +0200 +++ b/src/SbpOperators/stencil.jl Tue Feb 10 22:41:19 2026 +0100 @@ -1,18 +1,17 @@ -export CenteredStencil -export CenteredNestedStencil - struct Stencil{T,N} range::UnitRange{Int64} weights::NTuple{N,T} - function Stencil(range::UnitRange,weights::NTuple{N,T}) where {T, N} + function Stencil(range::UnitRange,weights::NTuple{N,Any}) where N + T = eltype(weights) + @assert length(range) == N new{T,N}(range,weights) end end """ - Stencil(weights::NTuple; center::Int) + Stencil(weights...; center::Int) Create a stencil with the given weights with element `center` as the center of the stencil. """ @@ -69,42 +68,67 @@ end Base.@propagate_inbounds @inline function apply_stencil(s::Stencil, v::AbstractVector, i::Int) - w = zero(promote_type(eltype(s),eltype(v))) - @simd for k ∈ 1:length(s) - w += s.weights[k]*v[i + s.range[k]] + return sum(enumerate(s.weights)) do (k,w) #TBD: Which optimizations are needed here? + w*v[i + @inbounds s.range[k]] end - - return w end Base.@propagate_inbounds @inline function apply_stencil_backwards(s::Stencil, v::AbstractVector, i::Int) - w = zero(promote_type(eltype(s),eltype(v))) - @simd for k ∈ length(s):-1:1 - w += s.weights[k]*v[i - s.range[k]] + return sum(enumerate(s.weights)) do (k,w) #TBD: Which optimizations are needed here? + w*v[i - @inbounds s.range[k]] end - return w end +# There are many options for the implementation of `apply_stencil` and +# `apply_stencil_backwards`. Some alternatives were tried on the branch +# bugfix/sbp_operators/stencil_return_type and can be found at the following +# revision: +# +# * 237b980ffb91 (baseline) +# * a72bab15228e (mapreduce) +# * ffd735354d54 (multiplication) +# * b5abd5191f2c (promote_op) +# * 8d56846185fc (return_type) +# + +function left_pad(s::Stencil, N) + weights = LazyTensors.left_pad_tuple(s.weights, zero(eltype(s)), N) + range = (first(s.range) - (N - length(s.weights))):last(s.range) + + return Stencil(range, weights) +end + +function right_pad(s::Stencil, N) + weights = LazyTensors.right_pad_tuple(s.weights, zero(eltype(s)), N) + range = first(s.range):(last(s.range) + (N - length(s.weights))) + + return Stencil(range, weights) +end + + struct NestedStencil{T,N,M} s::Stencil{Stencil{T,N},M} end +NestedStencil(;center) = NestedStencil(Stencil(;center)) +CenteredNestedStencil() = NestedStencil(CenteredStencil()) + # Stencil input NestedStencil(s::Vararg{Stencil}; center) = NestedStencil(Stencil(s... ; center)) CenteredNestedStencil(s::Vararg{Stencil}) = NestedStencil(CenteredStencil(s...)) # Tuple input -function NestedStencil(weights::Vararg{NTuple{N,Any}}; center) where N +function NestedStencil(weights::Vararg{NTuple{N,Any} where N}; center) inner_stencils = map(w -> Stencil(w...; center), weights) return NestedStencil(Stencil(inner_stencils... ; center)) end -function CenteredNestedStencil(weights::Vararg{NTuple{N,Any}}) where N + +function CenteredNestedStencil(weights::Vararg{NTuple{N,Any} where N}) inner_stencils = map(w->CenteredStencil(w...), weights) return CenteredNestedStencil(inner_stencils...) end - # Conversion function NestedStencil{T,N,M}(ns::NestedStencil{S,N,M}) where {T,S,N,M} return NestedStencil(Stencil{Stencil{T}}(ns.s)) @@ -117,7 +141,7 @@ function Base.convert(::Type{NestedStencil{T,N,M}}, s::NestedStencil{S,N,M}) where {T,S,N,M} return NestedStencil{T,N,M}(s) end -Base.convert(::Type{NestedStencil{T}}, stencil) where T = NestedStencil{T}(stencil) +Base.convert(::Type{NestedStencil{T}}, stencil::NestedStencil) where T = NestedStencil{T}(stencil) function Base.promote_rule(::Type{NestedStencil{T,N,M}}, ::Type{NestedStencil{S,N,M}}) where {T,S,N,M} return NestedStencil{promote_type(T,S),N,M}
