PWG1/AliKFParticleTest.h

   1 //---------------------------------------------------------------------------------
   2 // The example of usage of AliKFParticle & AliKFVertex classes for V0 analysis
   3 // .
   4 // @author  S.Gorbunov, I.Kisel
   5 // @version 1.0
   6 // @since   13.05.07
   7 //
   8 // The AliKFParticleTest macro contains a toy V0 finder for ESD tracks.
   9 // At the first step, event primary vertex is reconstructed.
  10 // At the second step, ideal PID hypothesis are assigned to all the particles.
  11 // At the third step, V0 candidates are constructed for each pair
  12 // of positive and negative particles.
  13 // V0 candidate considered as good when it passes Chi^2 cut and
  14 // it is placed >= 3 Sigma away from the primary vertex.
  15 // Invariant mass distribution for all good V0 candidates is plotted.
  16 //
  17 //  -= Copyright &copy ALICE HLT Group =-
  18 //_________________________________________________________________________________
  19
  20
  21 #ifndef ALIKFPARTICLETEST_H
  22 #define ALIKFPARTICLETEST_H
  23
  24 #include "TH1.h"
  25 #include "TCanvas.h"
  26 #include "AliESD.h"
  27 #include "AliAnalysisTaskRL.h"
  28
  29 class AliKFParticleTest: public AliAnalysisTaskRL {
  30  public:
  31   AliKFParticleTest(const char *name);
  32   virtual ~AliKFParticleTest() {}
  33
  34   virtual void   ConnectInputData(Option_t *);
  35   virtual void   CreateOutputObjects();
  36   virtual void   Exec(Option_t *option);
  37   virtual void   Terminate(Option_t *);
  38
  39  private:
  40   AliESD *fESD; //ESD object
  41   TH1D   *fHistoMassAll;
  42   TH1D   *fHistoMassSignal;
  43   TCanvas *fCanvas;
  44   void DrawV0();
  45
  46   TObjArray * fOutputContainer; // ! output data container
  47
  48   ClassDef(AliKFParticleTest, 0); // example of analysis
  49 };
  50
  51
  52 #include "Riostream.h"
  53 #include "TChain.h"
  54 #include "TSystem.h"
  55 #include "TROOT.h"
  56 #include "TParticle.h"
  57 #include "TRandom.h"
  58 #include "AliESD.h"
  59 #include "AliLog.h"
  60 #include "AliStack.h"
  61 #include "AliTracker.h"
  62 #include "TLatex.h"
  63 #include "TDatabasePDG.h"
  64 #include "AliKFParticle.h"
  65 #include "AliKFVertex.h"
  66
  67
  68 ClassImp(AliKFParticleTest);
  69
  70
  71 //________________________________________________________________________
  72 AliKFParticleTest::AliKFParticleTest(const char *name) :AliAnalysisTaskRL(name,""), fESD(0) {
  73   // Constructor.
  74   // Input slot #0 works with an Ntuple
  75   DefineInput(0, TChain::Class());
  76  // Output slot #0 writes into a TH1 container
  77   DefineOutput(0, TObjArray::Class());
  78 }
  79
  80 //___________________________________________________________________________
  81 void AliKFParticleTest::ConnectInputData(Option_t *) {
  82   // Initialize branches.
  83   printf("   ConnectInputData of task %s\n", GetName());
  84   if (!fESD) {
  85     char ** address = (char **)GetBranchAddress(0, "ESD");
  86     if (address) fESD = (AliESD*)(*address);
  87     if (!fESD) {
  88       fESD = new AliESD();
  89       SetBranchAddress(0, "ESD", &fESD);
  90     }
  91   }
  92
  93 }
  94
  95 //___________________________________________________________________________
  96 void AliKFParticleTest::CreateOutputObjects() {
  97   printf("   CreateOutputObjects of task %s\n", GetName());
  98
  99   if ( !gROOT->IsBatch() ) fCanvas = new TCanvas();
 100   else fCanvas = 0;
 101
 102   fHistoMassAll = new TH1D("massAll","AliKFParticle test", 500,0,3);
 103   fHistoMassAll->SetXTitle("V^{0} invariant mass [GeV]");
 104   fHistoMassAll->SetLineColor(kGreen);
 105   fHistoMassAll->SetFillColor(kGreen);
 106
 107   fHistoMassSignal = new TH1D("massSignal","V^{0} signal", 500,0,3);
 108   fHistoMassSignal->SetXTitle("V^{0} invariant mass [GeV]");
 109   fHistoMassSignal->SetLineColor(kBlue);
 110   fHistoMassSignal->SetFillColor(kBlue);
 111
 112   fOutputContainer = new TObjArray(1);
 113   fOutputContainer->SetName(GetName()) ;
 114   fOutputContainer->Add(fHistoMassAll);
 115   fOutputContainer->Add(fHistoMassSignal);
 116 }
 117
 118
 119 void AliKFParticleTest::DrawV0()
 120 {
 121   //* Draw the invariant mass histogram
 122
 123   if ( gROOT->IsBatch() ) return;
 124
 125   const Int_t histoPDG [4]= {22,310,3122,421};
 126   const Char_t* histoName[4]= {"#gamma","K^{0}_{s}","#Lambda","D^{0}"};
 127   TLatex histoText[4];
 128
 129   if( !fCanvas ) return;
 130   fCanvas->Clear();
 131   fCanvas->cd();
 132
 133   for( Int_t i=0; i<4; i++ ){
 134     histoText[i].SetTextColor(kBlue);
 135   }
 136
 137   fHistoMassAll->Draw();
 138   fHistoMassSignal->Draw("same");
 139
 140   for( Int_t i=0; i<4; i++ ){
 141     Double_t mass = TDatabasePDG::Instance()->GetParticle(histoPDG[i])->Mass();
 142     Int_t bin = fHistoMassSignal->FindBin(mass) -5;
 143     if( bin<0 ) bin =0;
 144     Double_t max = 0;
 145     for( Int_t j=bin; j<bin+10; j++ )
 146     if( max<fHistoMassSignal->GetBinContent(j) ) max = fHistoMassSignal->GetBinContent(j);
 147     if(max>0) histoText[i].DrawLatex( mass+.05, max, histoName[i] );
 148   }
 149   fCanvas->Update();
 150 }
 151
 152 //________________________________________________________________________
 153 void AliKFParticleTest::Exec(Option_t *) {
 154
 155   static Int_t iEvent = 0;
 156   if( ++iEvent%100 ==0 ) cout<<"Event "<<iEvent<<endl;
 157
 158   // Get input data
 159   TTree *tinput = (TTree*)GetInputData(0);
 160   Long64_t ientry = tinput->GetReadEntry();
 161   if (AliAnalysisTaskRL::GetEntry(ientry) == kFALSE) {
 162     printf("Couldn't get event from the runLoader\n");
 163     return;
 164   }
 165   if (!fESD) return;
 166
 167   Double_t Bz = fESD->GetMagneticField();
 168   AliKFParticle::SetField( Bz );
 169
 170   if (ientry==0) Notify ();
 171
 172  //cout <<"Event "<<ientry<<endl;
 173
 174   AliStack* stack = GetStack();
 175   if (!stack) {
 176     AliDebug(AliLog::kError, "Stack not available");
 177     return;
 178   }
 179
 180
 181   class TESDTrackInfo
 182   {
 183   public:
 184     TESDTrackInfo(){}
 185     AliKFParticle fParticle; //* assigned KFParticle
 186     Bool_t fPrimUsedFlag;    //* flag shows that the particle was used for primary vertex fit
 187     Bool_t fOK;              //* is the track good enough
 188     Int_t mcPDG;             //* Monte Carlo PDG code
 189     Int_t mcMotherID;        //* Monte Carlo ID of its mother
 190   };
 191
 192   Int_t nESDTracks=fESD->GetNumberOfTracks();
 193   if( nESDTracks>1000 ) nESDTracks=1000;
 194
 195   TESDTrackInfo ESDTrackInfo[1000]; //* parallel to ESD tracks
 196
 197   //* Fill ESDTrackInfo array
 198
 199   for (Int_t iTr=0; iTr<nESDTracks; iTr++)
 200     {
 201       TESDTrackInfo &info = ESDTrackInfo[iTr];
 202       info.fOK = 0;
 203       info.fPrimUsedFlag = 0;
 204       info.mcPDG = -1;
 205       info.mcMotherID = -1;
 206
 207       //* track quality check
 208
 209       AliESDtrack *pTrack = fESD->GetTrack(iTr);
 210       if( !pTrack  ) continue;
 211       if (pTrack->GetKinkIndex(0)>0) continue;
 212       if ( !( pTrack->GetStatus()&AliESDtrack::kITSrefit ) ) continue;
 213       Int_t indi[12];
 214       if( pTrack->GetITSclusters(indi) <5 ) continue;
 215       Int_t PDG = ( pTrack->Get1Pt() <0 ) ?321 :211;
 216
 217       //* take MC PDG
 218       {
 219         Int_t mcID = TMath::Abs(pTrack->GetLabel());
 220         TParticle * part = stack->Particle(TMath::Abs(mcID));
 221         if( part ){
 222           info.mcPDG = part->GetPdgCode();
 223           PDG = info.mcPDG;
 224           if( mcID>=0 ) info.mcMotherID = part->GetFirstMother();
 225         }
 226       }
 227
 228       //* Construct KFParticle for the track
 229
 230       info.fParticle = AliKFParticle( *pTrack, PDG );
 231       info.fOK = 1;
 232     }
 233
 234   //* Find event primary vertex
 235
 236   AliKFVertex primVtx;
 237   {
 238     const AliKFParticle * vSelected[1000]; //* Selected particles for vertex fit
 239     Int_t vIndex[1000];                    //* Indices of selected particles
 240     Bool_t vFlag[1000];                    //* Flags returned by the vertex finder
 241
 242     Int_t nSelected = 0;
 243     for( Int_t i = 0; i<nESDTracks; i++){
 244       if(ESDTrackInfo[i].fOK ){
 245         vSelected[nSelected] = &(ESDTrackInfo[i].fParticle);
 246         vIndex[nSelected] = i;
 247         nSelected++;
 248       }
 249     }
 250     primVtx.ConstructPrimaryVertex( vSelected, nSelected, vFlag, 3. );
 251     for( Int_t i = 0; i<nSelected; i++){
 252       if( vFlag[i] ) ESDTrackInfo[vIndex[i]].fPrimUsedFlag = 1;
 253     }
 254     if( primVtx.GetNDF() <1 ) return; //* Less then two tracks in primary vertex
 255   }
 256
 257   //* V0 finder
 258
 259   for( Int_t iTr = 0; iTr<nESDTracks; iTr++ ){ //* first daughter
 260
 261     TESDTrackInfo &info = ESDTrackInfo[iTr];
 262     if( !info.fOK ) continue;
 263
 264     for( Int_t jTr = iTr+1; jTr<nESDTracks; jTr++ ){  //* second daughter
 265       TESDTrackInfo &jnfo = ESDTrackInfo[jTr];
 266       if( !jnfo.fOK ) continue;
 267
 268       //* check for different charge
 269
 270       if( info.fParticle.GetQ() == jnfo.fParticle.GetQ() ) continue;
 271
 272       //* construct V0 mother
 273
 274       AliKFParticle V0( info.fParticle, jnfo.fParticle );
 275
 276       //* check V0 Chi^2
 277
 278       if( V0.GetNDF()<1 ) continue;
 279       if( TMath::Sqrt(TMath::Abs(V0.GetChi2()/V0.GetNDF())) >3. ) continue;
 280
 281       //* subtruct daughters from primary vertex
 282
 283       AliKFVertex primVtxCopy = primVtx;
 284
 285       if( info.fPrimUsedFlag ) primVtxCopy -= info.fParticle;
 286       if( jnfo.fPrimUsedFlag ) primVtxCopy -= jnfo.fParticle;
 287
 288       //* Check V0 Chi^2 deviation from primary vertex
 289
 290       if( V0.GetDeviationFromVertex( primVtxCopy ) >3. ) continue;
 291
 292       //* Add V0 to primary vertex to improve the primary vertex resolution
 293
 294       primVtxCopy += V0;
 295
 296       //* Set production vertex for V0
 297
 298       V0.SetProductionVertex( primVtxCopy );
 299
 300       //* Check chi^2 for a case
 301
 302       if( TMath::Sqrt( TMath::Abs(V0.GetChi2()/V0.GetNDF()) >3. )) continue;
 303
 304       //* Get V0 decay length with estimated error
 305
 306       Double_t length, sigmaLength;
 307       if( V0.GetDecayLength( length, sigmaLength ) ) continue;
 308
 309       //* Reject V0 if it decays too close to the primary vertex
 310
 311       if( length  <3.*sigmaLength ) continue;
 312
 313       //* Get V0 invariant mass and plot it
 314
 315       Double_t mass, sigmaMass;
 316       if( V0.GetMass( mass, sigmaMass ) ) continue;
 317       fHistoMassAll->Fill(mass);
 318
 319       //* Fill signal histograms using MC information
 320
 321       if( info.mcMotherID==jnfo.mcMotherID && info.mcMotherID>=0 ){
 322         TParticle *mother = stack->Particle(info.mcMotherID);
 323         if( mother && TMath::Abs(mother->GetPdgCode())!=21 ){
 324           if( mother->GetNDaughters()==2 ){
 325             fHistoMassSignal->Fill(mass);
 326           }
 327           cout<<"PDG V0,pi,pj, ndaughters, mc mass, reco mass = "<<mother->GetPdgCode()<<","<<info.mcPDG<<","<<jnfo.mcPDG<<", "
 328               << mother->GetNDaughters()<<", "<<mother->GetMass()<<", "<<mass<<endl;
 329         }
 330       }
 331     }
 332   }
 333   if( iEvent %1000 == 0 || (iEvent ==200)) DrawV0();
 334
 335   // Post final data. It will be written to a file with option "RECREATE"
 336   PostData(0, fOutputContainer);
 337 }
 338
 339
 340 //________________________________________________________________________
 341 void AliKFParticleTest::Terminate(Option_t *)
 342 {
 343   DrawV0();
 344   fOutputContainer = (TObjArray*)GetOutputData(0);
 345   //fHistoMassAll=(TH1D*)fOutputContainer->At(0) ;
 346 }
 347
 348 #endif