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