I improved (by a factor 2.5) the speed of the MatchToMC method

[u/mrichter/AliRoot.git] / STEER / AliAODRecoDecay.cxx

diff --git a/STEER/AliAODRecoDecay.cxx b/STEER/AliAODRecoDecay.cxx
index fb93e84c38948e0d1c3adf34281d4c0362bc609c..0c5c4fc17fb1d75f451eb9c57485577be4e9776c 100644 (file)
--- a/STEER/AliAODRecoDecay.cxx
+++ b/STEER/AliAODRecoDecay.cxx
@@ -22,8 +22,10 @@
 
 #include <TDatabasePDG.h>
 #include <TVector3.h>
+#include <TClonesArray.h>
 #include "AliLog.h"
 #include "AliVTrack.h"
+#include "AliAODMCParticle.h"
 #include "AliAODRecoDecay.h"
 
 ClassImp(AliAODRecoDecay)
@@ -472,6 +474,126 @@ Double_t AliAODRecoDecay::InvMass2Prongs(Int_t ip1,Int_t ip2,
 
   return mass;
 }
+//----------------------------------------------------------------------------
+Int_t AliAODRecoDecay::MatchToMC(Int_t pdgabs,TClonesArray *mcArray) const
+{
+  //
+  // Check if this candidate is matched to a MC signal
+  // If no, return -1
+  // If yes, return label (>=0) of the AliAODMCParticle
+  // 
+
+  Int_t ndg=GetNDaughters();
+  if(!ndg) {
+    AliError("No daughters available");
+    return -1;
+  }
+  if(ndg>10) {
+    AliError("Only decays with <10 daughters supported");
+    return -1;
+  }
+  
+  Int_t dgLabels[10];
+
+  // loop on daughters and write the labels
+  for(Int_t i=0; i<ndg; i++) {
+    AliAODTrack *trk = (AliAODTrack*)GetDaughter(i);
+    dgLabels[i] = trk->GetLabel();
+  }
+
+  return MatchToMC(pdgabs,mcArray,dgLabels,ndg);
+}
+//----------------------------------------------------------------------------
+Int_t AliAODRecoDecay::MatchToMC(Int_t pdgabs,TClonesArray *mcArray,
+                                Int_t dgLabels[10],Int_t ndg) const
+{
+  //
+  // Check if this candidate is matched to a MC signal
+  // If no, return -1
+  // If yes, return label (>=0) of the AliAODMCParticle
+  // 
+
+  Int_t labMom[10];
+  Int_t i,lab,labMother,pdgMother;
+  AliAODMCParticle *part=0;
+  AliAODMCParticle *mother=0;
+  Double_t pxSumDgs=0.,pySumDgs=0.,pzSumDgs=0.;
+
+  // loop on daughter labels
+  for(i=0; i<ndg; i++) {
+    labMom[i]=-1;
+    lab = dgLabels[i];
+    if(lab<0) {
+      printf("daughter with negative label %d\n",lab);
+      return -1;
+    }
+    part = (AliAODMCParticle*)mcArray->At(lab);
+    if(!part) { 
+      printf("no MC particle\n");
+      return -1;
+    }
+
+    // for the J/psi, check that the daughters are electrons
+    if(pdgabs==443) {
+      if(TMath::Abs(part->GetPdgCode())!=11) return -1;
+    }
+
+    mother = part;
+    while(mother->GetMother()>=0) {
+      labMother=mother->GetMother();
+      mother = (AliAODMCParticle*)mcArray->At(labMother);
+      if(!mother) {
+       printf("no MC mother particle\n");
+       break;
+      }
+      pdgMother = TMath::Abs(mother->GetPdgCode());
+      if(pdgMother==pdgabs) {
+       labMom[i]=labMother;
+       // keep sum of daughters' momenta, to check for mom conservation
+       pxSumDgs += part->Px();
+       pySumDgs += part->Py();
+       pzSumDgs += part->Pz();
+       break;
+      } else if(pdgMother>pdgabs || pdgMother<10) {
+       break;
+      }
+    }
+    if(labMom[i]==-1) return -1; // mother PDG not ok for this daughter
+  } // end loop on daughters
+
+  // check if the candidate is signal
+  labMother=labMom[0];
+  // all labels have to be the same and !=-1
+  for(i=0; i<ndg; i++) {
+    if(labMom[i]==-1)        return -1;
+    if(labMom[i]!=labMother) return -1;
+  }
+
+  
+  /*
+  // check that the mother decayed in <GetNDaughters()> prongs
+  Int_t ndg2 = TMath::Abs(mother->GetDaughter(1)-mother->GetDaughter(0))+1;
+  printf("  MC daughters %d\n",ndg2);
+  //if(ndg!=GetNDaughters()) return -1;
+  AliAODMCParticle* p1=(AliAODMCParticle*)(mcArray->At(mother->GetDaughter(1)));
+  AliAODMCParticle* p0=(AliAODMCParticle*)(mcArray->At(mother->GetDaughter(0)));
+  printf(">>>>>>>> pdg %d  %d %d %d   %d %d\n",mother->GetDaughter(0),mother->GetDaughter(1),dgLabels[0],dgLabels[1],p0->GetPdgCode(),p1->GetPdgCode());
+  */
+
+  // the above works only for non-resonant decays,
+  // it's better to check for mom conservation
+  mother = (AliAODMCParticle*)mcArray->At(labMother);
+  Double_t pxMother = mother->Px();
+  Double_t pyMother = mother->Py();
+  Double_t pzMother = mother->Pz();
+  // within 0.1%
+  if((TMath::Abs(pxMother-pxSumDgs)/(TMath::Abs(pxMother)+1.e-13)) > 0.001 &&
+     (TMath::Abs(pyMother-pySumDgs)/(TMath::Abs(pyMother)+1.e-13)) > 0.001 &&
+     (TMath::Abs(pzMother-pzSumDgs)/(TMath::Abs(pzMother)+1.e-13)) > 0.001) 
+    return -1;
+ 
+  return labMother;
+}
 //---------------------------------------------------------------------------
 void AliAODRecoDecay::Print(Option_t* /*option*/) const 
 {