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