add vertex information to final state hadron and partons

config/jetscape_main.xml

@@ -46,9 +46,11 @@
     <!-- -9999 corresponds to skipping no status code. It needs to have some entry to be avoid a segfault w/ empty values in tinyxml -->
     <FinalStateHadrons>
        <statusToSkip>-9999</statusToSkip>
+       <headerVersion>2</headerVersion>
     </FinalStateHadrons>
     <FinalStatePartons>
        <statusToSkip>-9999</statusToSkip>
+       <headerVersion>2</headerVersion>
     </FinalStatePartons>
   </Writer>
 

examples/FinalStateHadrons.cc

@@ -48,12 +48,13 @@ int main(int argc, char** argv)
   // //If you want to suppress it: use SetVerboseLevel(0) or max  SetVerboseLevel(9) or 10
   JetScapeLogger::Instance()->SetVerboseLevel(0);
 
-  // Whether to write the new header (ie. v2), including xsec info.
-  // To enable, pass anything as the third argument to enable this option.
-  // Default: version 1.
-  int headerVersion = 1;
+  // Whether to write a particular header version (eg. v2), including xsec info.
+  // To enable, pass the desired version (just the number) as the third argument.
+  // Default: v1
+  unsigned int headerVersion = 1;
   if (argc > 3) {
     headerVersion = std::atoi(argv[3]);
+    std::cout << "NOTE: Writing header v" << headerVersion << ", and final cross section and error at EOF.\n";
   }
   std::cout << "NOTE: Writing with output version v" << headerVersion << "\n";
 
