Просмотр исходного кода

First implementation [model-root-data]->[AnalysisTree ver. 2]

PeterParfenov 1 год назад
Родитель
Сommit
3891f49349

+ 2 - 0
CMakeLists.txt

@@ -78,6 +78,7 @@ set(MATC_LIBRARY_h_files
   ${CMAKE_CURRENT_SOURCE_DIR}/readers/qaReader_mcpico.h
   ${CMAKE_CURRENT_SOURCE_DIR}/readers/qaReader_manager.h
   ${CMAKE_CURRENT_SOURCE_DIR}/writers/qaWriter_manager.h
+  ${CMAKE_CURRENT_SOURCE_DIR}/writers/qaWriter_analysistree2.h
   ${CMAKE_CURRENT_SOURCE_DIR}/utility/Utility.h
 )
 
@@ -88,6 +89,7 @@ set(MATC_LIBRARY_cxx_files
   ${CMAKE_CURRENT_SOURCE_DIR}/readers/qaReader_mcpico.cxx
   ${CMAKE_CURRENT_SOURCE_DIR}/readers/qaReader_manager.cxx
   ${CMAKE_CURRENT_SOURCE_DIR}/writers/qaWriter_manager.cxx
+  ${CMAKE_CURRENT_SOURCE_DIR}/writers/qaWriter_analysistree2.cxx
   ${CMAKE_CURRENT_SOURCE_DIR}/utility/Utility.cxx
 )
 

+ 213 - 0
bin/main.cpp

