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