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

Small addition for CreateExtendedAT.C (11)

PeterParfenov 1 год назад
Родитель
Сommit
cb058f8d02
1 измененных файлов с 96 добавлено и 12 удалено
  1. 96 12
      CreateExtendedAT.C

+ 96 - 12
CreateExtendedAT.C

@@ -34,6 +34,7 @@ TVector3 GetFHCalPos(int iModule);
 Bool_t IsCharged(Int_t pdg);
 Float_t GetRapidity(Float_t p, Float_t pz, Int_t pid_type);
 Float_t GetRapidityPDG(Float_t p, Float_t pz, Int_t pdg);
+Float_t GetBCent(Float_t bimp);
 
 // Reads base AnalysisTree and creates additional fields that will be used in the QnAnalysis
 void CreateExtendedAT(TString iFileName, TString oFileName)
@@ -43,24 +44,23 @@ void CreateExtendedAT(TString iFileName, TString oFileName)
   const Long64_t nevents = (Long64_t)chain->GetEntries();
 
   // Initialize pointers to the branches that needed to be read
-  chain->InitPointersToBranches({"RecoEvent", "TpcTracks", "FHCalModules", "McTracks"});
+  chain->InitPointersToBranches({"RecoEvent", "McEvent", "TpcTracks", "FHCalModules", "McTracks"});
 
   // Read and print general info from DataHeader and Configuration
   auto *in_config = chain->GetConfiguration();
   auto *in_data_header = chain->GetDataHeader();
 
   // in_data_header->Print();
-  // in_config->Print();
 
   // Get branches from the AnalysisTree
   auto in_reco_event = chain->GetBranch("RecoEvent");
+  auto in_mc_event = chain->GetBranch("McEvent");
   auto in_tpc_tracks = chain->GetBranch("TpcTracks");
   auto in_mc_tracks = chain->GetBranch("McTracks");
   auto in_fhcal_mods = chain->GetBranch("FHCalModules");
   auto in_tpcmc_match = chain->GetMatching("TpcTracks", "McTracks");
 
   // Set up output AnalysisTree with extended information
-  // Set up output dst
   TFile *outFile = new TFile(oFileName.Data(), "RECREATE");
   TTree *outTree = new TTree("bTree", "Extended AnalysisTree Dst at MPD");
 
@@ -68,10 +68,10 @@ void CreateExtendedAT(TString iFileName, TString oFileName)
   AnalysisTree::Configuration *out_config = new AnalysisTree::Configuration;
 
   // Set up AnalysisTree configuration
-  std::string str_reco_event_branch = "RecoEventExt";
-  std::string str_tpc_tracks_branch = "TpcTracksExt";
-  std::string str_fhcal_branch = "FHCalModulesExt";
-  std::string str_mc_tracks_branch = "McTracksExt";
+  std::string str_reco_event_branch    = "RecoEventExt";
+  std::string str_tpc_tracks_branch    = "TpcTracksExt";
+  std::string str_fhcal_branch         = "FHCalModulesExt";
+  std::string str_mc_tracks_branch     = "McTracksExt";
   std::string str_tpc2mc_tracks_branch = "TpcTracksExt2McTracksExt";
 
   AnalysisTree::BranchConfig reco_event_branch(str_reco_event_branch.c_str(), AnalysisTree::DetType::kEventHeader); 
@@ -80,25 +80,32 @@ void CreateExtendedAT(TString iFileName, TString oFileName)
   AnalysisTree::BranchConfig mc_tracks_branch(str_mc_tracks_branch.c_str(), AnalysisTree::DetType::kParticle);
 
   // Create new branches which will store extended data for further analysis
+  AnalysisTree::Branch out_reco_event(reco_event_branch);
   AnalysisTree::Branch out_tpc_tracks(tpc_tracks_branch);
   AnalysisTree::Branch out_mc_tracks(mc_tracks_branch);
   AnalysisTree::Branch out_fhcal_mods(fhcal_branch);
 
   // Enable SetMutable so these output branches could be modified
+  out_reco_event.SetMutable();
   out_tpc_tracks.SetMutable();
   out_mc_tracks.SetMutable();
   out_fhcal_mods.SetMutable();
 
   // Enable Freeze so input branches cannot be modified
+  in_reco_event.Freeze();
   in_tpc_tracks.Freeze();
   in_mc_tracks.Freeze();
   in_fhcal_mods.Freeze();
 
   // Copy variables from the input - not required for branches with even-wise information (i.e. EventHeader)
+  out_reco_event.CloneVariables(in_reco_event.GetConfig());
   out_tpc_tracks.CloneVariables(in_tpc_tracks.GetConfig());
   out_mc_tracks.CloneVariables(in_mc_tracks.GetConfig());
   out_fhcal_mods.CloneVariables(in_fhcal_mods.GetConfig());
 
+  // Add extended variables to reco event
+  auto ireco_evt_bcent = out_reco_event.NewVariable("bcent", "centrality based on impact parameter (0%-100%)", Analysis::Types::kFloat);
+
   // Add extended variables to reco tracks
   auto itpc_eta_sign = out_tpc_tracks.NewVariable("eta_sign", "sign of track's pseudorapidity", AnalysisTree::Types::kInteger);
   auto itpc_mc_pdg = out_tpc_tracks.NewVariable("mc_pid", "PDG code from associated mc track", AnalysisTree::Types::kInteger);
@@ -118,10 +125,22 @@ void CreateExtendedAT(TString iFileName, TString oFileName)
   auto ifhcal_phi = out_fhcal_mods.NewVariable("phi", "Azimuthal angle of the center of FHCal module", AnalysisTree::Types::kFloat);
   auto ifhcal_signal_eta_signed = out_fhcal_mods.NewVariable("signal_eta_signed", "Energy*sign(eta) from the FHCal module", AnalysisTree::Types::kFloat);
 
