Intrepid2_TensorData.hpp

Go to the documentation of this file.

// @HEADER
// *****************************************************************************
//                           Intrepid2 Package
//
// Copyright 2007 NTESS and the Intrepid2 contributors.
// SPDX-License-Identifier: BSD-3-Clause
// *****************************************************************************
// @HEADER

#ifndef Intrepid2_TensorData_h
#define Intrepid2_TensorData_h

#include "Intrepid2_Data.hpp"

#include "Intrepid2_ScalarView.hpp"
#include "Intrepid2_Types.hpp"

namespace Intrepid2
{
  template<class Scalar, typename DeviceType>
  class TensorData {
  public:
    using value_type = Scalar;
    using execution_space = typename DeviceType::execution_space;
  protected:
    Kokkos::Array< Data<Scalar,DeviceType>, Parameters::MaxTensorComponents> tensorComponents_;
    Kokkos::Array<ordinal_type, 7> extents_;
    Kokkos::Array<Kokkos::Array<ordinal_type, Parameters::MaxTensorComponents>, 7> entryModulus_;
    ordinal_type rank_;
    bool separateFirstComponent_ = false; // true supported only for rank 1 components; uses a rank-2 operator() in that case, where the first argument corresponds to the index in the first component, while the second argument corresponds to the tensorially-multiplied remaining components
    ordinal_type numTensorComponents_ = 0;
    
    void initialize()
    {
      ordinal_type maxComponentRank = -1;
      for (ordinal_type r=0; r<numTensorComponents_; r++)
      {
        const ordinal_type componentRank = tensorComponents_[r].rank();
        maxComponentRank = (maxComponentRank > componentRank) ? maxComponentRank : componentRank;
      }
      INTREPID2_TEST_FOR_EXCEPTION_DEVICE_SAFE(separateFirstComponent_ && (maxComponentRank != 1), std::invalid_argument, "separateFirstComponent = true only supported if all components have rank 1");
      ordinal_type rd_start; // where to begin the extents_ and entryModulus_ loops below
      if ((maxComponentRank == 1) && separateFirstComponent_)
      {
        rank_ = 2;
        rd_start = 1;
        extents_[0] = tensorComponents_[0].extent_int(0);
        entryModulus_[0][0] = extents_[0]; // should not be used
      }
      else
      {
        rank_ = maxComponentRank;
        rd_start = 0;
      }
      
      for (ordinal_type d=rd_start; d<7; d++)
      {
        extents_[d] = 1;
        for (ordinal_type r=rd_start; r<numTensorComponents_; r++)
        {
          extents_[d] *= tensorComponents_[r].extent_int(d-rd_start);
        }
        ordinal_type entryModulus = extents_[d];
        for (ordinal_type r=rd_start; r<numTensorComponents_; r++)
        {
          entryModulus /= tensorComponents_[r].extent_int(d-rd_start);
          entryModulus_[d][r] = entryModulus;
        }
      }
    }
  public:
    template<size_t numTensorComponents>
    TensorData(Kokkos::Array< Data<Scalar,DeviceType>, numTensorComponents> tensorComponents, bool separateFirstComponent = false)
    :
    separateFirstComponent_(separateFirstComponent),
    numTensorComponents_(numTensorComponents)
    {
      for (size_t r=0; r< numTensorComponents; r++)
      {
        tensorComponents_[r] = tensorComponents[r];
      }
      
      initialize();
    }
    
    TensorData(std::vector< Data<Scalar,DeviceType> > tensorComponents, bool separateFirstComponent = false)
    :
    separateFirstComponent_(separateFirstComponent),
    numTensorComponents_(tensorComponents.size())
    {
      for (ordinal_type r=0; r<numTensorComponents_; r++)
      {
        tensorComponents_[r] = tensorComponents[r];
      }
      
      initialize();
    }
    
    TensorData(const TensorData &first, const TensorData &second, bool separateFirstComponent = false)
    :
    separateFirstComponent_(separateFirstComponent),
    numTensorComponents_(first.numTensorComponents() + second.numTensorComponents())
    {
      ordinal_type r = 0;
      for (ordinal_type r1=0; r1<first.numTensorComponents(); r1++, r++)
      {
        tensorComponents_[r] = first.getTensorComponent(r1);
      }
      for (ordinal_type r2=0; r2<second.numTensorComponents(); r2++, r++)
      {
        tensorComponents_[r] = second.getTensorComponent(r2);
      }
      
      initialize();
    }
    
