123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188 |
- #include "UserTaskWrite.hpp"
- #include "TaskManager.hpp"
- using namespace AnalysisTree;
- void UserTaskWrite::Init() { // Preparations before reading events
- // Read input chain with AnalysisTree
- auto man = TaskManager::GetInstance();
- auto chain = man->GetChain();
- AddInputBranch("RecoEvent");
- reco_evt_ = chain->GetBranch("RecoEvent");
- reco_evt_.Freeze();
- AddInputBranch("McEvent");
- mc_evt_ = chain->GetBranch("McEvent");
- mc_evt_.Freeze();
- AddInputBranch("TpcTracks");
- tracks_ = chain->GetBranch("TpcTracks");
- tracks_.Freeze();
- AddInputBranch("McTracks");
- particles_ = chain->GetBranch("McTracks");
- particles_.Freeze();
- match_ = chain->GetMatching("TpcTracks", "McTracks");
- // Set up branch configurations (copy them from the input and then modify)
- auto br_conf_reco_evt = chain->GetConfiguration()->GetBranchConfig("RecoEvent");
- auto br_conf_mc_evt = chain->GetConfiguration()->GetBranchConfig("McEvent");
- auto new_conf_reco_evt = br_conf_reco_evt.Clone("RecoEventExt", DetType::kEventHeader);
- new_conf_reco_evt.AddFields<float>({"bcent"}, "Extended parameter");
- auto br_conf_tracks = chain->GetConfiguration()->GetBranchConfig("TpcTracks");
- auto new_conf_tracks = br_conf_tracks.Clone("TpcTracksExt", DetType::kTrack);
- new_conf_tracks.AddFields<float>({"mc_E", "mc_pT", "mc_eta", "mc_phi", "mc_y", "rapidity"}, "Extended parameter");
- new_conf_tracks.AddFields<int>({"eta_sign", "mc_pid", "mc_mother_id"}, "Extended parameter");
- auto br_conf_particles = chain->GetConfiguration()->GetBranchConfig("McTracks");
- auto new_conf_particles = br_conf_particles.Clone("McTracksExt", DetType::kParticle);
- new_conf_particles.AddFields<bool>({"is_charged"}, "Extended parameter");
- new_conf_particles.AddFields<int>({"eta_sign"}, "Extended parameter");
- // Prepare output branches and then add them to the output file via task manager
- new_reco_evt_ = Branch(new_conf_reco_evt);
- new_reco_evt_.SetMutable();
- new_tracks_ = Branch(new_conf_tracks);
- new_tracks_.SetMutable();
- new_particles_ = Branch(new_conf_particles);
- new_particles_.SetMutable();
- man->AddBranch(&new_reco_evt_);
- man->AddBranch(&new_tracks_);
- man->AddBranch(&new_particles_);
- }
- void UserTaskWrite::Exec() { // Procedures to do iside loop over events
- // Get links to variables from the input mc event header and output reco event header
- auto imc_event_b = mc_evt_.GetField("B");
- auto ireco_evt_bcent = new_reco_evt_.GetField("bcent");
- // Get links to variables from the input and output reco tracks
- auto itpc_pt = tracks_.GetField("pT");
- auto itpc_p = tracks_.GetField("p");
- auto itpc_pz = tracks_.GetField("pz");
- auto itpc_eta = tracks_.GetField("eta");
- auto itpc_track_mc_E = new_tracks_.GetField("mc_E");
- auto itpc_track_mc_pT = new_tracks_.GetField("mc_pT");
- auto itpc_track_mc_eta = new_tracks_.GetField("mc_eta");
- auto itpc_track_mc_phi = new_tracks_.GetField("mc_phi");
- auto itpc_track_mc_y = new_tracks_.GetField("mc_y");
- auto itpc_track_rapidity = new_tracks_.GetField("rapidity");
- auto itpc_track_eta_sign = new_tracks_.GetField("eta_sign");
- auto itpc_track_mc_pid = new_tracks_.GetField("mc_pid");
- auto itpc_track_mc_mother_id = new_tracks_.GetField("mc_mother_id");
- // Get links to variables from the input and output mc tracks
- auto imc_mass = particles_.GetField("mass");
- auto imc_p = particles_.GetField("p");
- auto imc_pT = particles_.GetField("pT");
- auto imc_eta = particles_.GetField("eta");
- auto imc_phi = particles_.GetField("phi");
- auto imc_rapidity = particles_.GetField("rapidity");
- auto imc_pid = particles_.GetField("pid");
- auto imc_mother_id = particles_.GetField("mother_id");
- auto imc_particle_is_charged = new_particles_.GetField("is_charged");
- auto imc_particle_eta_sign = new_particles_.GetField("eta_sign");
- // Processing reco & mc event
- for(size_t i=0; i<1; ++i){
- auto mc_evt_var = mc_evt_[i];
- auto new_reco_evt_var = new_reco_evt_[i];
- float bcent = GetBCent(mc_evt_var.Value(imc_event_b));
- new_reco_evt_var.SetValue(ireco_evt_bcent, bcent);
- }
- // Processing reco tracks
- new_tracks_.ClearChannels();
- for(size_t i=0; i<tracks_.size(); ++i){ // Loop over reco tracks
- auto track = tracks_[i];
- auto match_id = match_->GetMatch(track.GetId());
- auto new_track = new_tracks_.NewChannel();
- new_track.CopyContent(track);
- int tpc_eta_sign = (track.Value(itpc_eta) < 0) ? -1 : 1;
- new_track.SetValue(itpc_track_eta_sign, tpc_eta_sign);
- // Fill extended fields with info from associated mc tracks
- if(match_id >= 0){
- auto particle = particles_[match_id];
- new_track.SetValue(itpc_track_mc_pid, particle.Value(imc_pid));
- new_track.SetValue(itpc_track_mc_pT, particle.Value(imc_pT));
- new_track.SetValue(itpc_track_mc_eta, particle.Value(imc_eta));
- new_track.SetValue(itpc_track_mc_phi, particle.Value(imc_phi));
- new_track.SetValue(itpc_track_mc_y, particle.Value(imc_rapidity));
- new_track.SetValue(itpc_track_mc_mother_id, particle.Value(imc_mother_id));
- double Ene = sqrt(pow(particle.Value(imc_p),2) + pow(particle.Value(imc_mass),2));
- new_track.SetValue(itpc_track_mc_E, Ene);
- double rapidityPDG = GetRapidityPDG((float)track.Value(itpc_p), (float)track.Value(itpc_pz), (int)particle.Value(imc_pid));
- new_track.SetValue(itpc_track_rapidity, rapidityPDG);
- }
- }
- // Processing mc tracks
- new_particles_.ClearChannels();
- for(size_t i=0; i<particles_.size(); ++i){ // Loop over mc tracks
- auto particle = particles_[i];
- auto new_particle = new_particles_.NewChannel();
- new_particle.CopyContent(particle);
- int mc_eta_sign = (particle.Value(itpc_eta) < 0) ? -1 : 1;
- bool is_charged = IsCharged(particle.Value(imc_pid));
- new_particle.SetValue(imc_particle_eta_sign, mc_eta_sign);
- new_particle.SetValue(imc_particle_is_charged, is_charged);
- }
- }
- UserTaskWrite::~UserTaskWrite() { // Destructor. Make sure to delete branch pointers for each output branch
- delete new_reco_evt_ptr_;
- delete new_particles_ptr_;
- delete new_tracks_ptr_;
- }
- bool UserTaskWrite::IsCharged(Int_t pdg) { // Additional function returning bool value wheather the particle is charged or not
- auto particle = (TParticlePDG *)TDatabasePDG::Instance()->GetParticle(pdg);
- if (!particle)
- return false;
- return (particle->Charge() != 0);
- }
- float UserTaskWrite::GetRapidityPDG(Float_t p, Float_t pz, Int_t pdg) { // Additional function returning rapidity value based on particle's pdg code
- Float_t mass;
- auto pdgParticle = (TParticlePDG *)TDatabasePDG::Instance()->GetParticle(pdg);
- if (!pdgParticle)
- return -999.;
- mass = pdgParticle->Mass();
- Float_t E = TMath::Sqrt(p * p + mass * mass);
- return 0.5 * TMath::Log((E + pz) / (E - pz));
- }
- float UserTaskWrite::GetBCent(Float_t bimp) { // Additional function returning centrality class based on fixed impact parameter ranges
- // 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;
- }
|