Intrepid2_TensorArgumentIterator.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_TensorArgumentIterator_h
#define Intrepid2_TensorArgumentIterator_h

#include "Intrepid2_TensorData.hpp"
#include "Intrepid2_Types.hpp"

namespace Intrepid2
{
  class TensorArgumentIterator {
    Kokkos::Array<ordinal_type,Parameters::MaxTensorComponents> arguments_;
    Kokkos::Array<ordinal_type,Parameters::MaxTensorComponents> bounds_;
    ordinal_type numTensorComponents_;
  public:
    template<class Scalar, typename ExecSpaceType>
    KOKKOS_INLINE_FUNCTION
    TensorArgumentIterator(const TensorData<Scalar,ExecSpaceType> &tensorData, const ordinal_type argumentOrdinal)
    :
    numTensorComponents_(tensorData.numTensorComponents())
    {
      for (ordinal_type r=0; r<numTensorComponents_; r++)
      {
        arguments_[r] = 0;
        bounds_[r]    = tensorData.getTensorComponent(r).extent_int(argumentOrdinal);
      }
    }
    
    template<class Scalar, typename ExecSpaceType>
    KOKKOS_INLINE_FUNCTION
    TensorArgumentIterator(const TensorData<Scalar,ExecSpaceType> &tensorData, const ordinal_type argumentOrdinal, const ordinal_type numTensorComponents)
    :
    numTensorComponents_(numTensorComponents)
    {
      for (ordinal_type r=0; r<numTensorComponents_; r++)
      {
        arguments_[r] = 0;
        bounds_[r]    = tensorData.getTensorComponent(r).extent_int(argumentOrdinal);
      }
    }
    
    TensorArgumentIterator(const std::vector<ordinal_type> tensorComponentBounds)
    :
    numTensorComponents_(tensorComponentBounds.size())
    {
      for (ordinal_type r=0; r<numTensorComponents_; r++)
      {
        arguments_[r] = 0;
        bounds_[r]    = tensorComponentBounds[r];
      }
    }
    
    template<size_t rank>
    KOKKOS_INLINE_FUNCTION
    TensorArgumentIterator(const Kokkos::Array<ordinal_type,rank> &tensorComponentBounds)
    :
    numTensorComponents_(rank)
    {
      for (ordinal_type r=0; r<rank; r++)
      {
        arguments_[r] = 0;
        bounds_[r]    = tensorComponentBounds[r];
      }
    }
    
    KOKKOS_INLINE_FUNCTION ordinal_type increment()
    {
      ordinal_type r = numTensorComponents_ - 1;
      while (arguments_[r] + 1 >= bounds_[r])
      {
        arguments_[r] = 0; // reset
        r--;
        if (r < 0) break;
      }
      if (r >= 0) ++arguments_[r];
      return r;
    }
    
    KOKKOS_INLINE_FUNCTION
    ordinal_type nextIncrementResult() const
    {
      ordinal_type r = numTensorComponents_ - 1;
      while (arguments_[r] + 1 >= bounds_[r])
      {
        r--;
        if (r < 0) break;
      }
      return r;
    }
    
    KOKKOS_INLINE_FUNCTION const ordinal_type & argument(const ordinal_type &r) const
    {
      return arguments_[r];
    }
    
    KOKKOS_INLINE_FUNCTION ordinal_type enumerationIndex() const
    {
      // commented-out code belongs to implementation with rightmost argument as the fastest-moving.  We may want to support this as an option.
//      ordinal_type i = 0;
//      for (ordinal_type r=0; r<numTensorComponents_-1; r++)
//      {
//        i += arguments_[r];
//        i *= bounds_[r+1];
//      }
//      i += arguments_[numTensorComponents_-1];
//      return i;
      
      // TensorData's numbering has the leftmost argument as the fastest-moving
      // We return that numbering here.
      ordinal_type i = 0;
      for (ordinal_type r=numTensorComponents_-1; r>0; r--)
      {
        i += arguments_[r];
        i *= bounds_[r-1];
      }
      i += arguments_[0];
      return i;
    }
    
    KOKKOS_INLINE_FUNCTION ordinal_type relativeEnumerationIndex(const ordinal_type &startingComponent) const
    {
      // commented-out code belongs to implementation with rightmost argument as the fastest-moving.  We may want to support this as an option.
//      ordinal_type i = 0;
//      for (ordinal_type r=startingComponent; r<numTensorComponents_-1; r++)
//      {
//        i += arguments_[r];
//        i *= bounds_[r+1];
//      }
//      i += arguments_[numTensorComponents_-1];
//      return i;
      
      // TensorData's numbering has the leftmost argument as the fastest-moving
      // We return that numbering here.
      ordinal_type i = 0;
      for (ordinal_type r=numTensorComponents_-1; r>startingComponent; r--)
      {
        i += arguments_[r];
        i *= bounds_[r-1];
      }
      i += arguments_[startingComponent];
      return i;
    }
    
    KOKKOS_INLINE_FUNCTION ordinal_type relativeEnumerationSpan(const ordinal_type &startingComponent) const
    {
      ordinal_type i = 1;
      for (ordinal_type r=startingComponent; r<numTensorComponents_; r++)
      {
        i *= bounds_[r];
      }
      return i;
    }
    
    KOKKOS_INLINE_FUNCTION
    void reset(ordinal_type from_component_number=0)
    {
      for (ordinal_type r=from_component_number; r<numTensorComponents_; r++)
      {
        arguments_[r] = 0;
      }
    }
    
    KOKKOS_INLINE_FUNCTION
    void setArgumentForComponent(const ordinal_type &r, const ordinal_type &i)
    {
      arguments_[r] = i;
    }
    
    KOKKOS_INLINE_FUNCTION
    void setEnumerationIndex(const ordinal_type &enumerationIndex)
    {
      ordinal_type remainder = enumerationIndex;
      for (ordinal_type d=0; d<numTensorComponents_; d++)
      {
        arguments_[d] = remainder % bounds_[d];
        remainder  /= bounds_[d];
      }
    }
    
    KOKKOS_INLINE_FUNCTION
    void copyArguments(TensorArgumentIterator &otherArgumentIterator, const ordinal_type &r0_from, const ordinal_type &r0_to, const ordinal_type &numArguments)
    {
      for (ordinal_type i=0; i<numArguments; i++)
      {
        arguments_[r0_to + i] = otherArgumentIterator.argument(r0_from + i);
      }
    }
    
  };
} // namespace Intrepid2

#endif /* Intrepid2_TensorArgumentIterator_h */