    TensorData(Data<Scalar,DeviceType> tensorComponent)
    :
    TensorData(Kokkos::Array< Data<Scalar,DeviceType>, 1>({tensorComponent}), false)
    {}
    
    TensorData()
    :
    extents_({0,0,0,0,0,0,0}),
    rank_(0)
    {}
    
    TensorData(TensorData otherTensorData, std::vector<int> whichComps)
    :
    numTensorComponents_(whichComps.size())
    {
      int r = 0;
      for (const auto & componentOrdinal : whichComps)
      {
        tensorComponents_[r++] = otherTensorData.getTensorComponent(componentOrdinal);
      }
      
      initialize();
    }
    
    template<typename OtherDeviceType, class = typename std::enable_if< std::is_same<typename DeviceType::memory_space, typename OtherDeviceType::memory_space>::value>::type,
                                       class = typename std::enable_if<!std::is_same<DeviceType,OtherDeviceType>::value>::type>
    TensorData(const TensorData<Scalar,OtherDeviceType> &tensorData)
    {
      if (tensorData.isValid())
      {
        numTensorComponents_ = tensorData.numTensorComponents();
        for (ordinal_type r=0; r<numTensorComponents_; r++)
        {
          Data<Scalar,OtherDeviceType> otherTensorComponent = tensorData.getTensorComponent(r);
          tensorComponents_[r] = Data<Scalar,DeviceType>(otherTensorComponent);
        }
        initialize();
      }
      else
      {
        extents_ = Kokkos::Array<ordinal_type,7>{0,0,0,0,0,0,0};
        rank_    = 0;
      }
    }
    
    template<typename OtherDeviceType, class = typename std::enable_if<!std::is_same<typename DeviceType::memory_space, typename OtherDeviceType::memory_space>::value>::type>
    TensorData(const TensorData<Scalar,OtherDeviceType> &tensorData)
    {
      if (tensorData.isValid())
      {
        numTensorComponents_ = tensorData.numTensorComponents();
        for (ordinal_type r=0; r<numTensorComponents_; r++)
        {
          Data<Scalar,OtherDeviceType> otherTensorComponent = tensorData.getTensorComponent(r);
          tensorComponents_[r] = Data<Scalar,DeviceType>(otherTensorComponent);
        }
        initialize();
      }
      else
      {
        extents_ = Kokkos::Array<ordinal_type,7>{0,0,0,0,0,0,0};
        rank_    = 0;
      }
    }
    
    KOKKOS_INLINE_FUNCTION
    const Data<Scalar,DeviceType> & getTensorComponent(const ordinal_type &r) const
    {
      return tensorComponents_[r];
    }
    
    template <typename iType0>
    KOKKOS_INLINE_FUNCTION typename std::enable_if<std::is_integral<iType0>::value, Scalar>::type
    operator()(const iType0& tensorEntryIndex) const {
#ifdef HAVE_INTREPID2_DEBUG
      INTREPID2_TEST_FOR_EXCEPTION_DEVICE_SAFE(rank_ != 1, std::invalid_argument, "This method is only valid for rank 1 containers.");
#endif
      Scalar value = 1.0;
      iType0 remainingEntryOrdinal = tensorEntryIndex;
      for (ordinal_type r=0; r<numTensorComponents_; r++)
      {
        const ordinal_type componentEntryCount   = tensorComponents_[r].extent_int(0);
        const ordinal_type componentEntryOrdinal = remainingEntryOrdinal % componentEntryCount;
        remainingEntryOrdinal /= componentEntryCount;
        
        value *= tensorComponents_[r](componentEntryOrdinal);
      }
      
      return value;
    }
    
    template <typename iType0, ordinal_type numTensorComponents>
    KOKKOS_INLINE_FUNCTION typename std::enable_if<std::is_integral<iType0>::value, Scalar>::type
    operator()(const Kokkos::Array<iType0,numTensorComponents>& entryComponents) const {
#ifdef HAVE_INTREPID2_DEBUG
      INTREPID2_TEST_FOR_EXCEPTION_DEVICE_SAFE(rank_ != 1, std::invalid_argument, "This method is only valid for rank 1 containers.");
      INTREPID2_TEST_FOR_EXCEPTION_DEVICE_SAFE(numTensorComponents_ != numTensorComponents, std::invalid_argument, "Tensorial component count mismatch");
#endif
      Scalar value = 1.0;
      for (ordinal_type r=0; r<numTensorComponents; r++)
      {
        value *= tensorComponents_[r](entryComponents[r]);
      }
      return value;
    }
    
