From ba974960e6f65835b8216eca72c2b6120016993f Mon Sep 17 00:00:00 2001
From: Christopher Kelly <ckelly@phys.columbia.edu>
Date: Mon, 14 Feb 2022 08:09:01 -0800
Subject: [PATCH] Added an HMC checkpoint start option that loads the fields
 and then reseeds the RNGs, suitable for creating new evolution streams Added
 option to choose RNG seeds in 40ID main binary

---
 Grid/qcd/hmc/GenericHMCrunner.h        | 12 ++++++-
 Grid/qcd/hmc/HMCModules.h              |  2 ++
 HMC/Mobius2p1fIDSDRGparityEOFA_40ID.cc | 45 ++++++++++++++++++++++++--
 3 files changed, 55 insertions(+), 4 deletions(-)

diff --git a/Grid/qcd/hmc/GenericHMCrunner.h b/Grid/qcd/hmc/GenericHMCrunner.h
index 73381351..13811265 100644
--- a/Grid/qcd/hmc/GenericHMCrunner.h
+++ b/Grid/qcd/hmc/GenericHMCrunner.h
@@ -151,12 +151,22 @@ public:
       Resources.GetCheckPointer()->CheckpointRestore(Parameters.StartTrajectory, U,
 						     Resources.GetSerialRNG(),
 						     Resources.GetParallelRNG());
+    } else if (Parameters.StartingType == "CheckpointStartReseed") {
+      // Same as CheckpointRestart but reseed the RNGs using the fixed integer seeding used for ColdStart and HotStart
+      // Useful for creating new evolution streams from an existing stream
+      
+      // WARNING: Unfortunately because the checkpointer doesn't presently allow us to separately restore the RNG and gauge fields we have to load
+      // an existing RNG checkpoint first; make sure one is available and named correctly
+      Resources.GetCheckPointer()->CheckpointRestore(Parameters.StartTrajectory, U,
+						     Resources.GetSerialRNG(),
+						     Resources.GetParallelRNG());
+      Resources.SeedFixedIntegers();      
     } else {
       // others
       std::cout << GridLogError << "Unrecognized StartingType\n";
       std::cout
 	<< GridLogError
-	<< "Valid [HotStart, ColdStart, TepidStart, CheckpointStart]\n";
+	<< "Valid [HotStart, ColdStart, TepidStart, CheckpointStart, CheckpointStartReseed]\n";
       exit(1);
     }
   }
diff --git a/Grid/qcd/hmc/HMCModules.h b/Grid/qcd/hmc/HMCModules.h
index 4c61a006..cf0edd26 100644
--- a/Grid/qcd/hmc/HMCModules.h
+++ b/Grid/qcd/hmc/HMCModules.h
@@ -80,7 +80,9 @@ public:
       std::cout << GridLogError << "Seeds not initialized" << std::endl;
       exit(1);
     }
+    std::cout << GridLogMessage << "Reseeding serial RNG with seed vector " << SerialSeeds << std::endl;
     sRNG_.SeedFixedIntegers(SerialSeeds);
+    std::cout << GridLogMessage << "Reseeding parallel RNG with seed vector " << ParallelSeeds << std::endl;
     pRNG_->SeedFixedIntegers(ParallelSeeds);
   }
 };
diff --git a/HMC/Mobius2p1fIDSDRGparityEOFA_40ID.cc b/HMC/Mobius2p1fIDSDRGparityEOFA_40ID.cc
index 9ba622e5..6de636bd 100644
--- a/HMC/Mobius2p1fIDSDRGparityEOFA_40ID.cc
+++ b/HMC/Mobius2p1fIDSDRGparityEOFA_40ID.cc
@@ -499,16 +499,47 @@ int main(int argc, char **argv) {
 
   std::string param_file = "params.xml";
   bool file_load_check = false;
-  for(int i=1;i<argc;i++){
+
+  std::string serial_seeds = "1 2 3 4 5";
+  std::string parallel_seeds = "6 7 8 9 10";
+
+  int i=1;
+  while(i < argc){
     std::string sarg(argv[i]);
     if(sarg == "--param_file"){
       assert(i!=argc-1);
       param_file = argv[i+1];
+      i+=2;
     }else if(sarg == "--read_check"){ //check the fields load correctly and pass checksum/plaquette repro
       file_load_check = true;
+      i++;
+    }else if(sarg == "--set_seeds"){ //set the rng seeds. Expects two vector args, e.g.  --set_seeds 1.2.3.4 5.6.7.8
+      assert(i < argc-2);
+      std::vector<int> tmp;
+      GridCmdOptionIntVector(argv[i+1],tmp);
+      {
+	std::stringstream ss;
+	for(int j=0;j<tmp.size()-1;j++) ss << tmp[j] << " ";
+	ss << tmp.back();
+	serial_seeds = ss.str();
+      }
+      GridCmdOptionIntVector(argv[i+2],tmp);
+      {
+	std::stringstream ss;
+	for(int j=0;j<tmp.size()-1;j++) ss << tmp[j] << " ";
+	ss << tmp.back();
+	parallel_seeds = ss.str();
+      }
+      i+=3;
+      std::cout << GridLogMessage << "Set serial seeds to " << serial_seeds << std::endl;
+      std::cout << GridLogMessage << "Set parallel seeds to " << parallel_seeds << std::endl;
+      
+    }else{
+      i++;
     }
   }
 
+  
   //Read the user parameters
   EvolParameters user_params;
   
@@ -587,8 +618,8 @@ int main(int argc, char **argv) {
 
   //Note that checkpointing saves the RNG state so that this initialization is required only for the very first configuration
   RNGModuleParameters RNGpar;
-  RNGpar.serial_seeds = "1 2 3 4 5";
-  RNGpar.parallel_seeds = "6 7 8 9 10";
+  RNGpar.serial_seeds = serial_seeds;
+  RNGpar.parallel_seeds = parallel_seeds;
   TheHMC.Resources.SetRNGSeeds(RNGpar);
 
   typedef PlaquetteMod<GaugeImplPolicy> PlaqObs;
@@ -614,6 +645,14 @@ int main(int argc, char **argv) {
   RealD mob_b = (mobius_scale + mob_bmc)/2.;
   RealD mob_c = (mobius_scale - mob_bmc)/2.;
 
+  std::cout << GridLogMessage
+	    << "Ensemble parameters:" << std::endl
+	    << "Ls=" << Ls << std::endl
+	    << "beta=" << beta << std::endl
+	    << "light_mass=" << light_mass << std::endl
+	    << "strange_mass=" << strange_mass << std::endl
+	    << "mobius_scale=" << mobius_scale << std::endl;
+  
   //Setup the Grids
   auto UGridD   = TheHMC.Resources.GetCartesian();
   auto UrbGridD = TheHMC.Resources.GetRBCartesian();