@@ -0,0 +1,213 @@
+#include <iostream>
+#include <vector>
+
+#include <TString.h>
+#include <TH1D.h>
+#include <TH1I.h>
+#include <TH2D.h>
+#include <TProfile.h>
+#include <TProfile2D.h>
+#include <TFile.h>
+#include <TStopwatch.h>
+
+#include <qaParticle.h>
+#include <qaParticleLight.h>
+#include <qaEvent.h>
+#include <qaReader_manager.h>
+#include <qaReader_smash_root.h>
+#include <qaReader_mcpico.h>
+#include <Utility.h>
+#include <qaWriter_manager.h>
+#include <qaWriter_analysistree2.h>
+
+#ifdef _MCINI_
+#include <qaReader_mcini.h>
+#endif
+#ifdef _PHQMD_
+#include <qaReader_phqmd.h>
+#endif
+#ifdef _HSD_ROOT_
+#include <qaReader_hsd_root.h>
+#endif
+
+// AnalysisTree headers
+#include <AnalysisTree/Configuration.hpp>
+#include <AnalysisTree/DataHeader.hpp>
+#include <AnalysisTree/EventHeader.hpp>
+#include <AnalysisTree/Detector.hpp>
+#include <AnalysisTree/Matching.hpp>
+
+int main(int argc, char **argv)
+{
+  TString iFileName, oFileName;
+  TString outputFormat = "analysistree_v2";
+
+  if (argc < 7)
+  {
+    std::cerr << "./matc -i input.list -o output.root -input_format [INPUT_FORMAT] [OPTIONAL: -output_format [OUTPUT_FORMAT]]" << std::endl;
+    std::cerr << "Available input formats:" << std::endl;
+    std::cerr << "\tmcpico - simple custom ROOT format to store model data" << std::endl;
+    std::cerr << "\tparticle - ROOT format that is used by the SMASH model" << std::endl;
+#ifdef _MCINI_
+    std::cerr << "\tmcini - custom ROOT format to store both initial state and final state (UniGen data format) model data" << std::endl;
+#endif
+#ifdef _PHQMD_
+    std::cerr << "\tphqmd - custom ROOT format to store PHQMD (with MST) model data" << std::endl;
+#endif
+#ifdef _HSD_ROOT_
+    std::cerr << "\thsd - custom ROOT format to store HSD model data" << std::endl;
+#endif
+    std::cerr << "Available output formats:" << std::endl;
+    std::cerr << "\tanalysistree_v2 - AnalysisTree ver. 2 data format" << std::endl;
+    return 1;
+  }
+  for (int i = 1; i < argc; i++)
+  {
+    if (std::string(argv[i]) != "-i" &&
+        std::string(argv[i]) != "-o" &&
+        std::string(argv[i]) != "-input_format" &&
+        std::string(argv[i]) != "-output_format")
+    {
+      std::cerr << "\n[ERROR]: Unknown parameter " << i << ": " << argv[i] << std::endl;
+      return 2;
+    }
+    else
+    {
+      if (std::string(argv[i]) == "-i" && i != argc - 1)
+      {
+        iFileName = argv[++i];
+        continue;
+      }
+      if (std::string(argv[i]) == "-i" && i == argc - 1)
+      {
+        std::cerr << "\n[ERROR]: Input file name was not specified " << std::endl;
+        return 1;
+      }
+      if (std::string(argv[i]) == "-o" && i != argc - 1)
+      {
+        oFileName = argv[++i];
+        continue;
+      }
+      if (std::string(argv[i]) == "-o" && i == argc - 1)
+      {
+        std::cerr << "\n[ERROR]: Output file name was not specified " << std::endl;
+        return 1;
+      }
+      if (std::string(argv[i]) == "-input_format" && i != argc - 1)
+      {
+        qaUtility::GetInstance()->format = argv[++i];
+        continue;
+      }
+      if (std::string(argv[i]) == "-input_format" && i == argc - 1)
+      {
+        std::cerr << "\n[ERROR]: Input file format was not specified " << std::endl;
+        return 1;
+      }
+      if (std::string(argv[i]) == "-output_format" && i != argc - 1)
+      {
+        outputFormat = argv[++i];
+        continue;
+      }
+      if (std::string(argv[i]) == "-output_format" && i == argc - 1)
+      {
+        std::cerr << "\n[ERROR]: Output file format was not specified " << std::endl;
+        return 1;
+      }
+    }
+  }
+
+  TStopwatch timer;
+  timer.Start();
+
+  qaReader_manager *readerManager;
+  if (qaUtility::GetInstance()->format == "mcpico")
+  {
+    readerManager = new qaReader_mcpico();
+  }
+#ifdef _MCINI_
+  if (qaUtility::GetInstance()->format == "mcini")
+  {
+    readerManager = new qaReader_mcini();
+  }
+#endif
+#ifdef _PHQMD_
+  if (qaUtility::GetInstance()->format == "phqmd")
+  {
+    readerManager = new qaReader_phqmd();
+  }
+#endif
+#ifdef _HSD_ROOT_
+  if (qaUtility::GetInstance()->format == "hsd")
+  {
+    readerManager = new qaReader_hsd_root();
+  }
+#endif
+  if (qaUtility::GetInstance()->format == "particles")
+  {
+    readerManager = new qaReader_smash_root();
+  }
+
+  if (!readerManager)
+  {
+    std::cerr << "This input format is not found!" << std::endl;
+    return 20;
+  }
+
+  qaWriter_manager *writerManager;
+  if (outputFormat == "analysistree_v2")
+  {
+    writerManager = new qaWriter_analysistree2();
+  }
+
+  if (!writerManager)
+  {
+    std::cerr << "This output format is not found!" << std::endl;
+    return 30;
+  }
+
+  readerManager->SetChain(iFileName.Data());
+  Long64_t Nentries = readerManager->GetEntries();
+  Int_t Nparticles;
+
+  qaEvent *event = nullptr;
+  qaParticle *particle = nullptr;
+
+  writerManager->Init(oFileName.Data(), "modelTree");
+
+  for (Long64_t iev = 0; iev < Nentries; iev++)
+  {
+    if (iev % 1000 == 0)
+      std::cout << "Event [" << Minbias_counter << "/" << Nentries << "]" << std::endl;
+
+    event = (qaEvent *)readerManager->ReadEvent(iev);
+    if (!event)
+      continue;
+
+    writerManager->WriteEvent(event);
+
+    Nparticles = event->GetNparticles();
+
+    for (int iparticle = 0; iparticle < Nparticles; iparticle++)
+    {
+      particle = readerManager->ReadParticle(iparticle);
+
+      if (!particle)
+        continue;
+
+      writerManager->WriteParticle(particle);
+      
+      delete particle;
+    }
+
+    writerManager->FillTree();
+
+    delete event;
+  }
+
+  writerManager->WriteTree();
+
+  timer.Stop();
+  timer.Print();
+
+  return 0;
+}

