Shaken out the peekIndex support.

Hardwire constants "SpinIndex, ColourIndex" and LorentzIndex in Grid_QCD.h

lib/Grid_Lattice.h

@@ -505,6 +505,34 @@ public:
         }
         return ret;
     };
+    ////////////////////////////////////////////////////////////////////////////////////////////////////
+    // Poke internal indices of a Lattice object
+    ////////////////////////////////////////////////////////////////////////////////////////////////////
+    template<int Index,class vobj> inline
+    void pokeIndex(Lattice<vobj> &lhs,const Lattice<decltype(peekIndex<Index>(lhs._odata[0]))> & rhs)
+    {
+#pragma omp parallel for
+        for(int ss=0;ss<lhs._grid->oSites();ss++){
+	  pokeIndex<Index>(lhs._odata[ss],rhs._odata[ss]);
+	}      
+    }
+    template<int Index,class vobj> inline
+    void pokeIndex(Lattice<vobj> &lhs,const Lattice<decltype(peekIndex<Index>(lhs._odata[0],0))> & rhs,int i)
+    {
+#pragma omp parallel for
+        for(int ss=0;ss<lhs._grid->oSites();ss++){
+	  pokeIndex<Index>(lhs._odata[ss],rhs._odata[ss],i);
+	}      
+    }
+    template<int Index,class vobj> inline
+    void pokeIndex(Lattice<vobj> &lhs,const Lattice<decltype(peekIndex<Index>(lhs._odata[0],0,0))> & rhs,int i,int j)
+    {
+#pragma omp parallel for
+        for(int ss=0;ss<lhs._grid->oSites();ss++){
+	  pokeIndex<Index>(lhs._odata[ss],rhs._odata[ss],i,j);
+	}      
+    }
+
     ////////////////////////////////////////////////////////////////////////////////////////////////////
     // Reduction operations
     ////////////////////////////////////////////////////////////////////////////////////////////////////

lib/Grid_QCD.h