@@ -77,13 +78,20 @@ int main(int argc, char** argv)
     if (hadrons.size() > 0)
     {
       ++SN;
-      if (headerVersion == 2) {
+      if (headerVersion > 1) {
         // NOTE: Needs consistent "\t" between all entries to simplify parsing later.
         dist_output << "#"
             << "\t" << "Event\t" << SN
             << "\t" << "weight\t" << reader->GetEventWeight()
             << "\t" << "EPangle\t" << reader->GetEventPlaneAngle()
-            << "\t" << "N_hadrons\t" << hadrons.size()
+            << "\t" << "N_hadrons\t" << hadrons.size();
+        if (headerVersion == 3) {
+          dist_output
+              << "\t" << "vertex_x\t" << reader->GetVertexX()
+              << "\t" << "vertex_y\t" << reader->GetVertexY()
+              << "\t" << "vertex_z\t" << reader->GetVertexZ();
+        }
+        dist_output
             << "\t" << "|"  // As a delimiter
             << "\t" << "N"
             << "\t" << "pid"

examples/FinalStatePartons.cc

@@ -49,12 +49,13 @@ int main(int argc, char** argv)
   // //If you want to suppress it: use SetVerboseLevle(0) or max  SetVerboseLevel(9) or 10
   JetScapeLogger::Instance()->SetVerboseLevel(0);
 
-  // Whether to write the new header (ie. v2), including xsec info.
-  // To enable, pass anything as the third argument to enable this option.
-  // Default: version 1.
-  int headerVersion = 1;
+  // Whether to write a particular header version (eg. v2), including xsec info.
+  // To enable, pass the desired version (just the number) as the third argument.
+  // Default: v1
+  unsigned int headerVersion = 1;
   if (argc > 3) {
     headerVersion = std::atoi(argv[3]);
+    std::cout << "NOTE: Writing header v" << headerVersion << ", and final cross section and error at EOF.\n";
   }
   std::cout << "NOTE: Writing with output version v" << headerVersion << "\n";
 
@@ -83,13 +84,20 @@ int main(int argc, char** argv)
     if(TotalPartons > 0)
     {
       ++SN;
-      if (headerVersion == 2) {
+      if (headerVersion > 1) {
         // NOTE: Needs consistent "\t" between all entries to simplify parsing later.
         dist_output << "#"
             << "\t" << "Event\t" << SN
             << "\t" << "weight\t" << reader->GetEventWeight()
             << "\t" << "EPangle\t" << reader->GetEventPlaneAngle()
-            << "\t" << "N_partons\t" << TotalPartons
+            << "\t" << "N_partons\t" << TotalPartons;
+        if (headerVersion == 3) {
+          dist_output
+              << "\t" << "vertex_x\t" << reader->GetVertexX()
+              << "\t" << "vertex_y\t" << reader->GetVertexY()
+              << "\t" << "vertex_z\t" << reader->GetVertexZ();
+        }
+        dist_output
             << "\t" << "|"  // As a delimiter
             << "\t" << "N"
             << "\t" << "pid"

src/framework/HardProcess.cc

@@ -64,6 +64,7 @@ void HardProcess::Init() {
         printer = GetXMLElementText({"PartonPrinter","FileName"});
         JSINFO << BOLDYELLOW << "Extra parton info goes to " << printer ;
     }
+
   InitTask();
   InitTasks();
 }
@@ -122,6 +123,9 @@ void HardProcess::CollectHeader(weak_ptr<JetScapeWriter> w) {
     header.SetSigmaErr(GetSigmaErr());
     header.SetPtHat(GetPtHat());
     header.SetEventWeight(GetEventWeight());
+    header.SetVertexX(hp_list[0]->x_in().x());
+    header.SetVertexY(hp_list[0]->x_in().y());
+    header.SetVertexZ(hp_list[0]->x_in().z());
   }
 }
 

src/framework/JetScapeEventHeader.h

@@ -58,6 +58,18 @@ class JetScapeEventHeader {
   double GetEventWeight() { return EventWeight; };
   /// Initial Hard Process: Set additionally created weight (e.g. pythia.event().weight())
   void SetEventWeight(double d) { EventWeight = d; };
+  /// Initial Hard Process: Get vertex X position
+  double GetVertexX() { return vX; };
+  /// Initial Hard Process: Set vertex X position
+  void SetVertexX(double d) { vX = d; };
+  /// Initial Hard Process: Get vertex Y position
+  double GetVertexY() { return vY; };
+  /// Initial Hard Process: Set vertex Y position
+  void SetVertexY(double d) { vY = d; };
+  /// Initial Hard Process: Get vertex Z position
+  double GetVertexZ() { return vZ; };
+  /// Initial Hard Process: Set vertex Z position
+  void SetVertexZ(double d) { vZ = d; };
 
   // ============================ Initial State =================================
   /// Initial State: Get number of participants
@@ -92,6 +104,9 @@ class JetScapeEventHeader {
   double SigmaErr = -1;
   double PtHat = -1;
   double EventWeight = 1;
+  double vX = -999;
+  double vY = -999;
+  double vZ = -999;
 
   // ============================ Initial State =================================
   double Npart = -1; // could be int, but using double to allow averaged values

src/framework/JetScapeWriterFinalStateStream.cc

@@ -64,7 +64,9 @@ JetScapeWriterFinalStateStream<T>::JetScapeWriterFinalStateStream(string m_file_
   particles{},
   writeCentrality{false},
   writePtHat{false},
-  particleStatusToSkip{}
+  particleStatusToSkip{},
+  output_file{},
+  headerVersion{2}
 {
   SetOutputFileName(m_file_name_out);
 }
@@ -130,7 +132,14 @@ template <class T> void JetScapeWriterFinalStateStream<T>::WriteEvent() {
       << "\t" << "Event\t" << GetCurrentEvent() + 1  // +1 to index the event count from 1
       << "\t" << "weight\t" << std::setprecision(15) << GetHeader().GetEventWeight() << std::setprecision(6)
       << "\t" << "EPangle\t" << (GetHeader().GetEventPlaneAngle() > -999 ? GetHeader().GetEventPlaneAngle() : 0)
-      << "\t" << "N_" << GetName() << "\t" << ipart
+      << "\t" << "N_" << GetName() << "\t" << ipart;
+  if (headerVersion == 3) {
+    output_file
+        << "\t" << "vertex_x\t" << GetHeader().GetVertexX()
+        << "\t" << "vertex_y\t" << GetHeader().GetVertexY()
+        << "\t" << "vertex_z\t" << GetHeader().GetVertexZ();
+  }
+  output_file
       << centrality_text
       << pt_hat_text
       <<  "\n";
@@ -161,14 +170,23 @@ template <class T> void JetScapeWriterFinalStateStream<T>::InitTask() {
     particleStatusToSkip.erase(std::remove(particleStatusToSkip.begin(), particleStatusToSkip.end(), -9999), particleStatusToSkip.end());
   }
   if (GetActive()) {
-    JSINFO << "JetScape Final State " << name << " Stream Writer initialized with output file = "
+    // We need the header version to determine how to write.
+    // NOTE: Don't require the version. If not specified, defaults to v2.
+    int result = GetXMLElementInt({"Writer", (std::string("FinalState") + name).c_str(), "headerVersion"}, false);
+    // If it fails to retrieve the value, it will return 0. The header version must be >= 2,
+    // so only assign if the value is actually set.
+    if (result) {
+      headerVersion = static_cast<unsigned int>(result);
+    }
+
+    JSINFO << "JetScape Final State " << name << " Stream Writer v" << headerVersion << " initialized with output file = "
            << GetOutputFileName();
     output_file.open(GetOutputFileName().c_str());
     // NOTE: This header will only be printed once at the beginning on the file.
     output_file << "#"
         // The specifics the version number. For consistency in parsing, the string
         // will always be "v<number>"
-        << "\t" << "JETSCAPE_FINAL_STATE\t" << "v2"
+        << "\t" << "JETSCAPE_FINAL_STATE\t" << "v" << headerVersion
         << "\t" << "|"  // As a delimiter
         << "\t" << "N"
         << "\t" << "pid"

src/framework/JetScapeWriterFinalStateStream.h

@@ -64,6 +64,7 @@ class JetScapeWriterFinalStateStream : public JetScapeWriter {
 
 protected:
   T output_file; //!< Output file
+  unsigned int headerVersion;
   std::vector<std::shared_ptr<JetScapeParticleBase>> particles;
   bool writeCentrality;
   bool writePtHat;

src/reader/JetScapeReader.cc

@@ -22,6 +22,9 @@ template <class T> JetScapeReader<T>::JetScapeReader():
   currentEvent{-1}
   , sigmaGen{-1}
   , sigmaErr{-1}
+  , vertexX{-999}
+  , vertexY{-999}
+  , vertexZ{-999}
   , eventWeight{-1}
   , EventPlaneAngle{0.0}
 {
@@ -39,6 +42,9 @@ template <class T> void JetScapeReader<T>::ClearTask() {
 
   sigmaGen = -1;
   sigmaErr = -1;
+  vertexX = -999;
+  vertexY = -999;
+  vertexZ = -999;
   eventWeight = -1;
   EventPlaneAngle = 0.0;
 }
@@ -151,6 +157,16 @@ template <class T> void JetScapeReader<T>::Next() {
         data >> dummy >> dummy >> dummy >> EventPlaneAngle;
         JSDEBUG << " EventPlaneAngle=" << EventPlaneAngle;
       }
+      // vertex position of hard scattering
+      if (line.find("HardProcess") != std::string::npos) {
+        getline(inFile, line); // get next line to get vertex position
+        std::stringstream data(line);
+        double dummy;
+        data >> dummy >> dummy >> dummy >> dummy >> dummy >> dummy >> dummy >> vertexX >> vertexY >> vertexZ;
+        JSDEBUG << " vertexX=" << vertexX;
+        JSDEBUG << " vertexY=" << vertexY;
+        JSDEBUG << " vertexZ=" << vertexZ;
+      }
       continue;
     }
 

src/reader/JetScapeReader.h

@@ -64,6 +64,9 @@ template <class T> class JetScapeReader {
   double GetSigmaErr() const { return sigmaErr; }
   double GetEventWeight() const { return eventWeight; }
   double GetEventPlaneAngle() const { return EventPlaneAngle; }
+  double GetVertexX() const { return vertexX; }
+  double GetVertexY() const { return vertexY; }
+  double GetVertexZ() const { return vertexZ; }
 
 private:
   StringTokenizer strT;
@@ -89,6 +92,9 @@ template <class T> class JetScapeReader {
   double sigmaErr;
   double eventWeight;
   double EventPlaneAngle;
+  double vertexX;
+  double vertexY;
+  double vertexZ;
 };
 
 typedef JetScapeReader<ifstream> JetScapeReaderAscii;