Branch data Line data Source code
1 : : //
2 : : // File IpplOperations.h
3 : : // Expression Templates operations.
4 : : //
5 : : #ifndef IPPL_OPERATIONS_H
6 : : #define IPPL_OPERATIONS_H
7 : :
8 : : #include <Kokkos_MathematicalFunctions.hpp>
9 : : #include <tuple>
10 : :
11 : : namespace ippl {
12 : : /*!
13 : : * @file IpplOperations.h
14 : : */
15 : :
16 : : /*!
17 : : * Utility function for apply (see its docstring)
18 : : * @tparam Idx... indices of the elements to take (in practice, always the sequence of natural
19 : : * numbers up to the dimension of the view)
20 : : */
21 : : template <typename View, typename Coords, size_t... Idx>
22 : 133629100 : KOKKOS_INLINE_FUNCTION constexpr decltype(auto) apply_impl(const View& view,
23 : : const Coords& coords,
24 : : const std::index_sequence<Idx...>&) {
25 : 219138966 : return view(coords[Idx]...);
26 : : }
27 : :
28 : : /*!
29 : : * Extracts the mathematical rank of an expression (i.e. the number of dimensions) based on
30 : : * its type
31 : : */
32 : : struct ExtractExpressionRank {
33 : : /*!
34 : : * Extracts the extent of an array-like expression
35 : : * @tparam Coords the array type
36 : : * @return The array's rank
37 : : */
38 : : template <typename Coords, std::enable_if_t<std::is_array_v<Coords>, int> = 0>
39 : : KOKKOS_INLINE_FUNCTION constexpr static unsigned getRank() {
40 : : return std::extent_v<Coords>;
41 : : }
42 : :
43 : : /*!
44 : : * Extracts the rank of an expression type
45 : : * @tparam Coords the expression type that evaluates to a set of coordinates
46 : : * @return The number of dimensions in the expression
47 : : */
48 : : template <typename Coords, std::enable_if_t<std::is_class_v<Coords>, int> = 0>
49 : : KOKKOS_INLINE_FUNCTION constexpr static unsigned getRank() {
50 : : return Coords::dim;
51 : : }
52 : : };
53 : :
54 : : /*!
55 : : * Accesses the element of a view at the indices contained in an array-like structure
56 : : * instead of having the indices being separate arguments
57 : : * @tparam View the view type
58 : : * @tparam Coords an array-like container of indices
59 : : * @param view the view to access
60 : : * @param coords the indices
61 : : * @return The element in the view at the given location
62 : : */
63 : : template <typename View, typename Coords>
64 : 133629100 : KOKKOS_INLINE_FUNCTION constexpr decltype(auto) apply(const View& view, const Coords& coords) {
65 : : using Indices = std::make_index_sequence<ExtractExpressionRank::getRank<Coords>()>;
66 [ + - ]: 133629100 : return apply_impl(view, coords, Indices{});
67 : : }
68 : :
69 : : #define DefineUnaryOperation(fun, name, op1, op2) \
70 : : template <typename E> \
71 : : struct fun : public detail::Expression<fun<E>, sizeof(E)> { \
72 : : constexpr static unsigned dim = E::dim; \
73 : : using value_type = typename E::value_type; \
74 : : \
75 : : KOKKOS_FUNCTION \
76 : : fun(const E& u) \
77 : : : u_m(u) {} \
78 : : \
79 : : KOKKOS_INLINE_FUNCTION auto operator[](size_t i) const { return op1; } \
80 : : \
81 : : template <typename... Args> \
82 : : KOKKOS_INLINE_FUNCTION auto operator()(Args... args) const { \
83 : : return op2; \
84 : : } \
85 : : \
86 : : private: \
87 : : const E u_m; \
88 : : }; \
89 : : \
90 : : template <typename E, size_t N> \
91 : : KOKKOS_INLINE_FUNCTION fun<E> name(const detail::Expression<E, N>& u) { \
92 : : return fun<E>(*static_cast<const E*>(&u)); \
93 : : }
94 : :
95 : : /// @cond
96 : :
97 : : // clang-format off
98 [ # # # # ]: 264 : DefineUnaryOperation(UnaryMinus, operator-, -u_m[i], -u_m(args...))
[ # # ]
99 : : DefineUnaryOperation(UnaryPlus, operator+, +u_m[i], +u_m(args...))
100 : : DefineUnaryOperation(BitwiseNot, operator~, ~u_m[i], ~u_m(args...))
101 : : DefineUnaryOperation(Not, operator!, !u_m[i], !u_m(args...))
102 : :
103 : 238680 : DefineUnaryOperation(ArcCos, acos, Kokkos::acos(u_m[i]), Kokkos::acos(u_m(args...)))
104 : 238680 : DefineUnaryOperation(ArcSin, asin, Kokkos::asin(u_m[i]), Kokkos::asin(u_m(args...)))
105 : 238680 : DefineUnaryOperation(ArcTan, atan, Kokkos::atan(u_m[i]), Kokkos::atan(u_m(args...)))
106 : 238680 : DefineUnaryOperation(Ceil, ceil, Kokkos::ceil(u_m[i]), Kokkos::ceil(u_m(args...)))
107 : 238680 : DefineUnaryOperation(Cos, cos, Kokkos::cos(u_m[i]), Kokkos::cos(u_m(args...)))
108 : 238680 : DefineUnaryOperation(HypCos, cosh, Kokkos::cosh(u_m[i]), Kokkos::cosh(u_m(args...)))
109 : 238680 : DefineUnaryOperation(Exp, exp, Kokkos::exp(u_m[i]), Kokkos::exp(u_m(args...)))
110 : 238680 : DefineUnaryOperation(Fabs, fabs, Kokkos::fabs(u_m[i]), Kokkos::fabs(u_m(args...)))
111 : 238680 : DefineUnaryOperation(Floor, floor, Kokkos::floor(u_m[i]), Kokkos::floor(u_m(args...)))
112 : 238680 : DefineUnaryOperation(Log, log, Kokkos::log(u_m[i]), Kokkos::log(u_m(args...)))
113 : 238680 : DefineUnaryOperation(Log10, log10, Kokkos::log10(u_m[i]), Kokkos::log10(u_m(args...)))
114 : 238680 : DefineUnaryOperation(Sin, sin, Kokkos::sin(u_m[i]), Kokkos::sin(u_m(args...)))
115 : 238680 : DefineUnaryOperation(HypSin, sinh, Kokkos::sinh(u_m[i]), Kokkos::sinh(u_m(args...)))
116 : 238680 : DefineUnaryOperation(Sqrt, sqrt, Kokkos::sqrt(u_m[i]), Kokkos::sqrt(u_m(args...)))
117 : 238680 : DefineUnaryOperation(Tan, tan, Kokkos::tan(u_m[i]), Kokkos::tan(u_m(args...)))
118 : 238680 : DefineUnaryOperation(HypTan, tanh, Kokkos::tanh(u_m[i]), Kokkos::tanh(u_m(args...)))
119 : 238680 : DefineUnaryOperation(Erf, erf, Kokkos::erf(u_m[i]), Kokkos::erf(u_m(args...)))
120 : : // clang-format on
121 : : /// @endcond
122 : :
123 : : /*!
124 : : * Macro to overload C++ operators for the Scalar, BareField and Vector class.
125 : : * @param fun name of the expression template function
126 : : * @param name overloaded operator
127 : : * @param op1 operation for single index access
128 : : * @param op2 operation for multiple indices access
129 : : */
130 : : #define DefineBinaryOperation(fun, name, op1, op2) \
131 : : template <typename E1, typename E2> \
132 : : struct fun : public detail::Expression<fun<E1, E2>, sizeof(E1) + sizeof(E2)> { \
133 : : constexpr static unsigned dim = std::max(E1::dim, E2::dim); \
134 : : using value_type = typename E1::value_type; \
135 : : \
136 : : KOKKOS_FUNCTION \
137 : : fun(const E1& u, const E2& v) \
138 : : : u_m(u) \
139 : : , v_m(v) {} \
140 : : \
141 : : KOKKOS_INLINE_FUNCTION auto operator[](size_t i) const { return op1; } \
142 : : \
143 : : template <typename... Args> \
144 : : KOKKOS_INLINE_FUNCTION auto operator()(Args... args) const { \
145 : : static_assert(sizeof...(Args) == dim || dim == 0); \
146 : : return op2; \
147 : : } \
148 : : \
149 : : private: \
150 : : const E1 u_m; \
151 : : const E2 v_m; \
152 : : }; \
153 : : \
154 : : template <typename E1, size_t N1, typename E2, size_t N2> \
155 : : KOKKOS_INLINE_FUNCTION fun<E1, E2> name(const detail::Expression<E1, N1>& u, \
156 : : const detail::Expression<E2, N2>& v) { \
157 : : return fun<E1, E2>(*static_cast<const E1*>(&u), *static_cast<const E2*>(&v)); \
158 : : } \
159 : : \
160 : : template <typename E, size_t N, typename T, \
161 : : typename = std::enable_if_t<std::is_scalar<T>::value>> \
162 : : KOKKOS_INLINE_FUNCTION fun<E, detail::Scalar<T>> name(const detail::Expression<E, N>& u, \
163 : : const T& v) { \
164 : : return fun<E, detail::Scalar<T>>(*static_cast<const E*>(&u), v); \
165 : : } \
166 : : \
167 : : template <typename E, size_t N, typename T, \
168 : : typename = std::enable_if_t<std::is_scalar<T>::value>> \
169 : : KOKKOS_INLINE_FUNCTION fun<detail::Scalar<T>, E> name(const T& u, \
170 : : const detail::Expression<E, N>& v) { \
171 : : return fun<detail::Scalar<T>, E>(u, *static_cast<const E*>(&v)); \
172 : : }
173 : :
174 : : /// @cond
175 : : // clang-format off
176 [ + - ][ # # : 96888444 : DefineBinaryOperation(Add, operator+, u_m[i] + v_m[i], u_m(args...) + v_m(args...))
# # # # ]
177 [ + - ][ # # : 24585486 : DefineBinaryOperation(Subtract, operator-, u_m[i] - v_m[i], u_m(args...) - v_m(args...))
# # # # ]
178 [ + - ]: 92140198 : DefineBinaryOperation(Multiply, operator*, u_m[i] * v_m[i], u_m(args...) * v_m(args...))
[ # # # # ]
179 [ + - ]: 53227026 : DefineBinaryOperation(Divide, operator/, u_m[i] / v_m[i], u_m(args...) / v_m(args...))
180 : : DefineBinaryOperation(Mod, operator%, u_m[i] % v_m[i], u_m(args...) % v_m(args...))
181 : : DefineBinaryOperation(LT, operator<, u_m[i] < v_m[i], u_m(args...) < v_m(args...))
182 : : DefineBinaryOperation(LE, operator<=, u_m[i] <= v_m[i], u_m(args...) <= v_m(args...))
183 : : DefineBinaryOperation(GT, operator>, u_m[i] > v_m[i], u_m(args...) > v_m(args...))
184 : : DefineBinaryOperation(GE, operator>=, u_m[i] >= v_m[i], u_m(args...) >= v_m(args...))
185 : : DefineBinaryOperation(EQ, operator==, u_m[i] == v_m[i], u_m(args...) == v_m(args...))
186 : : DefineBinaryOperation(NEQ, operator!=, u_m[i] != v_m[i], u_m(args...) != v_m(args...))
187 : : DefineBinaryOperation(And, operator&&, u_m[i] && v_m[i], u_m(args...) && v_m(args...))
188 : : DefineBinaryOperation(Or, operator||, u_m[i] || v_m[i], u_m(args...) || v_m(args...))
189 : :
190 : : DefineBinaryOperation(BitwiseAnd, operator&, u_m[i] & v_m[i], u_m(args...) & v_m(args...))
191 : : DefineBinaryOperation(BitwiseOr, operator|, u_m[i] | v_m[i], u_m(args...) | v_m(args...))
192 : : DefineBinaryOperation(BitwiseXor, operator^, u_m[i] ^ v_m[i], u_m(args...) ^ v_m(args...))
193 : : // clang-format on
194 : :
195 : : DefineBinaryOperation(Copysign, copysign, Kokkos::copysign(u_m[i], v_m[i]),
196 : : Kokkos::copysign(u_m(args...), v_m(args...)))
197 : : // ldexp not provided by Kokkos
198 : : DefineBinaryOperation(Ldexp, ldexp, ldexp(u_m[i], v_m[i]), ldexp(u_m(args...), v_m(args...)))
199 : : DefineBinaryOperation(Fmod, fmod, Kokkos::fmod(u_m[i], v_m[i]),
200 : : Kokkos::fmod(u_m(args...), v_m(args...)))
201 [ # # ]: 0 : DefineBinaryOperation(Pow, pow, Kokkos::pow(u_m[i], v_m[i]),
202 : : Kokkos::pow(u_m(args...), v_m(args...)))
203 : : DefineBinaryOperation(ArcTan2, atan2, Kokkos::atan2(u_m[i], v_m[i]),
204 : : Kokkos::atan2(u_m(args...), v_m(args...)))
205 : : /// @endcond
206 : :
207 : : namespace detail {
208 : : /*!
209 : : * Meta function of cross product. This function is only supported for 3-dimensional
210 : : * vectors.
211 : : */
212 : : template <typename E1, typename E2>
213 : : struct meta_cross : public detail::Expression<meta_cross<E1, E2>, sizeof(E1) + sizeof(E2)> {
214 : : constexpr static unsigned dim = E1::dim;
215 : : static_assert(E1::dim == E2::dim);
216 : :
217 : : KOKKOS_FUNCTION
218 : 0 : meta_cross(const E1& u, const E2& v)
219 : 0 : : u_m(u)
220 : 0 : , v_m(v) {}
221 : :
222 : : /*
223 : : * Vector::cross
224 : : */
225 : 0 : KOKKOS_INLINE_FUNCTION auto operator[](size_t i) const {
226 : 0 : const size_t j = (i + 1) % 3;
227 : 0 : const size_t k = (i + 2) % 3;
228 : 0 : return u_m[j] * v_m[k] - u_m[k] * v_m[j];
229 : : }
230 : :
231 : : /*
232 : : * This is required for BareField::cross
233 : : */
234 : : template <typename... Args>
235 : : KOKKOS_INLINE_FUNCTION auto operator()(Args... args) const {
236 : : return cross(u_m(args...), v_m(args...));
237 : : }
238 : :
239 : : private:
240 : : const E1 u_m;
241 : : const E2 v_m;
242 : : };
243 : : } // namespace detail
244 : :
245 : : template <typename E1, size_t N1, typename E2, size_t N2>
246 : 0 : KOKKOS_INLINE_FUNCTION detail::meta_cross<E1, E2> cross(const detail::Expression<E1, N1>& u,
247 : : const detail::Expression<E2, N2>& v) {
248 : 0 : return detail::meta_cross<E1, E2>(*static_cast<const E1*>(&u), *static_cast<const E2*>(&v));
249 : : }
250 : :
251 : : namespace detail {
252 : : /*!
253 : : * Meta function of dot product.
254 : : */
255 : : template <typename E1, typename E2>
256 : : struct meta_dot : public Expression<meta_dot<E1, E2>, sizeof(E1) + sizeof(E2)> {
257 : : constexpr static unsigned dim = E1::dim;
258 : : static_assert(E1::dim == E2::dim);
259 : :
260 : : KOKKOS_FUNCTION
261 : 1510700 : meta_dot(const E1& u, const E2& v)
262 : 1510700 : : u_m(u)
263 [ + - ]: 1510700 : , v_m(v) {}
264 : :
265 : : /*
266 : : * Vector::dot
267 : : */
268 : 1510688 : KOKKOS_INLINE_FUNCTION auto apply() const {
269 : 1510688 : typename E1::value_type res = 0.0;
270 : : // Equivalent computation in 3D:
271 : : // u_m[0] * v_m[0] + u_m[1] * v_m[1] + u_m[2] * v_m[2]
272 [ + + ]: 9741888 : for (size_t i = 0; i < E1::dim; ++i) {
273 : 8231200 : res += u_m[i] * v_m[i];
274 : : }
275 : 1510688 : return res;
276 : : }
277 : :
278 : : /*
279 : : * This is required for BareField::dot
280 : : */
281 : : template <typename... Args>
282 : 238656 : KOKKOS_INLINE_FUNCTION auto operator()(Args... args) const {
283 [ + - + - : 238656 : return dot(u_m(args...), v_m(args...)).apply();
+ - + - ]
284 : : }
285 : :
286 : : private:
287 : : const E1 u_m;
288 : : const E2 v_m;
289 : : };
290 : : } // namespace detail
291 : :
292 : : template <typename E1, size_t N1, typename E2, size_t N2>
293 : 1510700 : KOKKOS_INLINE_FUNCTION detail::meta_dot<E1, E2> dot(const detail::Expression<E1, N1>& u,
294 : : const detail::Expression<E2, N2>& v) {
295 : 1510700 : return detail::meta_dot<E1, E2>(*static_cast<const E1*>(&u), *static_cast<const E2*>(&v));
296 : : }
297 : :
298 : : namespace detail {
299 : : /*!
300 : : * Meta function of gradient
301 : : */
302 : :
303 : : template <typename E>
304 : : struct meta_grad
305 : : : public Expression<
306 : : meta_grad<E>,
307 : : sizeof(E) + sizeof(typename E::Mesh_t::vector_type[E::Mesh_t::Dimension])> {
308 : : constexpr static unsigned dim = E::dim;
309 : : using value_type = typename E::value_type;
310 : :
311 : : KOKKOS_FUNCTION
312 : 12 : meta_grad(const E& u, const typename E::Mesh_t::vector_type vectors[])
313 [ # # ][ + - : 54 : : u_m(u) {
+ + - - -
- ]
314 [ + + ]: 54 : for (unsigned d = 0; d < E::Mesh_t::Dimension; d++) {
315 : 42 : vectors_m[d] = vectors[d];
316 : : }
317 : 12 : }
318 : :
319 : : /*
320 : : * n-dimensional grad
321 : : */
322 : : template <typename... Idx>
323 : 238656 : KOKKOS_INLINE_FUNCTION auto operator()(const Idx... args) const {
324 : : using index_type = std::tuple_element_t<0, std::tuple<Idx...>>;
325 : :
326 : : /*
327 : : * Equivalent computation in 3D:
328 : : * xvector_m * (u_m(i + 1, j, k) - u_m(i - 1, j, k))
329 : : * + yvector_m * (u_m(i, j + 1, k) - u_m(i, j - 1, k))
330 : : * + zvector_m * (u_m(i, j, k + 1) - u_m(i, j, k - 1))
331 : : */
332 : :
333 [ + - ]: 238656 : vector_type res(0);
334 [ + + ]: 1510528 : for (unsigned d = 0; d < dim; d++) {
335 : 1271872 : index_type coords[dim] = {args...};
336 : :
337 : 1271872 : coords[d] += 1;
338 [ + - ]: 1271872 : auto&& right = apply(u_m, coords);
339 : :
340 : 1271872 : coords[d] -= 2;
341 [ + - ]: 1271872 : auto&& left = apply(u_m, coords);
342 : :
343 [ + - + - ]: 1271872 : res += vectors_m[d] * (right - left);
[ # # ]
344 : : }
345 : 238656 : return res;
346 : 0 : }
347 : :
348 : : private:
349 : : using Mesh_t = typename E::Mesh_t;
350 : : using vector_type = typename Mesh_t::vector_type;
351 : : const E u_m;
352 : : vector_type vectors_m[dim];
353 : : };
354 : : } // namespace detail
355 : :
356 : : namespace detail {
357 : :
358 : : /*!
359 : : * Meta function of divergence
360 : : */
361 : : template <typename E>
362 : : struct meta_div
363 : : : public Expression<
364 : : meta_div<E>,
365 : : sizeof(E) + sizeof(typename E::Mesh_t::vector_type[E::Mesh_t::Dimension])> {
366 : : constexpr static unsigned dim = E::dim;
367 : :
368 : : KOKKOS_FUNCTION
369 : 12 : meta_div(const E& u, const typename E::Mesh_t::vector_type vectors[])
370 [ + - + + : 54 : : u_m(u) {
- - - - ]
371 [ + + ]: 54 : for (unsigned d = 0; d < E::Mesh_t::Dimension; d++) {
372 : 42 : vectors_m[d] = vectors[d];
373 : : }
374 : 12 : }
375 : :
376 : : /*
377 : : * n-dimensional div
378 : : */
379 : : template <typename... Idx>
380 : 238656 : KOKKOS_INLINE_FUNCTION auto operator()(const Idx... args) const {
381 : : using index_type = std::tuple_element_t<0, std::tuple<Idx...>>;
382 : :
383 : : /*
384 : : * Equivalent computation in 3D:
385 : : * dot(xvector_m, (u_m(i + 1, j, k) - u_m(i - 1, j, k))).apply()
386 : : * + dot(yvector_m, (u_m(i, j + 1, k) - u_m(i, j - 1, k))).apply()
387 : : * + dot(zvector_m, (u_m(i, j, k + 1) - u_m(i, j, k - 1))).apply()
388 : : */
389 : 238656 : typename E::Mesh_t::value_type res = 0;
390 [ + + ]: 1510528 : for (unsigned d = 0; d < dim; d++) {
391 : 1271872 : index_type coords[dim] = {args...};
392 : :
393 : 1271872 : coords[d] += 1;
394 [ + - ]: 1271872 : auto&& right = apply(u_m, coords);
395 : :
396 : 1271872 : coords[d] -= 2;
397 [ + - ]: 1271872 : auto&& left = apply(u_m, coords);
398 : :
399 [ + - + - : 1271872 : res += dot(vectors_m[d], right - left).apply();
+ - ]
400 : : }
401 : 238656 : return res;
402 : : }
403 : :
404 : : private:
405 : : using Mesh_t = typename E::Mesh_t;
406 : : using vector_type = typename Mesh_t::vector_type;
407 : : const E u_m;
408 : : vector_type vectors_m[dim];
409 : : };
410 : :
411 : : /*!
412 : : * Meta function of Laplacian
413 : : */
414 : : template <typename E>
415 : : struct meta_laplace
416 : : : public Expression<meta_laplace<E>,
417 : : sizeof(E) + sizeof(typename E::Mesh_t::vector_type)> {
418 : : constexpr static unsigned dim = E::dim;
419 : : using value_type = typename E::value_type;
420 : :
421 : : KOKKOS_FUNCTION
422 : 12 : meta_laplace(const E& u, const typename E::Mesh_t::vector_type& hvector)
423 : 12 : : u_m(u)
424 : 12 : , hvector_m(hvector) {}
425 : :
426 : : /*
427 : : * n-dimensional Laplacian
428 : : */
429 : : template <typename... Idx>
430 : 238656 : KOKKOS_INLINE_FUNCTION auto operator()(const Idx... args) const {
431 : : using index_type = std::tuple_element_t<0, std::tuple<Idx...>>;
432 : : using T = typename E::Mesh_t::value_type;
433 : :
434 : : /*
435 : : * Equivalent computation in 3D:
436 : : * hvector_m[0] * (u_m(i+1, j, k) - 2 * u_m(i, j, k) + u_m(i-1, j, k ))
437 : : * + hvector_m[1] * (u_m(i , j+1, k) - 2 * u_m(i, j, k) + u_m(i , j-1, k ))
438 : : * + hvector_m[2] * (u_m(i , j , k+1) - 2 * u_m(i, j, k) + u_m(i , j , k-1))
439 : : */
440 : 238656 : T res = 0;
441 [ + + ]: 1510528 : for (unsigned d = 0; d < dim; d++) {
442 : 1271872 : index_type coords[dim] = {args...};
443 [ + - ]: 1271872 : auto&& center = apply(u_m, coords);
444 : :
445 : 1271872 : coords[d] -= 1;
446 [ + - ]: 1271872 : auto&& left = apply(u_m, coords);
447 : :
448 : 1271872 : coords[d] += 2;
449 [ + - ]: 1271872 : auto&& right = apply(u_m, coords);
450 : :
451 : 1271872 : res += hvector_m[d] * (left - 2 * center + right);
452 : : }
453 : 238656 : return res;
454 : : }
455 : :
456 : : private:
457 : : using Mesh_t = typename E::Mesh_t;
458 : : using vector_type = typename Mesh_t::vector_type;
459 : : const E u_m;
460 : : const vector_type hvector_m;
461 : : };
462 : : } // namespace detail
463 : :
464 : : namespace detail {
465 : : /*!
466 : : * Meta function of curl
467 : : */
468 : :
469 : : template <typename E>
470 : : struct meta_curl
471 : : : public Expression<meta_curl<E>,
472 : : sizeof(E) + 4 * sizeof(typename E::Mesh_t::vector_type)> {
473 : : constexpr static unsigned dim = E::dim;
474 : :
475 : : KOKKOS_FUNCTION
476 : 2 : meta_curl(const E& u, const typename E::Mesh_t::vector_type& xvector,
477 : : const typename E::Mesh_t::vector_type& yvector,
478 : : const typename E::Mesh_t::vector_type& zvector,
479 : : const typename E::Mesh_t::vector_type& hvector)
480 : 2 : : u_m(u)
481 : 2 : , xvector_m(xvector)
482 : 2 : , yvector_m(yvector)
483 : 2 : , zvector_m(zvector)
484 : 2 : , hvector_m(hvector) {}
485 : :
486 : : /*
487 : : * 3-dimensional curl
488 : : */
489 : 8192 : KOKKOS_INLINE_FUNCTION auto operator()(size_t i, size_t j, size_t k) const {
490 : 8192 : return xvector_m
491 [ + - + - ]: 24576 : * ((u_m(i, j + 1, k)[2] - u_m(i, j - 1, k)[2]) / (2 * hvector_m[1])
492 [ + - + - : 16384 : - (u_m(i, j, k + 1)[1] - u_m(i, j, k - 1)[1]) / (2 * hvector_m[2]))
+ - ]
493 [ + - ]: 16384 : + yvector_m
494 [ + - + - ]: 24576 : * ((u_m(i, j, k + 1)[0] - u_m(i, j, k - 1)[0]) / (2 * hvector_m[2])
495 [ + - + - : 16384 : - (u_m(i + 1, j, k)[2] - u_m(i - 1, j, k)[2]) / (2 * hvector_m[0]))
+ - ]
496 : 8192 : + zvector_m
497 [ + - + - ]: 24576 : * ((u_m(i + 1, j, k)[1] - u_m(i - 1, j, k)[1]) / (2 * hvector_m[0])
498 [ + - + - : 32768 : - (u_m(i, j + 1, k)[0] - u_m(i, j - 1, k)[0]) / (2 * hvector_m[1]));
+ - + - ]
499 : : }
500 : :
501 : : private:
502 : : using Mesh_t = typename E::Mesh_t;
503 : : using vector_type = typename Mesh_t::vector_type;
504 : : const E u_m;
505 : : const vector_type xvector_m;
506 : : const vector_type yvector_m;
507 : : const vector_type zvector_m;
508 : : const vector_type hvector_m;
509 : : };
510 : : } // namespace detail
511 : :
512 : : namespace detail {
513 : :
514 : : /*!
515 : : * Meta function of Hessian
516 : : */
517 : : template <typename E>
518 : : struct meta_hess
519 : : : public Expression<meta_hess<E>,
520 : : sizeof(E)
521 : : + sizeof(typename E::Mesh_t::vector_type[E::Mesh_t::Dimension])
522 : : + sizeof(typename E::Mesh_t::vector_type)> {
523 : : constexpr static unsigned dim = E::dim;
524 : :
525 : : KOKKOS_FUNCTION
526 : 12 : meta_hess(const E& u, const typename E::Mesh_t::vector_type vectors[],
527 : : const typename E::Mesh_t::vector_type& hvector)
528 : 12 : : u_m(u)
529 [ # # ][ + - : 54 : , hvector_m(hvector) {
+ + - - -
- ]
530 [ + + ]: 54 : for (unsigned d = 0; d < E::Mesh_t::Dimension; d++) {
531 : 42 : vectors_m[d] = vectors[d];
532 : : }
533 : 12 : }
534 : :
535 : : /*
536 : : * n-dimensional hessian (return Vector<Vector<T,n>,n>)
537 : : */
538 : : template <typename... Idx>
539 : 238656 : KOKKOS_INLINE_FUNCTION auto operator()(const Idx... args) const {
540 : 238656 : matrix_type hessian;
541 [ + - ]: 238656 : computeHessian(std::make_index_sequence<dim>{}, hessian, args...);
542 : 238656 : return hessian;
543 : 0 : }
544 : :
545 : : private:
546 : : using Mesh_t = typename E::Mesh_t;
547 : : using vector_type = typename Mesh_t::vector_type;
548 : : using matrix_type = typename Mesh_t::matrix_type;
549 : :
550 : : const E u_m;
551 : : vector_type vectors_m[dim];
552 : : const vector_type hvector_m;
553 : :
554 : : /*!
555 : : * Utility function for computing the Hessian. Computes the rows of the matrix
556 : : * one by one via fold expression.
557 : : * @tparam row... the row indices (in practice, the sequence 0...Dim - 1)
558 : : * @tparam Idx... the indices at which to access the field view
559 : : * @param is index sequence (reused for row computation)
560 : : * @param hessian matrix in which to store the Hessian
561 : : * @param args... the indices
562 : : */
563 : : template <size_t... row, typename... Idx>
564 : 238656 : KOKKOS_INLINE_FUNCTION constexpr void computeHessian(
565 : : const std::index_sequence<row...>& is, matrix_type& hessian,
566 : : const Idx... args) const {
567 : : // The comma operator forces left-to-right evaluation order, which reduces
568 : : // performance; therefore we apply a dummy operation to dummy values and discard the
569 : : // result
570 : 238656 : [[maybe_unused]] auto _ = (hessianRow<row>(is, hessian, args...) + ...);
571 : 238656 : }
572 : :
573 : : /*!
574 : : * Utility function for computing the Hessian. Computes the entries in a single
575 : : * row of the matrix via fold expression.
576 : : * @tparam row the row index
577 : : * @tparam col... the column indices (in practice, the sequence 0...Dim - 1)
578 : : * @tparam Idx... the indices at which to access the field view
579 : : * @param hessian matrix in which to store the hessian
580 : : * @param args... the indices
581 : : * @return An unused dummy value (required to allow use of a more performant fold
582 : : * expression)
583 : : */
584 : : template <size_t row, size_t... col, typename... Idx>
585 : 1271872 : KOKKOS_INLINE_FUNCTION constexpr int hessianRow(const std::index_sequence<col...>&,
586 : : matrix_type& hessian,
587 : : const Idx... args) const {
588 [ # # # # : 1271872 : hessian[row] = (hessianEntry<row, col>(args...) + ...);
# # # # #
# ][ + - +
- + - + -
+ - + - +
- + - + -
+ - + - ]
589 : 1271872 : return 0;
590 : : }
591 : :
592 : : /*!
593 : : * Utility function for computing the Hessian. Computes a single entry
594 : : * of the matrix
595 : : * @tparam row the row index
596 : : * @tparam col the column index
597 : : * @tparam Idx... the indices at which to access the field view
598 : : * @param args... the indices
599 : : * @return The entry of the Hessian at the given row and column
600 : : */
601 : : template <size_t row, size_t col, typename... Idx>
602 : 6959168 : KOKKOS_INLINE_FUNCTION constexpr vector_type hessianEntry(const Idx... args) const {
603 : : using index_type = std::tuple_element_t<0, std::tuple<Idx...>>;
604 : 6959168 : index_type coords[dim] = {args...};
605 : : if constexpr (row == col) {
606 [ + - ]: 1271872 : auto&& center = apply(u_m, coords);
607 : :
608 : 1271872 : coords[row] += 1;
609 [ + - ]: 1271872 : auto&& right = apply(u_m, coords);
610 : :
611 : 1271872 : coords[row] -= 2;
612 [ + - ]: 1271872 : auto&& left = apply(u_m, coords);
613 : :
614 : : // The diagonal elements correspond to second derivatives w.r.t. a single
615 : : // variable
616 : 1271872 : return vectors_m[row] * (right - 2. * center + left)
617 [ + - + - : 2543744 : / (hvector_m[row] * hvector_m[row]);
+ - ]
618 : : } else {
619 : 5687296 : coords[row] += 1;
620 : 5687296 : coords[col] += 1;
621 [ + - ]: 5687296 : auto&& uu = apply(u_m, coords);
622 : :
623 : 5687296 : coords[col] -= 2;
624 [ + - ]: 5687296 : auto&& ud = apply(u_m, coords);
625 : :
626 : 5687296 : coords[row] -= 2;
627 [ + - ]: 5687296 : auto&& dd = apply(u_m, coords);
628 : :
629 : 5687296 : coords[col] += 2;
630 [ + - ]: 5687296 : auto&& du = apply(u_m, coords);
631 : :
632 : : // The non-diagonal elements are mixed derivatives, whose finite difference form
633 : : // is slightly different from above
634 : 5687296 : return vectors_m[col] * (uu - du - ud + dd)
635 [ + - + - : 11374592 : / (4. * hvector_m[row] * hvector_m[col]);
+ - ]
636 : : }
637 : : // Silences incorrect nvcc warning: missing return statement at end of non-void
638 : : // function
639 : : return vector_type{};
640 : : }
641 : : };
642 : : } // namespace detail
643 : : } // namespace ippl
644 : :
645 : : #endif
|
