kPyMbDefault back to what it was.
[u/mrichter/AliRoot.git] / EVE / alice-macros / kine_tracks.C
CommitLineData
d810d0de 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 *
51346b82 7 * full copyright notice. *
d810d0de 8 **************************************************************************/
16718cdc 9
1b337638 10// Import tracks from kinematics-tree / particle-stack.
11// Preliminary/minimal solution.
12
3c2901eb 13#include <TEveManager.h>
14#include <TEveTrackPropagator.h>
15
16#include <EveBase/AliEveEventManager.h>
17#include <EveBase/AliEveMagField.h>
18#include <EveBase/AliEveTrack.h>
19#include <EveBase/AliEveKineTools.h>
20
21#include <AliRunLoader.h>
22#include <AliStack.h>
23#include <AliMagF.h>
24
fc9514f5 25#include "TParticlePDG.h"
a7867742 26
6336e5d7 27// Use magnetic-field as retrieved from GRP.
28Bool_t g_kine_tracks_true_field = kTRUE;
29
30// Use Runge-Kutta track stepper.
31Bool_t g_kine_tracks_rk_stepper = kFALSE;
32
3c2901eb 33//==============================================================================
34
35void kine_track_propagator_setup(TEveTrackPropagator* trkProp);
36
37TEveTrackList*
38kine_tracks(Double_t min_pt = 0, Double_t min_p = 0,
39 Bool_t pdg_col = kTRUE, Bool_t recurse = kTRUE,
40 Bool_t use_track_refs = kTRUE);
41
42void kine_daughters(AliEveTrack* parent, AliStack* stack,
43 Double_t min_pt, Double_t min_p,
44 Bool_t pdg_col, Bool_t recurse);
45
46void set_track_color(AliEveTrack* t, Bool_t pdg_col);
47Color_t get_pdg_color(Int_t pdg);
48
49TEveElement*
50kine_track(Int_t label,
51 Bool_t import_mother = kTRUE, Bool_t import_daughters = kTRUE,
52 Bool_t pdg_col = kTRUE, Bool_t recurse = kTRUE,
53 TEveElement* cont = 0);
54
55void kine_hide_neutrals(TEveElement* el=0, Int_t level=0);
6336e5d7 56
57//==============================================================================
58
59void kine_track_propagator_setup(TEveTrackPropagator* trkProp)
60{
61 AliMagF* fld = AliEveEventManager::AssertMagField();
62
63 if (g_kine_tracks_true_field)
64 {
65 trkProp->SetMagFieldObj(new AliEveMagField(fld));
66 }
67 else
68 {
69 trkProp->SetMagField(-0.1*fld->SolenoidField());
70 }
71 if (g_kine_tracks_rk_stepper)
72 {
73 trkProp->SetStepper(TEveTrackPropagator::kRungeKutta);
74 }
75}
76
77//==============================================================================
78
84aff7a4 79TEveTrackList*
3c2901eb 80kine_tracks(Double_t min_pt, Double_t min_p,
81 Bool_t pdg_col, Bool_t recurse,
82 Bool_t use_track_refs)
1b337638 83{
d810d0de 84 AliRunLoader* rl = AliEveEventManager::AssertRunLoader();
1b337638 85 rl->LoadKinematics();
86 AliStack* stack = rl->Stack();
8661a211 87 if (!stack)
88 {
1b337638 89 Error("kine_tracks.C", "can not get kinematics.");
90 return 0;
91 }
92
84aff7a4 93 gEve->DisableRedraw();
51346b82 94
95 TEveTrackList* cont = new TEveTrackList("Kine Tracks");
fbc350a3 96 cont->SetMainColor(3);
8661a211 97 TEveTrackPropagator* trkProp = cont->GetPropagator();
98
6336e5d7 99 kine_track_propagator_setup(trkProp);
1b337638 100
84aff7a4 101 gEve->AddElement(cont);
1b337638 102 Int_t count = 0;
92bd3b8a 103 Int_t Np = stack->GetNprimary();
104 for (Int_t i = 0; i < Np; ++i)
43d68fc4 105 {
92bd3b8a 106 TParticle* p = stack->Particle(i);
107 if (p->GetStatusCode() <= 1)
43d68fc4 108 {
ef9d2241 109 if (p->Pt() < min_pt && p->P() < min_p) continue;
43d68fc4 110
111 ++count;
0e33c639 112 AliEveTrack* track = new AliEveTrack(p, i, trkProp);
51346b82 113
43d68fc4 114 //PH The line below is replaced waiting for a fix in Root
115 //PH which permits to use variable siza arguments in CINT
116 //PH on some platforms (alphalinuxgcc, solariscc5, etc.)
117 //PH track->SetName(Form("%s [%d]", p->GetName(), i));
118 char form[1000];
119 sprintf(form,"%s [%d]", p->GetName(), i);
120 track->SetName(form);
fc9514f5 121 track->SetStdTitle();
a77bb3bc 122 Int_t ml = p->GetMother(0);
123 if (ml != -1)
124 {
125 track->SetTitle(Form("%s\nMother label=%d\nMother Pdg=%d",
126 track->GetElementTitle(),
127 ml, stack->Particle(ml)->GetPdgCode()));
128 }
fc9514f5 129 set_track_color(track, pdg_col);
130
84aff7a4 131 gEve->AddElement(track, cont);
ef9d2241 132
133 if (recurse)
134 kine_daughters(track, stack, min_pt, min_p, pdg_col, recurse);
43d68fc4 135 }
1b337638 136 }
ef9d2241 137
43d68fc4 138 // set path marks
51346b82 139 AliEveKineTools kt;
ef9d2241 140 kt.SetDaughterPathMarks(cont, stack, recurse);
32e219c2 141 if (use_track_refs && rl->LoadTrackRefs() == 0)
9a040a65 142 {
143 kt.SetTrackReferences(cont, rl->TreeTR(), recurse);
8661a211 144 trkProp->SetEditPathMarks(kTRUE);
9a040a65 145 }
146 kt.SortPathMarks(cont, recurse);
43d68fc4 147
ef9d2241 148 //PH const Text_t* tooltip = Form("min pT=%.2lf, min P=%.2lf), N=%d", min_pt, min_p, count);
7be1e8cd 149 char tooltip[1000];
ef9d2241 150 sprintf(tooltip,"min pT=%.2lf, min P=%.2lf), N=%d", min_pt, min_p, count);
1b337638 151 cont->SetTitle(tooltip); // Not broadcasted automatically ...
1b337638 152
ef9d2241 153 cont->MakeTracks(recurse);
84aff7a4 154 gEve->EnableRedraw();
155 gEve->Redraw3D();
1b337638 156
157 return cont;
158}
a7867742 159
0e33c639 160void kine_daughters(AliEveTrack* parent, AliStack* stack,
ef9d2241 161 Double_t min_pt, Double_t min_p,
162 Bool_t pdg_col, Bool_t recurse)
163{
164 TParticle *p = stack->Particle(parent->GetLabel());
51346b82 165 if (p->GetNDaughters() > 0)
ef9d2241 166 {
84aff7a4 167 TEveTrackPropagator* rs = parent->GetPropagator();
51346b82 168 for (int d=p->GetFirstDaughter(); d>0 && d<=p->GetLastDaughter(); ++d)
169 {
ef9d2241 170 TParticle* dp = stack->Particle(d);
171 if (dp->Pt() < min_pt && dp->P() < min_p) continue;
172
0e33c639 173 AliEveTrack* dtrack = new AliEveTrack(dp, d, rs);
ef9d2241 174 char form[1000];
175 sprintf(form,"%s [%d]", dp->GetName(), d);
176 dtrack->SetName(form);
fc9514f5 177 dtrack->SetStdTitle();
178 set_track_color(dtrack, pdg_col);
179
84aff7a4 180 gEve->AddElement(dtrack, parent);
ef9d2241 181
182 if (recurse)
183 kine_daughters(dtrack, stack, min_pt, min_p, pdg_col, recurse);
184 }
185 }
186}
a7867742 187
3c2901eb 188//==============================================================================
189
190void set_track_color(AliEveTrack* t, Bool_t pdg_col)
fc9514f5 191{
192 if (pdg_col)
193 t->SetMainColor(get_pdg_color(t->GetPdg()));
194 else
fbc350a3 195 t->SetMainColor(30);
fc9514f5 196}
197
539d6798 198Color_t get_pdg_color(Int_t pdg)
199{
fc9514f5 200 // PDG color indices
201 static const Color_t DefCol = 30;
202 static const Color_t ECol = 5;
203 static const Color_t MuCol = 6;
51346b82 204 static const Color_t GammaCol = 7;
fc9514f5 205 static const Color_t MesCol1 = 3;
206 static const Color_t MesCol2 = 38;
207 static const Color_t BarCol = 10;
208
209 Int_t pdga = TMath::Abs(pdg);
210 Color_t col = DefCol;
a7867742 211
212 // elementary particles
213 if (pdga < 100) {
214 switch (pdga) {
51346b82 215 case 11:
216 col = ECol; break;
a7867742 217 case 12:
218 col = MuCol; break;
219 case 22:
220 col = GammaCol; break;
221 }
222 }
fc9514f5 223 // mesons and barions
224 else if (pdga < 100000) {
a7867742 225 Int_t i = pdga;
3c2901eb 226 // Int_t i0 = i%10; // Not used at the moment.
227 i /= 10;
51346b82 228 Int_t i1 = i%10; i /= 10;
229 Int_t i2 = i%10; i /= 10;
230 Int_t i3 = i%10; i /= 10;
a7867742 231 Int_t i4 = i%10;
232 //printf("pdg(%d) quark indices (%d,%d,%d,%d,%d) \n",pdg, i4,i3,i2, i1, i0);
233 // meson
234 if ((i3 == 0) && ( i4 < 2)){
235 col = MesCol1; // quarks: i1,i2 (spin = i0)
236 if(i1 == 3 || i2 == 3)
237 col = MesCol2;
238 } // barion
239 else if ( i2 >= i1 && i3 >= i2 ) {
240 col = BarCol; // quarks: i1,i2, i3 (spin = i0))
241 }
242 }
fc9514f5 243
a7867742 244 return col;
245}
c7b57c0b 246
3c2901eb 247//==============================================================================
c7b57c0b 248
84aff7a4 249TEveElement*
b65b3044 250kine_track(Int_t label,
3c2901eb 251 Bool_t import_mother, Bool_t import_daughters,
252 Bool_t pdg_col, Bool_t recurse,
253 TEveElement* cont)
c7b57c0b 254{
fc9514f5 255 // Create mother and daughters tracks with given label.
256 // mother -> particle with label
257 // daughters -> daughters of label
258
c7b57c0b 259 if (label < 0) {
260 Warning("kine_track", "label not set.");
261 return 0;
262 }
51346b82 263
d810d0de 264 AliRunLoader* rl = AliEveEventManager::AssertRunLoader();
c7b57c0b 265 rl->LoadKinematics();
266 AliStack* stack = rl->Stack();
f1b2cd0b 267 if (label >= stack->GetNtrack())
268 {
269 Warning("kine_track", "label out of range.");
270 return 0;
271 }
272
c7b57c0b 273 TParticle* p = stack->Particle(label);
274
b65b3044 275 if (import_mother || (import_daughters && p->GetNDaughters()))
c7b57c0b 276 {
84aff7a4 277 TEveTrackPropagator* rs = 0;
2caed564 278
b65b3044 279 if (cont == 0)
280 {
84aff7a4 281 TEveTrackList* tlist = new TEveTrackList
ef9d2241 282 (Form("Kinematics of %d", label, p->GetNDaughters()));
283 cont = tlist;
c7b57c0b 284
8661a211 285 TEveTrackPropagator* trkProp = tlist->GetPropagator();
286
6336e5d7 287 kine_track_propagator_setup(trkProp);
8661a211 288
c7b57c0b 289 char tooltip[1000];
b65b3044 290 sprintf(tooltip,"Ndaughters=%d", p->GetNDaughters());
ef9d2241 291 tlist->SetTitle(tooltip);
8661a211 292 trkProp->fMaxOrbs = 2;
293 trkProp->SetEditPathMarks(kTRUE);
fc9514f5 294
84aff7a4 295 gEve->AddElement(cont);
296 rs = tlist->GetPropagator();
2caed564 297 }
ef9d2241 298 else
299 {
0e33c639 300 // check if container is TEveTrackList or AliEveTrack (has rnr-style)
301 AliEveTrack* t = dynamic_cast<AliEveTrack*>(cont);
fc9514f5 302 if (t) {
84aff7a4 303 rs = t->GetPropagator();
fc9514f5 304 } else {
84aff7a4 305 TEveTrackList* l = dynamic_cast<TEveTrackList*>(cont);
fc9514f5 306 if (l)
84aff7a4 307 rs = l->GetPropagator();
ef9d2241 308 else
2caed564 309 Error("kine_tracks.C", "TrackRenderStyle not set.");
2caed564 310 }
c7b57c0b 311 }
2caed564 312
b65b3044 313 if (import_mother)
314 {
0e33c639 315 AliEveTrack* track = new AliEveTrack(p, label, rs);
c7b57c0b 316 char form[1000];
317 sprintf(form,"%s [%d]", p->GetName(), label);
318 track->SetName(form);
fc9514f5 319 track->SetStdTitle();
320 set_track_color(track, pdg_col);
2caed564 321
fc9514f5 322 track->MakeTrack();
84aff7a4 323 gEve->AddElement(track, cont);
fc9514f5 324 cont = track;
c7b57c0b 325 }
326
51346b82 327 if (import_daughters && p->GetNDaughters())
c7b57c0b 328 {
51346b82 329 for (int d=p->GetFirstDaughter(); d>0 && d<=p->GetLastDaughter(); ++d)
330 {
c7b57c0b 331 TParticle* dp = stack->Particle(d);
0e33c639 332 AliEveTrack* track = new AliEveTrack(dp, d, rs);
c7b57c0b 333 char form[1000];
334 sprintf(form,"%s [%d]", dp->GetName(), d);
335 track->SetName(form);
fc9514f5 336 track->SetStdTitle();
337 set_track_color(track, pdg_col);
338
2caed564 339 track->MakeTrack();
84aff7a4 340 gEve->AddElement(track, cont);
fc9514f5 341
342 if (recurse)
343 kine_daughters(track, stack, 0, 0, pdg_col, recurse);
c7b57c0b 344 }
345 }
346 }
347
84aff7a4 348 gEve->Redraw3D();
c7b57c0b 349 return cont;
350}
2caed564 351
3c2901eb 352//==============================================================================
353
354void kine_hide_neutrals(TEveElement* el, Int_t level)
355{
356 if (el == 0)
357 {
358 el = gEve->GetCurrentEvent()->FindChild("Kine Tracks");
359 if (!el)
360 return;
361 }
362
363 TEveTrack* t = dynamic_cast<TEveTrack*>(el);
364 if (t && t->GetCharge() == 0)
365 t->SetRnrSelf(kFALSE);
366
367 for (TEveElement::List_i i = el->BeginChildren(); i != el->EndChildren(); ++i)
368 {
369 kine_hide_neutrals(*i, level + 1);
370 }
371
372 if (level == 0)
373 gEve->Redraw3D();
374}