New test macros for the analysis framework
[u/mrichter/AliRoot.git] / ANALYSIS / testEvent.C
1 // #####    TO RUN THIS MACRO:
2 // bash$ aliroot         (due to AliInfo, AliError, ...)
3 // root[0] gSystem->Load("libANALYSIS_NEW");
4 // IN case you do not have the include path set to AliRoot includes:
5 // root[1] gSystem->AddIncludePath("-I\"$ALICE_ROOT/include\"");
6 // root[2] .L testEvent.C+;
7 // root[3] generate();
8 // root[4] filter_reco();
9
10 #include "TClonesArray.h"
11 #include "TChain.h"
12 #include "TH1.h"
13 #include "TCanvas.h"
14 #include "testEvent.h"   
15
16 //============= First step: generate events
17 void generate()
18 {
19 // Simple event generation
20    AliAnalysisManager *mgr = new AliAnalysisManager();
21    TaskGenerate *task = new TaskGenerate("gener");
22    mgr->AddTask(task);
23
24    if (mgr->InitAnalysis()) {
25       mgr->PrintStatus();
26       mgr->ExecAnalysis();
27    }   
28    delete mgr;
29 }   
30
31 //============= Second step: filter gammas; use TSelector functionality
32 void filter_reco()
33 {
34 // Filter the input events having more than 100 gammas. Reconstruct pi0's
35 // From gammas coming from the same vertex.
36    // Get the input data as a chain
37    TChain *chain = new TChain("T");
38    chain->Add("event02000.root");
39    chain->Add("event04000.root");
40    chain->Add("event06000.root");
41    chain->Add("event08000.root");
42    chain->Add("event10000.root");
43    // Create an analysis manager
44    AliAnalysisManager *mgr = new AliAnalysisManager();
45    // Create a filter task and register it
46    TaskFilter *task1 = new TaskFilter("TaskFilter");
47    mgr->AddTask(task1);
48    // Create a reco task and register it
49    TaskRecoPi0 *task2 = new TaskRecoPi0("TaskRecoPi0");
50    mgr->AddTask(task2);
51    // Create containers for input/output
52    AliAnalysisDataContainer *cinput = mgr->CreateContainer("input0", 
53                       TTree::Class(), AliAnalysisManager::kInputContainer);
54    AliAnalysisDataContainer *coutput1 = mgr->CreateContainer("output0", 
55                       TTree::Class(), AliAnalysisManager::kOutputContainer);
56    AliAnalysisDataContainer *coutput2 = mgr->CreateContainer("output1", 
57                       TList::Class(), AliAnalysisManager::kOutputContainer);
58    AliAnalysisDataContainer *coutput = mgr->CreateContainer("output", 
59                       TH1::Class(), AliAnalysisManager::kOutputContainer);
60    // Connect containers to the task input/outputs
61    mgr->ConnectInput(task1,0,cinput);
62    mgr->ConnectOutput(task1,0,coutput1);
63    mgr->ConnectOutput(task1,1,coutput2);
64    mgr->ConnectInput(task2,0,coutput1);
65    mgr->ConnectOutput(task2,0,coutput);
66    // Connect input data
67    cinput->SetData(chain);
68    // Init analysis and start event loop
69    if (mgr->InitAnalysis()) {
70       mgr->PrintStatus();
71       chain->Process(mgr);
72    }
73 //   delete mgr;   
74 }
75    
76 ClassImp(TaskGenerate)
77
78 //______________________________________________________________________________
79 void TaskGenerate::Exec(Option_t *)
80 {
81 // Generate 10k events in 5 files
82    AnaEvent::CreateEvents(2000, "event02000.root");
83    AnaEvent::CreateEvents(2000, "event04000.root");
84    AnaEvent::CreateEvents(2000, "event06000.root");
85    AnaEvent::CreateEvents(2000, "event08000.root");
86    AnaEvent::CreateEvents(2000, "event10000.root");
87 }
88
89 ClassImp(TaskFilter)
90
91 //______________________________________________________________________________
92 TaskFilter::TaskFilter(const char *name) 
93            :AliAnalysisTask(name,""), fEvent(0), fOutput(0), fList(0), fHist1(0), fHist2(0)
94 {
95 // Constructor
96    // Input slot #0 works with a tree of events
97    DefineInput(0, TTree::Class());
98
99    // Output slot #0 writes into a tree, #1 produces a hisogram
100    DefineOutput(0, TTree::Class());
101    DefineOutput(1, TList::Class());
102 }
103
104 //______________________________________________________________________________
105 void TaskFilter::Init(Option_t *)
106 {
107 // Initialize branches.
108    printf("   Init %s\n", GetName());
109    if (!fEvent) {
110       // One should first check if the branch address was taken by some other task
111       char ** address = (char **)GetBranchAddress(0, "event");
112       if (address) fEvent = (AnaEvent*)(*address);
113       if (!fEvent) {
114          fEvent = new AnaEvent();
115          SetBranchAddress(0, "event", &fEvent);
116       }
117       // The output tree will be written to gammas.root
118       TDirectory *dirsav = gDirectory;
119       // Open a file for output #0
120       OpenFile(0, "gammas.root", "RECREATE");
121       fOutput = new TTree("TGAM", "gammas");
122       TBranch *branch = fOutput->Branch("event", &fEvent, 18000,1);
123       branch->SetAutoDelete(kFALSE);
124       if (dirsav) dirsav->cd();
125    } 
126    if (!fList) {
127       fList = new TList();
128       fHist1 = new TH1I("ntracks", "Number of tracks per event", 100, 0, 1000);
129       fHist1->SetLineColor(kRed);
130       fHist2 = new TH1I("ngammas", "Number of gammas per event", 100, 0, 1000);
131       fHist2->SetLineColor(kBlue);
132       fList->Add(fHist1);
133       fList->Add(fHist2);
134    }   
135 }
136
137 //______________________________________________________________________________
138 void TaskFilter::Exec(Option_t *)
139 {
140 // Filtering.
141    TTree *tinput = (TTree*)GetInputData(0);
142    Long64_t ientry = tinput->GetReadEntry();
143    // In this case fEvent address is already connected to the input
144    if (!fEvent) return;
145    // First check multiplicity
146    Int_t ntracks = fEvent->GetNtracks();
147    // Loop tracks and get rid of non-gammas
148    AnaTrack *track;
149    Int_t igamma = 0;
150    for (Int_t i=0; i<ntracks; i++) {
151       track = fEvent->GetTrack(i);
152       if (track->GetMass() < 1.e-3) igamma++;
153    }
154    if (ientry%100 == 0) printf("TaskFilter -> Event %lld: %i tracks filtered %i gammas\n", ientry, ntracks, igamma);
155    fHist1->Fill(ntracks);
156    fHist2->Fill(igamma);
157    if (igamma > 100) {
158       fOutput->Fill();
159       PostData(0, fOutput);
160    }   
161    PostData(1, fList);
162 }
163
164 //______________________________________________________________________________
165 void TaskFilter::Terminate(Option_t *)
166 {
167 // Draw some histogram at the end.
168    if (!gROOT->IsBatch()) {
169       fHist1->SetMaximum(2500);
170       fHist1->Draw();
171       fHist2->Draw("SAME");
172    }   
173 }
174
175 ClassImp(TaskRecoPi0)
176
177 //______________________________________________________________________________
178 TaskRecoPi0::TaskRecoPi0(const char *name) 
179            :AliAnalysisTask(name,""), fEvent(0), fGammas(0), fPions(0), fHist(0)
180 {
181 // Constructor
182    // Input slot #0 works with a tree of events
183    DefineInput(0, TTree::Class());
184
185    // Output slot #1 produces a hisogram
186    DefineOutput(0, TH1::Class());
187 }
188
189 //______________________________________________________________________________
190 TaskRecoPi0::~TaskRecoPi0()
191 {
192 // Dtor.
193    if (fEvent) delete fEvent;
194    if (fGammas) delete fGammas;
195    if (fPions) delete fPions;
196 }   
197
198 //______________________________________________________________________________
199 void TaskRecoPi0::Init(Option_t *)
200 {
201 // Initialize branches.
202    printf("   Init %s\n", GetName());
203    if (!fEvent) {
204       // One should first check if the branch address was taken by some other task
205       char ** address = (char **)GetBranchAddress(0, "event");
206       if (address) fEvent = (AnaEvent*)(*address);
207       if (!fEvent) {
208          fEvent = new AnaEvent();
209          SetBranchAddress(0, "event", &fEvent);
210       }
211       fGammas = new TObjArray();
212       fPions  = new TObjArray();
213    } 
214    if (!fHist) {
215       fHist = new TH1F("Pt_pi0", "Pt distribution for pi0's", 100, 0., 10.);
216       fHist->SetLineColor(kRed);
217    }   
218 }
219
220 //______________________________________________________________________________
221 void TaskRecoPi0::Exec(Option_t *)
222 {
223 // Reconstruct Pi0's for one event
224    AnaTrack *track = 0;
225    Int_t ntracks = fEvent->GetNtracks();
226    Int_t ngamma = 0;
227    Int_t i,j;
228    // Clear containers
229    fGammas->Clear();
230    fPions->Delete();
231    fPions->Clear();
232    // Loop tracks and move gammas to gamma container
233    for (i=0; i<ntracks; i++) {
234       track = fEvent->GetTrack(i);
235       if (track->GetMass() < 1.e-3) {
236          fGammas->Add(track);
237          ngamma++;
238       }
239    }
240    printf("TaskRecoPi0 -> Tracks %i \n", ntracks);
241
242    // Loop gammas and check vertex position
243    Double_t v1[3], v2[3];
244    AnaTrack *tracko = 0;
245    Double_t cutoff = 0.001;
246    for (i=0; i<ngamma-1; i++) {
247       track = (AnaTrack*)fGammas->At(i);
248       v1[0] = track->GetVertex(0);
249       v1[1] = track->GetVertex(1);
250       v1[2] = track->GetVertex(2);
251       for (j=i+1; j<ngamma; j++) {
252          tracko = (AnaTrack*)fGammas->At(j);
253          v2[0] = tracko->GetVertex(0);
254          v2[1] = tracko->GetVertex(1);
255          v2[2] = tracko->GetVertex(2);
256          Double_t dist2 = (v2[0]-v1[0])*(v2[0]-v1[0])+
257                           (v2[1]-v1[1])*(v2[1]-v1[1])+
258                           (v2[2]-v1[2])*(v2[2]-v1[2]);
259          if (dist2>cutoff*cutoff) continue;
260          // We have found a pair candidate
261          Double_t px = track->GetPx()+tracko->GetPx();
262          Double_t py = track->GetPy()+tracko->GetPy();
263          Double_t pz = track->GetPz()+tracko->GetPz();
264          track = new AnaTrack(px,py,pz,0.135,0,v1[0],v1[1],v1[2]);
265          fPions->Add(track);
266          fHist->Fill(track->GetPt());
267          break;
268       }   
269    }   
270    PostData(0,fHist);         
271 }
272
273 //______________________________________________________________________________
274 void TaskRecoPi0::Terminate(Option_t *)
275 {
276 // Draw some histogram at the end.
277    if (!gROOT->IsBatch()) {
278       new TCanvas("pi0", "Pt for pi0's", 800,600);
279       fHist->Draw();
280    }   
281 }
282
283
284 ClassImp(AnaTrack)
285
286 //______________________________________________________________________________
287 AnaTrack::AnaTrack(Double_t random, Double_t *vertex) 
288 {
289 // Constructor
290    Int_t itype;
291    if (random<0.3) itype = 1;      // pi+
292    else if (random<0.6) itype = 2; // pi-
293    else if (random<0.9) itype = 3; // p
294    else if (random<0.95) itype = 4; // pi0
295    else itype = 5;                 // gamma
296    gRandom->Rannor(fPx, fPy);
297    Double_t vert_width = 0.1;
298    
299    switch (itype) {
300       case 1:
301          fMass = 0.13957 + gRandom->Gaus(0.,0.001);
302          fCharge = 1;
303          break;
304       case 2:
305          fMass = 0.13957 + gRandom->Gaus(0.,0.001);
306          fCharge = -1;
307          break;
308       case 3:
309          fMass = 0.938 + gRandom->Gaus(0.,0.002);
310          fCharge = 1;
311          fPx *= 0.15;
312          fPy *= 0.15;
313          break;
314       case 4:
315          fMass = 0.135 + gRandom->Gaus(0.,0.001);
316          fCharge = 0;
317          fPx *= 0.8;
318          fPy *= 0.8;
319          vert_width = 10.;
320          break;
321       case 5:
322          fMass = 0.;
323          fCharge = 0;
324          fPx *= 0.5;
325          fPy *= 1.5;
326          break;
327    };
328    fPz = gRandom->Gaus(4., 2.);
329    if (vertex) {
330       fVertex[0] = vertex[0];
331       fVertex[1] = vertex[1];
332       fVertex[2] = vertex[2];
333    } else {   
334       fVertex[0] = gRandom->Gaus(0,vert_width);
335       fVertex[1] = gRandom->Gaus(0,vert_width);
336       fVertex[2] = gRandom->Gaus(0,0.01);
337    }   
338 }
339
340 //______________________________________________________________________________
341 Bool_t AnaTrack::Decay(Double_t &px1, Double_t &py1, Double_t &pz1, 
342                        Double_t &px2, Double_t &py2, Double_t &pz2)
343 {
344 // Decay a pi0 in 2 gammas.
345    if (fCharge != 0) return kFALSE;
346    if (fMass<0.132 || fMass>0.138) return kFALSE;
347    Double_t phi1 = 2.*TMath::Pi()*gRandom->Rndm(); // [0,2*pi]
348    Double_t phi2 = phi1 + TMath::Pi();
349    if (phi2 > 2.*TMath::Pi()) phi2 -= 2.*TMath::Pi();
350    Double_t r2 = gRandom->Rndm();
351    Double_t theta1 = TMath::ACos(1.-2.*r2);
352    Double_t p0 = GetP();
353    Double_t m0 = 0.135;
354    Double_t p1 = 0.5*(p0*(1-2*r2)+TMath::Sqrt(p0*p0*(1-2*r2)*(1-2*r2)+2*m0*m0));
355    Double_t p2 = TMath::Sqrt(p0*p0+m0*m0-p1*p1);
356    Double_t theta2 = TMath::ACos((p0-p1*(1-2*r2))/p2);
357    // Px, Py and Pz in the reference frame of the pion
358    px1 = p1 * TMath::Sin(theta1)*TMath::Cos(phi1);
359    px2 = p2 * TMath::Sin(theta2)*TMath::Cos(phi2);
360    py1 = p1 * TMath::Sin(theta1)*TMath::Sin(phi1);
361    py2 = p2 * TMath::Sin(theta2)*TMath::Sin(phi2);
362    pz1 = p1 * TMath::Cos(theta1);
363    pz2 = p2 * TMath::Cos(theta2);
364    Double_t phi = TMath::ATan2(fPy, fPx) * TMath::RadToDeg();
365    Double_t theta = TMath::ACos(GetPt()/p0) * TMath::RadToDeg();
366    TGeoRotation r("rot", phi+90., theta, 0.);
367    Double_t loc[3], vect[3];
368    p0 = TMath::Sqrt(px1*px1+py1*py1+pz1*pz1);
369    loc[0] = px1/p0;
370    loc[1] = py1/p0;
371    loc[2] = pz1/p0;
372    r.LocalToMasterVect(loc, vect);
373    px1 = vect[0]*p0;
374    py1 = vect[1]*p0;
375    pz1 = vect[2]*p0;
376 //   t1 = new AnaTrack(1., fVertex);
377    p0 = TMath::Sqrt(px2*px2+py2*py2+pz2*pz2);
378    loc[0] = px2/p0;
379    loc[1] = py2/p0;
380    loc[2] = pz2/p0;
381    r.LocalToMasterVect(loc, vect);
382    px2 = vect[0]*p0;
383    py2 = vect[1]*p0;
384    pz2 = vect[2]*p0;
385 //   t2 = new AnaTrack(1., fVertex);
386    return kTRUE;
387 }
388
389 //______________________________________________________________________________
390 ClassImp(AnaEvent)
391
392 TClonesArray *AnaEvent::fgTracks = 0;
393
394 //______________________________________________________________________________
395 AnaEvent::AnaEvent()
396 {
397 // Ctor
398    if (!fgTracks) fgTracks = new TClonesArray("AnaTrack", 2000);
399    fTracks = fgTracks;
400    fEventNumber = 0;
401    fNtracks = 0;
402 }
403
404 //______________________________________________________________________________
405 AnaTrack *AnaEvent::AddTrack(Double_t rnd, Double_t *vert)
406 {
407 // Add a random track
408 //   printf("track %d\n", fNtracks);
409    TClonesArray &tracks = *fTracks;
410    AnaTrack *track = new(tracks[fNtracks++]) AnaTrack(rnd, vert);
411    return track;
412 }   
413
414 //______________________________________________________________________________
415 void AnaEvent::Clear(Option_t *)
416 {
417 // Clears current event.
418    fTracks->Clear();
419    fNtracks = 0;
420    fEventNumber = 0;
421 }   
422
423 //______________________________________________________________________________
424 Int_t AnaEvent::Build(Int_t ev)
425 {
426 // Create a random event
427    Clear();
428    fEventNumber = ev;
429    Int_t ntracks = Int_t(gRandom->Gaus(500., 100.));
430    if (ntracks < 1) ntracks = 1;
431    if (ntracks>1000) ntracks = 1000;
432    AnaTrack *track, *track0;
433    for (Int_t i=0; i<ntracks; i++) {
434       Double_t rnd = gRandom->Rndm();
435       if (rnd>=0.90 && rnd<0.95) {
436       // Pi0 -> decay the track in 2 gammas
437          track0 = new AnaTrack(rnd);
438          Double_t vert[3];
439          vert[0] = track0->GetVertex(0);
440          vert[1] = track0->GetVertex(1);
441          vert[2] = track0->GetVertex(2);
442          Double_t px1,py1,pz1,px2,py2,pz2;
443          if (track0->Decay(px1,py1,pz1,px2,py2,pz2)) {
444             track = AddTrack(1.,vert);
445             track->SetPx(px1);
446             track->SetPy(py1);
447             track->SetPz(pz1);
448             track = AddTrack(1.,vert);
449             track->SetPx(px2);
450             track->SetPy(py2);
451             track->SetPz(pz2);
452          }
453          delete track0;
454       } else {   
455          track = AddTrack(rnd);
456       }   
457    }
458    return fNtracks;
459 }   
460
461 //______________________________________________________________________________
462 void AnaEvent::CreateEvents(Int_t nevents, const char *filename)
463 {
464 // Create nevents in one tree and put them in filename.
465    TFile *hfile = new TFile(filename, "RECREATE", "Some AnaEvents...");
466    TTree *tree  = new TTree("T", "Tree of AnaEvents");
467    tree->SetAutoSave(1000000000);  // autosave when 1 Gbyte written
468    Int_t bufsize = 16000;
469    AnaEvent *event = new AnaEvent();
470    TBranch *branch = tree->Branch("event", &event, bufsize,1);
471    branch->SetAutoDelete(kFALSE);
472    
473    for (Int_t ev=0; ev<nevents; ev++) {
474       Int_t ntracks = event->Build(ev);
475       if (ev%100 == 0) printf("event: %d  ntracks=%d\n", ev, ntracks);
476       tree->Fill();
477    }   
478    hfile->Write();
479    tree->Print();
480    hfile->Close();
481 }  
482