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