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 120435732 : KOKKOS_INLINE_FUNCTION constexpr decltype(auto) apply_impl(const View& view,
23 : const Coords& coords,
24 : const std::index_sequence<Idx...>&) {
25 192308652 : 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 120435732 : KOKKOS_INLINE_FUNCTION constexpr decltype(auto) apply(const View& view, const Coords& coords) {
65 : using Indices = std::make_index_sequence<ExtractExpressionRank::getRank<Coords>()>;
66 120435732 : 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 96842654 : DefineBinaryOperation(Add, operator+, u_m[i] + v_m[i], u_m(args...) + v_m(args...))
177 24496734 : DefineBinaryOperation(Subtract, operator-, u_m[i] - v_m[i], u_m(args...) - v_m(args...))
178 92202602 : DefineBinaryOperation(Multiply, operator*, u_m[i] * v_m[i], u_m(args...) * v_m(args...))
179 53226828 : 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 1519276 : meta_dot(const E1& u, const E2& v)
262 1519276 : : u_m(u)
263 1519276 : , v_m(v) {}
264 :
265 : /*
266 : * Vector::dot
267 : */
268 1519264 : KOKKOS_INLINE_FUNCTION auto apply() const {
269 1519264 : 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 9775936 : for (size_t i = 0; i < E1::dim; ++i) {
273 8256672 : res += u_m[i] * v_m[i];
274 : }
275 1519264 : 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 1519276 : KOKKOS_INLINE_FUNCTION detail::meta_dot<E1, E2> dot(const detail::Expression<E1, N1>& u,
294 : const detail::Expression<E2, N2>& v) {
295 1519276 : 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
|