    template <typename iType0, typename iType1>
    KOKKOS_INLINE_FUNCTION typename std::enable_if<
        (std::is_integral<iType0>::value && std::is_integral<iType1>::value),
        Scalar>::type
    operator()(const iType0& tensorEntryIndex0, const iType1& tensorEntryIndex1) const {
#ifdef HAVE_INTREPID2_DEBUG
      INTREPID2_TEST_FOR_EXCEPTION_DEVICE_SAFE(rank_ != 2, std::invalid_argument, "This method is only valid for rank 2 containers.");
#endif
      
      if (numTensorComponents_ == 1)
      {
        return tensorComponents_[0](tensorEntryIndex0,tensorEntryIndex1);
      }
      
      if (!separateFirstComponent_)
      {
        Scalar value = 1.0;
        iType0 remainingEntryOrdinal0 = tensorEntryIndex0;
        iType1 remainingEntryOrdinal1 = tensorEntryIndex1;
        for (ordinal_type r=0; r<numTensorComponents_; r++)
        {
          auto & component = tensorComponents_[r];
          const ordinal_type componentEntryCount0   = component.extent_int(0);
          const ordinal_type componentEntryCount1   = component.extent_int(1);
          const iType0 componentEntryOrdinal0 = remainingEntryOrdinal0 % componentEntryCount0;
          const iType1 componentEntryOrdinal1 = remainingEntryOrdinal1 % componentEntryCount1;
          remainingEntryOrdinal0 /= componentEntryCount0;
          remainingEntryOrdinal1 /= componentEntryCount1;
          
          const ordinal_type componentRank = component.rank();
          
          if (componentRank == 2)
          {
            value *= component(componentEntryOrdinal0,componentEntryOrdinal1);
          }
          else if (componentRank == 1)
          {
            value *= component(componentEntryOrdinal0);
          }
          else
          {
            INTREPID2_TEST_FOR_EXCEPTION_DEVICE_SAFE(true, std::logic_error, "unsupported component rank encountered");
          }
        }
        
        return value;
      }
      else
      {
        Scalar value = tensorComponents_[0](tensorEntryIndex0);
        iType0 remainingEntryOrdinal = tensorEntryIndex1;
        for (ordinal_type r=1; r<numTensorComponents_; r++)
        {
          const ordinal_type componentEntryCount   = tensorComponents_[r].extent_int(0);
          const ordinal_type componentEntryOrdinal = remainingEntryOrdinal % componentEntryCount;
          remainingEntryOrdinal /= componentEntryCount;
          
          value *= tensorComponents_[r](componentEntryOrdinal);
        }
        return value;
      }
    }
    
    KOKKOS_INLINE_FUNCTION
    ordinal_type getTensorComponentIndex(const ordinal_type &tensorComponent, const ordinal_type &dim, const ordinal_type &enumerationIndex) const
    {
      ordinal_type remainingEntryOrdinal = enumerationIndex;
      for (ordinal_type r=0; r<tensorComponent; r++)
      {
        const auto & component = tensorComponents_[r];
        const ordinal_type & componentEntryCount = component.extent_int(dim);
        
        remainingEntryOrdinal /= componentEntryCount;
      }
      return remainingEntryOrdinal % tensorComponents_[tensorComponent].extent_int(dim);
    }
    
