Line data    Source code

       1            0 : // Distributed under the MIT License.
       2              : // See LICENSE.txt for details.
       3              : 
       4              : #pragma once
       5              : 
       6              : #include <concepts>
       7              : #include <cstddef>
       8              : 
       9              : #include "DataStructures/DataBox/Tag.hpp"
      10              : #include "DataStructures/Tensor/Metafunctions.hpp"
      11              : #include "DataStructures/Tensor/Tensor.hpp"
      12              : #include "Utilities/Kokkos/KokkosCore.hpp"
      13              : 
      14              : /*!
      15              :  * \brief Get the `Tensor` at a specific grid point index.
      16              :  *
      17              :  * Works with a `Tensor` of `Kokkos::View`s. Can also be made to work with
      18              :  * `DataVector` if needed (use `operator[]` instead of `operator()`).
      19              :  */
      20              : template <
      21              :     typename TensorType, std::integral... Is,
      22              :     typename ValueType = typename TensorType::type::value_type,
      23              :     typename ResultType = TensorMetafunctions::swap_type<ValueType, TensorType>>
      24            1 : KOKKOS_FUNCTION ResultType make_at_index(const TensorType& tensor,
      25              :                                          const Is&... i) {
      26              :   ResultType result{};
      27              :   for (size_t component = 0; component < TensorType::size(); ++component) {
      28              :     result[component] = tensor[component](i...);
      29              :   }
      30              :   return result;
      31              : }
      32              : 
      33              : /*!
      34              :  * \brief Set the `Tensor` at a specific grid point index.
      35              :  *
      36              :  * Works with a `Tensor` of `Kokkos::View`s. Can also be made to work with
      37              :  * `DataVector` if needed (use `operator[]` instead of `operator()`).
      38              :  */
      39              : template <typename TensorType, std::integral... Is,
      40              :           typename ValueType = typename TensorType::type::value_type>
      41            1 : KOKKOS_FUNCTION void set_at_index(
      42              :     const gsl::not_null<TensorType*> tensor,
      43              :     const TensorMetafunctions::swap_type<ValueType, TensorType>& value,
      44              :     const Is&... i) {
      45              :   for (size_t component = 0; component < value.size(); ++component) {
      46              :     (*tensor)[component](i...) = value[component];
      47              :   }
      48              : }
      49              : 
      50              : namespace Tags {
      51              : 
      52              : /*!
      53              :  * \brief Tag representing a specific grid point index of a tensor.
      54              :  *
      55              :  * This tag replaces the `Tensor` data type with `double` so that it can be
      56              :  * used with pointwise operations, e.g. in a `Kokkos::parallel_for` kernel.
      57              :  */
      58              : template <typename Tag>
      59            1 : struct AtIndex : db::PrefixTag {
      60              :  private:
      61            0 :   using TensorType = typename Tag::type;
      62            0 :   using ValueType = typename TensorType::type::value_type;
      63              : 
      64              :  public:
      65            0 :   using tag = Tag;
      66            0 :   using type = TensorMetafunctions::swap_type<ValueType, TensorType>;
      67              : };
      68              : 
      69              : }  // namespace Tags