@@ -1,6 +1,7 @@
 #ifndef GRID_QCD_H
 #define GRID_QCD_H
 namespace Grid{
+
 namespace QCD {
 
     static const int Nc=3;
@@ -13,12 +14,18 @@ namespace QCD {
     //////////////////////////////////////////////////////////////////////////////
     // QCD iMatrix types
     // Index conventions:                            Lorentz x Spin x Colour
-    //
+    //////////////////////////////////////////////////////////////////////////////
+    static const int ColourIndex = 1;
+    static const int SpinIndex   = 2;
+    static const int LorentzIndex= 3;
+
     // ChrisK very keen to add extra space for Gparity doubling.
     //
     // Also add domain wall index, in a way where Wilson operator 
     // naturally distributes across the 5th dimensions.
-    //////////////////////////////////////////////////////////////////////////////
+    //
+    // That probably makes for GridRedBlack4dCartesian grid.
+
     template<typename vtype> using iSinglet          = iScalar<iScalar<iScalar<vtype> > >;
     template<typename vtype> using iSpinMatrix       = iScalar<iMatrix<iScalar<vtype>, Ns> >;
     template<typename vtype> using iSpinColourMatrix = iScalar<iMatrix<iMatrix<vtype, Nc>, Ns> >;

lib/Grid_math_types.h

@@ -1485,6 +1485,125 @@ template<int Level,class vtype,int N> inline
 }
 
 
+//////////////////////////////////////////////////////////////////////////////
+// Poke a specific index; 
+//////////////////////////////////////////////////////////////////////////////
+
+// Scalar poke
+template<int Level,class vtype> inline 
+  void pokeIndex(iScalar<vtype> &ret, 
+		 const typename std::enable_if<matchGridTensorIndex<iScalar<vtype>,Level>::value,iScalar<vtype> >::type &arg)
+{
+  ret._internal = arg._internal;
+}
+// Vector poke, one index
+template<int Level,class vtype,int N> inline 
+  void pokeIndex(iVector<vtype,N> &ret, 
+		 const typename std::enable_if<matchGridTensorIndex<iVector<vtype,N>,Level>::value,iScalar<vtype> >::type &arg,int i)
+{
+  ret._internal[i] = arg._internal;
+}
+// Vector poke, two indices
+template<int Level,class vtype,int N> inline 
+  void pokeIndex(iMatrix<vtype,N> &ret, 
+		 const typename std::enable_if<matchGridTensorIndex<iMatrix<vtype,N>,Level>::value,iScalar<vtype> >::type &arg,int i,int j)
+{
+  ret._internal[i][j] = arg._internal;
+}
+
+/////////////
+// No match poke for scalar,vector,matrix must forward on either 0,1,2 args. Must have 9 routines with notvalue
+/////////////
+// scalar
+template<int Level,class vtype> inline 
+  void pokeIndex(iScalar<vtype> &ret,
+		 const typename std::enable_if<matchGridTensorIndex<iScalar<vtype>,Level>::notvalue,iScalar<decltype(peekIndex<Level>(ret._internal))> >::type &arg)
+		 
+{
+  pokeIndex<Level>(ret._internal,arg._internal);
+}
+template<int Level,class vtype> inline 
+  void pokeIndex(iScalar<vtype> &ret,
+		 const typename std::enable_if<matchGridTensorIndex<iScalar<vtype>,Level>::notvalue,iScalar<decltype(peekIndex<Level>(ret._internal,0))> >::type &arg,
+		 int i)
+		 
+{
+  pokeIndex<Level>(ret._internal,arg._internal,i);
+}
+template<int Level,class vtype> inline 
+  void pokeIndex(iScalar<vtype> &ret,
+		 const typename std::enable_if<matchGridTensorIndex<iScalar<vtype>,Level>::notvalue,iScalar<decltype(peekIndex<Level>(ret._internal,0,0))> >::type &arg,
+		 int i,int j)
+		 
+{
+  pokeIndex<Level>(ret._internal,arg._internal,i,j);
+}
+
+// Vector
+template<int Level,class vtype,int N> inline 
+  void pokeIndex(iVector<vtype,N> &ret,
+		 const typename std::enable_if<matchGridTensorIndex<iVector<vtype,N>,Level>::notvalue,iVector<decltype(peekIndex<Level>(ret._internal)),N> >::type &arg)
+		 
+{
+  for(int ii=0;ii<N;ii++){
+    pokeIndex<Level>(ret._internal[ii],arg._internal[ii]);
+  }
+}
+template<int Level,class vtype,int N> inline 
+  void pokeIndex(iVector<vtype,N> &ret,
+		 const typename std::enable_if<matchGridTensorIndex<iVector<vtype,N>,Level>::notvalue,iVector<decltype(peekIndex<Level>(ret._internal,0)),N> >::type &arg,
+		 int i)
+		 
+{
+  for(int ii=0;ii<N;ii++){
+    pokeIndex<Level>(ret._internal[ii],arg._internal[ii],i);
+  }
+}
+template<int Level,class vtype,int N> inline 
+  void pokeIndex(iVector<vtype,N> &ret,
+		 const typename std::enable_if<matchGridTensorIndex<iVector<vtype,N>,Level>::notvalue,iVector<decltype(peekIndex<Level>(ret._internal,0,0)),N> >::type &arg,
+		 int i,int j)
+		 
+{
+  for(int ii=0;ii<N;ii++){
+    pokeIndex<Level>(ret._internal[ii],arg._internal[ii],i,j);
+  }
+}
+
+// Matrix
+template<int Level,class vtype,int N> inline 
+  void pokeIndex(iMatrix<vtype,N> &ret,
+		 const typename std::enable_if<matchGridTensorIndex<iMatrix<vtype,N>,Level>::notvalue,iMatrix<decltype(peekIndex<Level>(ret._internal)),N> >::type &arg)
+		 
+{
+  for(int ii=0;ii<N;ii++){
+  for(int jj=0;jj<N;jj++){
+    pokeIndex<Level>(ret._internal[ii][jj],arg._internal[ii][jj]);
+  }}
+}
+template<int Level,class vtype,int N> inline 
+  void pokeIndex(iMatrix<vtype,N> &ret,
+		 const typename std::enable_if<matchGridTensorIndex<iMatrix<vtype,N>,Level>::notvalue,iMatrix<decltype(peekIndex<Level>(ret._internal,0)),N> >::type &arg,
+		 int i)
+		 
+{
+  for(int ii=0;ii<N;ii++){
+  for(int jj=0;jj<N;jj++){
+    pokeIndex<Level>(ret._internal[ii][jj],arg._internal[ii][jj],i);
+  }}
+}
+template<int Level,class vtype,int N> inline 
+  void pokeIndex(iMatrix<vtype,N> &ret,
+		 const typename std::enable_if<matchGridTensorIndex<iMatrix<vtype,N>,Level>::notvalue,iMatrix<decltype(peekIndex<Level>(ret._internal,0,0)),N> >::type &arg,
+		 int i,int j)
+		 
+{
+  for(int ii=0;ii<N;ii++){
+  for(int jj=0;jj<N;jj++){
+    pokeIndex<Level>(ret._internal[ii][jj],arg._internal[ii][jj],i,j);
+  }}
+}
+
 /////////////////////////////////////////////////////////////////
 // Can only take the real/imag part of scalar objects, since
 // lattice objects of different complex nature are non-conformable.

tests/Grid_main.cc

@@ -131,11 +131,14 @@ int main (int argc, char ** argv)
     //    rscalar=real(scalar);
     //    iscalar=imag(scalar);
     //    scalar =cmplx(rscalar,iscalar);
+    pokeIndex<1>(cVec,scalar,1);
+
 
     scalar=transpose(scalar);
     scalar=transposeIndex<1>(scalar);
     scalar=traceIndex<1>(scalar);
     scalar=peekIndex<1>(cVec,0);
+
     scalar=trace(scalar);
     scalar=localInnerProduct(cVec,cVec);
     scalar=localNorm2(cVec);
@@ -205,7 +208,7 @@ int main (int argc, char ** argv)
     LatticeComplex trscMat(&Fine);
     trscMat = trace(scMat); // Trace
 
-    {
+    { // Peek-ology and Poke-ology, with a little app-ology
       TComplex      c;
       ColourMatrix c_m;   
       SpinMatrix   s_m;   
@@ -224,10 +227,15 @@ int main (int argc, char ** argv)
       printf("c. Level %d\n",c_m().TensorLevel);
       printf("c. Level %d\n",c_m()().TensorLevel);
       
-      c_m()() = scm()(0,0); //ColourComponents of CM <= ColourComponents of SpinColourMatrix
+      c_m()()    = scm()(0,0); //ColourComponents of CM <= ColourComponents of SpinColourMatrix
       scm()(1,1) = cm()();  //ColourComponents of CM <= ColourComponents of SpinColourMatrix
+      c          = scm()(1,1)(1,2);
+      scm()(1,1)(2,1) = c;
+
+      pokeIndex<1> (c_m,c,0,0);
     }
 
+    
     FooBar = Bar;
  
     /*