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