97e4f519810b9c5a47e94f24bcc3cdbf350fcdea
[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       if (!(InputEvent()->FindListObject(fOutTracksName)))
293         InputEvent()->AddObject(fOutTracks);
294     }
295   }
296
297   if (fNClusters > 0 && !fClusters && !fCaloName.IsNull()) {
298     fClusters = dynamic_cast<TClonesArray*>(InputEvent()->FindListObject(fCaloName));
299  
300     if (!fClusters) {
301       AliError(Form("%s: Couldn't retrieve clusters with name %s!", GetName(), fCaloName.Data()));
302       return;
303     }
304
305     if (!fClusters->GetClass()->GetBaseClass("AliVCluster")) {
306       AliError(Form("%s: Collection %s does not contain AliVCluster objects!", GetName(), fCaloName.Data())); 
307       fClusters = 0;
308       return;
309     }
310
311     if (!fOutClusters) {
312       fOutCaloName = fCaloName;
313       if (fCopyArray) {
314         fOutCaloName += fSuffix;
315         fOutClusters = new TClonesArray(fClusters->GetClass()->GetName(), fClusters->GetSize());
316         fOutClusters->SetName(fOutCaloName);
317       }
318       else {
319         fOutClusters = fClusters;
320       }
321     }
322
323     if (fCopyArray) {
324       if (!(InputEvent()->FindListObject(fOutCaloName)))
325         InputEvent()->AddObject(fOutClusters);
326     }
327
328     if (!fGeom) {
329       if (fGeomName.Length() > 0) {
330         fGeom = AliEMCALGeometry::GetInstance(fGeomName);
331         if (!fGeom)
332           AliError(Form("Could not get geometry with name %s!", fGeomName.Data()));
333       } else {
334         fGeom = AliEMCALGeometry::GetInstance();
335         if (!fGeom) 
336           AliError("Could not get default geometry!");
337       }
338     }
339     
340     const Double_t EmcalMinEta = fGeom->GetArm1EtaMin();
341     const Double_t EmcalMaxEta = fGeom->GetArm1EtaMax();
342     const Double_t EmcalMinPhi = fGeom->GetArm1PhiMin() * TMath::DegToRad();
343     const Double_t EmcalMaxPhi = fGeom->GetArm1PhiMax() * TMath::DegToRad();
344     
345     if (fEtaMax > EmcalMaxEta) fEtaMax = EmcalMaxEta;
346     if (fEtaMax < EmcalMinEta) fEtaMax = EmcalMinEta;
347     if (fEtaMin > EmcalMaxEta) fEtaMin = EmcalMaxEta;
348     if (fEtaMin < EmcalMinEta) fEtaMin = EmcalMinEta;
349   
350     if (fPhiMax > EmcalMaxPhi) fPhiMax = EmcalMaxPhi;
351     if (fPhiMax < EmcalMinPhi) fPhiMax = EmcalMinPhi;
352     if (fPhiMin > EmcalMaxPhi) fPhiMin = EmcalMaxPhi;
353     if (fPhiMin < EmcalMinPhi) fPhiMin = EmcalMinPhi;
354   }
355 }
356
357 //________________________________________________________________________
358 void AliJetModelBaseTask::GetRandomCell(Double_t &eta, Double_t &phi, Int_t &absId)
359 {
360   // Get random cell.
361
362   Int_t repeats = 0;
363   Double_t rndEta = eta;
364   Double_t rndPhi = phi;
365   do {
366     if (eta < -100)
367       rndEta = GetRandomEta();
368     if (phi < 0)
369       rndPhi = GetRandomPhi();
370     fGeom->GetAbsCellIdFromEtaPhi(rndEta, rndPhi, absId);  
371     repeats++;
372   } while (absId == -1 && repeats < 100);
373   
374   if (!(absId > -1)) {
375     AliWarning(Form("Could not extract random cluster! Random eta-phi extracted more than 100 times!\n"
376                     "eta [%f, %f], phi [%f, %f]\n", fEtaMin, fEtaMax, fPhiMin, fPhiMax));
377   }
378   else {
379     eta = rndEta;
380     phi = rndPhi;
381   }
382 }
383
384 //________________________________________________________________________
385 Double_t AliJetModelBaseTask::GetRandomEta()
386 {
387   // Get random eta.
388
389   return gRandom->Rndm() * (fEtaMax - fEtaMin) + fEtaMin;
390 }
391
392 //________________________________________________________________________
393 Double_t AliJetModelBaseTask::GetRandomPhi()
394 {
395   // Get random phi.
396
397   return gRandom->Rndm() * (fPhiMax - fPhiMin) + fPhiMin;
398 }
399
400 //________________________________________________________________________
401 Double_t AliJetModelBaseTask::GetRandomPt()
402 {
403   // Get random pt.
404
405   if (fPtSpectrum)
406     return fPtSpectrum->GetRandom();
407   else
408     return gRandom->Rndm() * (fPtMax - fPtMin) + fPtMin;
409 }
410
411 //________________________________________________________________________
412 void AliJetModelBaseTask::Run() 
413 {
414   // Run.
415 }