Grid/lib/algorithms/CoarsenedMatrix.h

#ifndef  GRID_ALGORITHM_COARSENED_MATRIX_H
#define  GRID_ALGORITHM_COARSENED_MATRIX_H

#include <Grid.h>

namespace Grid {

  class Geometry {
    //    int dimension;
  public:
    int npoint;
    std::vector<int> directions   ;
    std::vector<int> displacements;

    // FIXME -- don't like xposing the operator directions
    // as different to the geometrical dirs
    // Also don't like special casing five dim.. should pass an object in template
  Geometry(int _d)  {
  
      int base = (_d==5) ? 1:0;

      // make coarse grid stencil for 4d , not 5d
      if ( _d==5 ) _d=4;

      npoint = 2*_d+1;
      directions.resize(npoint);
      displacements.resize(npoint);
      for(int d=0;d<_d;d++){
	directions[2*d  ] = d+base;
	directions[2*d+1] = d+base;
	displacements[2*d  ] = +1;
	displacements[2*d+1] = -1;
      }
      directions   [2*_d]=0;
      displacements[2*_d]=0;
      
      //// report back
      std::cout<<"directions    :";
      for(int d=0;d<npoint;d++) std::cout<< directions[d]<< " ";
      std::cout <<std::endl;
      std::cout<<"displacements :";
      for(int d=0;d<npoint;d++) std::cout<< displacements[d]<< " ";
      std::cout <<std::endl;
    }
  
    /*
      // Original cleaner code
    Geometry(int _d) : dimension(_d), npoint(2*_d+1), directions(npoint), displacements(npoint) {
      for(int d=0;d<dimension;d++){
	directions[2*d  ] = d;
	directions[2*d+1] = d;
	displacements[2*d  ] = +1;
	displacements[2*d+1] = -1;
      }
      directions   [2*dimension]=0;
      displacements[2*dimension]=0;
    }
    std::vector<int> GetDelta(int point) {
      std::vector<int> delta(dimension,0);
      delta[directions[point]] = displacements[point];
      return delta;
    };
    */    

  };
  
  // Fine Object == (per site) type of fine field
  // nbasis      == number of deflation vectors
  template<class Fobj,class CComplex,int nbasis>
  class CoarsenedMatrix : public SparseMatrixBase<Lattice<iVector<CComplex,nbasis > > >  {
  public:
    
    typedef iVector<CComplex,nbasis >             siteVector;
    typedef Lattice<siteVector>                 CoarseVector;
    typedef Lattice<iMatrix<CComplex,nbasis > > CoarseMatrix;

    typedef Lattice< CComplex >   CoarseScalar; // used for inner products on fine field
    typedef Lattice<Fobj >        FineField;

    ////////////////////
    // Data members
    ////////////////////
    Geometry         geom;
    GridBase *       _grid; 
    CartesianStencil Stencil; 

    std::vector<CoarseMatrix> A;

    std::vector<siteVector,alignedAllocator<siteVector> >   comm_buf;
      
    ///////////////////////
    // Interface
    ///////////////////////
    GridBase * Grid(void)         { return _grid; };   // this is all the linalg routines need to know

    RealD M (const CoarseVector &in, CoarseVector &out){

      conformable(_grid,in._grid);
      conformable(in._grid,out._grid);

      SimpleCompressor<siteVector> compressor;
      Stencil.HaloExchange(in,comm_buf,compressor);

      //PARALLEL_FOR_LOOP
      for(int ss=0;ss<Grid()->oSites();ss++){
        siteVector res = zero;
	siteVector nbr;
	int offset,local,perm,ptype;

	for(int point=0;point<geom.npoint;point++){
	  offset = Stencil._offsets [point][ss];
	  local  = Stencil._is_local[point][ss];
	  perm   = Stencil._permute [point][ss];
	  ptype  = Stencil._permute_type[point];
	  
	  if(local&&perm) { 
	    permute(nbr,in._odata[offset],ptype);
	  } else if(local) { 
	    nbr = in._odata[offset];
	  } else {
	    nbr = comm_buf[offset];
	  }
	  res = res + A[point]._odata[ss]*nbr;
	}
	vstream(out._odata[ss],res);
      }
      return norm2(out);
    };

    RealD Mdag (const CoarseVector &in, CoarseVector &out){ 
      return M(in,out);
    };

    // Defer support for further coarsening for now
    void Mdiag    (const CoarseVector &in,  CoarseVector &out){};
    void Mdir     (const CoarseVector &in,  CoarseVector &out,int dir, int disp){};

    CoarsenedMatrix(GridCartesian &CoarseGrid) 	: 

      _grid(&CoarseGrid),
      geom(CoarseGrid._ndimension),
      Stencil(&CoarseGrid,geom.npoint,Even,geom.directions,geom.displacements),
      A(geom.npoint,&CoarseGrid)
    {
      comm_buf.resize(Stencil._unified_buffer_size);
    };

