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