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