From d5ab571a894682e36bc24488958f01a18983c7f7 Mon Sep 17 00:00:00 2001
From: Christopher Kelly <ckelly@bnl.gov>
Date: Tue, 23 Feb 2021 11:49:56 -0500
Subject: [PATCH] Added the ability to apply a custom "filter" to the conjugate
 momentum in the Integrator classes, applied both after refresh and after
 applying the forces Added a conjugate momentum "filter" that applies a phase
 to each site. With sites set to 1.0 or 0.0 this acts as a mask and enables,
 for example, the freezing of inactive gauge links in DDHMC Added
 tests/forces/Test_momentum_filter demonstrating the use of the filter to
 freeze boundary links

---
 Grid/qcd/hmc/integrators/Integrator.h     |  28 ++++
 Grid/qcd/hmc/integrators/MomentumFilter.h |  94 +++++++++++++
 tests/forces/Test_momentum_filter.cc      | 154 ++++++++++++++++++++++
 3 files changed, 276 insertions(+)
 create mode 100644 Grid/qcd/hmc/integrators/MomentumFilter.h
 create mode 100644 tests/forces/Test_momentum_filter.cc

diff --git a/Grid/qcd/hmc/integrators/Integrator.h b/Grid/qcd/hmc/integrators/Integrator.h
index 70055754..77b7de52 100644
--- a/Grid/qcd/hmc/integrators/Integrator.h
+++ b/Grid/qcd/hmc/integrators/Integrator.h
@@ -33,6 +33,7 @@ directory
 #define INTEGRATOR_INCLUDED
 
 #include <memory>
+#include "MomentumFilter.h"
 
 NAMESPACE_BEGIN(Grid);
 
@@ -78,8 +79,19 @@ protected:
   RepresentationPolicy Representations;
   IntegratorParameters Params;
 
+  //Filters allow the user to manipulate the conjugate momentum, for example to freeze links in DDHMC
+  //It is applied whenever the momentum is updated / refreshed
+  //The default filter does nothing
+  MomentumFilterBase<MomentaField> const* MomFilter;
+
   const ActionSet<Field, RepresentationPolicy> as;
 
+  //Get a pointer to a shared static instance of the "do-nothing" momentum filter to serve as a default
+  static MomentumFilterBase<MomentaField> const* getDefaultMomFilter(){ 
+    static MomentumFilterNone<MomentaField> filter;
+    return &filter;
+  }
+
   void update_P(Field& U, int level, double ep) 
   {
     t_P[level] += ep;
@@ -135,6 +147,8 @@ protected:
 
     // Force from the other representations
     as[level].apply(update_P_hireps, Representations, Mom, U, ep);
+
+    MomFilter->applyFilter(Mom);
   }
 
   void update_U(Field& U, double ep) 
@@ -174,11 +188,23 @@ public:
     t_P.resize(levels, 0.0);
     t_U = 0.0;
     // initialization of smearer delegated outside of Integrator
+
+    //Default the momentum filter to "do-nothing"
+    MomFilter = getDefaultMomFilter();
   };
 
   virtual ~Integrator() {}
 
   virtual std::string integrator_name() = 0;
+  
+  //Set the momentum filter allowing for manipulation of the conjugate momentum
+  void setMomentumFilter(const MomentumFilterBase<MomentaField> &filter){
+    MomFilter = &filter;
+  }
+
+  //Access the conjugate momentum
+  const MomentaField & getMomentum() const{ return P; }
+  
 
   void print_parameters()
   {
@@ -249,6 +275,8 @@ public:
       // Refresh the higher representation actions
       as[level].apply(refresh_hireps, Representations, pRNG);
     }
+
+    MomFilter->applyFilter(P);
   }
 
   // to be used by the actionlevel class to iterate
