bc964af5b364d8c19aafc85d2c87526c34272b5e
[u/mrichter/AliRoot.git] / PWGJE / EMCALJetTasks / AliJetModelBaseTask.cxx
1 // $Id$
2 //
3 // Jet modelling task.
4 //
5 // Author: S.Aiola, C.Loizides
6
7 #include "AliJetModelBaseTask.h"
8
9 #include <TClonesArray.h>
10 #include <TH1.h>
11 #include <TLorentzVector.h>
12 #include <TRandom3.h>
13
14 #include "AliAODCaloCluster.h"
15 #include "AliAnalysisManager.h"
16 #include "AliEMCALDigit.h"
17 #include "AliEMCALGeometry.h"
18 #include "AliEMCALRecPoint.h"
19 #include "AliESDCaloCluster.h"
20 #include "AliLog.h"
21 #include "AliPicoTrack.h"
22 #include "AliVCluster.h"
23 #include "AliVEvent.h"
24
25 ClassImp(AliJetModelBaseTask)
26
27 //________________________________________________________________________
28 AliJetModelBaseTask::AliJetModelBaseTask() : 
29   AliAnalysisTaskSE("AliJetModelBaseTask"),
30   fGeomName(),
31   fTracksName(),
32   fOutTracksName(),
33   fCaloName(),
34   fOutCaloName(),
35   fSuffix(),
36   fEtaMin(-1),
37   fEtaMax(1),
38   fPhiMin(0),
39   fPhiMax(TMath::Pi() * 2),
40   fPtMin(0),
41   fPtMax(0),
42   fCopyArray(kTRUE),
43   fNClusters(0),
44   fNTracks(0),
45   fMarkMC(kTRUE),
46   fPtSpectrum(0),
47   fIsInit(0),
48   fGeom(0),
49   fClusters(0),
50   fOutClusters(0),
51   fTracks(0),
52   fOutTracks(0)
53 {
54   // Default constructor.
55 }
56
57 //________________________________________________________________________
58 AliJetModelBaseTask::AliJetModelBaseTask(const char *name) : 
59   AliAnalysisTaskSE(name),
60   fGeomName(""),
61   fTracksName("PicoTracks"),
62   fOutTracksName("PicoTracksEmbedded"),
63   fCaloName("CaloClustersCorr"),
64   fOutCaloName("CaloClustersCorrEmbedded"),
65   fSuffix("Processed"),
66   fEtaMin(-1),
67   fEtaMax(1),
68   fPhiMin(0),
69   fPhiMax(TMath::Pi() * 2),
70   fPtMin(50),
71   fPtMax(60),
72   fCopyArray(kTRUE),
73   fNClusters(0),
74   fNTracks(1),
75   fMarkMC(kTRUE),
76   fPtSpectrum(0),
77   fIsInit(0),
78   fGeom(0),
79   fClusters(0),
80   fOutClusters(0),
81   fTracks(0),
82   fOutTracks(0)
83 {
84   // Standard constructor.
85 }
86
87 //________________________________________________________________________
88 AliJetModelBaseTask::~AliJetModelBaseTask()
89 {
90   // Destructor
91 }
92
93 //________________________________________________________________________
94 void AliJetModelBaseTask::UserExec(Option_t *) 
95 {
96   // Execute per event.
97
98   if (!fIsInit) {
99     ExecOnce();
100     fIsInit = 1;
101   }
102   Run();
103 }
104
105 //________________________________________________________________________
106 AliVCluster* AliJetModelBaseTask::AddCluster(Double_t e, Double_t eta, Double_t phi)
107 {
108   // Add a cluster to the event.
109
110   Int_t absId = 0;
111   if (eta < -100 || phi < 0) {
112     GetRandomCell(eta, phi, absId);
113   }
114   else {
115     fGeom->EtaPhiFromIndex(absId, eta, phi);
116   }
117
118   if (absId == -1) {
119     AliWarning(Form("Unable to embed cluster in eta = %f, phi = %f!"
120                     " Maybe the eta-phi range is not inside the EMCal acceptance (eta = [%f, %f], phi = [%f, %f])", 
121                     eta, phi, fEtaMin, fEtaMax, fPhiMin, fPhiMax));
122     return 0;
123   } 
124
125   if (e < 0) {
126     Double_t pt = GetRandomPt();
127     TLorentzVector nPart;
128     nPart.SetPtEtaPhiM(pt, eta, phi, 0);
129     e = nPart.E();
130   }
131
132   return AddCluster(e, absId);
133 }
134       
135 //________________________________________________________________________
136 AliVCluster* AliJetModelBaseTask::AddCluster(Double_t e, Int_t absId)
137 {
138   // Add a cluster to the event.
139
140   const Int_t nClusters = fOutClusters->GetEntriesFast();
141
142   TClonesArray digits("AliEMCALDigit", 1);
143
144   AliEMCALDigit *digit = static_cast<AliEMCALDigit*>(digits.New(0));
145   digit->SetId(absId);
146   digit->SetIndexInList(0);
147   digit->SetType(AliEMCALDigit::kHG);
148   digit->SetAmplitude(e);
149       
150   AliEMCALRecPoint recPoint("");
151   recPoint.AddDigit(*digit, e, kFALSE);
152   recPoint.EvalGlobalPosition(0, &digits);
153
154   TVector3 gpos;
155   recPoint.GetGlobalPosition(gpos);
156   Float_t g[3];
157   gpos.GetXYZ(g);
158       
159   AliVCluster *cluster = static_cast<AliVCluster*>(fOutClusters->New(nClusters));
160   cluster->SetType(AliVCluster::kEMCALClusterv1);
161   cluster->SetE(recPoint.GetEnergy());
162   cluster->SetPosition(g);
163   cluster->SetNCells(1);
164   UShort_t shortAbsId = absId;
165   cluster->SetCellsAbsId(&shortAbsId);
166   Double32_t fract = 1;
167   cluster->SetCellsAmplitudeFraction(&fract);
168   cluster->SetID(nClusters);
169   cluster->SetEmcCpvDistance(-1);
170   if (fMarkMC)
171     cluster->SetChi2(100); // MC flag!
172
173   return cluster;
174 }
175
176 //________________________________________________________________________
177 AliPicoTrack* AliJetModelBaseTask::AddTrack(Double_t pt, Double_t eta, Double_t phi)
178 {
179   // Add a track to the event.
180
181   const Int_t nTracks = fOutTracks->GetEntriesFast();
182   
183   if (pt < 0) 
184     pt = GetRandomPt();
185   if (eta < -100) 
186     eta = GetRandomEta();
187   if (phi < 0) 
188     phi = GetRandomPhi();
189
190   Int_t label = fMarkMC ? 100 : 0;
191
192   AliPicoTrack *track = new ((*fOutTracks)[nTracks]) AliPicoTrack(pt, 
193                                                                   eta, 
194                                                                   phi, 
195                                                                   1, 
196                                                                   label,    // MC flag!      
197                                                                   0, 
198                                                                   0, 
199                                                                   kFALSE);
200   return track;
201 }
202
203 //________________________________________________________________________
204 void AliJetModelBaseTask::CopyClusters()
205 {
206   // Copy all the clusters in the new collection
207
208   Bool_t esdMode = (Bool_t)(fClusters->GetClass()->GetBaseClass("AliESDCaloCluster") != 0);
209   const Int_t nClusters = fClusters->GetEntriesFast();
210   Int_t nCopiedClusters = 0;
211   
212   if (esdMode) {
213     for (Int_t i = 0; i < nClusters; ++i) {
214       AliESDCaloCluster *esdcluster = static_cast<AliESDCaloCluster*>(fClusters->At(i));
215       if (!esdcluster || !esdcluster->IsEMCAL())
216         continue;
217       new ((*fOutClusters)[nCopiedClusters]) AliESDCaloCluster(*esdcluster);
218       nCopiedClusters++;
219     }
220   }
221   else {
222     for (Int_t i = 0; i < nClusters; ++i) {
223       AliAODCaloCluster *aodcluster = static_cast<AliAODCaloCluster*>(fClusters->At(i));
224       if (!aodcluster || !aodcluster->IsEMCAL())
225         continue;
226       new ((*fOutClusters)[nCopiedClusters]) AliAODCaloCluster(*aodcluster);
227       nCopiedClusters++;
228     }
229   }
230 }
231
232 //________________________________________________________________________
233 void AliJetModelBaseTask::CopyTracks()
234 {
235   const Int_t nTracks = fTracks->GetEntriesFast();
236   Int_t nCopiedTracks = 0;
237   for (Int_t i = 0; i < nTracks; ++i) {
238     AliPicoTrack *track = static_cast<AliPicoTrack*>(fTracks->At(i));
239     if (!track)
240       continue;
241     new ((*fOutTracks)[nCopiedTracks]) AliPicoTrack(*track);
242     nCopiedTracks++;
243   }
244 }
245
246 //________________________________________________________________________
247 void AliJetModelBaseTask::ExecOnce()
248 {
249   // Init task.
250
251   if (fPtMax < fPtMin) {
252     AliWarning (Form("PtMax (%f) < PtMin (%f), setting PtMax = PtMin = %f", fPtMax, fPtMin, fPtMin));
253     fPtMax = fPtMin;
254   }
255
256   if (fEtaMax < fEtaMin) {
257     AliWarning (Form("EtaMax (%f) < EtaMin (%f), setting EtaMax = EtaMin = %f", fEtaMax, fEtaMin, fEtaMin));
258     fEtaMax = fEtaMin;
259   }
260
261   if (fPhiMax < fPhiMin) {
262     AliWarning (Form("PhiMax (%f) < PhiMin (%f), setting PhiMax = PhiMin = %f", fPhiMax, fPhiMin, fPhiMin));
263     fPhiMax = fPhiMin;
264   }
265
266   if (fNTracks > 0 && !fTracks && !fTracksName.IsNull()) {
267     fTracks = dynamic_cast<TClonesArray*>(InputEvent()->FindListObject(fTracksName));
268     if (!fTracks) {
269       AliError(Form("%s: Couldn't retrieve tracks with name %s!", GetName(), fTracksName.Data()));
270       return;
271     }
272     
273     if (!fTracks->GetClass()->GetBaseClass("AliPicoTrack")) {
274       AliError(Form("%s: Collection %s does not contain AliPicoTrack objects!", GetName(), fTracksName.Data())); 
275       fTracks = 0;
276       return;
277     }
278
279     if (!fOutTracks) {
280       fOutTracksName = fTracksName;
281       if (fCopyArray) {
282         fOutTracksName += fSuffix;
283         fOutTracks = new TClonesArray("AliPicoTrack", fTracks->GetSize());
284         fOutTracks->SetName(fOutTracksName);
285       }
286       else {
287         fOutTracks = fTracks;
288       }
289     }
290
291     if (fCopyArray) {
292       fOutTracks->Delete();
293       if (!(InputEvent()->FindListObject(fOutTracksName)))
294         InputEvent()->AddObject(fOutTracks);
295     }
296   }
297
298   if (fNClusters > 0 && !fClusters && !fCaloName.IsNull()) {
299     fClusters = dynamic_cast<TClonesArray*>(InputEvent()->FindListObject(fCaloName));
300  
301     if (!fClusters) {
302       AliError(Form("%s: Couldn't retrieve clusters with name %s!", GetName(), fCaloName.Data()));
303       return;
304     }
305
306     if (!fClusters->GetClass()->GetBaseClass("AliVCluster")) {
307       AliError(Form("%s: Collection %s does not contain AliVCluster objects!", GetName(), fCaloName.Data())); 
308       fClusters = 0;
309       return;
310     }
311
312     if (!fOutClusters) {
313       fOutCaloName = fCaloName;
314       if (fCopyArray) {
315         fOutCaloName += fSuffix;
316         fOutClusters = new TClonesArray(fClusters->GetClass()->GetName(), fClusters->GetSize());
317         fOutClusters->SetName(fOutCaloName);
318       }
319       else {
320         fOutClusters = fClusters;
321       }
322     }
323
324     if (fCopyArray) {
325       fOutClusters->Delete();
326       if (!(InputEvent()->FindListObject(fOutCaloName)))
327         InputEvent()->AddObject(fOutClusters);
328     }
329
330     if (!fGeom) {
331       if (fGeomName.Length() > 0) {
332         fGeom = AliEMCALGeometry::GetInstance(fGeomName);
333         if (!fGeom)
334           AliError(Form("Could not get geometry with name %s!", fGeomName.Data()));
335       } else {
336         fGeom = AliEMCALGeometry::GetInstance();
337         if (!fGeom) 
338           AliError("Could not get default geometry!");
339       }
340     }
341     
342     const Double_t EmcalMinEta = fGeom->GetArm1EtaMin();
343     const Double_t EmcalMaxEta = fGeom->GetArm1EtaMax();
344     const Double_t EmcalMinPhi = fGeom->GetArm1PhiMin() * TMath::DegToRad();
345     const Double_t EmcalMaxPhi = fGeom->GetArm1PhiMax() * TMath::DegToRad();
346     
347     if (fEtaMax > EmcalMaxEta) fEtaMax = EmcalMaxEta;
348     if (fEtaMax < EmcalMinEta) fEtaMax = EmcalMinEta;
349     if (fEtaMin > EmcalMaxEta) fEtaMin = EmcalMaxEta;
350     if (fEtaMin < EmcalMinEta) fEtaMin = EmcalMinEta;
351   
352     if (fPhiMax > EmcalMaxPhi) fPhiMax = EmcalMaxPhi;
353     if (fPhiMax < EmcalMinPhi) fPhiMax = EmcalMinPhi;
354     if (fPhiMin > EmcalMaxPhi) fPhiMin = EmcalMaxPhi;
355     if (fPhiMin < EmcalMinPhi) fPhiMin = EmcalMinPhi;
356   }
357 }
358
359 //________________________________________________________________________
360 void AliJetModelBaseTask::GetRandomCell(Double_t &eta, Double_t &phi, Int_t &absId)
361 {
362   // Get random cell.
363
364   Int_t repeats = 0;
365   Double_t rndEta = eta;
366   Double_t rndPhi = phi;
367   do {
368     if (eta < -100)
369       rndEta = GetRandomEta();
370     if (phi < 0)
371       rndPhi = GetRandomPhi();
372     fGeom->GetAbsCellIdFromEtaPhi(rndEta, rndPhi, absId);  
373     repeats++;
374   } while (absId == -1 && repeats < 100);
375   
376   if (!(absId > -1)) {
377     AliWarning(Form("Could not extract random cluster! Random eta-phi extracted more than 100 times!\n"
378                     "eta [%f, %f], phi [%f, %f]\n", fEtaMin, fEtaMax, fPhiMin, fPhiMax));
379   }
380   else {
381     eta = rndEta;
382     phi = rndPhi;
383   }
384 }
385
386 //________________________________________________________________________
387 Double_t AliJetModelBaseTask::GetRandomEta()
388 {
389   // Get random eta.
390
391   return gRandom->Rndm() * (fEtaMax - fEtaMin) + fEtaMin;
392 }
393
394 //________________________________________________________________________
395 Double_t AliJetModelBaseTask::GetRandomPhi()
396 {
397   // Get random phi.
398
399   return gRandom->Rndm() * (fPhiMax - fPhiMin) + fPhiMin;
400 }
401
402 //________________________________________________________________________
403 Double_t AliJetModelBaseTask::GetRandomPt()
404 {
405   // Get random pt.
406
407   if (fPtSpectrum)
408     return fPtSpectrum->GetRandom();
409   else
410     return gRandom->Rndm() * (fPtMax - fPtMin) + fPtMin;
411 }
412
413 //________________________________________________________________________
414 void AliJetModelBaseTask::Run() 
415 {
416   // Run.
417 }