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