-  std::cout << (short)in_tpc_tracks.GetBranchType() << std::endl;
-  in_tpc_tracks.GetConfig().Print();
-  std::cout << (short)out_tpc_tracks.GetBranchType() << std::endl;
-  out_tpc_tracks.GetConfig().Print();
+  out_config->AddBranchConfig(out_reco_event.GetConfig());
+  out_config->AddBranchConfig(out_tpc_tracks.GetConfig());
+  out_config->AddBranchConfig(out_mc_tracks.GetConfig());
+  out_config->AddBranchConfig(out_fhcal_mods.GetConfig());
+  
+  std::cout << std::endl;
+  std::cout << "Input base AnalysisTree configuration (aTree):" << std::endl;
+  in_config->Print();
+  std::cout << std::endl;
+  std::cout << "Output extended AnalysisTree configuration (bTree):" << std::endl;
+  out_config->Print();
+  std::cout << std::endl;
+
+
+  // Get links to variables from the input mc event header
+  auto imc_event_b = in_mc_event.GetField("B");
 
   // Get links to variables from the input reco tracks
   auto itpc_pt = in_tpc_tracks.GetField("pT");
@@ -139,13 +158,29 @@ void CreateExtendedAT(TString iFileName, TString oFileName)
   auto imc_pid       = in_mc_tracks.GetField("pid");
   auto imc_mother_id = in_mc_tracks.GetField("mother_id");
 
+  // Get links to variables from the input FHCal modules
+  auto ifhcal_signal = in_fhcal_mods.GetField("signal");
+  auto ifhcal_number = in_fhcal_mods.GetField("number");
+
   // Start loop over events
-  Int_t ntpc_tracks, nmc_tracks;
+  Int_t ntpc_tracks, nmc_tracks, nfhcal_mods;
   for (Long64_t ievent = 0; ievent < nevents; ievent++)
   {
     std::cout << "Event [" << ievent << "/" << nevents << "]" << std::endl;
     chain->GetEntry(ievent);
 
+    // Set extended reco event
+    out_reco_event.ClearChannels();
+
+    auto in_mc_event_var = in_mc_event[0];
+    auto in_reco_event_var = in_reco_event[0];
+    auto out_reco_event_var = out_reco_event.NewChannel();
+
+    out_reco_event_var.CopyContent(in_reco_event_var);
+
+    float bcent = GetBCent(in_mc_event_var.Value(imc_event_b));
+    out_reco_event_var.SetValue(ireco_evt_bcent, bcent);
+
     // Start loop over reco tracks
     ntpc_tracks = in_tpc_tracks.size();
     out_tpc_tracks.ClearChannels();
@@ -178,6 +213,36 @@ void CreateExtendedAT(TString iFileName, TString oFileName)
       }
 
     } // End loop over reco tracks
+
+    // Start loop over mc tracks
+    nmc_tracks = in_mc_tracks.size();
+    out_mc_tracks.ClearChannels();
+    for (int itrack = 0; itrack < nmc_tracks; itrack++)
+    {
+      auto in_mc_track = in_mc_tracks[itrack];
+      auto out_mc_track = out_mc_tracks.NewChannel();
+
+      bool is_charged = IsCharged((int)in_mc_track.Value(imc_pid));
+      out_mc_track.SetValue(imc_is_charged, is_charged);
+
+      int mc_eta_sign = (in_mc_track.Value(imc_eta) < 0) ? -1 : 1;
+      out_mc_track.SetValue(imc_eta_sign, mc_eta_sign);
+    } // End loop over mc tracks
+
+    // Start loop over FHCal modules
+    nfhcal_mods = in_fhcal_mods.size();
+    out_fhcal_mods.ClearChannels();
+    for (int imod = 0; imod < nfhcal_mods; imod++)
+    {
+      auto in_fhcal_mod = in_fhcal_mods[imod];
+      auto out_fhcal_mod = out_fhcal_mods.NewChannel();
+
+      double phi_mod = GetFHCalPhi((int)in_fhcal_mod.Value(ifhcal_number));
+      out_fhcal_mod.SetValue(ifhcal_phi, phi_mod);
+
+      float fhcal_sign = ((int)in_fhcal_mod.Value(ifhcal_number) < 45) ? -1. : 1.;
+      out_fhcal_mod.SetValue(ifhcal_signal_eta_signed, fhcal_sign*in_fhcal_mod.Value(ifhcal_signal));
+    } // End loop over FHCal modules
   }   // End loop over events
 }
 
@@ -187,6 +252,25 @@ float get_beamP(float sqrtSnn, float m_target, float m_beam)
   return sqrt(pow((pow(sqrtSnn, 2) - pow(m_target, 2) - pow(m_beam, 2)) / (2 * m_target), 2) - pow(m_beam, 2));
 }
 
+Float_t GetBCent(Float_t bimp)
+{
+  // Impact parameter values for each centrlity class
+  const std::vector<float> v_bimp = {0., 4.06, 5.84, 7.17, 8.27, 9.26, 10.17, 11.04, 11.99, 14.};
+  // Centrality class values
+  const std::vector<float>   v_cent = {0., 10., 20., 30., 40., 50., 60., 70., 80., 100.};
+
+  Float_t cent=-1;
+
+  for (int i = 0; i<v_bimp.size()-1; i++)
+  {
+    if (bimp >= v_bimp.at(i) && bimp < v_bimp.at(i+1))
+    {
+      cent = v_cent.at(i) + (v_cent.at(i+1) - v_cent.at(i))/2.;
+    }
+  }
+  return cent;
+}
+
 float GetFHCalPhi(int iModule)
 {
   const int Nmodules = 45;