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