Tensor.hpp

/*
 * ------------------------------------------------------------------------------------------------------------
 * SPDX-License-Identifier: LGPL-2.1-only
 *
 * Copyright (c) 2016-2024 Lawrence Livermore National Security LLC
 * Copyright (c) 2018-2024 TotalEnergies
 * Copyright (c) 2018-2024 The Board of Trustees of the Leland Stanford Junior University
 * Copyright (c) 2023-2024 Chevron
 * Copyright (c) 2019-     GEOS/GEOSX Contributors
 * All rights reserved
 *
 * See top level LICENSE, COPYRIGHT, CONTRIBUTORS, NOTICE, and ACKNOWLEDGEMENTS files for details.
 * ------------------------------------------------------------------------------------------------------------
 */
 
#ifndef GEOS_COMMON_TENSOR_HPP_
#define GEOS_COMMON_TENSOR_HPP_
 
#include "GeosxMacros.hpp"
 
namespace geos
{
 
template< typename T, int SIZE_TPARAM >
class Tensor
{
public:
 
  static_assert( SIZE_TPARAM > 0, "Tensor size must be a positive value" );
 
  using value_type = T;
 
  static constexpr int SIZE = SIZE_TPARAM;
 
  GEOS_HOST_DEVICE
  GEOS_FORCE_INLINE
  T const & operator[]( std::ptrdiff_t const i ) const
  {
    return data[i];
  }
 
  GEOS_HOST_DEVICE
  GEOS_FORCE_INLINE
  T & operator[]( std::ptrdiff_t const i )
  {
    return data[i];
  }
 
  template< typename U = T >
  GEOS_HOST_DEVICE
  GEOS_FORCE_INLINE
  std::enable_if_t< std::is_floating_point< U >::value, bool >
  operator==( Tensor< U, SIZE > const & rhs ) const
  {
    for( int i = 0; i < SIZE; ++i )
    {
      if( (data[i] > rhs.data[i]) || (data[i] < rhs.data[i]) )
      {
        return false;
      }
    }
    return true;
  }
 
  template< typename U = T >
  GEOS_HOST_DEVICE
  GEOS_FORCE_INLINE
  std::enable_if_t< !std::is_floating_point< U >::value, bool >
  operator==( Tensor< U, SIZE > const & rhs ) const
  {
    for( int i = 0; i < SIZE; ++i )
    {
      if( data[i] != rhs.data[i] )
      {
        return false;
      }
    }
    return true;
  }
 
  GEOS_HOST_DEVICE
  GEOS_FORCE_INLINE
  constexpr int size( int junk ) const
  {
    GEOS_UNUSED_VAR( junk );
    return SIZE;
  }
 
  GEOS_HOST_DEVICE Tensor< T, SIZE_TPARAM > & operator=( const double & rhs );
  GEOS_HOST_DEVICE Tensor< T, SIZE_TPARAM > & operator+=( const double & rhs );
  GEOS_HOST_DEVICE Tensor< T, SIZE_TPARAM > & operator+=( const Tensor< T, SIZE_TPARAM > & rhs );
 
  friend double GEOS_HOST_DEVICE operator*( const Tensor< T, SIZE_TPARAM > & lhs, const Tensor< T, SIZE_TPARAM > & rhs )
  {
    double result = 0;
    for( int i = 0; i < SIZE_TPARAM; ++i )
    {
      result += lhs.data[i] * rhs.data[i];
    }
    return result;
  };
 
  T data[SIZE] = {};
 
private:
 
  friend inline std::ostream & operator<<( std::ostream & os, Tensor< T, SIZE > const & t )
  {
    os << '{' << t.data[0];
    for( int i = 1; i < SIZE; ++i )
    {
      os << ',' << t.data[i];
    }
    os << '}';
    return os;
  }
};
 
// *****************************************************************************
// ****************************** END DECLARATION ******************************
// *****************************************************************************
 
template< typename T, int SIZE_TPARAM >
GEOS_HOST_DEVICE
GEOS_FORCE_INLINE Tensor< T, SIZE_TPARAM > &
Tensor< T, SIZE_TPARAM >::operator=( const double & rhs )
{
  for( int i = 0; i < SIZE_TPARAM; ++i )
  {
    data[i] = rhs;
  }
  return *this;
}
 
template< typename T, int SIZE_TPARAM >
GEOS_HOST_DEVICE
GEOS_FORCE_INLINE Tensor< T, SIZE_TPARAM > &
Tensor< T, SIZE_TPARAM >::operator+=( const double & rhs )
{
  for( int i = 0; i < SIZE_TPARAM; ++i )
  {
    data[i] += rhs;
  }
  return *this;
}
 
template< typename T, int SIZE_TPARAM >
GEOS_HOST_DEVICE
GEOS_FORCE_INLINE Tensor< T, SIZE_TPARAM > &
Tensor< T, SIZE_TPARAM >::operator+=( const Tensor< T, SIZE_TPARAM > & rhs )
{
  for( int i = 0; i < SIZE_TPARAM; ++i )
  {
    data[i] += rhs.data[i];
  }
  return *this;
}
 
}
 
#endif //GEOS_COMMON_TENSOR_HPP_