    template <typename iType0, typename iType1, typename iType2>
    KOKKOS_INLINE_FUNCTION typename std::enable_if<
        (std::is_integral<iType0>::value && std::is_integral<iType1>::value && std::is_integral<iType2>::value),
        Scalar>::type
    operator()(const iType0& tensorEntryIndex0, const iType1& tensorEntryIndex1, const iType2& tensorEntryIndex2) const
    {
#ifdef HAVE_INTREPID2_DEBUG
      INTREPID2_TEST_FOR_EXCEPTION_DEVICE_SAFE(rank_ != 3, std::invalid_argument, "This method is only valid for rank 3 containers.");
      INTREPID2_TEST_FOR_EXCEPTION_DEVICE_SAFE(separateFirstComponent_, std::logic_error, "This method should never be called when separateFirstComponent is true");
#endif
      
      Scalar value = 1.0;
      Kokkos::Array<ordinal_type,3> remainingEntryOrdinal {tensorEntryIndex0, tensorEntryIndex1, tensorEntryIndex2};
      for (ordinal_type r=0; r<numTensorComponents_; r++)
      {
        auto & component = tensorComponents_[r];
        const ordinal_type componentEntryCount0   = component.extent_int(0);
        const ordinal_type componentEntryCount1   = component.extent_int(1);
        const ordinal_type componentEntryCount2   = component.extent_int(2);
        const ordinal_type componentEntryOrdinal0 = remainingEntryOrdinal[0] % componentEntryCount0;
        const ordinal_type componentEntryOrdinal1 = remainingEntryOrdinal[1] % componentEntryCount1;
        const ordinal_type componentEntryOrdinal2 = remainingEntryOrdinal[2] % componentEntryCount2;
        remainingEntryOrdinal[0] /= componentEntryCount0;
        remainingEntryOrdinal[1] /= componentEntryCount1;
        remainingEntryOrdinal[2] /= componentEntryCount2;
        
        const ordinal_type componentRank = component.rank();
        
        if (componentRank == 3)
        {
          value *= component(componentEntryOrdinal0,componentEntryOrdinal1,componentEntryOrdinal2);
        }
        else if (componentRank == 2)
        {
          value *= component(componentEntryOrdinal0,componentEntryOrdinal1);
        }
        else if (componentRank == 1)
        {
          value *= component(componentEntryOrdinal0);
        }
        else
        {
          INTREPID2_TEST_FOR_EXCEPTION_DEVICE_SAFE(true, std::logic_error, "unsupported component rank encountered");
        }
      }
      return value;
    }
    
    template <typename iType0, typename iType1, ordinal_type numTensorComponents>
    KOKKOS_INLINE_FUNCTION typename std::enable_if<
        (std::is_integral<iType0>::value && std::is_integral<iType1>::value),
        Scalar>::type
    operator()(const Kokkos::Array<iType0,numTensorComponents>& entryComponents0, const Kokkos::Array<iType1,numTensorComponents>& entryComponents1) const {
#ifdef HAVE_INTREPID2_DEBUG
      INTREPID2_TEST_FOR_EXCEPTION_DEVICE_SAFE(numTensorComponents_ != numTensorComponents, std::invalid_argument, "Tensorial component count mismatch");
#endif
      Scalar value = 1.0;
      for (ordinal_type r=0; r<numTensorComponents; r++)
      {
        auto & component = tensorComponents_[r];
        const ordinal_type componentRank = component.rank();
        if (componentRank == 2)
        {
          value *= component(entryComponents0[r],entryComponents1[r]);
        }
        else if (componentRank == 1)
        {
          value *= component(entryComponents0[r]);
        }
        else
        {
          INTREPID2_TEST_FOR_EXCEPTION_DEVICE_SAFE(true, std::logic_error, "unsupported component rank encountered");
        }
      }
      return value;
    }
    
    template <typename iType>
    KOKKOS_INLINE_FUNCTION
    typename std::enable_if<std::is_integral<iType>::value, ordinal_type>::type
    extent_int(const iType& d) const {
      return extents_[d];
    }
    
    template <typename iType>
    KOKKOS_INLINE_FUNCTION constexpr
    typename std::enable_if<std::is_integral<iType>::value, size_t>::type
    extent(const iType& d) const {
      return extents_[d];
    }
    
    KOKKOS_INLINE_FUNCTION constexpr bool isValid() const
    {
      return extents_[0] > 0;
    }
    
    KOKKOS_INLINE_FUNCTION
    ordinal_type rank() const
    {
      return rank_;
    }
    
    KOKKOS_INLINE_FUNCTION
    ordinal_type numTensorComponents() const
    {
      return numTensorComponents_;
    }
    
    KOKKOS_INLINE_FUNCTION
    bool separateFirstComponent() const
    {
      return separateFirstComponent_;
    }
    
    void setFirstComponentExtentInDimension0(const ordinal_type &newExtent)
    {
      INTREPID2_TEST_FOR_EXCEPTION(!separateFirstComponent_ && (numTensorComponents_ != 1), std::invalid_argument, "setFirstComponentExtent() is only allowed when separateFirstComponent_ is true, or there is only one component");
      tensorComponents_[0].setExtent(0,newExtent);
      extents_[0] = newExtent;
    }
  };
}

#endif /* Intrepid2_TensorData_h */