reduce binning and ranges
[u/mrichter/AliRoot.git] / PWG1 / TRD / AliTRDpidRefMaker.cxx
CommitLineData
1ee39b3a 1#include "TPDGCode.h"
2#include "TFile.h"
3#include "TTree.h"
4#include "TEventList.h"
5#include "TObjArray.h"
6#include "TH2.h"
7
8#include "AliLog.h"
9#include "AliESDtrack.h"
10#include "AliTrackReference.h"
11
12#include "AliTRDReconstructor.h"
13#include "AliTRDtrackV1.h"
14#include "AliTRDseedV1.h"
15#include "AliTRDpidRefMaker.h"
16#include "info/AliTRDv0Info.h"
17
18
19// Defines and implements basic functionality for building reference data for TRD PID.
20//
21// Here is the list of functionality provided by this class
22// 1.
23// 2.
24//
25// Authors:
26// Alex Bercuci <A.Bercuci@gsi.de>
27// Alex Wilk <wilka@uni-muenster.de>
28// Markus Fasel <mfasel@gsi.de>
29// Markus Heide <mheide@uni-muenster.de>
30//
31
32ClassImp(AliTRDpidRefMaker)
33ClassImp(AliTRDpidRefMaker::AliTRDpidRefData)
34ClassImp(AliTRDpidRefMaker::AliTRDpidRefDataArray)
35
36//________________________________________________________________________
37AliTRDpidRefMaker::AliTRDpidRefDataArray::AliTRDpidRefDataArray() :
38 fNtracklets(0)
b91fdd71 39 ,fData(NULL)
1ee39b3a 40{
41 // Constructor of data array
42 fData = new AliTRDpidRefData[AliTRDgeometry::kNlayer];
43}
44
45//________________________________________________________________________
46AliTRDpidRefMaker::AliTRDpidRefDataArray::~AliTRDpidRefDataArray()
47{
48 // Destructor
49 delete [] fData;
50}
51
52//________________________________________________________________________
53void AliTRDpidRefMaker::AliTRDpidRefDataArray::PushBack(Int_t ly, Int_t p, Float_t *dedx)
54{
55// Add PID data to the end of the array
56 fData[fNtracklets].fPLbin= (ly<<4) | (p&0xf);
57 memcpy(fData[fNtracklets].fdEdx, dedx, 8*sizeof(Float_t));
58 fNtracklets++;
59}
60
61//________________________________________________________________________
62void AliTRDpidRefMaker::AliTRDpidRefDataArray::Reset()
63{
64// Reset content
65
66 if(!fNtracklets) return;
67 while(fNtracklets--){
68 fData[fNtracklets].fPLbin = 0xff;
69 memset(fData[fNtracklets].fdEdx, 0, 8*sizeof(Float_t));
70 }
71 fNtracklets=0;
72}
73
74//________________________________________________________________________
75AliTRDpidRefMaker::AliTRDpidRefMaker(const char *name, const char *title)
76 :AliTRDrecoTask(name, title)
b91fdd71 77 ,fReconstructor(NULL)
78 ,fV0s(NULL)
79 ,fData(NULL)
80 ,fInfo(NULL)
81 ,fPIDdataArray(NULL)
1ee39b3a 82 ,fRefPID(kMC)
83 ,fRefP(kMC)
84 ,fPIDbin(0xff)
85 ,fFreq(1.)
86 ,fP(-1.)
87 ,fPthreshold(0.5)
88{
89 //
90 // Default constructor
91 //
92
93 fReconstructor = new AliTRDReconstructor();
94 fReconstructor->SetRecoParam(AliTRDrecoParam::GetLowFluxParam());
95 memset(fdEdx, 0, 10*sizeof(Float_t));
96 memset(fPID, 0, AliPID::kSPECIES*sizeof(Float_t));
97
98 DefineInput(1, TObjArray::Class());
b91fdd71 99 DefineInput(2, TObjArray::Class());
1ee39b3a 100 DefineOutput(1, TTree::Class());
101}
102
103
104//________________________________________________________________________
105AliTRDpidRefMaker::~AliTRDpidRefMaker()
106{
107 if(fPIDdataArray) delete fPIDdataArray;
108 if(fReconstructor) delete fReconstructor;
109}
110
111//________________________________________________________________________
112Bool_t AliTRDpidRefMaker::CheckQuality(AliTRDseedV1* /*trklt*/)
113{
114// Place holder for checking tracklet quality for PID.
115 return kTRUE;
116}
117
118
119//________________________________________________________________________
120Float_t* AliTRDpidRefMaker::CookdEdx(AliTRDseedV1 *trklt)
121{
122 trklt->CookdEdx(AliTRDpidUtil::kNNslices);
123 memcpy(fdEdx, trklt->GetdEdx(), AliTRDpidUtil::kNNslices*sizeof(Float_t));
124 return fdEdx;
125}
126
127
128//________________________________________________________________________
129void AliTRDpidRefMaker::ConnectInputData(Option_t *opt)
130{
131 AliTRDrecoTask::ConnectInputData(opt);
b91fdd71 132 fV0s = dynamic_cast<TObjArray*>(GetInputData(1));
133 fInfo = dynamic_cast<TObjArray*>(GetInputData(2));
134}
1ee39b3a 135
136//________________________________________________________________________
f8f46e4d 137void AliTRDpidRefMaker::UserCreateOutputObjects()
1ee39b3a 138{
139 // Create histograms
140 // Called once
141
f8f46e4d 142 OpenFile(1, "RECREATE");
1ee39b3a 143 fContainer = new TObjArray();
144 fContainer->SetName(Form("Moni%s", GetName()));
145
146 TH2 *h2 = new TH2I("hPDG","Particle abundance", AliPID::kSPECIES, -0.5, 4.5, AliTRDCalPID::kNMom, -0.5, AliTRDCalPID::kNMom-0.5);
147 TAxis *ax = h2->GetXaxis();
148 ax->SetNdivisions(505);
149 ax->SetTitle("Particle species");
150 for(Int_t is=AliPID::kSPECIES; is--;) ax->SetBinLabel(is+1, AliPID::ParticleShortName(is));
151 h2->GetYaxis()->SetTitle("P bins");
152 h2->GetYaxis()->SetNdivisions(511);
153
154 fContainer->AddAt(h2, 0);
155
156 fData = new TTree(GetName(), Form("Reference data for %s", GetName()));
157 fData->Branch("s", &fPIDbin, "s/b");
158 fData->Branch("data", &fPIDdataArray);
159}
160
161//________________________________________________________________________
f8f46e4d 162void AliTRDpidRefMaker::UserExec(Option_t *)
1ee39b3a 163{
164 // Main loop
165 // Called for each event
166
b91fdd71 167 AliTRDtrackInfo *track = NULL;
168 AliTRDtrackV1 *trackTRD = NULL;
169 AliTrackReference *ref = NULL;
170 const AliTRDtrackInfo::AliESDinfo *infoESD = NULL;
1ee39b3a 171 for(Int_t itrk=0; itrk<fTracks->GetEntriesFast(); itrk++){
172 track = (AliTRDtrackInfo*)fTracks->UncheckedAt(itrk);
173 if(!track->HasESDtrack()) continue;
174 trackTRD = track->GetTrack();
175 infoESD = track->GetESDinfo();
176 Double32_t *infoPID = infoESD->GetSliceIter();
177 Int_t n = infoESD->GetNSlices() - AliTRDgeometry::kNlayer;
178 Double32_t *p = &infoPID[n];
4226db3e 179 AliDebug(4, Form("n[%d] p[GeV/c]{%6.2f %6.2f %6.2f %6.2f %6.2f %6.2f}", n, p[0], p[1], p[2], p[3], p[4], p[5]));
1ee39b3a 180
181 ULong_t status = track->GetStatus();
182 if(!(status&AliESDtrack::kTPCout)) continue;
183
184 // fill the pid information
185 SetRefPID(fRefPID, track, fPID);
186
187 // prepare PID data array
188 if(!fPIDdataArray){
189 fPIDdataArray = new AliTRDpidRefDataArray();
190 } else fPIDdataArray->Reset();
191
192 // fill PID information
193 for(Int_t ily = 0; ily < AliTRDgeometry::kNlayer; ily++){
194
195 // fill P & dE/dx information
196 if(HasFriends()){ // from TRD track
197 if(!trackTRD) continue;
b91fdd71 198 AliTRDseedV1 *trackletTRD(NULL);
1ee39b3a 199 trackTRD -> SetReconstructor(fReconstructor);
200 if(!(trackletTRD = trackTRD -> GetTracklet(ily))) continue;
201 if(!CheckQuality(trackletTRD)) continue;
202 if(!CookdEdx(trackletTRD)) continue;
203
204 // fill momentum information
205 fP = 0.;
206 switch(fRefP){
207 case kMC:
208 if(!(ref = track->GetTrackRef(trackletTRD))) continue;
209 fP = ref->P();
210 break;
211 case kRec:
212 fP = trackletTRD->GetMomentum();
213 break;
214 default: continue;
215 }
216 } else { // from ESD track
217 // fill momentum information
218 switch(fRefP){
219 case kMC:
220 if(!(ref = track->GetTrackRef(ily))) continue;
221 fP = ref->P();
222 break;
223 case kRec:
224 fP = p[ily];
225 break;
226 default: continue;
227 }
228 Double32_t *it = &infoPID[ily*AliTRDCalPID::kNSlicesNN];
229 for(Int_t is=AliTRDCalPID::kNSlicesNN; is--; it++) fdEdx[is] = (*it);
230 }
231
232 // momentum threshold
233 if(fP < fPthreshold) continue;
234
235 // store information
236 fPIDdataArray->PushBack(ily, AliTRDpidUtil::GetMomentumBin(fP), fdEdx);
237 }
238
239 Fill();
240 }
241
f8f46e4d 242 PostData(1, fContainer);
243 PostData(2, fData);
1ee39b3a 244}
245
246
247//________________________________________________________________________
248void AliTRDpidRefMaker::Fill()
249{
250// Fill data tree
251
252 if(!fPIDdataArray->fNtracklets) return;
253 fPIDbin = TMath::LocMax(AliPID::kSPECIES, fPID); // get particle type
254// Fill data tree
255 fData->Fill();
256
257
258 // fill monitor
259 for(Int_t ily=fPIDdataArray->fNtracklets; ily--;){
260 Int_t pBin = fPIDdataArray->fData[ily].fPLbin & 0xf;
261 ((TH2*)fContainer->At(0))->Fill(fPIDbin, pBin);
262 }
263}
264
265//________________________________________________________________________
266void AliTRDpidRefMaker::LinkPIDdata()
267{
268// Link data tree to data members
269 fData->SetBranchAddress("s", &fPIDbin);
270 fData->SetBranchAddress("data", &fPIDdataArray);
271}
272
273//________________________________________________________________________
274void AliTRDpidRefMaker::SetRefPID(ETRDpidRefMakerSource select, void *source, Float_t *pid)
275{
276// Fill the reference PID values "pid" from "source" object
277// according to the option "select". Possible options are
278// - kV0 - v0 based PID
279// - kMC - MC truth [default]
280// - kRec - outside detectors
281
282
283 memset(fPID, 0, AliPID::kSPECIES*sizeof(Float_t));
284 switch(select){
285 case kV0:
286 {
287 AliTRDtrackInfo *track = static_cast<AliTRDtrackInfo*>(source);
288 if(!track){
289 AliError("No trackInfo found");
290 return;
291 }
292
293 //Get V0 PID decisions from the AliTRDv0Info for all particle species (implemented so far : electrons from conversions, pions from K0s and protons from Lambdas) :
294 AliTRDv0Info v0;
295 for(Int_t is=AliPID::kSPECIES; is--;) fPID[is] = v0.GetV0PID(is, track);
296 }
297 break;
298 case kMC:
299 if(!HasMCdata()){
300 AliError("Could not retrive reference PID from MC");
301 return;
302 }
303 {
304 AliTRDtrackInfo *track = static_cast<AliTRDtrackInfo*>(source);
305 switch(track->GetPDG()){
306 case kElectron:
307 case kPositron:
308 fPID[AliPID::kElectron] = 1.;
309 break;
310 case kMuonPlus:
311 case kMuonMinus:
312 fPID[AliPID::kMuon] = 1.;
313 break;
314 case kPiPlus:
315 case kPiMinus:
316 fPID[AliPID::kPion] = 1.;
317 break;
318 case kKPlus:
319 case kKMinus:
320 fPID[AliPID::kKaon] = 1.;
321 break;
322 case kProton:
323 case kProtonBar:
324 fPID[AliPID::kProton] = 1.;
325 break;
326 }
327 }
328 break;
329 case kRec:
330 {
331 AliTRDtrackInfo *track = static_cast<AliTRDtrackInfo*>(source);
332 AliTRDtrackV1 *trackTRD = track->GetTrack();
333 trackTRD -> SetReconstructor(fReconstructor);
334 //fReconstructor -> SetOption("nn");
335 trackTRD -> CookPID();
336 for(Int_t iPart = 0; iPart < AliPID::kSPECIES; iPart++){
337 pid[iPart] = trackTRD -> GetPID(iPart);
338 AliDebug(4, Form("PDG is (in V0info) %d %f", iPart, pid[iPart]));
339 }
340 }
341 break;
342 default:
343 AliWarning("PID reference source not implemented");
344 return;
345 }
346 AliDebug(4, Form("Ref PID [%] : %s[%5.2f] %s[%5.2f] %s[%5.2f] %s[%5.2f] %s[%5.2f]"
347 ,AliPID::ParticleShortName(0), 1.e2*fPID[0]
348 ,AliPID::ParticleShortName(1), 1.e2*fPID[1]
349 ,AliPID::ParticleShortName(2), 1.e2*fPID[2]
350 ,AliPID::ParticleShortName(3), 1.e2*fPID[3]
351 ,AliPID::ParticleShortName(4), 1.e2*fPID[4]
352 ));
353}
354
355//________________________________________________________________________
356void AliTRDpidRefMaker::SetAbundance(Float_t train)
357{
358// Split data sample between trainning and testing
359
360 if(train<0. || train >1.){
361 AliWarning("The input data should be in the interval [0, 1]");
362 return;
363 }
364
365 fFreq = train;
366}
367