Writing of Eigen::Tensor of grid objects now works (for Hdf5)

Grid/serialisation/BaseIO.h

@@ -65,12 +65,20 @@ namespace Grid {
     template <typename ETensor>
     typename std::enable_if<std::is_base_of<Eigen::TensorBase<ETensor, Eigen::ReadOnlyAccessors>, ETensor>::value && (std::is_arithmetic<typename ETensor::Scalar>::value || Grid::is_complex<typename ETensor::Scalar>::value), void>::type
     write(const std::string &s, const ETensor &output);
-    template<typename U, int NumIndices_, int Options_, typename IndexType_>
+    template <typename ETensor/*, typename U, int N*/>
-    void write(const std::string &s, const Eigen::Tensor<iScalar<U>, NumIndices_, Options_, IndexType_> &output);
+    typename std::enable_if<std::is_base_of<Eigen::TensorBase<ETensor, Eigen::ReadOnlyAccessors>, ETensor>::value && !(std::is_arithmetic<typename ETensor::Scalar>::value || Grid::is_complex<typename ETensor::Scalar>::value)
-    template<typename U, int N, int NumIndices_, int Options_, typename IndexType_>
+                         /*&& ( std::is_base_of<typename ETensor::Scalar, iScalar<U>   >::value
-    void write(const std::string &s, const Eigen::Tensor<iVector<U, N>, NumIndices_, Options_, IndexType_> &output);
+                           || std::is_base_of<typename ETensor::Scalar, iVector<U, N>>::value
-    template<typename U, int N, int NumIndices_, int Options_, typename IndexType_>
+                           || std::is_base_of<typename ETensor::Scalar, iMatrix<U, N>>::value )*/, void>::type
-    void write(const std::string &s, const Eigen::Tensor<iMatrix<U, N>, NumIndices_, Options_, IndexType_> &output);
+    write(const std::string &s, const ETensor &output);
+    /*template <typename ETensor, typename U, int N>
+    typename std::enable_if<std::is_base_of<Eigen::TensorBase<ETensor, Eigen::ReadOnlyAccessors>, ETensor>::value
+                         && std::is_base_of<typename ETensor::Scalar, iVector<U, N>>::value, void>::type
+    write(const std::string &s, const ETensor &output);
+    template <typename ETensor, typename U, int N>
+    typename std::enable_if<std::is_base_of<Eigen::TensorBase<ETensor, Eigen::ReadOnlyAccessors>, ETensor>::value
+                         && std::is_base_of<typename ETensor::Scalar, iMatrix<U, N>>::value, void>::type
+    write(const std::string &s, const ETensor &output);*/
 
     void         scientificFormat(const bool set);
     bool         isScientific(void);
@@ -182,13 +190,12 @@ namespace Grid {
   typename std::enable_if<std::is_base_of<Eigen::TensorBase<ETensor, Eigen::ReadOnlyAccessors>, ETensor>::value && (std::is_arithmetic<typename ETensor::Scalar>::value || Grid::is_complex<typename ETensor::Scalar>::value), void>::type
   Writer<T>::write(const std::string &s, const ETensor &output)
   {
+    std::cout << "Eigen::Tensors of arithmetic/complex base type" << std::endl;
     const typename ETensor::Index NumElements{output.size()};
     assert( NumElements > 0 );
     if( NumElements == 1 )
       upcast->writeDefault(s, * output.data());
     else {
-      // Create a single, flat vector to hold all the data
-      std::vector<typename ETensor::Scalar> flat(NumElements);
       // We're not interested in trivial dimensions, i.e. dimensions = 1
       unsigned int TrivialDimCount{0};
       std::vector<size_t> ReducedDims;
@@ -203,8 +210,9 @@ namespace Grid {
           ReducedDims.push_back(s);
         }
       }
-      const unsigned int ReducedDimCount{output.NumDimensions - TrivialDimCount};
+      assert( output.NumDimensions > TrivialDimCount > 0 ); // NB: NumElements > 1 implies this is not a scalar, so some dims should be left
-      assert( ReducedDimCount > 0 ); // NB: NumElements > 1 implies this is not a scalar
+      // Create a single, flat vector to hold all the data
+      std::vector<typename ETensor::Scalar> flat(NumElements);
       // Now copy all the data to my flat vector
       // Regardless of the Eigen::Tensor storage order, the copy will be Row Major
       std::array<typename ETensor::Index, ETensor::NumIndices> MyIndex;
@@ -221,17 +229,71 @@ namespace Grid {
   
   // Eigen::Tensors of iScalar<U>
   template <typename T>
-  template<typename U, int NumIndices_, int Options_, typename IndexType_>
+  template <typename ETensor/*, typename U, int N*/>
-  void Writer<T>::write(const std::string &s, const Eigen::Tensor<iScalar<U>, NumIndices_, Options_, IndexType_> &output)
+  typename std::enable_if<std::is_base_of<Eigen::TensorBase<ETensor, Eigen::ReadOnlyAccessors>, ETensor>::value && !(std::is_arithmetic<typename ETensor::Scalar>::value || Grid::is_complex<typename ETensor::Scalar>::value)
+                     /*&& ( std::is_base_of<typename ETensor::Scalar, iScalar<U>   >::value
+                       || std::is_base_of<typename ETensor::Scalar, iVector<U, N>>::value
+                       || std::is_base_of<typename ETensor::Scalar, iMatrix<U, N>>::value )*/, void>::type
+  Writer<T>::write(const std::string &s, const ETensor &output)
   {
-    //upcast->writeDefault(s, tensorToVec(output));
+    std::cout << "Eigen::Tensors of iScalar<U>" << std::endl;
-    std::cout << "I really should add code to write Eigen::Tensor (iScalar) ..." << std::endl;
+    const typename ETensor::Index NumElements{output.size()};
+    assert( NumElements > 0 );
+    if( NumElements == 1 )
+      upcast->writeDefault(s, tensorToVec(* output.data()));
+    else {
+      // We're not interested in trivial dimensions, i.e. dimensions = 1
+      unsigned int TrivialDimCount{0};
+      std::vector<size_t> ReducedDims;
+      for(auto i = 0; i < output.NumDimensions; i++ ) {
+        auto dim = output.dimension(i);
+        if( dim <= 1 ) {
+          TrivialDimCount++;
+          assert( dim == 1 ); // Not expecting dimension to be <= 0
+        } else {
+          size_t s = static_cast<size_t>(dim);
+          assert( s == dim ); // check we didn't lose anything in the conversion
+          ReducedDims.push_back(s);
+        }
+      }
+      assert( output.NumDimensions > TrivialDimCount > 0 ); // NB: NumElements > 1 implies this is not a scalar, so some dims should be left
+      // Now add the extra dimensions, based on object zero
+      typename TensorToVec<typename ETensor::Scalar>::type ttv = tensorToVec(* output.data());
+      Flatten<typename TensorToVec<typename ETensor::Scalar>::type> f(ttv);
+      const std::vector<size_t> & ExtraDims{f.getDim()};
+      size_t ExtraCount{1};
+      for( auto i : ExtraDims ) {
+        assert( i > 0 );
+        ExtraCount *= i;
+        ReducedDims.push_back(i);
+      }
+      typedef typename ETensor::Scalar::scalar_type Scalar;
+      assert( sizeof( typename ETensor::Scalar ) == ExtraCount * sizeof( Scalar ) );
+      // Create a single, flat vector to hold all the data
+      const typename ETensor::Index TotalNumElements = NumElements * ExtraCount;
+      std::vector<Scalar> flat(TotalNumElements);
+      // Now copy all the data to my flat vector
+      // Regardless of the Eigen::Tensor storage order, the copy will be Row Major
+      std::array<typename ETensor::Index, ETensor::NumIndices> MyIndex;
+      for( int i = 0 ; i < output.NumDimensions ; i++ ) MyIndex[i] = 0;
+      for( typename ETensor::Index n = 0; n < TotalNumElements; ) {
+        const Scalar * p = reinterpret_cast<const Scalar *>( &output( MyIndex ));
+        for( auto j = 0; j < ExtraCount ; j++ )
+          flat[n++] =  * p++;
+        // Now increment the index
+        for( int i = output.NumDimensions - 1; i >= 0 && ++MyIndex[i] == output.dimension(i); i-- )
+          MyIndex[i] = 0;
+      }
+      upcast->template writeMultiDim<Scalar>(s, ReducedDims, flat);
+    }
   }
-  
+
   // Eigen::Tensors of iVector<U, N>
-  template <typename T>
+  /*template <typename T>
-  template<typename U, int N, int NumIndices_, int Options_, typename IndexType_>
+  template <typename ETensor, typename U, int N>
-  void Writer<T>::write(const std::string &s, const Eigen::Tensor<iVector<U, N>, NumIndices_, Options_, IndexType_> &output)
+  typename std::enable_if<std::is_base_of<Eigen::TensorBase<ETensor, Eigen::ReadOnlyAccessors>, ETensor>::value
+  && std::is_base_of<typename ETensor::Scalar, iVector<U, N>>::value, void>::type
+  Writer<T>::write(const std::string &s, const ETensor &output)
   {
     //upcast->writeDefault(s, tensorToVec(output));
     std::cout << "I really should add code to write Eigen::Tensor (iVector) ..." << std::endl;
@@ -239,12 +301,14 @@ namespace Grid {
   
   // Eigen::Tensors of iMatrix<U, N>
   template <typename T>
-  template<typename U, int N, int NumIndices_, int Options_, typename IndexType_>
+  template <typename ETensor, typename U, int N>
-  void Writer<T>::write(const std::string &s, const Eigen::Tensor<iMatrix<U, N>, NumIndices_, Options_, IndexType_> &output)
+  typename std::enable_if<std::is_base_of<Eigen::TensorBase<ETensor, Eigen::ReadOnlyAccessors>, ETensor>::value
+  && std::is_base_of<typename ETensor::Scalar, iMatrix<U, N>>::value, void>::type
+  Writer<T>::write(const std::string &s, const ETensor &output)
   {
     //upcast->writeDefault(s, tensorToVec(output));
     std::cout << "I really should add code to write Eigen::Tensor (iMatrix) ..." << std::endl;
-  }
+  }*/
 
   template <typename T>
   void Writer<T>::scientificFormat(const bool set)
@@ -398,6 +462,18 @@ namespace Grid {
       }
       return bResult;
     }
+
+    template <typename T>
+    static inline typename std::enable_if<!std::is_base_of<Eigen::TensorBase<T, Eigen::ReadOnlyAccessors>, T>::value, void>::type
+    WriteMember(std::ostream &os, const T &object) {
+      os << object;
+    }
+    
+    template <typename T>
+    static inline typename std::enable_if<std::is_base_of<Eigen::TensorBase<T, Eigen::ReadOnlyAccessors>, T>::value, void>::type
+    WriteMember(std::ostream &os, const T &object) {
+      os << "Eigen::Tensor";
+    }
   };
 
   // Generic writer interface //////////////////////////////////////////////////

Grid/serialisation/MacroMagic.h

@@ -110,7 +110,7 @@ THE SOFTWARE.
 
 #define GRID_MACRO_MEMBER(A,B)        A B;
 #define GRID_MACRO_COMP_MEMBER(A,B) result = (result and CompareMember(lhs. B, rhs. B));
-#define GRID_MACRO_OS_WRITE_MEMBER(A,B) os<< #A <<" " #B << " = " << obj. B << " ; " <<std::endl;
+#define GRID_MACRO_OS_WRITE_MEMBER(A,B) os<< #A <<" " #B << " = "; WriteMember( os, obj. B ); os << " ; " <<std::endl;
 #define GRID_MACRO_READ_MEMBER(A,B) Grid::read(RD,#B,obj. B);
 #define GRID_MACRO_WRITE_MEMBER(A,B) Grid::write(WR,#B,obj. B);
 

tests/hadrons/Test_hadrons_distil.cc

@@ -466,6 +466,8 @@ typedef std::complex<double> TestScalar;
 typedef Eigen::Tensor<TestScalar, 3> TestTensor;
 typedef Eigen::TensorFixedSize<TestScalar, Eigen::Sizes<9,4,2>> TestTensorFixed;
 typedef std::vector<TestTensorFixed> aTestTensorFixed;
+typedef Eigen::TensorFixedSize<SpinColourVector, Eigen::Sizes<11,3,2>> LSCTensor;
+typedef Eigen::TensorFixedSize<LorentzColourMatrix, Eigen::Sizes<5,7,2>> LCMTensor;
 // From Test_serialisation.cc
 class myclass: Serializable {
 public:
@@ -475,12 +477,15 @@ public:
                                   , TestTensor, Critter
                                   , TestTensorFixed, FixedCritter
                                   , aTestTensorFixed, aFixedCritter
+                                  , LSCTensor, MyLSCTensor
+                                  , LCMTensor, MyLCMTensor
                                   );
   myclass() : Critter(7,3,2), aFixedCritter(3) {}
 };
 
 bool DebugIOTest(void) {
-  SpinColourVector scv;
+  SpinColourVector scv, scv2;
+  scv2 = scv;
   ioTest<Hdf5Writer, Hdf5Reader, SpinColourVector>("iotest_vector.h5", scv, "SpinColourVector");
   SpinColourMatrix scm;
   ioTest<Hdf5Writer, Hdf5Reader, SpinColourMatrix>("iotest_matrix.h5", scm, "SpinColourMatrix");
@@ -512,6 +517,35 @@ bool DebugIOTest(void) {
   myclass o;
   ioTest<Hdf5Writer, Hdf5Reader, myclass>("iotest_object.h5", o, "myclass_object_instance_name");
 
+  // Tensor of spin colour
+  LSCTensor l;
+  Val = 0;
+  for( int i = 0 ; i < l.dimension(0) ; i++)
+    for( int j = 0 ; j < l.dimension(1) ; j++)
+      for( int k = 0 ; k < l.dimension(2) ; k++)
+        for( int s = 0 ; s < Ns ; s++ )
+          for( int c = 0 ; c < Nc ; c++ )
+          {
+            l(i,j,k)()(s)(c) = Val;
+            Val += Inc;
+          }
+  ioTest<Hdf5Writer, Hdf5Reader, LSCTensor>("iotest_LSCTensor.h5", l, "LSCTensor_object_instance_name");
+
+  // Tensor of spin colour
+  LCMTensor l2;
+  Val = 0;
+  for( int i = 0 ; i < l2.dimension(0) ; i++)
+    for( int j = 0 ; j < l2.dimension(1) ; j++)
+      for( int k = 0 ; k < l2.dimension(2) ; k++)
+        for( int l = 0 ; l < Ns ; l++ )
+          for( int c = 0 ; c < Nc ; c++ )
+            for( int c2 = 0 ; c2 < Nc ; c2++ )
+            {
+              l2(i,j,k)(l)()(c,c2) = Val;
+              Val += Inc;
+            }
+  ioTest<Hdf5Writer, Hdf5Reader, LCMTensor>("iotest_LCMTensor.h5", l2, "LCMTensor_object_instance_name");
+  
   std::cout << "Wow!" << std::endl;
   
   return true;