Belle II Software  release-05-02-19
PIDNtupleModule.cc
1 /**************************************************************************
2  * BASF2 (Belle Analysis Framework 2) *
3  * Copyright(C) 2010 - Belle II Collaboration *
4  * *
5  * Author: The Belle II Collaboration *
6  * Contributors: Marko Staric, Jake Bennett *
7  * *
8  * This software is provided "as is" without any warranty. *
9  **************************************************************************/
10 
11 // Own include
12 #include <reconstruction/modules/PIDNtuple/PIDNtupleModule.h>
13 
14 // framework - DataStore
15 #include <framework/datastore/StoreArray.h>
16 #include <framework/datastore/StoreObjPtr.h>
17 
18 // framework aux
19 #include <framework/gearbox/Const.h>
20 #include <framework/logging/Logger.h>
21 
22 // DataStore classes
23 #include <framework/dataobjects/EventMetaData.h>
24 #include <mdst/dataobjects/Track.h>
25 #include <mdst/dataobjects/MCParticle.h>
26 #include <mdst/dataobjects/PIDLikelihood.h>
27 
28 // ROOT
29 #include <TRandom3.h>
30 
31 using namespace Belle2;
32 
33 //-----------------------------------------------------------------
34 // Register the Module
35 //-----------------------------------------------------------------
36 REG_MODULE(PIDNtuple)
37 
38 //-----------------------------------------------------------------
39 // Implementation
40 //-----------------------------------------------------------------
41 
42 PIDNtupleModule::PIDNtupleModule() : Module(),
43  m_norm(0), m_value(0), m_file(0), m_tree(0)
44 
45 {
46  // set module description (e.g. insert text)
47  setDescription("Writes a flat ntuple of PIDLikelihoods with track info into a root file");
48  setPropertyFlags(c_ParallelProcessingCertified);
49 
50  // Add parameters
51  addParam("outputFileName", m_outputFileName, "Output file name",
52  std::string("PIDNtuple.root"));
53  addParam("makeFlat", m_makeFlat, "make momentum distribution flat up to pMax", false);
54  addParam("p1", m_p1, "parameter of momentum distribution", 0.90311E-01);
55  addParam("p2", m_p2, "parameter of momentum distribution", 0.56846);
56  addParam("pMax", m_pMax, "make distribution flat up to this momentum", 3.0);
57 }
58 
59 PIDNtupleModule::~PIDNtupleModule() { }
60 
61 void PIDNtupleModule::initialize()
62 {
63 
64  // required inputs
65  StoreArray<Track> tracks;
66  StoreArray<TrackFitResult> trackfitResults;
67  StoreArray<PIDLikelihood> PIDLikelihoods;
68 
69  tracks.isRequired();
70  trackfitResults.isRequired();
71  PIDLikelihoods.isRequired();
72 
73  // optional inputs
74  StoreArray<MCParticle> mcparticles;
75  mcparticles.isOptional();
76  tracks.optionalRelationTo(mcparticles);
77 
78  // Create and book the ROOT file and TTree
79  m_file = new TFile(m_outputFileName.c_str(), "RECREATE");
80  m_tree = new TTree("pid", "PID tree");
81 
82  m_tree->Branch("evt", &m_pid.evt, "evt/I");
83  m_tree->Branch("run", &m_pid.run, "run/I");
84  m_tree->Branch("p", &m_pid.p, "p/F");
85  m_tree->Branch("cth", &m_pid.cth, "cth/F");
86  m_tree->Branch("phi", &m_pid.phi, "phi/F");
87  m_tree->Branch("pValue", &m_pid.pValue, "pValue/F");
88  m_tree->Branch("PDG", &m_pid.PDG, "PDG/I");
89  m_tree->Branch("motherPDG", &m_pid.motherPDG, "motherPDG/I");
90  m_tree->Branch("primary", &m_pid.primary, "primary/S");
91  m_tree->Branch("rhoProd", &m_pid.rhoProd, "rhoProd/F");
92  m_tree->Branch("zProd", &m_pid.zProd, "zProd/F");
93  m_tree->Branch("phiProd", &m_pid.phiProd, "phiProd/F");
94  m_tree->Branch("rhoDec", &m_pid.rhoDec, "rhoDec/F");
95  m_tree->Branch("zDec", &m_pid.zDec, "zDec/F");
96  m_tree->Branch("phiDec", &m_pid.phiDec, "phiDec/F");
97  m_tree->Branch("cdcdedx", &m_pid.cdcdedx, "le/F:lmu:lpi:lk:lp:flag/S:seen/S");
98  m_tree->Branch("svddedx", &m_pid.svddedx, "le/F:lmu:lpi:lk:lp:flag/S:seen/S");
99  m_tree->Branch("top", &m_pid.top, "le/F:lmu:lpi:lk:lp:flag/S:seen/S");
100  m_tree->Branch("arich", &m_pid.arich, "le/F:lmu:lpi:lk:lp:flag/S:seen/S");
101  m_tree->Branch("ecl", &m_pid.ecl, "le/F:lmu:lpi:lk:lp:flag/S:seen/S");
102  m_tree->Branch("klm", &m_pid.klm, "le/F:lmu:lpi:lk:lp:flag/S:seen/S");
103 
104  m_norm = 1;
105  double p0 = m_p1 * log1p(4.0 * m_p2 / m_p1); // distribution peak
106  m_norm = 1.0 / momDistribution(p0);
107  m_value = momDistribution(m_pMax);
108 }
109 
110 void PIDNtupleModule::beginRun()
111 {
112 }
113 
114 void PIDNtupleModule::event()
115 {
116 
117  StoreObjPtr<EventMetaData> evtMetaData;
118  StoreArray<Track> tracks;
119 
120  // loop over tracks
121  for (const auto& track : tracks) {
122 
123  const TrackFitResult* trackFit = track.getTrackFitResultWithClosestMass(Const::pion);
124  if (!trackFit) {
125  B2WARNING("No track fit result... Skipping.");
126  continue;
127  }
128 
129  const PIDLikelihood* pid = track.getRelated<PIDLikelihood>();
130  if (!pid) {
131  B2WARNING("No track fit result... Skipping.");
132  continue;
133  }
134 
135  const MCParticle* mcParticle = track.getRelated<MCParticle>();
136  const MCParticle* mother = 0;
137  if (mcParticle) mother = mcParticle->getMother();
138 
139  m_pid.clear();
140  m_pid.evt = evtMetaData->getEvent();
141  m_pid.run = evtMetaData->getRun();
142  TVector3 momentum = trackFit->getMomentum();
143  m_pid.p = momentum.Mag();
144  if (m_makeFlat) {
145  if (gRandom->Rndm() * momDistribution(m_pid.p) > m_value) continue;
146  }
147  m_pid.cth = momentum.CosTheta();
148  m_pid.phi = momentum.Phi();
149  m_pid.pValue = trackFit->getPValue();
150  if (mcParticle) {
151  m_pid.PDG = mcParticle->getPDG();
152  if (mother) m_pid.motherPDG = mother->getPDG();
153  m_pid.primary = mcParticle->getStatus(MCParticle::c_PrimaryParticle);
154  TVector3 prodVertex = mcParticle->getProductionVertex();
155  m_pid.rhoProd = prodVertex.Perp();
156  m_pid.zProd = prodVertex.Z();
157  m_pid.phiProd = prodVertex.Phi();
158  TVector3 decVertex = mcParticle->getDecayVertex();
159  m_pid.rhoDec = decVertex.Perp();
160  m_pid.zDec = decVertex.Z();
161  m_pid.phiDec = decVertex.Phi();
162  }
163 
164  m_pid.cdcdedx.le = pid->getLogL(Const::electron, Const::CDC);
165  m_pid.cdcdedx.lmu = pid->getLogL(Const::muon, Const::CDC);
166  m_pid.cdcdedx.lpi = pid->getLogL(Const::pion, Const::CDC);
167  m_pid.cdcdedx.lk = pid->getLogL(Const::kaon, Const::CDC);
168  m_pid.cdcdedx.lp = pid->getLogL(Const::proton, Const::CDC);
169  m_pid.cdcdedx.flag = pid->isAvailable(Const::CDC);
170  if (mcParticle)
171  m_pid.cdcdedx.seen = mcParticle->hasSeenInDetector(Const::CDC);
172 
173  m_pid.svddedx.le = pid->getLogL(Const::electron, Const::SVD);
174  m_pid.svddedx.lmu = pid->getLogL(Const::muon, Const::SVD);
175  m_pid.svddedx.lpi = pid->getLogL(Const::pion, Const::SVD);
176  m_pid.svddedx.lk = pid->getLogL(Const::kaon, Const::SVD);
177  m_pid.svddedx.lp = pid->getLogL(Const::proton, Const::SVD);
178  m_pid.svddedx.flag = pid->isAvailable(Const::SVD);
179  if (mcParticle)
180  m_pid.svddedx.seen = mcParticle->hasSeenInDetector(Const::SVD);
181 
182  m_pid.top.le = pid->getLogL(Const::electron, Const::TOP);
183  m_pid.top.lmu = pid->getLogL(Const::muon, Const::TOP);
184  m_pid.top.lpi = pid->getLogL(Const::pion, Const::TOP);
185  m_pid.top.lk = pid->getLogL(Const::kaon, Const::TOP);
186  m_pid.top.lp = pid->getLogL(Const::proton, Const::TOP);
187  m_pid.top.flag = pid->isAvailable(Const::TOP);
188  if (mcParticle)
189  m_pid.top.seen = mcParticle->hasSeenInDetector(Const::TOP);
190 
191  m_pid.arich.le = pid->getLogL(Const::electron, Const::ARICH);
192  m_pid.arich.lmu = pid->getLogL(Const::muon, Const::ARICH);
193  m_pid.arich.lpi = pid->getLogL(Const::pion, Const::ARICH);
194  m_pid.arich.lk = pid->getLogL(Const::kaon, Const::ARICH);
195  m_pid.arich.lp = pid->getLogL(Const::proton, Const::ARICH);
196  m_pid.arich.flag = pid->isAvailable(Const::ARICH);
197  if (mcParticle)
198  m_pid.arich.seen = mcParticle->hasSeenInDetector(Const::ARICH);
199 
200  m_pid.ecl.le = pid->getLogL(Const::electron, Const::ECL);
201  m_pid.ecl.lmu = pid->getLogL(Const::muon, Const::ECL);
202  m_pid.ecl.lpi = pid->getLogL(Const::pion, Const::ECL);
203  m_pid.ecl.lk = pid->getLogL(Const::kaon, Const::ECL);
204  m_pid.ecl.lp = pid->getLogL(Const::proton, Const::ECL);
205  m_pid.ecl.flag = pid->isAvailable(Const::ECL);
206  if (mcParticle)
207  m_pid.ecl.seen = mcParticle->hasSeenInDetector(Const::ECL);
208 
209  m_pid.klm.le = pid->getLogL(Const::electron, Const::KLM);
210  m_pid.klm.lmu = pid->getLogL(Const::muon, Const::KLM);
211  m_pid.klm.lpi = pid->getLogL(Const::pion, Const::KLM);
212  m_pid.klm.lk = pid->getLogL(Const::kaon, Const::KLM);
213  m_pid.klm.lp = pid->getLogL(Const::proton, Const::KLM);
214  m_pid.klm.flag = pid->isAvailable(Const::KLM);
215  if (mcParticle)
216  m_pid.klm.seen = mcParticle->hasSeenInDetector(Const::KLM);
217 
218  m_tree->Fill();
219  }
220 
221 }
222 
223 
224 void PIDNtupleModule::endRun() { }
225 
226 void PIDNtupleModule::terminate()
227 {
228  m_file->cd();
229  m_tree->Write();
230  m_file->Close();
231 }
232 
233 void PIDNtupleModule::printModuleParams() const { }
Belle2::PID::LogLikelihoods::le
Float_t le
log likelihood for electron
Definition: PIDTree.h:35
Belle2::PIDNtupleModule::beginRun
virtual void beginRun() override
Called when entering a new run.
Definition: PIDNtupleModule.cc:110
Belle2::PID::PIDTree::klm
LogLikelihoods klm
log likelihoods from KLM
Definition: PIDTree.h:91
Belle2::PID::PIDTree::cdcdedx
LogLikelihoods cdcdedx
log likelihoods from CDC dE/dx
Definition: PIDTree.h:86
Belle2::PID::PIDTree::rhoDec
Float_t rhoDec
decay vertex (cylindrical coordinate r) of MCParticle
Definition: PIDTree.h:82
Belle2::PIDNtupleModule
makes PID ntuple from mdst input
Definition: PIDNtupleModule.h:36
Belle2::PID::LogLikelihoods::seen
Short_t seen
is seen in this component (from related MCParticle)
Definition: PIDTree.h:41
Belle2::PIDNtupleModule::m_outputFileName
std::string m_outputFileName
output file name
Definition: PIDNtupleModule.h:86
Belle2::PID::LogLikelihoods::lmu
Float_t lmu
log likelihood for muon
Definition: PIDTree.h:36
Belle2::TrackFitResult::getMomentum
TVector3 getMomentum() const
Getter for vector of momentum at closest approach of track in r/phi projection.
Definition: TrackFitResult.h:116
Belle2::PIDNtupleModule::m_pid
PID::PIDTree m_pid
PID tree structure.
Definition: PIDNtupleModule.h:97
Belle2::Const::electron
static const ChargedStable electron
electron particle
Definition: Const.h:533
REG_MODULE
#define REG_MODULE(moduleName)
Register the given module (without 'Module' suffix) with the framework.
Definition: Module.h:652
Belle2::PID::PIDTree::ecl
LogLikelihoods ecl
log likelihoods from ECL
Definition: PIDTree.h:90
Belle2::PID::LogLikelihoods::lp
Float_t lp
log likelihood for proton
Definition: PIDTree.h:39
Belle2::PID::PIDTree::clear
void clear()
Clear the structure: set elements to zero.
Definition: PIDTree.h:103
Belle2::PID::PIDTree::rhoProd
Float_t rhoProd
production vertex (cylindrical coordinate r) of MCParticle
Definition: PIDTree.h:79
Belle2::TrackFitResult::getPValue
double getPValue() const
Getter for Chi2 Probability of the track fit.
Definition: TrackFitResult.h:163
Belle2::PID::PIDTree::zProd
Float_t zProd
production vertex (cylindrical coordinate z) of MCParticle
Definition: PIDTree.h:80
Belle2::PIDNtupleModule::m_value
double m_value
distribution value at m_pMax
Definition: PIDNtupleModule.h:93
Belle2::PID::PIDTree::arich
LogLikelihoods arich
log likelihoods from ARICH
Definition: PIDTree.h:89
Belle2::MCParticle::hasSeenInDetector
bool hasSeenInDetector(Const::DetectorSet set) const
Return if the seen-in flag for a specific subdetector is set or not.
Definition: MCParticle.h:318
Belle2::PID::PIDTree::phiDec
Float_t phiDec
decay vertex (cylindrical coordinate phi) of MCParticle
Definition: PIDTree.h:84
Belle2::PIDLikelihood
Class to collect log likelihoods from TOP, ARICH, dEdx, ECL and KLM aimed for output to mdst includes...
Definition: PIDLikelihood.h:37
Belle2::PID::PIDTree::primary
Short_t primary
is a primary particle (from related MCParticle)
Definition: PIDTree.h:78
Belle2::PIDNtupleModule::~PIDNtupleModule
virtual ~PIDNtupleModule()
Destructor.
Definition: PIDNtupleModule.cc:59
Belle2::PID::PIDTree::pValue
Float_t pValue
p-value of Track fit
Definition: PIDTree.h:74
Belle2::TrackFitResult
Values of the result of a track fit with a given particle hypothesis.
Definition: TrackFitResult.h:59
Belle2::PID::PIDTree::phiProd
Float_t phiProd
production vertex (cylindrical coordinate phi) of MCParticle
Definition: PIDTree.h:81
Belle2::PIDNtupleModule::terminate
virtual void terminate() override
Termination action.
Definition: PIDNtupleModule.cc:226
Belle2::Const::kaon
static const ChargedStable kaon
charged kaon particle
Definition: Const.h:536
Belle2::MCParticle::getDecayVertex
TVector3 getDecayVertex() const
Return decay vertex.
Definition: MCParticle.h:230
Belle2::Module
Base class for Modules.
Definition: Module.h:74
Belle2::PID::PIDTree::motherPDG
Int_t motherPDG
PDG code of related mother MCParticle.
Definition: PIDTree.h:77
Belle2::Const::pion
static const ChargedStable pion
charged pion particle
Definition: Const.h:535
Belle2::MCParticle::getStatus
unsigned int getStatus(unsigned short int bitmask=USHRT_MAX) const
Return status code of particle.
Definition: MCParticle.h:133
Belle2::PID::LogLikelihoods::flag
Short_t flag
flag: information is available (1) or not (0)
Definition: PIDTree.h:40
Belle2::PID::PIDTree::top
LogLikelihoods top
log likelihoods from TOP
Definition: PIDTree.h:88
Belle2::PID::LogLikelihoods::lk
Float_t lk
log likelihood for kaon
Definition: PIDTree.h:38
Belle2
Abstract base class for different kinds of events.
Definition: MillepedeAlgorithm.h:19
Belle2::StoreObjPtr
Type-safe access to single objects in the data store.
Definition: ParticleList.h:33
Belle2::PID::PIDTree::p
Float_t p
momentum magnitude of Track
Definition: PIDTree.h:71
Belle2::PIDNtupleModule::m_p2
double m_p2
parameter of momentum distribution
Definition: PIDNtupleModule.h:89
Belle2::MCParticle::getProductionVertex
TVector3 getProductionVertex() const
Return production vertex position.
Definition: MCParticle.h:200
Belle2::MCParticle::getPDG
int getPDG() const
Return PDG code of particle.
Definition: MCParticle.h:123
Belle2::PIDNtupleModule::m_norm
double m_norm
distribution normalization
Definition: PIDNtupleModule.h:92
Belle2::PIDNtupleModule::endRun
virtual void endRun() override
End-of-run action.
Definition: PIDNtupleModule.cc:224
Belle2::MCParticle::getMother
MCParticle * getMother() const
Returns a pointer to the mother particle.
Definition: MCParticle.h:593
Belle2::PID::PIDTree::evt
Int_t evt
event number
Definition: PIDTree.h:68
Belle2::PIDNtupleModule::momDistribution
double momDistribution(double p) const
parameterized momentum distribution
Definition: PIDNtupleModule.h:105
Belle2::PIDNtupleModule::m_makeFlat
bool m_makeFlat
if true, make momentum distribution flat up to m_pMax
Definition: PIDNtupleModule.h:87
Belle2::PIDNtupleModule::initialize
virtual void initialize() override
Initialize the Module.
Definition: PIDNtupleModule.cc:61
Belle2::PID::PIDTree::run
Int_t run
run number
Definition: PIDTree.h:69
Belle2::Const::proton
static const ChargedStable proton
proton particle
Definition: Const.h:537
Belle2::PIDNtupleModule::printModuleParams
void printModuleParams() const
Prints module parameters.
Definition: PIDNtupleModule.cc:233
Belle2::PIDNtupleModule::m_p1
double m_p1
parameter of momentum distribution
Definition: PIDNtupleModule.h:88
Belle2::PID::PIDTree::phi
Float_t phi
azimuthal angle of Track
Definition: PIDTree.h:73
Belle2::PID::LogLikelihoods::lpi
Float_t lpi
log likelihood for pion
Definition: PIDTree.h:37
Belle2::Const::muon
static const ChargedStable muon
muon particle
Definition: Const.h:534
Belle2::MCParticle
A Class to store the Monte Carlo particle information.
Definition: MCParticle.h:43
Belle2::PID::PIDTree::svddedx
LogLikelihoods svddedx
log likelihoods from SVD dE/dx
Definition: PIDTree.h:87
Belle2::StoreArray
Accessor to arrays stored in the data store.
Definition: ECLMatchingPerformanceExpertModule.h:33
Belle2::PIDNtupleModule::m_tree
TTree * m_tree
TTree with PIDTree structure.
Definition: PIDNtupleModule.h:96
Belle2::PID::PIDTree::PDG
Int_t PDG
PDG code of related MCParticle.
Definition: PIDTree.h:76
Belle2::PIDNtupleModule::m_pMax
double m_pMax
flatten distribution up to this momentum
Definition: PIDNtupleModule.h:90
Belle2::PID::PIDTree::cth
Float_t cth
cosine of polar angle of Track
Definition: PIDTree.h:72
Belle2::PID::PIDTree::zDec
Float_t zDec
decay vertex (cylindrical coordinate z) of MCParticle
Definition: PIDTree.h:83
Belle2::MCParticle::c_PrimaryParticle
@ c_PrimaryParticle
bit 0: Particle is primary particle.
Definition: MCParticle.h:58
Belle2::PIDNtupleModule::m_file
TFile * m_file
TFile.
Definition: PIDNtupleModule.h:95
Belle2::PIDNtupleModule::event
virtual void event() override
Event processor.
Definition: PIDNtupleModule.cc:114