Bug correction (energy recalculation when adding new incarnation of the particle).
[u/mrichter/AliRoot.git] / HBTAN / AliHBTReaderInternal.cxx
CommitLineData
91923a38 1#include "AliHBTReaderInternal.h"
2
91923a38 3#include <TTree.h>
4#include <TFile.h>
5#include <TParticle.h>
f701e886 6#include <TError.h>
91923a38 7#include <AliRun.h>
8#include <AliMagF.h>
9
10#include "AliHBTRun.h"
11#include "AliHBTEvent.h"
12#include "AliHBTParticle.h"
13#include "AliHBTParticleCut.h"
14
bed069a4 15AliHBTReaderInternal test;
91923a38 16
91923a38 17ClassImp(AliHBTReaderInternal)
18/********************************************************************/
19
bed069a4 20AliHBTReaderInternal::AliHBTReaderInternal():
21 fFileName(),
22 fPartBranch(0x0),
23 fTrackBranch(0x0),
24 fTree(0x0),
25 fFile(0x0)
91923a38 26{
f701e886 27//Defalut constructor
91923a38 28}
29/********************************************************************/
30
bed069a4 31AliHBTReaderInternal::AliHBTReaderInternal(const char *filename):
32 fFileName(filename),
33 fPartBranch(0x0),
34 fTrackBranch(0x0),
35 fTree(0x0),
36 fFile(0x0)
91923a38 37{
f701e886 38//constructor
39//filename - name of file to open
91923a38 40}
41/********************************************************************/
f701e886 42
91923a38 43AliHBTReaderInternal::AliHBTReaderInternal(TObjArray* dirs, const char *filename):
bed069a4 44 AliHBTReader(dirs),
45 fFileName(filename),
46 fPartBranch(0x0),
47 fTrackBranch(0x0),
48 fTree(0x0),
49 fFile(0x0)
91923a38 50{
f701e886 51//ctor
52//dirs contains strings with directories to look data in
53//filename - name of file to open
91923a38 54}
f701e886 55/********************************************************************/
56
91923a38 57AliHBTReaderInternal::~AliHBTReaderInternal()
58 {
59 //desctructor
bed069a4 60 delete fTree;
61 delete fFile;
91923a38 62 }
63/********************************************************************/
bed069a4 64
65void AliHBTReaderInternal::Rewind()
66{
67 delete fTree;
68 fTree = 0x0;
69 delete fFile;
70 fFile = 0x0;
71 fCurrentDir = 0;
72 fNEventsRead= 0;
73}
91923a38 74/********************************************************************/
75
bed069a4 76Int_t AliHBTReaderInternal::ReadNext()
91923a38 77 {
78 //reads data and puts put to the particles and tracks objects
79 //reurns 0 if everything is OK
80 //
bed069a4 81 Info("ReadNext","");
91923a38 82 Int_t i; //iterator and some temprary values
f701e886 83 Int_t counter;
9d09c308 84 AliHBTParticle* tpart = 0x0, *ttrack = 0x0;
bed069a4 85
27bb2a54 86 TDatabasePDG* pdgdb = TDatabasePDG::Instance();
87 if (pdgdb == 0x0)
88 {
bed069a4 89 Error("ReadNext","Can not get PDG Particles Data Base");
27bb2a54 90 return 1;
91923a38 91 }
bed069a4 92
93 if (fParticlesEvent == 0x0) fParticlesEvent = new AliHBTEvent();
94 if (fTracksEvent == 0x0) fTracksEvent = new AliHBTEvent();
95
96 fParticlesEvent->Reset();
97 fTracksEvent->Reset();
d4c1d7b6 98
91923a38 99 do //do{}while; is OK even if 0 dirs specified. In that case we try to read from "./"
100 {
bed069a4 101 if (fTree == 0x0)
102 if( (i=OpenNextFile()) )
103 {
4d91c73a 104 Error("ReadNext","Skipping directory due to problems with opening files. Errorcode %d",i);
bed069a4 105 fCurrentDir++;
106 continue;
107 }
108 if (fCurrentEvent == (Int_t)fTree->GetEntries())
91923a38 109 {
bed069a4 110 delete fTree;
111 fTree = 0x0;
112 delete fFile;
113 fFile = 0x0;
114 fPartBranch = 0x0;
115 fTrackBranch= 0x0;
116 fCurrentDir++;
27bb2a54 117 continue;
91923a38 118 }
119 /***************************/
120 /***************************/
121 /***************************/
122
bed069a4 123
124 Info("ReadNext","Event %d",fCurrentEvent);
125 fTree->GetEvent(fCurrentEvent);
91923a38 126
bed069a4 127 counter = 0;
128 if (fPartBranch && fTrackBranch)
4d1d0fe9 129 {
bed069a4 130 for(i = 0; i < fPartBuffer->GetEntries(); i++)
131 {
132 tpart = dynamic_cast<AliHBTParticle*>(fPartBuffer->At(i));
133 ttrack = dynamic_cast<AliHBTParticle*>(fTrackBuffer->At(i));
134
135 if( tpart == 0x0 ) continue; //if returned pointer is NULL
4d1d0fe9 136 if( ttrack == 0x0 ) continue; //if returned pointer is NULL
bed069a4 137
138 if (ttrack->GetUID() != tpart->GetUID())
9d09c308 139 {
bed069a4 140 Error("ReadNext","Sth. is wrong: Track and Particle has different UID.");
141 Error("ReadNext","They probobly do not correspond to each other.");
142 }
9d09c308 143
4d1d0fe9 144 for (Int_t s = 0; s < ttrack->GetNumberOfPids(); s++)
bed069a4 145 {
4d1d0fe9 146 if( pdgdb->GetParticle(ttrack->GetNthPid(s)) == 0x0 ) continue; //if particle has crazy PDG code (not known to our database)
147 if( Pass(ttrack->GetNthPid(s)) ) continue; //check if we are intersted with particles of this type
bed069a4 148 //if not take next partilce
4d1d0fe9 149
150 AliHBTParticle* track = new AliHBTParticle(*ttrack);
151
152 //apart of setting PDG code of an incarnation
153 //it is necessary tu recalculate energy on the basis of
154 //new PDG code (mass) hypothesis
155 TParticlePDG* pdgp = pdgdb->GetParticle(ttrack->GetNthPid(s));
156 if (pdgp == 0x0)//PDG part corresponding to new incarnation
157 {
158 Error("ReadNext","Particle code unknown to PDG DB.");
159 continue;
160 }
161 Double_t mass = pdgp->Mass();//mass of new incarnation
162 Double_t tEtot = TMath::Sqrt( ttrack->Px()*ttrack->Px() +
163 ttrack->Py()*ttrack->Py() +
164 ttrack->Pz()*ttrack->Pz() +
165 mass*mass);//total energy of the new incarnation
166 //update Energy and Calc Mass
167 track->SetMomentum(ttrack->Px(),ttrack->Py(),ttrack->Pz(),tEtot);
168 track->SetCalcMass(mass);
169 track->SetPdgCode(ttrack->GetNthPid(s),ttrack->GetNthPidProb(s));
170
171 if( Pass(track) )
27bb2a54 172 {
4d1d0fe9 173 delete track;
bed069a4 174 continue;
27bb2a54 175 }
4d1d0fe9 176 AliHBTParticle* part = new AliHBTParticle(*tpart);//particle has only 1 incarnation (real)
bed069a4 177
178 fParticlesEvent->AddParticle(part);
179 fTracksEvent->AddParticle(track);
180
181 counter++;
182 }
183 }
184 Info("ReadNext"," Read: %d particles and tracks.",counter);
185 }
186 else
187 {
188 if (fPartBranch)
189 {
190 Info("ReadNext","Found %d particles in total.",fPartBuffer->GetEntries());
191 for(i = 0; i < fPartBuffer->GetEntries(); i++)
192 {
193 tpart = dynamic_cast<AliHBTParticle*>(fPartBuffer->At(i));
194 if(tpart == 0x0) continue; //if returned pointer is NULL
195
196 for (Int_t s = 0; s < tpart->GetNumberOfPids(); s++)
197 {
198 if( pdgdb->GetParticle(tpart->GetNthPid(s)) == 0x0 ) continue; //if particle has crazy PDG code (not known to our database)
199 if( Pass(tpart->GetNthPid(s)) ) continue; //check if we are intersted with particles of this type
bed069a4 200 AliHBTParticle* part = new AliHBTParticle(*tpart);
201 part->SetPdgCode(tpart->GetNthPid(s),tpart->GetNthPidProb(s));
202 if( Pass(part) )
27bb2a54 203 {
bed069a4 204 delete part;
205 continue;
27bb2a54 206 }
bed069a4 207 fParticlesEvent->AddParticle(part);
208 counter++;
209 }
210 }
211 Info("ReadNext"," Read: %d particles.",counter);
212 }
213 else Info("ReadNext"," Read: 0 particles.");
214
215 if (fTrackBranch)
216 {
217 for(i = 0; i < fTrackBuffer->GetEntries(); i++)
218 {
4d1d0fe9 219 ttrack = dynamic_cast<AliHBTParticle*>(fTrackBuffer->At(i));
220 if( ttrack == 0x0 ) continue; //if returned pointer is NULL
bed069a4 221
4d1d0fe9 222 for (Int_t s = 0; s < ttrack->GetNumberOfPids(); s++)
223 {
224 if( pdgdb->GetParticle(ttrack->GetNthPid(s)) == 0x0 ) continue; //if particle has crazy PDG code (not known to our database)
225 if( Pass(ttrack->GetNthPid(s)) ) continue; //check if we are intersted with particles of this type
226 //if not take next partilce
227 AliHBTParticle* track = new AliHBTParticle(*ttrack);
228
229 //apart of setting PDG code of an incarnation
230 //it is necessary tu recalculate energy on the basis of
231 //new PDG code (mass) hypothesis
232 TParticlePDG* pdgp = pdgdb->GetParticle(ttrack->GetNthPid(s));
233 if (pdgp == 0x0)//PDG part corresponding to new incarnation
234 {
235 Error("ReadNext","Particle code unknown to PDG DB.");
236 continue;
237 }
238 Double_t mass = pdgp->Mass();//mass of new incarnation
239 Double_t tEtot = TMath::Sqrt( ttrack->Px()*ttrack->Px() +
240 ttrack->Py()*ttrack->Py() +
241 ttrack->Pz()*ttrack->Pz() +
242 mass*mass);//total energy of the new incarnation
243 //update Energy and Calc Mass
244 track->SetMomentum(ttrack->Px(),ttrack->Py(),ttrack->Pz(),tEtot);
245 track->SetCalcMass(mass);
246 track->SetPdgCode(ttrack->GetNthPid(s),ttrack->GetNthPidProb(s));
247
248 if( Pass(track) )
27bb2a54 249 {
4d1d0fe9 250 delete track;
bed069a4 251 continue;
27bb2a54 252 }
4d1d0fe9 253 fTracksEvent->AddParticle(track);
254
bed069a4 255 counter++;
256 }
257 }
258 Info("ReadNext"," Read: %d tracks",counter);
259 }
260 else Info("ReadNext"," Read: 0 tracks.");
261 }
d7f24a3d 262
263 if (fPartBuffer) fPartBuffer->Delete();
264 if (fTrackBuffer) fTrackBuffer->Delete();
bed069a4 265
266 fCurrentEvent++;
267 fNEventsRead++;
268
269 return 0;
9d09c308 270
4d91c73a 271 }while(fCurrentDir < GetNumberOfDirs());
91923a38 272
bed069a4 273 return 1;//no more directories to read
91923a38 274}
91923a38 275/********************************************************************/
276
bed069a4 277Int_t AliHBTReaderInternal::OpenNextFile()
91923a38 278{
bed069a4 279 //open file in current directory
4d91c73a 280 const TString& dirname = GetDirName(fCurrentDir);
91923a38 281 if (dirname == "")
282 {
bed069a4 283 Error("OpenNextFile","Can not get directory name");
91923a38 284 return 4;
285 }
286
287 TString filename = dirname +"/"+ fFileName;
bed069a4 288 fFile = TFile::Open(filename.Data());
289 if ( fFile == 0x0 )
91923a38 290 {
bed069a4 291 Error("OpenNextFile","Can't open file named %s",filename.Data());
91923a38 292 return 1;
293 }
bed069a4 294 if (fFile->IsOpen() == kFALSE)
91923a38 295 {
bed069a4 296 Error("OpenNextFile","Can't open filenamed %s",filename.Data());
91923a38 297 return 1;
298 }
bed069a4 299
300 fTree = (TTree*)fFile->Get("data");
301 if (fTree == 0x0)
302 {
303 Error("OpenNextFile","Can not get the tree.");
304 return 1;
305 }
306
307
308 fTrackBranch = fTree->GetBranch("tracks");//get the branch with tracks
309 if (fTrackBranch == 0x0) ////check if we got the branch
310 {//if not return with error
311 Info("OpenNextFile","Can't find a branch with tracks !\n");
312 }
313 else
314 {
315 if (fTrackBuffer == 0x0)
316 fTrackBuffer = new TClonesArray("AliHBTParticle",15000);
317 fTrackBranch->SetAddress(&fTrackBuffer);
318 }
319
320 fPartBranch = fTree->GetBranch("particles");//get the branch with particles
321 if (fPartBranch == 0x0) ////check if we got the branch
322 {//if not return with error
323 Info("OpenNextFile","Can't find a branch with particles !\n");
324 }
325 else
326 {
327 if (fPartBuffer == 0x0)
328 fPartBuffer = new TClonesArray("AliHBTParticle",15000);
329 fPartBranch->SetAddress(&fPartBuffer);
330 }
331
332 Info("OpenNextFile","________________________________________________________");
333 Info("OpenNextFile","Found %d event(s) in directory %s",
334 (Int_t)fTree->GetEntries(),GetDirName(fCurrentEvent).Data());
335
336 fCurrentEvent = 0;
337
91923a38 338 return 0;
339}
f701e886 340/********************************************************************/
91923a38 341
27bb2a54 342Int_t AliHBTReaderInternal::Write(AliHBTReader* reader,const char* outfile, Bool_t multcheck)
91923a38 343 {
f701e886 344 //reads tracks from reader and writes runs to file
27bb2a54 345 //reader - provides data for writing in internal format
346 //name of output file
347 //multcheck - switches of checking if particle was stored with other incarnation
348 // usefull e.g. when using kine data, where all particles have 100% pid prob.and saves a lot of time
349
f701e886 350 Int_t i,j;
91923a38 351
f701e886 352 ::Info("AliHBTReaderInternal::Write","________________________________________________________");
353 ::Info("AliHBTReaderInternal::Write","________________________________________________________");
354 ::Info("AliHBTReaderInternal::Write","________________________________________________________");
9d09c308 355
91923a38 356 TFile *histoOutput = TFile::Open(outfile,"recreate");
357
358 if (!histoOutput->IsOpen())
359 {
f701e886 360 ::Error("AliHBTReaderInternal::Write","File is not opened");
91923a38 361 return 1;
362 }
363
91923a38 364 TTree *tracktree = new TTree("data","Tree with tracks");
365
366 TClonesArray* pbuffer = new TClonesArray("AliHBTParticle",15000);
367 TClonesArray* tbuffer = new TClonesArray("AliHBTParticle",15000);
91923a38 368
369 TClonesArray &particles = *pbuffer;
370 TClonesArray &tracks = *tbuffer;
371
372 TString name("Tracks");
91923a38 373
27bb2a54 374 Int_t nt = reader->GetNumberOfTrackEvents();
375 Int_t np = reader->GetNumberOfPartEvents();
91923a38 376
27bb2a54 377 if (AliHBTParticle::GetDebug() > 0)
378 ::Info("Write","Reader has %d track events and %d particles events.",nt,np);
379
380 Bool_t trck = (nt > 0) ? kTRUE : kFALSE;
381 Bool_t part = (np > 0) ? kTRUE : kFALSE;
91923a38 382
383 TBranch *trackbranch = 0x0, *partbranch = 0x0;
384
d4c1d7b6 385 if (trck) trackbranch = tracktree->Branch("tracks",&tbuffer,32000,0);
386 if (part) partbranch = tracktree->Branch("particles",&pbuffer,32000,0);
91923a38 387
27bb2a54 388 if ( (trck) && (part) && (np != nt))
91923a38 389 {
f701e886 390 ::Warning("AliHBTReaderInternal::Write","Number of track and particle events is different");
91923a38 391 }
392
27bb2a54 393 Int_t n;
394 if (nt >= np ) n = nt; else n = np;
395
396 if (AliHBTParticle::GetDebug() > 0)
397 ::Info("Write","Will loop over %d events",n);
91923a38 398
27bb2a54 399 for ( i =0;i< n; i++)
91923a38 400 {
f701e886 401 ::Info("AliHBTReaderInternal::Write","Event %d",i+1);
27bb2a54 402 Int_t counter = 0;
403 if (trck && (i<=nt))
404 {
91923a38 405 AliHBTEvent* trackev = reader->GetTrackEvent(i);
406 for ( j = 0; j< trackev->GetNumberOfParticles();j++)
407 {
27bb2a54 408 const AliHBTParticle& t = *(trackev->GetParticle(j));
409 if (multcheck)
410 {
411 if (FindIndex(tbuffer,t.GetUID()))
412 {
413 if (AliHBTParticle::GetDebug()>4)
414 {
415 ::Info("Write","Track with Event UID %d already stored",t.GetUID());
416 }
417 continue; //not to write the same particles with other incarnations
418 }
419 }
420 new (tracks[counter++]) AliHBTParticle(t);
91923a38 421 }
27bb2a54 422 ::Info("AliHBTReaderInternal::Write"," Tracks: %d",tracks.GetEntries());
f701e886 423 }else ::Info("AliHBTReaderInternal::Write","NO TRACKS");
91923a38 424
27bb2a54 425 counter = 0;
426 if (part && (i<=np))
91923a38 427 {
27bb2a54 428// ::Warning("AliHBTReaderInternal::Write","Find index switched off!!!");
429
91923a38 430 AliHBTEvent* partev = reader->GetParticleEvent(i);
91923a38 431 for ( j = 0; j< partev->GetNumberOfParticles();j++)
432 {
9d09c308 433 const AliHBTParticle& part= *(partev->GetParticle(j));
27bb2a54 434 if (multcheck)
435 {
436 if (FindIndex(pbuffer,part.GetUID()))
437 {
438 if (AliHBTParticle::GetDebug()>4)
439 {
440 ::Info("Write","Particle with Event UID %d already stored",part.GetUID());
441 }
442 continue; //not to write the same particles with other incarnations
443 }
444 }
445 new (particles[counter++]) AliHBTParticle(part);
91923a38 446 }
27bb2a54 447 ::Info("AliHBTReaderInternal::Write"," Particles: %d",particles.GetEntries());
f701e886 448 }else ::Info("AliHBTReaderInternal::Write","NO PARTICLES");
91923a38 449
91923a38 450 histoOutput->cd();
451 tracktree->Fill();
9d09c308 452 tracktree->AutoSave();
91923a38 453 tbuffer->Delete();
454 pbuffer->Delete();
91923a38 455 }
456
91923a38 457 histoOutput->cd();
91923a38 458 tracktree->Write(0,TObject::kOverwrite);
9d09c308 459 delete tracktree;
91923a38 460
9d09c308 461 tbuffer->SetOwner();
462 pbuffer->SetOwner();
463 delete pbuffer;
464 delete tbuffer;
465
466 histoOutput->Close();
91923a38 467 return 0;
468 }
27bb2a54 469/********************************************************************/
470
471Bool_t AliHBTReaderInternal::FindIndex(TClonesArray* arr,Int_t idx)
472{
473//Checks if in the array exists already partilce with Unique ID idx
474 if (arr == 0x0)
475 {
476 ::Error("FindIndex","Array is 0x0");
477 return kTRUE;
478 }
479 TIter next(arr);
480 AliHBTParticle* p;
481 while (( p = (AliHBTParticle*)next()))
482 {
483 if (p->GetUID() == idx) return kTRUE;
484 }
485 return kFALSE;
486}