+ 1 - 0
matc.LinkDef.h

@@ -23,6 +23,7 @@
 #endif
 #pragma link C++ class qaReader_manager+;
 #pragma link C++ class qaWriter_manager+;
+#pragma link C++ class qaWriter_analysistree2+;
 #pragma link C++ class qaUtility+;
 
 #endif

+ 7 - 0
utility/Utility.cxx

@@ -79,4 +79,11 @@ Int_t qaUtility::GetPdgId(Int_t pdg)
   if (pdg == vpdg.at(9))
     return 9;
   return -1;
+}
+
+Double_t qaUtility::GetBeamP(Double_t sqrtSnn)
+{
+  double m_target = 0.938;
+  double m_beam = 0.938;
+  return sqrt( pow((pow(sqrtSnn,2) - pow(m_target,2) - pow(m_beam,2))/(2*m_target), 2) - pow(m_beam,2) );
 }

+ 125 - 0
writers/qaWriter_analysistree2.cxx

@@ -0,0 +1,125 @@
+#include <qaWriter_analysistree2.h>
+
+ClassImp(qaWriter_analysistree2);
+
+qaWriter_analysistree2::qaWriter_analysistree2() : fFile(nullptr),
+                                                   fTree(nullptr),
+                                                   fDataHeader(nullptr),
+                                                   fConfig(nullptr),
+                                                   fEvent(nullptr),
+                                                   fParticles(nullptr),
+                                                   isInit(false),
+                                                   iB(qaUtility::GetInstance()->error_code),
+                                                   iPhiRp(qaUtility::GetInstance()->error_code),
+                                                   icharge(qaUtility::GetInstance()->error_code),
+{
+}
+
+qaWriter_analysistree2::~qaWriter_analysistree2()
+{
+}
+
+void qaWriter_analysistree2::Init(std::string filename, std::string treename, std::string system, float sqrtSnn)
+{
+  fFile = new TFile(filename.c_str(), "recreate");
+  fTree = new TTree(treename.c_str(), "AnalysisTree format, model data");
+
+  // Set up AnalysisTree data header
+  fDataHeader = new AnalysisTree::DataHeader;
+  if (system != "")
+    fDataHeader->SetSystem(system);
+  if (sqrtSnn > 0)
+    fDataHeader->SetBeamMomentum(qaUtility::GetInstance()->GetBeamP(sqrtSnn));
+
+  // Set up AnalysisTree configureation
+  fConfig = new AnalysisTree::Configuration;
+
+  // Set up AnalysisTree configuration
+  std::string str_event_branch = "Event";
+  std::string str_particles_branch = "Particles";
+
+  AnalysisTree::BranchConfig event_branch(str_event_branch.c_str(), AnalysisTree::DetType::kEventHeader);
+  AnalysisTree::BranchConfig particles_branch(str_particles_branch.c_str(), AnalysisTree::DetType::kParticle);
+
+  event_branch.AddField<float>("B");
+  event_branch.AddField<float>("PhiRp");
+
+  particles_branch.AddField<bool>("is_charged");
+
+  auto hasher = std::hash<std::string>();
+
+  fConfig->AddBranchConfig(event_branch);
+  fEvent = new AnalysisTree::EventHeader(Short_t(hasher(event_branch.GetName())));
+
+  fEvent->Init(event_branch);
+
+  fConfig->AddBranchConfig(particles_branch);
+  fParticles = new AnalysisTree::Particles(Short_t(hasher(particles_branch.GetName())));
+
+  iB = fConfig->GetBranchConfig(str_event_branch).GetFieldId("B");
+  iPhiRp = fConfig->GetBranchConfig(str_event_branch).GetFieldId("PhiRp");
+
+  icharge = fConfig->GetBranchConfig(str_particles_branch).GetFieldId("is_charged");
+
+  // Create branches in the output tree
+  fTree->Branch(str_event_branch.c_str(), "AnalysisTree::EventHeader", &fEvent, 32000, 99);
+  fTree->Branch(str_particles_branch.c_str(), "AnalysisTree::Particles", &fParticles, 256000, 99);
+
+  fConfig->Print();
+
+  isInit = true;
+}
+
+void qaWriter_analysistree2::WriteEvent(qaEvent *event)
+{
+  if (!event)
+    return;
+  if (!isInit)
+    return;
+
+  fParticles->ClearChannels();
+
+  fEvent->SetField(float(event->GetB()), iB);
+  fEvent->SetField(float(event->GetPhiRP()), iPhiRp);
+}
+
+void qaWriter_analysistree2::WriteParticle(qaParticle *particle)
+{
+  if (!particle)
+    return;
+  if (!isInit)
+    return;
+
+  auto *at_particle = fParticles->AddChannel();
+  at_particle->Init(fConfig->GetBranchConfig(str_particles_branch));
+
+  at_particle->SetMomentum(particle->GetPx(), particle->GetPy(), particle->GetPz());
+  at_particle->SetPid(particle->GetPdg());
+  auto particlePDG = (TParticlePDG *)TDatabasePDG::Instance()->GetParticle(particle->GetPdg());
+  if (!particlePDG)
+    at_particle->SetMass(qaUtility::GetInstance()->error_code);
+  if (particlePDG)
+    at_particle->SetMass(particlePDG->Mass());
+  bool is_charged = (qaUtility::GetInstance()->GetCharge(particle->GetPdg()) == 0 ||
+                     qaUtility::GetInstance()->GetCharge(particle->GetPdg()) == qaUtility::GetInstance()->error_code)
+                        ? false
+                        : true;
+  at_particle->SetField(is_charged, icharge);
+}
+
+void qaWriter_analysistree2::WriteTree()
+{
+  if (!isInit) return;
+  fFile->cd();
+  fTree->Print();
+  fTree->Write();
+
+  fDataHeader->Write("DataHeader");
+  fConfig->Write("Configuration");
+  fFile->Close();
+}
+
+void qaWriter_analysistree2::FillTree()
+{
+  fTree->Fill();
+}

