Catching up to NewIO -> Particle stores all passible PID and their probabilities
[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
16 ClassImp(AliHBTReaderInternal)
17 /********************************************************************/
18
19 AliHBTReaderInternal::AliHBTReaderInternal()
20 {
21 //Defalut constructor
22   fParticles = 0x0; 
23   fTracks = 0x0;
24   fIsRead = kFALSE;
25 }
26 /********************************************************************/
27
28 AliHBTReaderInternal::AliHBTReaderInternal(const char *filename):fFileName(filename)
29
30 //constructor 
31 //filename - name of file to open
32  fParticles = new AliHBTRun();
33  fTracks    = new AliHBTRun();
34  fIsRead = kFALSE;
35 }
36 /********************************************************************/
37
38 AliHBTReaderInternal::AliHBTReaderInternal(TObjArray* dirs, const char *filename):
39   AliHBTReader(dirs),fFileName(filename)
40
41 //ctor
42 //dirs contains strings with directories to look data in
43 //filename - name of file to open
44  fParticles = new AliHBTRun();
45  fTracks    = new AliHBTRun();
46  fIsRead = kFALSE;
47 }
48 /********************************************************************/
49
50 AliHBTReaderInternal::~AliHBTReaderInternal()
51  {
52  //desctructor
53    delete fParticles;
54    delete fTracks;
55  }
56 /********************************************************************/
57
58 AliHBTEvent* AliHBTReaderInternal::GetParticleEvent(Int_t n)
59  {
60  //returns Nth event with simulated particles
61    if (!fIsRead) 
62     if(Read(fParticles,fTracks))
63      {
64        Error("GetParticleEvent","Error in reading");
65        return 0x0;
66      }
67    return fParticles->GetEvent(n);
68  }
69 /********************************************************************/
70
71 AliHBTEvent* AliHBTReaderInternal::GetTrackEvent(Int_t n)
72  {
73  //returns Nth event with reconstructed tracks
74    if (!fIsRead) 
75     if(Read(fParticles,fTracks))
76      {
77        Error("GetTrackEvent","Error in reading");
78        return 0x0;
79      }
80    return fTracks->GetEvent(n);
81  }
82 /********************************************************************/
83
84 Int_t AliHBTReaderInternal::GetNumberOfPartEvents()
85  {
86  //returns number of events of particles
87    if (!fIsRead) 
88     if ( Read(fParticles,fTracks))
89      {
90        Error("GetNumberOfPartEvents","Error in reading");
91        return 0;
92      }
93    return fParticles->GetNumberOfEvents();
94  }
95
96 /********************************************************************/
97 Int_t AliHBTReaderInternal::GetNumberOfTrackEvents()
98  {
99  //returns number of events of tracks
100   if (!fIsRead)
101     if(Read(fParticles,fTracks))
102      {
103        Error("GetNumberOfTrackEvents","Error in reading");
104        return 0;
105      }
106   return fTracks->GetNumberOfEvents();
107  }
108 /********************************************************************/
109
110 Int_t AliHBTReaderInternal::Read(AliHBTRun* particles, AliHBTRun *tracks)
111  {
112  //reads data and puts put to the particles and tracks objects
113  //reurns 0 if everything is OK
114  //
115   Info("Read","");
116   Int_t i; //iterator and some temprary values
117   Int_t totalNevents = 0; //total number of read events
118   Int_t Nevents = 0;
119   Int_t currentdir = 0;
120   Int_t Ndirs;
121   Int_t counter;
122   TFile *aFile;//file with tracks
123   AliHBTParticle* tpart = 0x0, *ttrack = 0x0;
124   
125   if (!particles) //check if an object is instatiated
126    {
127      Error("Read"," particles object must instatiated before passing it to the reader");
128    }
129   if (!tracks)  //check if an object is instatiated
130    {
131      Error("Read"," tracks object must instatiated before passing it to the reader");
132    }
133   particles->Reset();//clear runs == delete all old events
134   tracks->Reset();
135  
136   if (fDirs) //if array with directories is supplied by user
137    {
138      Ndirs = fDirs->GetEntries(); //get the number if directories
139    }
140   else
141    {
142      Ndirs = 0; //if the array is not supplied read only from current directory
143    }
144
145   TClonesArray* pbuffer = new TClonesArray("AliHBTParticle",15000);
146   TClonesArray* tbuffer = new TClonesArray("AliHBTParticle",15000);
147
148   do  //do{}while; is OK even if 0 dirs specified. In that case we try to read from "./"
149    {
150     
151     if( (i=OpenFile(aFile, currentdir)) )
152      {
153        Error("Read","Exiting due to problems with opening files. Errorcode %d",i);
154        return i;
155      }
156    /***************************/
157    /***************************/
158    /***************************/
159     
160      TTree* tree = (TTree*)aFile->Get("data");
161      if (tree == 0x0)
162       {
163        Error("Read","Can not get the tree");
164        return 1;
165       }
166      
167      TBranch *trackbranch=tree->GetBranch("tracks");//get the branch with tracks
168      if (trackbranch == 0x0) ////check if we got the branch
169        {//if not return with error
170          Warning("Read","Can't find a branch with tracks !\n"); 
171        }
172      else
173       {
174         trackbranch->SetAddress(&tbuffer);
175       }
176
177      TBranch *partbranch=tree->GetBranch("particles");//get the branch with particles
178      if (partbranch == 0x0) ////check if we got the branch
179        {//if not return with error
180          Warning("Read","Can't find a branch with particles !\n"); 
181        }
182      else
183       {
184         partbranch->SetAddress(&pbuffer);
185       }
186      
187      Nevents = (Int_t)tree->GetEntries();
188      Info("Read","________________________________________________________");
189      Info("Read","Found %d event(s) in directory %s",Nevents,GetDirName(currentdir).Data());
190      
191      for (Int_t currentEvent =0; currentEvent<Nevents;currentEvent++)
192       {
193        Info("Read","Event %d",currentEvent);
194        tree->GetEvent(currentEvent);
195        
196        counter = 0;  
197        if (partbranch && trackbranch)
198         {
199            for(i = 0; i < pbuffer->GetEntries(); i++)
200              {
201                tpart = dynamic_cast<AliHBTParticle*>(pbuffer->At(i));
202                ttrack =  dynamic_cast<AliHBTParticle*>(tbuffer->At(i));
203
204                if( tpart == 0x0 ) continue; //if returned pointer is NULL
205                if( tpart->GetPDG()==0x0 ) continue; //if particle has crezy PDG code (not known to our database)
206                if( Pass(tpart) ) continue; //check if we are intersted with particles of this type 
207                                                           //if not take next partilce
208                AliHBTParticle* part = new AliHBTParticle(*tpart);
209                AliHBTParticle* track = new AliHBTParticle(*ttrack);
210                particles->AddParticle(totalNevents,part);//put track and particle on the run
211                tracks->AddParticle(totalNevents,track);
212                counter++;
213              }
214             Info("Read","   Read: %d particles and tracks",counter);
215         }
216        else
217         {  
218          if (partbranch)
219           {
220             for(i = 0; i < pbuffer->GetEntries(); i++)
221              {
222                tpart = dynamic_cast<AliHBTParticle*>(pbuffer->At(i));
223                if(tpart == 0x0) continue; //if returned pointer is NULL
224                if(tpart->GetPDG() == 0x0) continue; //if particle has crezy PDG code (not known to our database)
225                if(Pass(tpart)) continue; //check if we are intersted with particles of this type 
226                                          //if not take next partilce
227  
228                AliHBTParticle* part = new AliHBTParticle(*tpart);
229                particles->AddParticle(totalNevents,part);//put track and particle on the run
230                counter++;
231              }
232             Info("Read","   Read: %d particles and tracks",counter);
233           }
234          else Info("Read","   Read: 0 particles");
235          
236          if (trackbranch)
237           {
238             for(i = 0; i < tbuffer->GetEntries(); i++)
239              {
240                tpart = dynamic_cast<AliHBTParticle*>(tbuffer->At(i));
241                if(tpart == 0x0) continue; //if returned pointer is NULL
242                if(tpart->GetPDG() == 0x0) continue; //if particle has crezy PDG code (not known to our database)
243                if(Pass(tpart)) continue; //check if we are intersted with particles of this type 
244                                                    //if not take next partilce
245                AliHBTParticle* part = new AliHBTParticle(*tpart);
246                tracks->AddParticle(totalNevents,part);//put track and particle on the run
247                counter++;
248              }
249             Info("Read","   Read: %d tracks",counter);
250           }
251          else Info("Read","   Read: 0 tracks");
252         }
253         totalNevents++;
254        }
255
256     /***************************/
257     /***************************/
258     /***************************/
259     currentdir++;
260     delete tree;
261     aFile->Close();
262     delete aFile;
263     aFile = 0x0;
264
265    }while(currentdir < Ndirs);
266
267   delete pbuffer;
268   delete tbuffer;
269   fIsRead = kTRUE;
270   return 0;
271 }
272
273 /********************************************************************/
274
275 Int_t AliHBTReaderInternal::OpenFile(TFile*& aFile,Int_t event)
276 {
277
278    const TString& dirname = GetDirName(event); 
279    if (dirname == "")
280     {
281       Error("OpenFile","Can not get directory name");
282       return 4;
283     }
284    
285    TString filename = dirname +"/"+ fFileName;
286    aFile = TFile::Open(filename.Data());
287    if ( aFile  == 0x0 ) 
288      {
289        Error("OpenFiles","Can't open file with tacks named %s",filename.Data());
290        return 1;
291      }
292    if (!aFile->IsOpen())
293      {
294        Error("OpenFiles","Can't open file with tacks named %s",filename.Data());
295        return 1;
296      }
297    return 0; 
298 }
299 /********************************************************************/
300
301 Int_t AliHBTReaderInternal::Write(AliHBTReader* reader,const char* outfile)
302  {
303   //reads tracks from reader and writes runs to file
304   Int_t i,j;
305   
306   ::Info("AliHBTReaderInternal::Write","________________________________________________________");
307   ::Info("AliHBTReaderInternal::Write","________________________________________________________");
308   ::Info("AliHBTReaderInternal::Write","________________________________________________________");
309
310   TFile *histoOutput = TFile::Open(outfile,"recreate");
311   
312   if (!histoOutput->IsOpen())
313    {
314      ::Error("AliHBTReaderInternal::Write","File is not opened");
315      return 1;
316    }
317     
318   TTree *tracktree = new TTree("data","Tree with tracks");
319
320   TClonesArray* pbuffer = new TClonesArray("AliHBTParticle",15000);
321   TClonesArray* tbuffer = new TClonesArray("AliHBTParticle",15000);
322
323   TClonesArray &particles = *pbuffer;
324   TClonesArray &tracks = *tbuffer;
325     
326   TString name("Tracks");
327   
328   Int_t NT = reader->GetNumberOfTrackEvents();
329   Int_t NP = reader->GetNumberOfPartEvents();
330   
331   Bool_t trck = (NT > 0) ? kTRUE : kFALSE;
332   Bool_t part = (NP > 0) ? kTRUE : kFALSE;
333
334   TBranch *trackbranch = 0x0, *partbranch = 0x0;
335   
336   if (trck) trackbranch = tracktree->Branch("tracks","TClonesArray",&tbuffer);
337   if (part) partbranch = tracktree->Branch("particles","TClonesArray",&pbuffer);
338   
339   if ( (trck) && (part) && (NP != NT))
340    {
341      ::Warning("AliHBTReaderInternal::Write","Number of track and particle events is different");
342    }
343   
344   Int_t N;
345   if (NT >= NP ) N = NT; else N = NP;
346
347   for ( i =0;i< N; i++)
348     {
349       ::Info("AliHBTReaderInternal::Write","Event %d",i+1);
350       if (trck && (i<=NT))
351        {
352          AliHBTEvent* trackev = reader->GetTrackEvent(i);
353          for ( j = 0; j< trackev->GetNumberOfParticles();j++)
354           {
355             const AliHBTParticle& t= *(trackev->GetParticle(j));
356             new (tracks[j]) AliHBTParticle(t);
357           }
358          ::Info("AliHBTReaderInternal::Write","    Tracks: %d",j);
359        }else ::Info("AliHBTReaderInternal::Write","NO TRACKS");
360       
361       if (part && (i<=NP))
362        {
363         AliHBTEvent* partev = reader->GetParticleEvent(i);
364         for ( j = 0; j< partev->GetNumberOfParticles();j++)
365          {
366            const AliHBTParticle& part= *(partev->GetParticle(j));
367            new (particles[j]) AliHBTParticle(part);
368          }
369          ::Info("AliHBTReaderInternal::Write","    Particles: %d",j);
370        }else ::Info("AliHBTReaderInternal::Write","NO PARTICLES");
371
372       histoOutput->cd();
373       tracktree->Fill();
374       tracktree->AutoSave();
375       tbuffer->Delete();
376       pbuffer->Delete();
377      }
378
379   histoOutput->cd();
380   tracktree->Write(0,TObject::kOverwrite);
381   delete tracktree;
382
383   tbuffer->SetOwner();
384   pbuffer->SetOwner();
385   delete pbuffer;
386   delete tbuffer;
387
388   histoOutput->Close();
389   return 0;
390  }