Line data    Source code

       1           0 : // Distributed under the MIT License.
       2             : // See LICENSE.txt for details.
       3             : 
       4             : #pragma once
       5             : 
       6             : #include <array>
       7             : #include <cmath>
       8             : #include <cstddef>
       9             : #include <limits>
      10             : #include <type_traits>
      11             : #include <utility>
      12             : 
      13             : #include "DataStructures/Tensor/Expressions/NumberAsExpression.hpp"
      14             : #include "DataStructures/Tensor/Expressions/TensorExpression.hpp"
      15             : #include "DataStructures/Tensor/Expressions/TimeIndex.hpp"
      16             : #include "Utilities/ForceInline.hpp"
      17             : #include "Utilities/TMPL.hpp"
      18             : 
      19             : namespace tenex {
      20             : /// \ingroup TensorExpressionsGroup
      21             : /// \brief Defines the tensor expression representing the square root of a
      22             : /// tensor expression that evaluates to a rank 0 tensor
      23             : ///
      24             : /// \details The expression can have a non-zero number of indices as long as
      25             : /// all indices are concrete time indices, as this represents a rank 0 tensor.
      26             : ///
      27             : /// For details on aliases and members defined in this class, as well as general
      28             : /// `TensorExpression` terminology used in its members' documentation, see
      29             : /// documentation for `TensorExpression`.
      30             : ///
      31             : /// \tparam T the type of the tensor expression of which to take the square
      32             : /// root
      33             : /// \tparam Args the TensorIndexs of the expression
      34             : template <typename T, typename... Args>
      35           1 : struct SquareRoot
      36             :     : public TensorExpression<SquareRoot<T, Args...>, typename T::type,
      37             :                               typename T::symmetry, typename T::index_list,
      38             :                               tmpl::list<Args...>> {
      39             :   static_assert(
      40             :       (... and tt::is_time_index<Args>::value),
      41             :       "Can only take the square root of a tensor expression that evaluates to "
      42             :       "a rank 0 tensor.");
      43             : 
      44             :   // === Index properties ===
      45             :   /// The type of the data being stored in the result of the expression
      46           1 :   using type = typename T::type;
      47             :   /// The ::Symmetry of the result of the expression
      48           1 :   using symmetry = typename T::symmetry;
      49             :   /// The list of \ref SpacetimeIndex "TensorIndexType"s of the result of the
      50             :   /// expression
      51           1 :   using index_list = typename T::index_list;
      52             :   /// The list of generic `TensorIndex`s of the result of the expression
      53           1 :   using args_list = tmpl::list<Args...>;
      54             :   /// The number of tensor indices in the result of the expression
      55           1 :   static constexpr auto num_tensor_indices = sizeof...(Args);
      56             : 
      57             :   // === Expression subtree properties ===
      58             :   /// The number of arithmetic tensor operations done in the subtree for the
      59             :   /// left operand
      60           1 :   static constexpr size_t num_ops_left_child = T::num_ops_subtree;
      61             :   /// The number of arithmetic tensor operations done in the subtree for the
      62             :   /// right operand. This is 0 because this expression represents a unary
      63             :   /// operation.
      64           1 :   static constexpr size_t num_ops_right_child = 0;
      65             :   /// The total number of arithmetic tensor operations done in this expression's
      66             :   /// whole subtree
      67           1 :   static constexpr size_t num_ops_subtree = num_ops_left_child + 1;
      68             :   /// The height of this expression's node in the expression tree relative to
      69             :   /// the closest `TensorAsExpression` leaf in its subtree
      70           1 :   static constexpr size_t height_relative_to_closest_tensor_leaf_in_subtree =
      71             :       T::height_relative_to_closest_tensor_leaf_in_subtree !=
      72             :               std::numeric_limits<size_t>::max()
      73             :           ? T::height_relative_to_closest_tensor_leaf_in_subtree + 1
      74             :           : T::height_relative_to_closest_tensor_leaf_in_subtree;
      75             : 
      76             :   // === Properties for splitting up subexpressions along the primary path ===
      77             :   // These definitions only have meaning if this expression actually ends up
      78             :   // being along the primary path that is taken when evaluating the whole tree.
      79             :   // See documentation for `TensorExpression` for more details.
      80             :   /// If on the primary path, whether or not the expression is an ending point
      81             :   /// of a leg
      82           1 :   static constexpr bool is_primary_end = T::is_primary_start;
      83             :   /// If on the primary path, this is the remaining number of arithmetic tensor
      84             :   /// operations that need to be done in the subtree of the child along the
      85             :   /// primary path, given that we will have already computed the whole subtree
      86             :   /// at the next lowest leg's starting point.
      87           1 :   static constexpr size_t num_ops_to_evaluate_primary_left_child =
      88             :       is_primary_end ? 0 : T::num_ops_to_evaluate_primary_subtree;
      89             :   /// If on the primary path, this is the remaining number of arithmetic tensor
      90             :   /// operations that need to be done in the right operand's subtree. No
      91             :   /// splitting is currently done, so this is just `num_ops_right_child`.
      92           1 :   static constexpr size_t num_ops_to_evaluate_primary_right_child =
      93             :       num_ops_right_child;
      94             :   /// If on the primary path, this is the remaining number of arithmetic tensor
      95             :   /// operations that need to be done for this expression's subtree, given that
      96             :   /// we will have already computed the subtree at the next lowest leg's
      97             :   /// starting point
      98           1 :   static constexpr size_t num_ops_to_evaluate_primary_subtree =
      99             :       num_ops_to_evaluate_primary_left_child +
     100             :       num_ops_to_evaluate_primary_right_child + 1;
     101             :   /// If on the primary path, whether or not the expression is a starting point
     102             :   /// of a leg
     103           1 :   static constexpr bool is_primary_start =
     104             :       num_ops_to_evaluate_primary_subtree >=
     105             :       detail::max_num_ops_in_sub_expression<type>;
     106             :   /// If on the primary path, whether or not the expression's child along the
     107             :   /// primary path is a subtree that contains a starting point of a leg along
     108             :   /// the primary path
     109           1 :   static constexpr bool primary_child_subtree_contains_primary_start =
     110             :       T::primary_subtree_contains_primary_start;
     111             :   /// If on the primary path, whether or not this subtree contains a starting
     112             :   /// point of a leg along the primary path
     113           1 :   static constexpr bool primary_subtree_contains_primary_start =
     114             :       is_primary_start or primary_child_subtree_contains_primary_start;
     115             : 
     116           0 :   SquareRoot(T t) : t_(std::move(t)) {}
     117           0 :   ~SquareRoot() override = default;
     118             : 
     119             :   /// \brief Assert that the LHS tensor of the equation does not also appear in
     120             :   /// this expression's subtree
     121             :   template <typename LhsTensor>
     122           1 :   SPECTRE_ALWAYS_INLINE void assert_lhs_tensor_not_in_rhs_expression(
     123             :       const gsl::not_null<LhsTensor*> lhs_tensor) const {
     124             :     if constexpr (not std::is_base_of_v<MarkAsNumberAsExpression, T>) {
     125             :       t_.assert_lhs_tensor_not_in_rhs_expression(lhs_tensor);
     126             :     }
     127             :   }
     128             : 
     129             :   /// \brief Assert that each instance of the LHS tensor in the RHS tensor
     130             :   /// expression uses the same generic index order that the LHS uses
     131             :   ///
     132             :   /// \tparam LhsTensorIndices the list of generic `TensorIndex`s of the LHS
     133             :   /// result `Tensor` being computed
     134             :   /// \param lhs_tensor the LHS result `Tensor` being computed
     135             :   template <typename LhsTensorIndices, typename LhsTensor>
     136           1 :   SPECTRE_ALWAYS_INLINE void assert_lhs_tensorindices_same_in_rhs(
     137             :       const gsl::not_null<LhsTensor*> lhs_tensor) const {
     138             :     if constexpr (not std::is_base_of_v<MarkAsNumberAsExpression, T>) {
     139             :       t_.template assert_lhs_tensorindices_same_in_rhs<LhsTensorIndices>(
     140             :           lhs_tensor);
     141             :     }
     142             :   }
     143             : 
     144             :   /// \brief Get the size of a component from a `Tensor` in this expression's
     145             :   /// subtree of the RHS `TensorExpression`
     146             :   ///
     147             :   /// \return the size of a component from a `Tensor` in this expression's
     148             :   /// subtree of the RHS `TensorExpression`
     149           1 :   SPECTRE_ALWAYS_INLINE size_t get_rhs_tensor_component_size() const {
     150             :     return t_.get_rhs_tensor_component_size();
     151             :   }
     152             : 
     153             :   /// \brief Returns the square root of the component of the tensor evaluated
     154             :   /// from the contained tensor expression
     155             :   ///
     156             :   /// \details
     157             :   /// SquareRoot only supports tensor expressions that evaluate to a rank 0
     158             :   /// Tensor. This is why `multi_index` is always an array of size 0.
     159             :   ///
     160             :   /// \param multi_index the multi-index of the component of which to take the
     161             :   /// square root
     162             :   /// \return the square root of the component of the tensor evaluated from the
     163             :   /// contained tensor expression
     164           1 :   SPECTRE_ALWAYS_INLINE decltype(auto) get(
     165             :       const std::array<size_t, num_tensor_indices>& multi_index) const {
     166             :     return sqrt(t_.get(multi_index));
     167             :   }
     168             : 
     169             :   /// \brief Returns the square root of the component of the tensor evaluated
     170             :   /// from the contained tensor expression
     171             :   ///
     172             :   /// \details
     173             :   /// SquareRoot only supports tensor expressions that evaluate to a rank 0
     174             :   /// Tensor. This is why `multi_index` is always an array of size 0.
     175             :   ///
     176             :   /// This function differs from `get` in that it takes into account whether we
     177             :   /// have already computed part of the result component at a lower subtree.
     178             :   /// In recursively computing this square root, the current result component
     179             :   /// will be substituted in for the most recent (highest) subtree below it that
     180             :   /// has already been evaluated.
     181             :   ///
     182             :   /// \param result_component the LHS tensor component to evaluate
     183             :   /// \param multi_index the multi-index of the component of which to take the
     184             :   /// square root
     185             :   /// \return the square root of the component of the tensor evaluated from the
     186             :   /// contained tensor expression
     187             :   template <typename ResultType>
     188           1 :   SPECTRE_ALWAYS_INLINE decltype(auto) get_primary(
     189             :       const ResultType& result_component,
     190             :       const std::array<size_t, num_tensor_indices>& multi_index) const {
     191             :     if constexpr (is_primary_end) {
     192             :       (void)multi_index;
     193             :       // We've already computed the whole child subtree on the primary path, so
     194             :       // just return the square root of the current result component
     195             :       return sqrt(result_component);
     196             :     } else {
     197             :       // We haven't yet evaluated the whole subtree for this expression, so
     198             :       // return the square root of this expression's subtree
     199             :       return sqrt(t_.template get_primary(result_component, multi_index));
     200             :     }
     201             :   }
     202             : 
     203             :   /// \brief Successively evaluate the LHS Tensor's result component at each
     204             :   /// leg in this expression's subtree
     205             :   ///
     206             :   /// \details
     207             :   /// This function takes into account whether we have already computed part of
     208             :   /// the result component at a lower subtree. In recursively computing this
     209             :   /// square root, the current result component will be substituted in for the
     210             :   /// most recent (highest) subtree below it that has already been evaluated.
     211             :   ///
     212             :   /// \param result_component the LHS tensor component to evaluate
     213             :   /// \param multi_index the multi-index of the component of the result tensor
     214             :   /// to evaluate
     215             :   template <typename ResultType>
     216           1 :   SPECTRE_ALWAYS_INLINE void evaluate_primary_subtree(
     217             :       ResultType& result_component,
     218             :       const std::array<size_t, num_tensor_indices>& multi_index) const {
     219             :     if constexpr (primary_child_subtree_contains_primary_start) {
     220             :       // The primary child's subtree contains at least one leg, so recurse down
     221             :       // and evaluate that first
     222             :       t_.template evaluate_primary_subtree(result_component, multi_index);
     223             :     }
     224             : 
     225             :     if constexpr (is_primary_start) {
     226             :       // We want to evaluate the subtree for this expression
     227             :       result_component = get_primary(result_component, multi_index);
     228             :     }
     229             :   }
     230             : 
     231             :  private:
     232             :   /// Operand expression
     233           1 :   T t_;
     234             : };
     235             : }  // namespace tenex
     236             : 
     237             : /// \ingroup TensorExpressionsGroup
     238             : /// \brief Returns the tensor expression representing the square root of a
     239             : /// tensor expression that evaluates to a rank 0 tensor
     240             : ///
     241             : /// \details
     242             : /// `t` must be an expression that, when evaluated, would be a rank 0 tensor.
     243             : /// For example, if `R` and `S` are Tensors, here is a non-exhaustive list of
     244             : /// some of the acceptable forms that `t` could take:
     245             : /// - `R()`
     246             : /// - `R(ti::A, ti::a)`
     247             : /// - `(R(ti::A, ti::B) * S(ti::a, ti::b))`
     248             : /// - `R(ti::t, ti::t) + 1.0`
     249             : ///
     250             : /// \param t the tensor expression of which to take the square root
     251             : template <typename T, typename X, typename Symm, typename IndexList,
     252             :           typename... Args>
     253           1 : SPECTRE_ALWAYS_INLINE auto sqrt(
     254             :     const TensorExpression<T, X, Symm, IndexList, tmpl::list<Args...>>& t) {
     255             :   return tenex::SquareRoot<T, Args...>(~t);
     256             : }