Mercurial > repos > public > sbplib_julia

comparison src/SbpOperators/volumeops/derivatives/second_derivative_variable.jl @ 1384:851d1e4ab3de

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
Merge feature/variable_derivatives
author Jonatan Werpers <jonatan@werpers.com>
date Thu, 08 Jun 2023 15:52:22 +0200
parents e9dfc1998d31
children 3684db043add
comparison
equal deleted inserted replaced
1379:8e59c85a25c0 1384:851d1e4ab3de
1 """
2 second_derivative_variable(g, coeff ..., [direction])
3
4 The variable second derivative operator as a `LazyTensor` on the given grid.
5 `coeff` is a grid function of the variable coefficient.
6
7 Approximates the d/dξ c d/dξ on `g` along the coordinate dimension specified
8 by `direction`.
9 """
10 function second_derivative_variable end
11
12 function second_derivative_variable(g::TensorGrid, coeff, stencil_set, dir::Int)
13 inner_stencil = parse_nested_stencil(eltype(coeff), stencil_set["D2variable"]["inner_stencil"])
14 closure_stencils = parse_nested_stencil.(eltype(coeff), stencil_set["D2variable"]["closure_stencils"])
15
16 return second_derivative_variable(g, coeff, inner_stencil, closure_stencils, dir)
17 end
18
19 function second_derivative_variable(g::EquidistantGrid, coeff, stencil_set)
20 return second_derivative_variable(TensorGrid(g), coeff, stencil_set, 1)
21 end
22
23 function second_derivative_variable(g::TensorGrid, coeff, inner_stencil::NestedStencil, closure_stencils, dir)
24 check_coefficient(g, coeff)
25
26 Δxᵢ = spacing(g.grids[dir])
27 scaled_inner_stencil = scale(inner_stencil, 1/Δxᵢ^2)
28 scaled_closure_stencils = scale.(Tuple(closure_stencils), 1/Δxᵢ^2)
29 return SecondDerivativeVariable(coeff, scaled_inner_stencil, scaled_closure_stencils, dir)
30 end
31
32 function check_coefficient(g, coeff)
33 if ndims(g) != ndims(coeff)
34 throw(ArgumentError("The coefficient has dimension $(ndims(coeff)) while the grid is dimension $(ndims(g))"))
35 end
36
37 if size(g) != size(coeff)
38 throw(DimensionMismatch("the size $(size(coeff)) of the coefficient does not match the size $(size(g)) of the grid"))
39 end
40 end
41
42
43 """
44 SecondDerivativeVariable{Dir,T,D,...} <: LazyTensor{T,D,D}
45
46 A second derivative operator in direction `Dir` with a variable coefficient.
47 """
48 struct SecondDerivativeVariable{Dir,T,D,M,IStencil<:NestedStencil{T},CStencil<:NestedStencil{T},TArray<:AbstractArray} <: LazyTensor{T,D,D}
49 inner_stencil::IStencil
50 closure_stencils::NTuple{M,CStencil}
51 coefficient::TArray
52
53 function SecondDerivativeVariable(coefficient::AbstractArray, inner_stencil::NestedStencil{T}, closure_stencils::NTuple{M,NestedStencil{T}}, dir) where {T,M}
54 D = ndims(coefficient)
55 IStencil = typeof(inner_stencil)
56 CStencil = eltype(closure_stencils)
57 TArray = typeof(coefficient)
58 return new{dir,T,D,M,IStencil,CStencil,TArray}(inner_stencil, closure_stencils, coefficient)
59 end
60 end
61
62 derivative_direction(::SecondDerivativeVariable{Dir}) where {Dir} = Dir
63
64 closure_size(op::SecondDerivativeVariable) = length(op.closure_stencils)
65
66 LazyTensors.range_size(op::SecondDerivativeVariable) = size(op.coefficient)
67 LazyTensors.domain_size(op::SecondDerivativeVariable) = size(op.coefficient)
68
69
70 function derivative_view(op, a, I)
71 d = derivative_direction(op)
72
73 Iview = Base.setindex(I,:,d)
74 return @view a[Iview...]
75 end
76
77 function apply_lower(op::SecondDerivativeVariable, v, I...)
78 ṽ = derivative_view(op, v, I)
79 c̃ = derivative_view(op, op.coefficient, I)
80
81 i = I[derivative_direction(op)]
82 return @inbounds apply_stencil(op.closure_stencils[i], c̃, ṽ, i)
83 end
84
85 function apply_interior(op::SecondDerivativeVariable, v, I...)
86 ṽ = derivative_view(op, v, I)
87 c̃ = derivative_view(op, op.coefficient, I)
88
89 i = I[derivative_direction(op)]
90 return apply_stencil(op.inner_stencil, c̃, ṽ, i)
91 end
92
93 function apply_upper(op::SecondDerivativeVariable, v, I...)
94 ṽ = derivative_view(op, v, I)
95 c̃ = derivative_view(op, op.coefficient, I)
96
97 i = I[derivative_direction(op)]
98 sz = domain_size(op)[derivative_direction(op)]
99 stencil = op.closure_stencils[sz-i+1]
100 return @inbounds apply_stencil_backwards(stencil, c̃, ṽ, i)
101 end
102
103 function LazyTensors.apply(op::SecondDerivativeVariable, v::AbstractArray, I::Vararg{Index})
104 if I[derivative_direction(op)] isa Index{Lower}
105 return apply_lower(op, v, Int.(I)...)
106 elseif I[derivative_direction(op)] isa Index{Upper}
107 return apply_upper(op, v, Int.(I)...)
108 elseif I[derivative_direction(op)] isa Index{Interior}
109 return apply_interior(op, v, Int.(I)...)
110 else
111 error("Invalid region")
112 end
113 end
114
115 function LazyTensors.apply(op::SecondDerivativeVariable, v::AbstractArray, I...)
116 dir = derivative_direction(op)
117 sz = domain_size(op)[dir]
118
119 i = I[dir]
120
121 I = map(i->Index(i, Interior), I)
122 if 0 < i <= closure_size(op)
123 I = Base.setindex(I, Index(i, Lower), dir)
124 return LazyTensors.apply(op, v, I...)
125 elseif closure_size(op) < i <= sz-closure_size(op)
126 I = Base.setindex(I, Index(i, Interior), dir)
127 return LazyTensors.apply(op, v, I...)
128 elseif sz-closure_size(op) < i <= sz
129 I = Base.setindex(I, Index(i, Upper), dir)
130 return LazyTensors.apply(op, v, I...)
131 else
132 error("Bounds error") # This should be `throw(BoundsError())` but the type inference is so fragile that it doesn't work. Needs investigation. / Jonatan 2023-06-08
133 end
134 end
135
136
137 # 2D Specific implementations to avoid type instability
138 # TBD: Can this be solved by fixing the general methods instead?
139
140
141 ## x-direction
142 function apply_lower(op::SecondDerivativeVariable{1}, v, i, j)
143 ṽ = @view v[:,j]
144 c̃ = @view op.coefficient[:,j]
145
146 return @inbounds apply_stencil(op.closure_stencils[i], c̃, ṽ, i)
147 end
148
149 function apply_interior(op::SecondDerivativeVariable{1}, v, i, j)
150 ṽ = @view v[:,j]
151 c̃ = @view op.coefficient[:,j]
152
153 return @inbounds apply_stencil(op.inner_stencil, c̃, ṽ, i)
154 end
155
156 function apply_upper(op::SecondDerivativeVariable{1}, v, i, j)
157 ṽ = @view v[:,j]
158 c̃ = @view op.coefficient[:,j]
159
160 sz = domain_size(op)[derivative_direction(op)]
161 stencil = op.closure_stencils[sz-i+1]
162 return @inbounds apply_stencil_backwards(stencil, c̃, ṽ, i)
163 end
164
165
166 ## y-direction
167 function apply_lower(op::SecondDerivativeVariable{2}, v, i, j)
168 ṽ = @view v[i,:]
169 c̃ = @view op.coefficient[i,:]
170
171 return @inbounds apply_stencil(op.closure_stencils[j], c̃, ṽ, j)
172 end
173
174 function apply_interior(op::SecondDerivativeVariable{2}, v, i, j)
175 ṽ = @view v[i,:]
176 c̃ = @view op.coefficient[i,:]
177
178 return @inbounds apply_stencil(op.inner_stencil, c̃, ṽ, j)
179 end
180
181 function apply_upper(op::SecondDerivativeVariable{2}, v, i, j)
182 ṽ = @view v[i,:]
183 c̃ = @view op.coefficient[i,:]
184
185 sz = domain_size(op)[derivative_direction(op)]
186 stencil = op.closure_stencils[sz-j+1]
187 return @inbounds apply_stencil_backwards(stencil, c̃, ṽ, j)
188 end