kine_tracks.C
[u/mrichter/AliRoot.git] / EVE / alice-macros / kine_tracks.C
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 // Import tracks from kinematics-tree / particle-stack.
11 // Preliminary/minimal solution.
12
13 #include "TParticlePDG.h"
14
15 TEveTrackList*
16 kine_tracks(Double_t min_pt  = 0,     Double_t min_p   = 0,
17             Bool_t   pdg_col = kTRUE, Bool_t   recurse = kTRUE,
18             Bool_t   use_track_refs = kTRUE)
19 {
20   AliRunLoader* rl =  AliEveEventManager::AssertRunLoader();
21   rl->LoadKinematics();
22   AliStack* stack = rl->Stack();
23   if (!stack) {
24     Error("kine_tracks.C", "can not get kinematics.");
25     return 0;
26   }
27
28   gEve->DisableRedraw();
29
30   TEveTrackList* cont = new TEveTrackList("Kine Tracks");
31   cont->SetMainColor(3);
32   TEveTrackPropagator* rnrStyle = cont->GetPropagator();
33   // !!! Watch the '-', apparently different sign convention then for ESD.
34   rnrStyle->SetMagField( -0.1*gAlice->Field()->SolenoidField() );
35
36   gEve->AddElement(cont);
37   Int_t count = 0;
38   Int_t N = stack->GetNtrack();
39   Int_t Np = stack->GetNprimary();
40   for (Int_t i = 0; i < Np; ++i)
41   {
42     TParticle* p = stack->Particle(i);
43     if (p->GetStatusCode() <= 1)
44     {
45       if (p->Pt() < min_pt && p->P() < min_p) continue;
46
47       ++count;
48       TEveTrack* track = new TEveTrack(p, i, rnrStyle);
49
50       //PH The line below is replaced waiting for a fix in Root
51       //PH which permits to use variable siza arguments in CINT
52       //PH on some platforms (alphalinuxgcc, solariscc5, etc.)
53       //PH    track->SetName(Form("%s [%d]", p->GetName(), i));
54       char form[1000];
55       sprintf(form,"%s [%d]", p->GetName(), i);
56       track->SetName(form);
57       track->SetStdTitle();
58       Int_t ml = p->GetMother(0);
59       if (ml != -1)
60       {
61         track->SetTitle(Form("%s\nMother label=%d\nMother Pdg=%d",
62                              track->GetElementTitle(),
63                              ml, stack->Particle(ml)->GetPdgCode()));
64       }
65       set_track_color(track, pdg_col);
66
67       gEve->AddElement(track, cont);
68
69       if (recurse)
70         kine_daughters(track, stack, min_pt, min_p, pdg_col, recurse);
71     }
72   }
73
74   // set path marks
75   AliEveKineTools kt;
76   kt.SetDaughterPathMarks(cont, stack, recurse);
77   if (use_track_refs && rl->LoadTrackRefs() == 0)
78   {
79     kt.SetTrackReferences(cont, rl->TreeTR(), recurse);
80     rnrStyle->SetEditPathMarks(kTRUE);
81   }
82   kt.SortPathMarks(cont, recurse);
83
84   //PH  const Text_t* tooltip = Form("min pT=%.2lf, min P=%.2lf), N=%d", min_pt, min_p, count);
85   char tooltip[1000];
86   sprintf(tooltip,"min pT=%.2lf, min P=%.2lf), N=%d", min_pt, min_p, count);
87   cont->SetTitle(tooltip); // Not broadcasted automatically ...
88
89   cont->MakeTracks(recurse);
90   gEve->EnableRedraw();
91   gEve->Redraw3D();
92
93   return cont;
94 }
95
96 void kine_daughters(TEveTrack* parent,  AliStack* stack,
97                     Double_t     min_pt,  Double_t  min_p,
98                     Bool_t       pdg_col, Bool_t    recurse)
99 {
100   TParticle *p = stack->Particle(parent->GetLabel());
101   if (p->GetNDaughters() > 0)
102   {
103     TEveTrackPropagator* rs = parent->GetPropagator();
104     for (int d=p->GetFirstDaughter(); d>0 && d<=p->GetLastDaughter(); ++d)
105     {
106       TParticle* dp = stack->Particle(d);
107       if (dp->Pt() < min_pt && dp->P() < min_p) continue;
108
109       TEveTrack* dtrack = new TEveTrack(dp, d, rs);
110       char form[1000];
111       sprintf(form,"%s [%d]", dp->GetName(), d);
112       dtrack->SetName(form);
113       dtrack->SetStdTitle();
114       set_track_color(dtrack, pdg_col);
115
116       gEve->AddElement(dtrack, parent);
117
118       if (recurse)
119         kine_daughters(dtrack, stack, min_pt, min_p, pdg_col, recurse);
120     }
121   }
122 }
123
124 Color_t set_track_color(TEveTrack* t, Bool_t pdg_col)
125 {
126   if (pdg_col)
127     t->SetMainColor(get_pdg_color(t->GetPdg()));
128   else
129     t->SetMainColor(30);
130 }
131
132 Color_t get_pdg_color(Int_t pdg)
133 {
134   // PDG color indices
135   static const Color_t DefCol   = 30;
136   static const Color_t ECol     = 5;
137   static const Color_t MuCol    = 6;
138   static const Color_t GammaCol = 7;
139   static const Color_t MesCol1  = 3;
140   static const Color_t MesCol2  = 38;
141   static const Color_t BarCol   = 10;
142
143   Int_t pdga  = TMath::Abs(pdg);
144   Color_t col = DefCol;
145
146   // elementary  particles
147   if (pdga < 100) {
148     switch (pdga) {
149       case 11:
150         col = ECol; break;
151       case 12:
152         col = MuCol; break;
153       case 22:
154         col = GammaCol; break;
155     }
156   }
157   // mesons and barions
158   else if (pdga < 100000) {
159     Int_t i  = pdga;
160     Int_t i0 = i%10; i /= 10;
161     Int_t i1 = i%10; i /= 10;
162     Int_t i2 = i%10; i /= 10;
163     Int_t i3 = i%10; i /= 10;
164     Int_t i4 = i%10;
165     //printf("pdg(%d) quark indices (%d,%d,%d,%d,%d) \n",pdg, i4,i3,i2, i1, i0);
166     // meson
167     if ((i3 == 0) && ( i4 < 2)){
168       col = MesCol1; // quarks: i1,i2 (spin = i0)
169       if(i1 == 3 || i2 == 3)
170         col = MesCol2;
171     } // barion
172     else if ( i2 >= i1 && i3 >= i2 ) {
173       col = BarCol; // quarks: i1,i2, i3 (spin = i0))
174     }
175   }
176
177   return col;
178 }
179
180 /******************************************************************************/
181
182 TEveElement*
183 kine_track(Int_t  label,
184            Bool_t import_mother    = kTRUE, Bool_t import_daughters = kTRUE,
185            Bool_t pdg_col          = kTRUE, Bool_t recurse          = kTRUE,
186            TEveElement* cont = 0)
187
188 {
189   // Create mother and daughters tracks with given label.
190   // mother     -> particle with label
191   // daughters  -> daughters of label
192
193   if (label < 0) {
194     Warning("kine_track", "label not set.");
195     return 0;
196   }
197
198   AliRunLoader* rl =  AliEveEventManager::AssertRunLoader();
199   rl->LoadKinematics();
200   AliStack* stack = rl->Stack();
201   if (label >= stack->GetNtrack())
202   {
203     Warning("kine_track", "label out of range.");
204     return 0;
205   }
206
207   TParticle* p = stack->Particle(label);
208
209   if (import_mother || (import_daughters && p->GetNDaughters()))
210   {
211     TEveTrack* toptrack = 0;
212     TEveTrackList* tracklist = 0;
213     TEveTrackPropagator* rs = 0;
214
215     if (cont == 0)
216     {
217       TEveTrackList* tlist = new TEveTrackList
218         (Form("Kinematics of %d", label, p->GetNDaughters()));
219       cont = tlist;
220
221       TEveTrackPropagator* rnrStyle = tlist->GetPropagator();
222       // !!! Watch the '-', apparently different sign convention then for ESD.
223       rnrStyle->SetMagField( -0.1*gAlice->Field()->SolenoidField() );
224       char tooltip[1000];
225       sprintf(tooltip,"Ndaughters=%d", p->GetNDaughters());
226       tlist->SetTitle(tooltip);
227       rnrStyle->fMaxOrbs = 2;
228       rnrStyle->SetEditPathMarks(kTRUE);
229
230       gEve->AddElement(cont);
231       rs = tlist->GetPropagator();
232     }
233     else
234     {
235       // check if container is TEveTrackList or TEveTrack (has rnr-style)
236       TEveTrack* t = dynamic_cast<TEveTrack*>(cont);
237       if (t) {
238         rs = t->GetPropagator();
239       } else {
240         TEveTrackList* l = dynamic_cast<TEveTrackList*>(cont);
241         if (l)
242           rs = l->GetPropagator();
243         else
244           Error("kine_tracks.C", "TrackRenderStyle not set.");
245       }
246     }
247
248     if (import_mother)
249     {
250       TEveTrack* track = new TEveTrack(p, label, rs);
251       char form[1000];
252       sprintf(form,"%s [%d]", p->GetName(), label);
253       track->SetName(form);
254       track->SetStdTitle();
255       set_track_color(track, pdg_col);
256
257       track->MakeTrack();
258       gEve->AddElement(track, cont);
259       cont = track;
260     }
261
262     if (import_daughters && p->GetNDaughters())
263     {
264       for (int d=p->GetFirstDaughter(); d>0 && d<=p->GetLastDaughter(); ++d)
265       {
266         TParticle* dp = stack->Particle(d);
267         TEveTrack* track = new TEveTrack(dp, d, rs);
268         char form[1000];
269         sprintf(form,"%s [%d]", dp->GetName(), d);
270         track->SetName(form);
271         track->SetStdTitle();
272         set_track_color(track, pdg_col);
273
274         track->MakeTrack();
275         gEve->AddElement(track, cont);
276
277         if (recurse)
278           kine_daughters(track, stack, 0, 0, pdg_col, recurse);
279       }
280     }
281   }
282
283   gEve->Redraw3D();
284   return cont;
285 }
286