604d8aa6c07998d960a506ad590aa3409f495e57
[u/mrichter/AliRoot.git] / EVE / EveBase / AliEveKineTools.cxx
1 // $Id$
2 // Main authors: Matevz Tadel & Alja Mrak-Tadel: 2006, 2007
3
4 /**************************************************************************
5  * Copyright(c) 1998-2008, ALICE Experiment at CERN, all rights reserved. *
6  * See http://aliceinfo.cern.ch/Offline/AliRoot/License.html for          *
7  * full copyright notice.                                                 *
8  **************************************************************************/
9
10 #include "AliEveKineTools.h"
11
12 #include <TObject.h>
13 #include <TTree.h>
14 #include <TBranchElement.h>
15 #include <TClonesArray.h>
16
17 #include <AliStack.h>
18 #include <AliTrackReference.h>
19
20 #include <TEveTrack.h>
21 #include <TEveElement.h>
22
23 #include <algorithm>
24 #include <map>
25
26 //______________________________________________________________________________
27 // AliEveKineTools
28 //
29
30 using namespace std;
31
32 ClassImp(AliEveKineTools)
33
34 AliEveKineTools::AliEveKineTools()
35 {}
36
37 /******************************************************************************/
38
39 void AliEveKineTools::SetDaughterPathMarks(TEveElement* cont, AliStack* stack, Bool_t recurse)
40 {
41   // Import daughters birth points.
42
43   TEveElement::List_i  iter = cont->BeginChildren();
44
45   while(iter != cont->EndChildren())
46   {
47     TEveTrack* track = dynamic_cast<TEveTrack*>(*iter);
48     TParticle* p = stack->Particle(track->GetLabel());
49     if(p->GetNDaughters()) {
50       Int_t d0 = p->GetDaughter(0), d1 = p->GetDaughter(1);
51       for(int d=d0; d>0 && d<=d1; ++d)
52       {
53         TParticle* dp = stack->Particle(d);
54         TEvePathMark* pm = new TEvePathMark(TEvePathMark::kDaughter);
55         pm->fV.Set(dp->Vx(), dp->Vy(), dp->Vz());
56         pm->fP.Set(dp->Px(), dp->Py(), dp->Pz());
57         pm->fTime = dp->T();
58         track->AddPathMark(pm);
59       }
60       if (recurse)
61         SetDaughterPathMarks(track, stack, recurse);
62     }
63     ++iter;
64   }
65 }
66
67 /******************************************************************************/
68
69 namespace {
70 struct cmp_pathmark
71 {
72   bool operator()(TEvePathMark* const & a, TEvePathMark* const & b)
73   { return a->fTime < b->fTime; }
74 };
75
76 void slurp_tracks(std::map<Int_t, TEveTrack*>& tracks, TEveElement* cont, Bool_t recurse)
77 {
78   TEveElement::List_i citer = cont->BeginChildren();
79   while(citer != cont->EndChildren())
80   {
81     TEveTrack* track = dynamic_cast<TEveTrack*>(*citer);
82     tracks[track->GetLabel()] = track;
83     if (recurse)
84       slurp_tracks(tracks, track, recurse);
85     ++citer;
86   }
87 }
88
89 }
90
91 void AliEveKineTools::SetTrackReferences(TEveElement* cont, TTree* treeTR, Bool_t recurse)
92 {
93   // set decay and reference points
94
95   static const TEveException eH("AliEveKineTools::ImportPathMarks");
96
97   // Fill map
98   std::map<Int_t, TEveTrack*> tracks;
99   slurp_tracks(tracks, cont, recurse);
100
101   Int_t nPrimaries = (Int_t) treeTR->GetEntries();
102   TIter next(treeTR->GetListOfBranches());
103   TBranchElement* el;
104   Bool_t isRef = kTRUE;
105
106   while ((el = (TBranchElement*) next()))
107   {
108     if (strcmp("AliRun",el->GetName()) == 0)
109       isRef = kFALSE;
110
111     TClonesArray* arr = 0;
112     el->SetAddress(&arr);
113     for (Int_t iPrimPart = 0; iPrimPart<nPrimaries; iPrimPart++)
114     {
115       el->GetEntry(iPrimPart);
116
117       Int_t last_label = -1;
118       std::map<Int_t, TEveTrack*>::iterator iter = tracks.end();
119       Int_t Nent =  arr->GetEntriesFast();
120       for (Int_t iTrackRef = 0; iTrackRef < Nent; iTrackRef++)
121       {
122         AliTrackReference* atr = (AliTrackReference*)arr->UncheckedAt(iTrackRef);
123
124         Int_t label = atr->GetTrack();
125         if (label < 0)
126           throw(eH + Form("negative label for entry %d in branch %s.",
127                           iTrackRef, el->GetName()));
128
129         if(label != last_label) {
130           iter = tracks.find(label);
131           last_label = label;
132         }
133
134         if (iter != tracks.end()) {
135           TEvePathMark* pm = new TEvePathMark(isRef ? TEvePathMark::kReference : TEvePathMark::kDecay);
136           pm->fV.Set(atr->X(),atr->Y(), atr->Z());
137           pm->fP.Set(atr->Px(),atr->Py(), atr->Pz());
138           pm->fTime = atr->GetTime();
139           TEveTrack* track  = iter->second;
140           track->AddPathMark(pm);
141         }
142       } // loop track refs
143     } // loop primaries, clones arrays
144     delete arr;
145   } // end loop through top branches
146 }
147
148 void AliEveKineTools::SortPathMarks(TEveElement* cont, Bool_t recurse)
149 {
150   // Sort path-marks for all tracks by time.
151
152   // Fill map
153   std::map<Int_t, TEveTrack*> tracks;
154   slurp_tracks(tracks, cont, recurse);
155
156   // sort
157   for(std::map<Int_t, TEveTrack*>::iterator j=tracks.begin(); j!=tracks.end(); ++j)
158   {
159     j->second->SortPathMarksByTime();
160   }
161 }