diff --git a/Grid/qcd/hmc/integrators/MomentumFilter.h b/Grid/qcd/hmc/integrators/MomentumFilter.h
new file mode 100644
index 00000000..2a15d80c
--- /dev/null
+++ b/Grid/qcd/hmc/integrators/MomentumFilter.h
@@ -0,0 +1,94 @@
+/*************************************************************************************
+
+Grid physics library, www.github.com/paboyle/Grid
+
+Source file: ./lib/qcd/hmc/integrators/MomentumFilter.h
+
+Copyright (C) 2015
+
+Author: Christopher Kelly <ckelly@bnl.gov>
+Author: Peter Boyle <paboyle@ph.ed.ac.uk>
+
+This program is free software; you can redistribute it and/or modify
+it under the terms of the GNU General Public License as published by
+the Free Software Foundation; either version 2 of the License, or
+(at your option) any later version.
+
+This program is distributed in the hope that it will be useful,
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+GNU General Public License for more details.
+
+You should have received a copy of the GNU General Public License along
+with this program; if not, write to the Free Software Foundation, Inc.,
+51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
+
+See the full license in the file "LICENSE" in the top level distribution
+directory
+*************************************************************************************/
+/*  END LEGAL */
+//--------------------------------------------------------------------
+#ifndef MOMENTUM_FILTER
+#define MOMENTUM_FILTER
+
+NAMESPACE_BEGIN(Grid);
+
+//These filter objects allow the user to manipulate the conjugate momentum as part of the update / refresh
+
+template<typename MomentaField>
+struct MomentumFilterBase{
+  virtual void applyFilter(MomentaField &P) const;
+};
+
+//Do nothing
+template<typename MomentaField>
+struct MomentumFilterNone: public MomentumFilterBase<MomentaField>{
+  void applyFilter(MomentaField &P) const override{}
+};
+
+//Multiply each site/direction by a Lorentz vector complex number field
+//Can be used to implement a mask, zeroing out sites
+template<typename MomentaField>
+struct MomentumFilterApplyPhase: public MomentumFilterBase<MomentaField>{
+  typedef typename MomentaField::vector_type vector_type; //SIMD-vectorized complex type
+  typedef typename MomentaField::scalar_type scalar_type; //scalar complex type
+  typedef iVector<iScalar<iScalar<vector_type> >, Nd > LorentzScalarType; //complex phase for each site/direction
+  typedef Lattice<LorentzScalarType> LatticeLorentzScalarType;
+  
+  LatticeLorentzScalarType phase;
+ 
+  MomentumFilterApplyPhase(const LatticeLorentzScalarType _phase): phase(_phase){}
+
+  //Default to uniform field of (1,0)
+  MomentumFilterApplyPhase(GridBase* _grid): phase(_grid){
+    LorentzScalarType one;
+    for(int mu=0;mu<Nd;mu++)
+      one(mu)()() = scalar_type(1.);
+    
+    phase = one;
+  }
+
+  void applyFilter(MomentaField &P) const override{
+    conformable(P,phase);
+    autoView( P_v , P, AcceleratorWrite);
+    autoView( phase_v , phase, AcceleratorRead);
+
+    accelerator_for(ss,P_v.size(),MomentaField::vector_type::Nsimd(),{
+    	auto site_mom = P_v(ss);
+    	auto site_phase = phase_v(ss);
+	for(int mu=0;mu<Nd;mu++)
+	  site_mom(mu) = site_mom(mu) * site_phase(mu);
+    	coalescedWrite(P_v[ss], site_mom);
+      });
+    
+  }
+
+
+};
+
+
+
+
+NAMESPACE_END(Grid);
+
+#endif
diff --git a/tests/forces/Test_momentum_filter.cc b/tests/forces/Test_momentum_filter.cc
new file mode 100644
index 00000000..856ea0f2
--- /dev/null
+++ b/tests/forces/Test_momentum_filter.cc
@@ -0,0 +1,154 @@
+    /*************************************************************************************
+
+    Grid physics library, www.github.com/paboyle/Grid 
+
+    Source file: ./tests/Test_wilson_force.cc
+
+    Copyright (C) 2015
+
+Author: Christopher Kelly <ckelly@bnl.gov>
+
+    This program is free software; you can redistribute it and/or modify
+    it under the terms of the GNU General Public License as published by
+    the Free Software Foundation; either version 2 of the License, or
+    (at your option) any later version.
+
+    This program is distributed in the hope that it will be useful,
+    but WITHOUT ANY WARRANTY; without even the implied warranty of
+    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+    GNU General Public License for more details.
+
+    You should have received a copy of the GNU General Public License along
+    with this program; if not, write to the Free Software Foundation, Inc.,
+    51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
+
+    See the full license in the file "LICENSE" in the top level distribution directory
+    *************************************************************************************/
+    /*  END LEGAL */
+#include <Grid/Grid.h>
+
+using namespace std;
+using namespace Grid;
+
+//Get the mu-directected links on the upper boundary and the bulk remainder
+template<typename Field>
+void getLinksBoundaryBulk(Field &bound, Field &bulk,  Field &from, const Coordinate &latt_size){
+  bound = Zero(); bulk = Zero();
+  for(int mu=0;mu<Nd;mu++){
+    LatticeInteger mucoor(bound.Grid());
+    LatticeCoordinate(mucoor, mu);
+
+    bound = where( mucoor == (Integer)(latt_size[mu] - 1), from, bound );
+    bulk = where( mucoor != (Integer)(latt_size[mu] - 1), from, bulk );
+  }
+}
+
+
+int main (int argc, char ** argv)
+{
+  Grid_init(&argc,&argv);
+
+  Coordinate latt_size   = GridDefaultLatt();
+  Coordinate simd_layout = GridDefaultSimd(Nd,vComplex::Nsimd());
+  Coordinate mpi_layout  = GridDefaultMpi();
+
+  GridCartesian               Grid(latt_size,simd_layout,mpi_layout);
+  GridRedBlackCartesian     RBGrid(&Grid);
+
+  int threads = GridThread::GetThreads();
+  std::cout<<GridLogMessage << "Grid is setup to use "<<threads<<" threads"<<std::endl;
+
+  std::vector<int> seeds({1,2,3,4});
+
+  GridParallelRNG          pRNG(&Grid);
+  pRNG.SeedFixedIntegers(seeds);
+
+  typedef PeriodicGimplR Gimpl;
+  typedef WilsonGaugeAction<Gimpl> GaugeAction;
+  typedef NoHirep Representation; //fundamental
+  typedef NoSmearing<Gimpl> Smearing;
+  typedef MinimumNorm2<Gimpl, Smearing> Omelyan;
+  typedef Gimpl::Field Field;
+  typedef MomentumFilterApplyPhase<Field> Filter;
+  Filter filter(&Grid);
+  
+  //Setup a filter that disables link update on links passing through the global lattice boundary
+  typedef Filter::LatticeLorentzScalarType MaskType;
+  typedef Filter::LorentzScalarType MaskSiteType;
+
+  MaskSiteType zero, one;
+  for(int mu=0;mu<Nd;mu++){
+    zero(mu)()() = 0.;
+    one(mu)()() = 1.;
+  }
+  MaskType zeroField(&Grid), oneField(&Grid);
+  zeroField = zero;
+  oneField = one;
+
+  
+  filter.phase = oneField; //make every site 1.0
+
+  //Zero mu-directed links at upper boundary
+  for(int mu=0;mu<Nd;mu++){
+    LatticeInteger mucoor(&Grid);
+    LatticeCoordinate(mucoor, mu);
+    
+    filter.phase = where( mucoor == (Integer)(latt_size[mu] - 1) , zeroField, filter.phase );
+  }
+
+  //Start with a random gauge field
+  Field U(&Grid);
+  SU<Nc>::HotConfiguration(pRNG,U);
+
+  //Get the original links on the bulk and boundary for later use
+  Field Ubnd_orig(&Grid), Ubulk_orig(&Grid);
+  getLinksBoundaryBulk(Ubnd_orig, Ubulk_orig, U, latt_size);
+
+  ActionSet<Field,Representation> actions(1);  
+  double beta=6;
+  GaugeAction gauge_action(beta);
+  actions[0].push_back(&gauge_action);
+
+  Smearing smear;
+  IntegratorParameters params(1,1.); //1 MD step
+  Omelyan integrator(&Grid, params, actions, smear);
+  
+  integrator.setMomentumFilter(filter);
+
+  integrator.refresh(U, pRNG); //doesn't actually change the gauge field
+
+  //Check the momentum is zero on the boundary
+  const auto &P = integrator.getMomentum();
+  Field Pbnd(&Grid), Pbulk(&Grid);
+  getLinksBoundaryBulk(Pbnd, Pbulk, const_cast<Field&>(P), latt_size);
+
+  RealD Pbnd_nrm = norm2(Pbnd); //expect zero
+  std::cout << GridLogMessage << "After refresh, norm2 of mu-directed conjugate momentum on boundary is: " << Pbnd_nrm << " (expect 0)" << std::endl;
+  RealD Pbulk_nrm = norm2(Pbulk); //expect non-zero
+  std::cout << GridLogMessage << "After refresh, norm2 of bulk conjugate momentum is: " << Pbulk_nrm << " (expect non-zero)" << std::endl;
+
+  //Evolve the gauge field
+  integrator.integrate(U);
+
+  //Check momentum is still zero on boundary
+  getLinksBoundaryBulk(Pbnd, Pbulk, const_cast<Field&>(P), latt_size);
+  
+  Pbnd_nrm = norm2(Pbnd); //expect zero
+  std::cout << GridLogMessage << "After integrate, norm2 of mu-directed conjugate momentum on boundary is: " << Pbnd_nrm << " (expect 0)" << std::endl;
+  Pbulk_nrm = norm2(Pbulk); //expect non-zero
+  std::cout << GridLogMessage << "After integrate, norm2 of bulk conjugate momentum is: " << Pbulk_nrm << " (expect non-zero)" << std::endl;
+
+  //Get the new bulk and bound links
+  Field Ubnd_new(&Grid), Ubulk_new(&Grid);
+  getLinksBoundaryBulk(Ubnd_new, Ubulk_new, U, latt_size);
+
+  Field Ubnd_diff = Ubnd_new - Ubnd_orig;
+  Field Ubulk_diff = Ubulk_new - Ubulk_orig;
+
+  RealD Ubnd_change = norm2( Ubnd_diff );
+  RealD Ubulk_change = norm2( Ubulk_diff );
+  std::cout << GridLogMessage << "After integrate, norm2 of change in mu-directed boundary links is : " << Ubnd_change << " (expect 0)" << std::endl;
+  std::cout << GridLogMessage << "After integrate, norm2 of change in bulk links is : " << Ubulk_change << " (expect non-zero)" << std::endl;
+
+  Grid_finalize();
+}