+ 46 - 0
writers/qaWriter_analysistree2.h

@@ -0,0 +1,46 @@
+#ifndef QATOOLS_WRITERS_ANALYSISTREE2_H
+#define QATOOLS_WRITERS_ANALYSISTREE2_H
+
+#include <Rtypes.h>
+#include <TString.h>
+#include <TFile.h>
+#include <TTree.h>
+
+#include <qaEvent.h>
+#include <qaParticle.h>
+#include <Utility.h>
+
+// AnalysisTree headers
+#include <AnalysisTree/Configuration.hpp>
+#include <AnalysisTree/DataHeader.hpp>
+#include <AnalysisTree/EventHeader.hpp>
+#include <AnalysisTree/Detector.hpp>
+#include <AnalysisTree/Matching.hpp>
+
+class qaWriter_analysistree2
+{
+private:
+  /* data */
+  TFile *fFile;
+  TTree *fTree;
+  AnalysisTree::DataHeader *fDataHeader;
+  AnalysisTree::Configuration *fConfig;
+  AnalysisTree::EventHeader *fEvent;
+  AnalysisTree::Particles *fParticles;
+
+  int iB, iPhiRp, icharge;
+
+  bool isInit;
+public:
+  virtual ~qaWriter_analysistree2();
+
+  virtual void Init(std::string filename, std::string treename);
+  virtual void WriteEvent(qaEvent* event);
+  virtual void WriteParticle(qaParticle* particle);
+  virtual void WriteTree();
+  virtual void FillTree();
+
+  ClassDef(qaWriter_analysistree2,0);
+};
+
+#endif

+ 1 - 0
writers/qaWriter_manager.h

@@ -19,6 +19,7 @@ public:
   virtual void WriteEvent(qaEvent* event) = 0;
   virtual void WriteParticle(qaParticle* particle) = 0;
   virtual void WriteTree() = 0;
+  virtual void FillTree() = 0;
 
   ClassDef(qaWriter_manager,0);
 };