020aa2079dcf700c937c23b6ed7301b155c2bce7
[u/mrichter/AliRoot.git] / EVE / Alieve / KineTools.cxx
1 // $Header$
2
3 #include "KineTools.h"
4
5 #include <TObject.h>
6 #include <TTree.h>
7 #include <TBranchElement.h>
8 #include <TClonesArray.h>
9
10 #include <AliStack.h>
11 #include <AliTrackReference.h>
12
13 #include "Reve/Track.h"
14 #include "Reve/RenderElement.h"
15
16 #include <algorithm>
17 #include <map>
18
19 //______________________________________________________________________
20 // KineTools
21 //
22
23 using namespace Reve;
24 using namespace Alieve;
25 using namespace std;
26
27 ClassImp(KineTools)
28
29 KineTools::KineTools()
30 {}
31
32 /**************************************************************************/
33
34 void KineTools::SetDaughterPathMarks(TrackList* cont,  AliStack* stack)
35 {
36   // Import daughters birth points.
37
38   RenderElement::List_i  iter = cont->BeginChildren();
39
40   while(iter != cont->EndChildren())
41   {
42     Track* track = dynamic_cast<Track*>(*iter); 
43     TParticle* p = stack->Particle(track->GetLabel());
44     if(p->GetNDaughters()) {
45       Int_t d0 = p->GetDaughter(0), d1 = p->GetDaughter(1);
46       for(int d=d0; d>0 && d<=d1;++d) 
47       { 
48         TParticle* dp = stack->Particle(d);
49         Reve::PathMark* pm = new PathMark( PathMark::Daughter);
50         pm->V.Set(dp->Vx(),dp->Vy(), dp->Vz());
51         pm->P.Set(dp->Px(),dp->Py(), dp->Pz()); 
52         pm->time = dp->T();
53         track->AddPathMark(pm);
54       }
55     }
56     ++iter;
57   }
58 }
59
60 /**************************************************************************/
61
62 namespace {
63 struct cmp_pathmark {
64   bool operator()(PathMark* const & a, PathMark* const & b)
65   { return a->time < b->time; }
66 };
67 }
68
69 void KineTools::SetPathMarks(TrackList* cont, AliStack* stack , TTree* treeTR)
70 {
71   // set decay and reference points
72
73   static const Exc_t eH("KineTools::ImportPathMarks");
74
75   if(treeTR == 0) {
76     SetDaughterPathMarks(cont, stack);
77     return;
78   }
79
80   map<Int_t, Track::vpPathMark_t > refs;
81
82   Int_t nPrimaries = (Int_t) treeTR->GetEntries();
83   TIter next(treeTR->GetListOfBranches());
84   TBranchElement* el;
85   Bool_t isRef = kTRUE;
86   treeTR->SetBranchStatus("*", 0);
87
88   while ((el = (TBranchElement*) next()))
89   {
90     if (strcmp("AliRun",el->GetName()) == 0)
91       isRef = kFALSE;
92
93     treeTR->SetBranchStatus(Form("%s*", el->GetName()), 1);
94     for (Int_t iPrimPart = 0; iPrimPart<nPrimaries; iPrimPart++) 
95     {
96      
97       TClonesArray* arr = 0;
98       treeTR->SetBranchAddress(el->GetName(), &arr);
99       treeTR->GetEntry(iPrimPart);
100     
101       for (Int_t iTrackRef = 0; iTrackRef < arr->GetEntriesFast(); iTrackRef++) 
102       {
103         AliTrackReference* atr = (AliTrackReference*)arr->At(iTrackRef);
104         Int_t track = atr->GetTrack();
105         if(atr->TestBit(TObject::kNotDeleted)) {
106           if(track > 0)
107           { 
108             PathMark* pm;
109             if(isRef) 
110               pm = new PathMark(PathMark::Reference);
111             else
112               pm = new PathMark(PathMark::Decay);
113               
114             pm->V.Set(atr->X(),atr->Y(), atr->Z());
115             pm->P.Set(atr->Px(),atr->Py(), atr->Pz());  
116             pm->time = atr->GetTime();
117
118             Track::vpPathMark_t& v = refs[track];
119             v.push_back(pm);
120           }
121           else
122             throw(eH + Form("negative label for entry %d in branch %s.",
123                             iTrackRef, el->GetName()));
124         }
125       } // loop track refs 
126       treeTR->SetBranchAddress(el->GetName(), 0);
127     } // loop primaries, clones arrays
128     treeTR->SetBranchStatus(Form("%s*", el->GetName()), 0);
129   } // end loop through top branches
130   treeTR->SetBranchStatus("*", 1);
131
132   // sort references and add it to tracks
133   RenderElement::List_i  cit = cont->BeginChildren();
134   while(cit != cont->EndChildren())
135   {
136     Track* track = dynamic_cast<Track*>(*cit);
137
138     // add daughters path marks in the map 
139     TParticle* p = stack->Particle(track->GetLabel());
140     if(p->GetNDaughters()) {
141       for(int d=p->GetFirstDaughter(); d>0 && d<=p->GetLastDaughter();++d) 
142       { 
143         TParticle* dp = stack->Particle(d);
144         Reve::PathMark* pm = new PathMark( PathMark::Daughter);
145         pm->V.Set(dp->Vx(),dp->Vy(), dp->Vz());
146         pm->P.Set(dp->Px(),dp->Py(), dp->Pz()); 
147         pm->time = dp->T();
148         Track::vpPathMark_t& v = refs[track->GetLabel()];
149         v.push_back(pm);
150       }
151     }
152     
153     map<Int_t, Track::vpPathMark_t > ::iterator mi = refs.find(track->GetLabel());
154     if(mi != refs.end()) {
155       Track::vpPathMark_t& v = mi->second;
156       sort(v.begin(), v.end(), cmp_pathmark());
157       for(Track::vpPathMark_i i=v.begin(); i!=v.end(); ++i){
158         track->AddPathMark(*i);
159       }
160     }
161     ++cit;
162   }
163 }