aad455e6a1b12aad9844f3f48fc27c6335d8d0d8
[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    aFile->Close();
261    aFile = 0x0;
262
263    }while(currentdir < Ndirs);
264
265   fIsRead = kTRUE;
266   return 0;
267 }
268
269 /********************************************************************/
270
271 Int_t AliHBTReaderInternal::OpenFile(TFile*& aFile,Int_t event)
272 {
273
274    const TString& dirname = GetDirName(event); 
275    if (dirname == "")
276     {
277       Error("OpenFile","Can not get directory name");
278       return 4;
279     }
280    
281    TString filename = dirname +"/"+ fFileName;
282    aFile = TFile::Open(filename.Data());
283    if ( aFile  == 0x0 ) 
284      {
285        Error("OpenFiles","Can't open file with tacks named %s",filename.Data());
286        return 1;
287      }
288    if (!aFile->IsOpen())
289      {
290        Error("OpenFiles","Can't open file with tacks named %s",filename.Data());
291        return 1;
292      }
293    return 0; 
294 }
295 /********************************************************************/
296
297 Int_t AliHBTReaderInternal::Write(AliHBTReader* reader,const char* outfile)
298  {
299   //reads tracks from reader and writes runs to file
300   Int_t i,j;
301   
302   ::Info("AliHBTReaderInternal::Write","________________________________________________________");
303   ::Info("AliHBTReaderInternal::Write","________________________________________________________");
304   ::Info("AliHBTReaderInternal::Write","________________________________________________________");
305
306   TFile *histoOutput = TFile::Open(outfile,"recreate");
307   
308   if (!histoOutput->IsOpen())
309    {
310      ::Error("AliHBTReaderInternal::Write","File is not opened");
311      return 1;
312    }
313     
314   TTree *tracktree = new TTree("data","Tree with tracks");
315
316   TClonesArray* pbuffer = new TClonesArray("AliHBTParticle",15000);
317   TClonesArray* tbuffer = new TClonesArray("AliHBTParticle",15000);
318
319   TClonesArray &particles = *pbuffer;
320   TClonesArray &tracks = *tbuffer;
321     
322   TString name("Tracks");
323   
324   Int_t NT = reader->GetNumberOfTrackEvents();
325   Int_t NP = reader->GetNumberOfPartEvents();
326   
327   Bool_t trck = (NT > 0) ? kTRUE : kFALSE;
328   Bool_t part = (NP > 0) ? kTRUE : kFALSE;
329
330   TBranch *trackbranch = 0x0, *partbranch = 0x0;
331   
332   if (trck) trackbranch = tracktree->Branch("tracks","TClonesArray",&tbuffer);
333   if (part) partbranch = tracktree->Branch("particles","TClonesArray",&pbuffer);
334   
335   if ( (trck) && (part) && (NP != NT))
336    {
337      ::Warning("AliHBTReaderInternal::Write","Number of track and particle events is different");
338    }
339   
340   Int_t N;
341   if (NT >= NP ) N = NT; else N = NP;
342
343   for ( i =0;i< N; i++)
344     {
345       ::Info("AliHBTReaderInternal::Write","Event %d",i+1);
346       if (trck && (i<=NT))
347        {
348          AliHBTEvent* trackev = reader->GetTrackEvent(i);
349          for ( j = 0; j< trackev->GetNumberOfParticles();j++)
350           {
351             const AliHBTParticle& t= *(trackev->GetParticle(j));
352             new (tracks[j]) AliHBTParticle(t);
353           }
354          ::Info("AliHBTReaderInternal::Write","    Tracks: %d",j);
355        }else ::Info("AliHBTReaderInternal::Write","NO TRACKS");
356       
357       if (part && (i<=NP))
358        {
359         AliHBTEvent* partev = reader->GetParticleEvent(i);
360         for ( j = 0; j< partev->GetNumberOfParticles();j++)
361          {
362            const AliHBTParticle& part= *(partev->GetParticle(j));
363            new (particles[j]) AliHBTParticle(part);
364          }
365          ::Info("AliHBTReaderInternal::Write","    Particles: %d",j);
366        }else ::Info("AliHBTReaderInternal::Write","NO PARTICLES");
367
368       histoOutput->cd();
369       tracktree->Fill();
370       tracktree->AutoSave();
371       tbuffer->Delete();
372       pbuffer->Delete();
373      }
374
375   histoOutput->cd();
376   tracktree->Write(0,TObject::kOverwrite);
377   delete tracktree;
378
379   tbuffer->SetOwner();
380   pbuffer->SetOwner();
381   delete pbuffer;
382   delete tbuffer;
383
384   histoOutput->Close();
385   return 0;
386  }