Belle II Software  light-2212-foldex
TreeFitterModule.cc
1 /**************************************************************************
2  * basf2 (Belle II Analysis Software Framework) *
3  * Author: The Belle II Collaboration *
4  * *
5  * See git log for contributors and copyright holders. *
6  * This file is licensed under LGPL-3.0, see LICENSE.md. *
7  **************************************************************************/
8 
9 // Implementation of Decay Tree Fitter based on arXiv:physics/0503191
10 // Main module implementation
11 
12 #include <analysis/modules/TreeFitter/TreeFitterModule.h>
13 #include <analysis/VertexFitting/TreeFitter/FitManager.h>
14 
15 #include <framework/datastore/StoreArray.h>
16 #include <framework/particledb/EvtGenDatabasePDG.h>
17 #include <framework/database/DBObjPtr.h>
18 
19 #include <analysis/utility/ParticleCopy.h>
20 
21 #include <analysis/VertexFitting/TreeFitter/FitParameterDimensionException.h>
22 
23 using namespace Belle2;
24 
25 REG_MODULE(TreeFitter);
26 
27 TreeFitterModule::TreeFitterModule() : Module(), m_nCandidatesBeforeFit(-1), m_nCandidatesAfter(-1)
28 {
29  setDescription("Tree Fitter module. Performs simultaneous fit of all vertices in a decay chain. Can also be used to just fit a single vertex.");
31  //
32  addParam("particleList", m_particleList,
33  "Type::[string]. Input mother of the decay tree to fit. For example 'B0:myB0particleList'.");
34  addParam("confidenceLevel", m_confidenceLevel,
35  "Type::[double]. Confidence level to accept fitted decay tree. Candidates with < confidenceLevel will be removed from the particle list! Typical Values: -1: keep all particle candidates, 0: remove all that fail the fit, 0.001: standard cut, 0.1: (too) tight cut. Optimise using a figure of merit (for example S/(sqrt{S+B}) ) for your analysis.",
36  0.0);
37  addParam("convergencePrecision", m_precision,
38  "Type::[double]. Fractional upper limit for chi2 fluctuations to accept result. Larger value = less signal rejection but also less background rejection. Optimized for FOM on different topologies - don't touch unless you REALLY want this.",
39  0.01);
40  addParam("massConstraintList", m_massConstraintList,
41  "Type::[int]. List of particles to mass constrain with int = pdg code. Note that the variables 'M': fit result for the particle and 'InvM': calculated from the daughter momenta, will look different (especially if you don't update the daughters!).", {});
42  addParam("massConstraintListParticlename", m_massConstraintListParticlename,
43  "Type::[string]. List of particles to mass constrain with string = particle name.", {});
44 
45 
46  addParam("geoConstraintList", m_geoConstraintListPDG,
47  "Type::[int], if 'autoSetGeoConstraintAndMergeVertices==False' you can manually set the particles that will be geometrically constrained here.", {});
48  addParam("sharedVertexList", m_fixedToMotherVertexListPDG,
49  "Type::[int], if 'autoSetGeoConstraintAndMergeVertices==False' you can manually set the particles that share the vertex with their mother here.", {});
50  addParam("autoSetGeoConstraintAndMergeVertices", m_automatic_vertex_constraining,
51  "Type::bool, shall vertices of strong resonance be merged with their mothers? Can the particles vertex be constraint geometrically?",
52  true);
53 
54  addParam("customOriginVertex", m_customOriginVertex,
55  "Type::[double]. List of vertex coordinates to be used in the custom origin constraint.", {0.001, 0, 0.0116});
56  addParam("customOriginCovariance", m_customOriginCovariance,
57  "Type::[double]. List vertex covariance elements used in the custom origin constraint (as a vector). Default is meant for B0 decays and is taken from 100k generated B0 to mumu events.",
58  {
59  0.0048, 0, 0,
60  0, 0.003567, 0,
61  0, 0, 0.0400
62  }
63  );
64  addParam("customOriginConstraint", m_customOrigin,
65  "Type::[bool]. Use a custom vertex as the production point of the highest hierarchy particle (register this as the mother of the list you specify). Like the beam constraint but you can specify the position its covariance yourself. ",
66  false);
67  addParam("ipConstraint", m_ipConstraint,
68  "Type::[bool]. Use the IP as the origin of the tree. This registers an internal IP particle as the mother of the list you give. Or in other words forces the PRODUCTION vertex of your particle to be the IP and its covariance as specified in the database.",
69  false);
70  addParam("originDimension", m_originDimension,
71  "Type int, default 3. If origin or ip constraint used, specify the dimension of the constraint 3->x,y,z; 2->x,y. This also changes the dimension of the geometric constraints! So you might want to turn them off for some particles. (That means turn auto off and manually on for the ones you want to constrain)",
72  3);
73  addParam("updateAllDaughters", m_updateDaughters,
74  "Type::[bool]. Update all daughters (vertex position and momenta) in the tree. If not set only the 4-momenta for the head of the tree will be updated. We also update the vertex position of the daughters regardless of what you put here, because otherwise the default when the particle list is created is {0,0,0}.",
75  false);
76  addParam("expertBeamConstraintPDG", m_beamConstraintPDG,
77  "Type int, default 0. The 4-momentum of particles with the given PDG will be constrained to the 4-momentum of the initial e+e- system.",
78  0);
79  addParam("expertMassConstraintType", m_massConstraintType,
80  "Type::[int]. False(0): use particles parameters in mass constraint; True: use sum of daughter parameters for mass constraint. WAARNING not even guaranteed that it works.",
81  0);
82  addParam("expertRemoveConstraintList", m_removeConstraintList,
83  "Type::[string]. List of constraints that you do not want to be used in the fit. WARNING don't use if you don't know exactly what it does.", {});
84  addParam("expertUseReferencing", m_useReferencing,
85  "Type::[bool]. Use the Extended Kalman Filter. This implementation linearises around the previous state vector which gives smoother convergence.",
86  true);
87  addParam("inflationFactorCovZ", m_inflationFactorCovZ,
88  "Inflate the covariance of the beamspot by this number so that the 3d beam constraint becomes weaker in Z.And: thisnumber->infinity : dim(beamspot constr) 3d->2d.",
89  1);
90  addParam("treatAsInvisible", m_treatAsInvisible,
91  "Type::[string]. Decay string to select one particle that will be ignored in the fit.", {});
92 }
93 
95 {
96  m_plist.isRequired(m_particleList);
97  m_particles.isRequired();
100 
101  if ((m_massConstraintList.size()) == 0 && (m_massConstraintListParticlename.size()) > 0) {
102  for (auto& containedParticle : m_massConstraintListParticlename) {
103  TParticlePDG* particletemp = TDatabasePDG::Instance()->GetParticle((containedParticle).c_str());
104  m_massConstraintList.push_back(particletemp->PdgCode());
105  }
106  }
107 
108  if (!m_treatAsInvisible.empty()) {
110  if (!valid)
111  B2ERROR("TreeFitterModule::initialize Invalid Decay Descriptor: " << m_treatAsInvisible);
112  else if (m_pDDescriptorInvisibles.getSelectionPDGCodes().size() != 1)
113  B2ERROR("TreeFitterModule::initialize Please select exactly one particle to ignore: " << m_treatAsInvisible);
114  }
115 }
116 
118 {
119  if (!m_beamparams.isValid())
120  B2FATAL("BeamParameters are not available!");
121 
122  const ROOT::Math::PxPyPzEVector& her = m_beamparams->getHER();
123  const ROOT::Math::PxPyPzEVector& ler = m_beamparams->getLER();
124  const ROOT::Math::PxPyPzEVector& cms = her + ler;
125 
126  m_beamMomE(0) = cms.X();
127  m_beamMomE(1) = cms.Y();
128  m_beamMomE(2) = cms.Z();
129  m_beamMomE(3) = cms.E();
130 
131  const TMatrixDSym& HERcoma = m_beamparams->getCovHER();
132  const TMatrixDSym& LERcoma = m_beamparams->getCovLER();
133 
134  m_beamCovariance = Eigen::Matrix4d::Zero();
135  const double covE = (HERcoma(0, 0) + LERcoma(0, 0));
136  for (size_t i = 0; i < 4; ++i) {
137  m_beamCovariance(i, i) =
138  covE;
139  // TODO Currently, we do not get a full covariance matrix from beamparams, and the py value is zero, which means there is no constraint on py. Therefore, we approximate it by a diagonal matrix using the energy value for all components. This is based on the assumption that the components of the beam four-momentum are independent and of comparable size.
140  }
141 }
142 
144 {
145  if (!m_plist) {
146  B2ERROR("ParticleList " << m_particleList << " not found");
147  return;
148  }
149 
150  std::vector<unsigned int> toRemove;
151  const unsigned int nParticles = m_plist->getListSize();
152  m_nCandidatesBeforeFit += nParticles;
153 
154  TMatrixFSym dummyCovMatrix(7);
155  for (int row = 0; row < 7; ++row) { //diag
156  dummyCovMatrix(row, row) = 10000;
157  }
158 
159  for (unsigned iPart = 0; iPart < nParticles; iPart++) {
160  Belle2::Particle* particle = m_plist->getParticle(iPart);
161 
162  if (m_updateDaughters == true) {
163  ParticleCopy::copyDaughters(particle);
164  }
165 
166  if (!m_treatAsInvisible.empty()) {
167  std::vector<const Particle*> selParticlesTarget = m_pDDescriptorInvisibles.getSelectionParticles(particle);
168  Particle* targetD = m_particles[selParticlesTarget[0]->getArrayIndex()];
169  Particle* daughterCopy = Belle2::ParticleCopy::copyParticle(targetD);
170  daughterCopy->writeExtraInfo("treeFitterTreatMeAsInvisible", 1);
171  daughterCopy->setMomentumVertexErrorMatrix(dummyCovMatrix);
172  bool isReplaced = particle->replaceDaughterRecursively(targetD, daughterCopy);
173  if (!isReplaced)
174  B2ERROR("TreeFitterModule::event No target particle found for " << m_treatAsInvisible);
175  }
176 
177  try {
178  const bool ok = fitTree(particle);
179  if (!ok) { particle->setPValue(-1); }
180  } catch (TreeFitter::FitParameterDimensionException const& e) {
181  B2ERROR(e.what());
182  }
183 
184  if (particle->getPValue() < m_confidenceLevel) {
185  toRemove.push_back(particle->getArrayIndex());
186  }
187 
188  }
189  m_plist->removeParticles(toRemove);
190  m_nCandidatesAfter += m_plist->getListSize();
191 }
192 
193 
195 {
196  if (m_nCandidatesAfter > 0) {
197  plotFancyASCII();
198  } else {
199  B2WARNING("Not a single candidate survived the fit. Candidates before fit: " << m_nCandidatesBeforeFit << " after: " <<
201  }
202 }
203 
205 {
206  const TreeFitter::ConstraintConfiguration constrConfig(
219  m_beamMomE,
222  );
223 
224  std::unique_ptr<TreeFitter::FitManager> TreeFitter(
226  head,
227  constrConfig,
228  m_precision,
231  )
232  );
233  bool rc = TreeFitter->fit();
234  return rc;
235 }
236 
238 {
239  B2INFO("\033[1;35m================================================================================\033[0m");
240  B2INFO("\033[40;97m ,., \033[0m");
241  B2INFO("\033[40;97m ;%&M%;_ ,.., \033[0m");
242  B2INFO("\033[40;97m \"_ \"__\" % M % M %; , ..., , \033[0m");
243  B2INFO("\033[40;97m , ..., __.\" --\" , ., _ - \" % &W % WM %; \033[0m");
244  B2INFO("\033[40;97m ; % M&$M % \"___ \"_._ %M%\"_.\"\" _ \"\"\"\"\"\" \033[0m");
245  B2INFO("\033[40;97m \"\"\"\"\" \"\" , \\_. \"_. .\" \033[0m");
246  B2INFO("\033[40;97m , ., _\"__ \\__./ .\" \033[0m");
247  B2INFO("\033[40;97m ___ __ | y , .., \033[97;40mThank you for using TreeFitter. \033[0m");
248  B2INFO("\033[40;97m /)'\\ ''''| u \\ %W%W%%; \033[0m");
249  B2INFO("\033[40;97m ___)/ \"---\\_ \\ |____\" \033[97;40mPlease cite both: \033[0m");
250  B2INFO("\033[40;97m ;&&%%; (|__.|)./ ,.., \033[97;40m10.1016/j.nima.2020.164269 \033[0m");
251  B2INFO("\033[40;97m ,.., ___\\ |/ &&\" \033[97;40m 10.1016/j.nima.2005.06.078 \033[0m");
252  B2INFO("\033[40;97m &&%%& (| Uo / '\" \033[97;40mEmail: \033[0m");
253  B2INFO("\033[40;97m '''' \\ 7 \\ \033[97;40mfrancesco.tenchini@desy.de \033[0m");
254  B2INFO("\033[40;97m ._______________.-'____\"\"-.____. \033[97;40mjo-frederik.krohn@desy.de \033[0m");
255  B2INFO("\033[40;97m \\ / \033[0m");
256  B2INFO("\033[40;97m \\ \033[0m\033[32;40mTREEFITTER\033[0m\033[40;97m / \033[0m");
257  B2INFO("\033[40;97m \\_______________________/ \033[0m");
258  B2INFO("\033[40;97m (_) (_) \033[0m");
259  B2INFO("\033[40;97m \033[0m");
260  B2INFO("\033[1;35m============= TREEFIT STATISTICS ===============================================\033[0m");
261  B2INFO("\033[1;39mTarget particle list: " << m_particleList << "\033[0m");
262  B2INFO("\033[1;39mCandidates before fit: " << m_nCandidatesBeforeFit << "\033[0m");
263  B2INFO("\033[1;39mCandidates after fit: " << m_nCandidatesAfter << "\033[0m");
264  B2INFO("\033[1;39mA total of " << m_nCandidatesBeforeFit - m_nCandidatesAfter <<
265  " candidates were removed during the fit.\033[0m");
266  B2INFO("\033[1;39m" << (double)m_nCandidatesAfter / (double)m_nCandidatesBeforeFit * 100.0 <<
267  "% of candidates survived the fit.\033[0m");
268  B2INFO("\033[1;39m" << 100. - (double)m_nCandidatesAfter / (double)m_nCandidatesBeforeFit * 100.0 <<
269  "% of candidates did not.\033[0m");
270  B2INFO("\033[1;39mYou chose to drop all candidates with pValue < " << m_confidenceLevel << ".\033[0m");
271  B2INFO("\033[1;35m================================================================================\033[0m");
272 }
bool init(const std::string &str)
Initialise the DecayDescriptor from given string.
std::vector< int > getSelectionPDGCodes()
Return list of PDG codes of selected particles.
std::vector< const Particle * > getSelectionParticles(const Particle *particle)
Get a vector of pointers with selected daughters in the decay tree.
Base class for Modules.
Definition: Module.h:72
void setDescription(const std::string &description)
Sets the description of the module.
Definition: Module.cc:214
void setPropertyFlags(unsigned int propertyFlags)
Sets the flags for the module properties.
Definition: Module.cc:208
@ c_ParallelProcessingCertified
This module can be run in parallel processing mode safely (All I/O must be done through the data stor...
Definition: Module.h:80
Class to store reconstructed particles.
Definition: Particle.h:74
void writeExtraInfo(const std::string &name, const double value)
Sets the user defined extraInfo.
Definition: Particle.cc:1335
double getPValue() const
Returns chi^2 probability of fit if done or -1.
Definition: Particle.h:640
void setMomentumVertexErrorMatrix(const TMatrixFSym &errMatrix)
Sets 7x7 error matrix.
Definition: Particle.cc:425
void setPValue(double pValue)
Sets chi^2 probability of fit.
Definition: Particle.h:353
bool replaceDaughterRecursively(const Particle *oldDaughter, const Particle *newDaughter)
Apply replaceDaughter to all Particles in the decay tree by looping recursively through it,...
Definition: Particle.cc:753
int getArrayIndex() const
Returns this object's array index (in StoreArray), or -1 if not found.
void plotFancyASCII()
plot ascii art and statistics
std::vector< double > m_customOriginVertex
vertex coordinates of the custom origin
DBObjPtr< BeamParameters > m_beamparams
Pointer to the BeamParameters.
virtual void initialize() override
initialize
virtual void event() override
performed for each event
Eigen::Matrix< double, 4, 1 > m_beamMomE
beam four-momentum
bool m_updateDaughters
flag if you want to update all particle momenta in the decay tree.
DecayDescriptor m_pDDescriptorInvisibles
Decay descriptor of the invisible particles.
std::vector< std::string > m_massConstraintListParticlename
vector carrying the names of the particles to be mass constraint
unsigned int m_nCandidatesAfter
after the fit
std::vector< int > m_massConstraintList
vector carrying the PDG codes of the particles to be mass constraint
std::vector< int > m_fixedToMotherVertexListPDG
list of pdg codes of particles where we use the same vertex for production and decay which is the ver...
StoreArray< Particle > m_particles
StoreArray of Particles.
virtual void terminate() override
stuff at the end
bool m_automatic_vertex_constraining
should the vertex be joined with the mother and should it be geometrically constrained?...
virtual void beginRun() override
performed at the start of run
bool m_ipConstraint
Use x-y-z beamspot constraint.
int m_originDimension
dimension to use for beam/origin constraint
int m_inflationFactorCovZ
inflate beamspot covariance of z by this number
std::vector< std::string > m_removeConstraintList
list of constraints not to be applied in tree fit WARNING only use if you know what you are doing
int m_massConstraintType
type of the mass constraint false: use normal one.
std::vector< int > m_geoConstraintListPDG
list of pdg codes of particles to use a geo constraint for
std::vector< double > m_customOriginCovariance
covariance of the custom origin
std::string m_particleList
name of the particle list fed to the fitter
double m_confidenceLevel
minimum confidence level to accept fit calculated as f(chiSquared, NDF) -2: accept all 0: only accept...
bool m_customOrigin
use a custom vertex as the production vertex of the highest hierarchy particle
std::string m_treatAsInvisible
decay string to select one particle that will be treated as invisible
Eigen::Matrix< double, 4, 4 > m_beamCovariance
beam covariance matrix
int m_beamConstraintPDG
PDG code of particle to be constrained to the beam 4-momentum.
bool m_useReferencing
linearise around a previous state of the Kalman Filter
bool fitTree(Particle *head)
this fits all particle candidates contained in the m_particleList
double m_precision
convergence precision for the newton method When the delta chiSquared between 2 iterations divided by...
StoreObjPtr< ParticleList > m_plist
input particle list
unsigned int m_nCandidatesBeforeFit
before the fit
exception template, runtime_error implements what()
REG_MODULE(B2BIIConvertBeamParams)
Register the module.
void addParam(const std::string &name, T &paramVariable, const std::string &description, const T &defaultValue)
Adds a new parameter to the module.
Definition: Module.h:560
void copyDaughters(Particle *mother)
Function copies all (grand-)^n-daughter particles of the argument mother Particle.
Definition: ParticleCopy.cc:55
Particle * copyParticle(const Particle *original)
Function takes argument Particle and creates a copy of it and copies of all its (grand-)^n-daughters.
Definition: ParticleCopy.cc:18
Abstract base class for different kinds of events.
Definition: ClusterUtils.h:23