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