    void CoarsenOperator(GridBase *FineGrid,LinearOperatorBase<Lattice<Fobj> > &linop,std::vector<Lattice<Fobj> > & subspace){

      FineField iblock(FineGrid); // contributions from within this block
      FineField oblock(FineGrid); // contributions from outwith this block

      FineField     phi(FineGrid);
      FineField     tmp(FineGrid);
      FineField     zz(FineGrid); zz=zero;
      FineField    Mphi(FineGrid);

      Lattice<iScalar<vInteger> > coor(FineGrid);

      CoarseVector iProj(Grid()); 
      CoarseVector oProj(Grid()); 
      CoarseScalar InnerProd(Grid()); 

      // Orthogonalise the subblocks over the basis
      blockOrthogonalise(InnerProd,subspace);
      blockProject(iProj,subspace[0],subspace);

      // Compute the matrix elements of linop between this orthonormal
      // set of vectors.
      int self_stencil=-1;
      for(int p=0;p<geom.npoint;p++){ 
	A[p]=zero;
	if( geom.displacements[p]==0){
	  self_stencil=p;
	}
      }
      assert(self_stencil!=-1);

      for(int i=0;i<nbasis;i++){
	phi=subspace[i];
	for(int p=0;p<geom.npoint;p++){ 

	  int dir   = geom.directions[p];
	  int disp  = geom.displacements[p];

	  Integer block=(FineGrid->_rdimensions[dir])/(Grid()->_rdimensions[dir]);

	  LatticeCoordinate(coor,dir);

	  if ( disp==0 ){
	    linop.OpDiag(phi,Mphi);
	  }
	  else  {
	    linop.OpDir(phi,Mphi,dir,disp); 
	  }

	  ////////////////////////////////////////////////////////////////////////
	  // Pick out contributions coming from this cell and neighbour cell
	  ////////////////////////////////////////////////////////////////////////
	  if ( disp==0 ) {
	    iblock = Mphi;
	    oblock = zero;
	  } else if ( disp==1 ) {
	    oblock = where(mod(coor,block)==(block-1),Mphi,zz);
	    iblock = where(mod(coor,block)!=(block-1),Mphi,zz);
	  } else if ( disp==-1 ) {
	    oblock = where(mod(coor,block)==(Integer)0,Mphi,zz);
	    iblock = where(mod(coor,block)!=(Integer)0,Mphi,zz);
	  } else {
	    assert(0);
	  }

	  blockProject(iProj,iblock,subspace);
	  blockProject(oProj,oblock,subspace);
	  for(int ss=0;ss<Grid()->oSites();ss++){
	    for(int j=0;j<nbasis;j++){
	      if( disp!= 0 ) {
		A[p]._odata[ss](j,i) = oProj._odata[ss](j);
	      }
	      A[self_stencil]._odata[ss](j,i) =	A[self_stencil]._odata[ss](j,i) + iProj._odata[ss](j);
	    }
	  }
	}
      }

#if 0
      ///////////////////////////
      // test code worth preserving in if block
      ///////////////////////////
      std::cout<< " Computed matrix elements "<< self_stencil <<std::endl;
      for(int p=0;p<geom.npoint;p++){
	std::cout<< "A["<<p<<"]" << std::endl;
	std::cout<< A[p] << std::endl;
      }
      std::cout<< " picking by block0 "<< self_stencil <<std::endl;

      phi=subspace[0];
      std::vector<int> bc(FineGrid->_ndimension,0);

      blockPick(Grid(),phi,tmp,bc);      // Pick out a block
      linop.Op(tmp,Mphi);                // Apply big dop
      blockProject(iProj,Mphi,subspace); // project it and print it
      std::cout<< " Computed matrix elements from block zero only "<<std::endl;
      std::cout<< iProj <<std::endl;
      std::cout<<"Computed Coarse Operator"<<std::endl;
#endif
      //      AssertHermitian();
      //      ForceHermitian();
      //      ForceDiagonal();
    }
    void ForceDiagonal(void) {


      std::cout<<"**************************************************"<<std::endl;
      std::cout<<"****   Forcing coarse operator to be diagonal ****"<<std::endl;
      std::cout<<"**************************************************"<<std::endl;
      for(int p=0;p<8;p++){
	A[p]=zero;
      }

      GridParallelRNG  RNG(Grid()); RNG.SeedRandomDevice();
      Lattice<iScalar<CComplex> > val(Grid()); random(RNG,val);

      Complex one(1.0);

      iMatrix<Complex,nbasis> ident;  ident=one;

      val = val*adj(val);
      val = val + 1.0;

      A[8] = val*ident;

      //      for(int s=0;s<Grid()->oSites();s++) {
      //	A[8]._odata[s]=val._odata[s];
      //      }
    }
    void ForceHermitian(void) {
      for(int d=0;d<4;d++){
	int dd=d+1;
	A[2*d] = adj(Cshift(A[2*d+1],dd,1));
      }
      A[8] = 0.5*(A[8] + adj(A[8]));
    }
    void AssertHermitian(void) {
      CoarseMatrix AA    (Grid());
      CoarseMatrix AAc   (Grid());
      CoarseMatrix Diff  (Grid());
      for(int d=0;d<4;d++){
	
	int dd=d+1;
	AAc = Cshift(A[2*d+1],dd,1);
	AA  = A[2*d];
	
	Diff = AA - adj(AAc);

	std::cout<<"Norm diff dim "<<d<<" "<< norm2(Diff)<<std::endl;
	std::cout<<"Norm dim "<<d<<" "<< norm2(AA)<<std::endl;
	  
      }
      Diff = A[8] - adj(A[8]);
      std::cout<<"Norm diff local "<< norm2(Diff)<<std::endl;
      std::cout<<"Norm local "<< norm2(A[8])<<std::endl;
    }
    
  };

}
#endif