Bug correction (energy recalculation when adding new incarnation of the particle).
[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            if( ttrack == 0x0 ) continue; //if returned pointer is NULL
137
138            if (ttrack->GetUID() != tpart->GetUID())
139              {
140                Error("ReadNext","Sth. is wrong: Track and Particle has different UID.");
141                Error("ReadNext","They probobly do not correspond to each other.");
142              }
143
144            for (Int_t s = 0; s < ttrack->GetNumberOfPids(); s++)
145             {
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
148                                           //if not take next partilce
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) )
172                 {
173                   delete track;
174                   continue; 
175                 }
176               AliHBTParticle* part = new AliHBTParticle(*tpart);//particle has only 1 incarnation (real)
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
200               AliHBTParticle* part = new AliHBTParticle(*tpart);
201               part->SetPdgCode(tpart->GetNthPid(s),tpart->GetNthPidProb(s));
202               if( Pass(part) )
203                 {
204                   delete part;
205                   continue; 
206                 }
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          {
219            ttrack =  dynamic_cast<AliHBTParticle*>(fTrackBuffer->At(i));
220            if( ttrack == 0x0 ) continue; //if returned pointer is NULL
221
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) )
249                 {
250                   delete track;
251                   continue; 
252                 }
253               fTracksEvent->AddParticle(track);
254
255               counter++;
256             }
257          }
258         Info("ReadNext","   Read: %d tracks",counter);
259       }
260      else Info("ReadNext","   Read: 0 tracks.");
261     }
262     
263     if (fPartBuffer) fPartBuffer->Delete();
264     if (fTrackBuffer) fTrackBuffer->Delete();
265     
266     fCurrentEvent++;
267     fNEventsRead++;
268     
269     return 0;
270
271    }while(fCurrentDir < GetNumberOfDirs());
272
273   return 1;//no more directories to read
274 }
275 /********************************************************************/
276
277 Int_t AliHBTReaderInternal::OpenNextFile()
278 {
279   //open file in current directory
280    const TString& dirname = GetDirName(fCurrentDir);
281    if (dirname == "")
282     {
283       Error("OpenNextFile","Can not get directory name");
284       return 4;
285     }
286    
287    TString filename = dirname +"/"+ fFileName;
288    fFile = TFile::Open(filename.Data());
289    if ( fFile  == 0x0 ) 
290      {
291        Error("OpenNextFile","Can't open file named %s",filename.Data());
292        return 1;
293      }
294    if (fFile->IsOpen() == kFALSE)
295      {
296        Error("OpenNextFile","Can't open filenamed %s",filename.Data());
297        return 1;
298      }
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
338    return 0; 
339 }
340 /********************************************************************/
341
342 Int_t AliHBTReaderInternal::Write(AliHBTReader* reader,const char* outfile, Bool_t multcheck)
343  {
344   //reads tracks from reader and writes runs to file
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   
350   Int_t i,j;
351   
352   ::Info("AliHBTReaderInternal::Write","________________________________________________________");
353   ::Info("AliHBTReaderInternal::Write","________________________________________________________");
354   ::Info("AliHBTReaderInternal::Write","________________________________________________________");
355
356   TFile *histoOutput = TFile::Open(outfile,"recreate");
357   
358   if (!histoOutput->IsOpen())
359    {
360      ::Error("AliHBTReaderInternal::Write","File is not opened");
361      return 1;
362    }
363     
364   TTree *tracktree = new TTree("data","Tree with tracks");
365
366   TClonesArray* pbuffer = new TClonesArray("AliHBTParticle",15000);
367   TClonesArray* tbuffer = new TClonesArray("AliHBTParticle",15000);
368
369   TClonesArray &particles = *pbuffer;
370   TClonesArray &tracks = *tbuffer;
371     
372   TString name("Tracks");
373   
374   Int_t nt = reader->GetNumberOfTrackEvents();
375   Int_t np = reader->GetNumberOfPartEvents();
376   
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;
382
383   TBranch *trackbranch = 0x0, *partbranch = 0x0;
384   
385   if (trck) trackbranch = tracktree->Branch("tracks",&tbuffer,32000,0);
386   if (part) partbranch = tracktree->Branch("particles",&pbuffer,32000,0);
387   
388   if ( (trck) && (part) && (np != nt))
389    {
390      ::Warning("AliHBTReaderInternal::Write","Number of track and particle events is different");
391    }
392   
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);
398
399   for ( i =0;i< n; i++)
400     {
401       ::Info("AliHBTReaderInternal::Write","Event %d",i+1);
402       Int_t counter = 0;
403       if (trck && (i<=nt))
404        { 
405          AliHBTEvent* trackev = reader->GetTrackEvent(i);
406          for ( j = 0; j< trackev->GetNumberOfParticles();j++)
407           {
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);
421           }
422          ::Info("AliHBTReaderInternal::Write","    Tracks: %d",tracks.GetEntries());
423        }else ::Info("AliHBTReaderInternal::Write","NO TRACKS");
424       
425       counter = 0;
426       if (part && (i<=np))
427        {
428 //        ::Warning("AliHBTReaderInternal::Write","Find index switched off!!!");
429
430         AliHBTEvent* partev = reader->GetParticleEvent(i);
431         for ( j = 0; j< partev->GetNumberOfParticles();j++)
432          {
433            const AliHBTParticle& part= *(partev->GetParticle(j));
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);
446          }
447          ::Info("AliHBTReaderInternal::Write","    Particles: %d",particles.GetEntries());
448        }else ::Info("AliHBTReaderInternal::Write","NO PARTICLES");
449
450       histoOutput->cd();
451       tracktree->Fill();
452       tracktree->AutoSave();
453       tbuffer->Delete();
454       pbuffer->Delete();
455      }
456
457   histoOutput->cd();
458   tracktree->Write(0,TObject::kOverwrite);
459   delete tracktree;
460
461   tbuffer->SetOwner();
462   pbuffer->SetOwner();
463   delete pbuffer;
464   delete tbuffer;
465
466   histoOutput->Close();
467   return 0;
468  }
469 /********************************************************************/
470
471 Bool_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 }