4c1e5e167885316f762981187f9bd504d6e94062
[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 <map>
24
25 //______________________________________________________________________________
26 // AliEveKineTools
27 //
28 // Tools for import of kinematics. Preliminary version.
29 //
30
31 using namespace std;
32
33 ClassImp(AliEveKineTools)
34
35
36 namespace {
37
38   typedef std::map<Int_t, TEveTrack*> TrackMap_t;
39
40   void MapTracks(TrackMap_t& map, TEveElement* cont, Bool_t recurse)
41   {
42     TEveElement::List_i i = cont->BeginChildren();
43     while (i != cont->EndChildren()) {
44       TEveTrack* track = dynamic_cast<TEveTrack*>(*i);
45       map[track->GetLabel()] = track;
46       if (recurse)
47         MapTracks(map, track, recurse);
48       ++i;
49     }
50   }
51 }
52
53 /**************************************************************************/
54
55 void AliEveKineTools::SetDaughterPathMarks(TEveElement* cont, AliStack* stack, Bool_t recurse)
56 {
57   // Import daughters birth points.
58
59   TEveElement::List_i  iter = cont->BeginChildren();
60   while(iter != cont->EndChildren())
61   {
62     TEveTrack* track = dynamic_cast<TEveTrack*>(*iter);
63     TParticle* p = stack->Particle(track->GetLabel());
64     if(p->GetNDaughters()) {
65       Int_t d0 = p->GetDaughter(0), d1 = p->GetDaughter(1);
66       for(int d=d0; d>0 && d<=d1; ++d)
67       {
68         TParticle* dp = stack->Particle(d);
69         TEvePathMark* pm = new TEvePathMark(TEvePathMark::kDaughter);
70         pm->fV.Set(dp->Vx(), dp->Vy(), dp->Vz());
71         pm->fP.Set(dp->Px(), dp->Py(), dp->Pz());
72         pm->fTime = dp->T();
73         track->AddPathMark(pm);
74       }
75       if (recurse)
76         SetDaughterPathMarks(track, stack, recurse);
77     }
78     ++iter;
79   }
80 }
81
82 /**************************************************************************/
83
84 void AliEveKineTools::SetTrackReferences(TEveElement* cont, TTree* treeTR, Bool_t recurse)
85 {
86   // Set decay and track reference path-marks.
87
88   static const TEveException eH("AliEveKineTools::ImportPathMarks");
89
90   TrackMap_t map;
91   MapTracks(map, cont, recurse);
92
93   Int_t nPrimaries = (Int_t) treeTR->GetEntries();
94   TIter next(treeTR->GetListOfBranches());
95   TBranchElement* el;
96   Bool_t isRef = kTRUE;
97
98   while ((el = (TBranchElement*) next()))
99   {
100     if (strcmp("AliRun",el->GetName()) == 0)
101       isRef = kFALSE;
102
103     TClonesArray* arr = 0;
104     el->SetAddress(&arr);
105     for (Int_t iPrimPart = 0; iPrimPart<nPrimaries; iPrimPart++)
106     {
107       el->GetEntry(iPrimPart);
108
109       Int_t last_label = -1;
110       TrackMap_t::iterator iter = map.end();
111       Int_t Nent =  arr->GetEntriesFast();
112       for (Int_t iTrackRef = 0; iTrackRef < Nent; iTrackRef++)
113       {
114         AliTrackReference* atr = (AliTrackReference*)arr->UncheckedAt(iTrackRef);
115
116         Int_t label = atr->GetTrack();
117         if (label < 0)
118           throw(eH + Form("negative label for entry %d in branch %s.",
119                           iTrackRef, el->GetName()));
120
121         if(label != last_label) {
122           iter = map.find(label);
123           last_label = label;
124         }
125
126         if (iter != map.end()) {
127           TEvePathMark* pm = new TEvePathMark(isRef ? TEvePathMark::kReference : TEvePathMark::kDecay);
128           pm->fV.Set(atr->X(),atr->Y(), atr->Z());
129           pm->fP.Set(atr->Px(),atr->Py(), atr->Pz());
130           pm->fTime = atr->GetTime();
131           TEveTrack* track  = iter->second;
132           track->AddPathMark(pm);
133         }
134       } // loop track refs
135     } // loop primaries in clones arrays
136     delete arr;
137   } // end loop through top branches
138 }
139
140 /**************************************************************************/
141
142 void AliEveKineTools::SortPathMarks(TEveElement* el, Bool_t recurse)
143 {
144   // Sort path-marks for all tracks by time.
145
146   TEveTrack* track = dynamic_cast<TEveTrack*>(el);
147   if(track) track->SortPathMarksByTime();
148
149   TEveElement::List_i i = el->BeginChildren();
150   while (i != el->EndChildren() && recurse) {
151     track = dynamic_cast<TEveTrack*>(el);
152     if (track) track->SortPathMarksByTime();
153     i++;
154   }
155 }