[u/mrichter/AliRoot.git] / EVE / alice-macros / kine_tracks.C

// $Id$
// Main authors: Matevz Tadel & Alja Mrak-Tadel: 2006, 2007

/**************************************************************************
 * Copyright(c) 1998-2008, ALICE Experiment at CERN, all rights reserved. *
 * See http://aliceinfo.cern.ch/Offline/AliRoot/License.html for          *
 * full copyright notice.                                                 *
 **************************************************************************/

// Import tracks from kinematics-tree / particle-stack.
// Preliminary/minimal solution.

#include "TParticlePDG.h"

TEveTrackList*
kine_tracks(Double_t min_pt  = 0,     Double_t min_p   = 0,
	    Bool_t   pdg_col = kTRUE, Bool_t   recurse = kTRUE,
	    Bool_t   use_track_refs = kTRUE)
{
  AliRunLoader* rl =  AliEveEventManager::AssertRunLoader();
  rl->LoadKinematics();
  AliStack* stack = rl->Stack();
  if (!stack) {
    Error("kine_tracks.C", "can not get kinematics.");
    return 0;
  }

  gEve->DisableRedraw();

  TEveTrackList* cont = new TEveTrackList("Kine Tracks");
  cont->SetMainColor(3);
  TEveTrackPropagator* rnrStyle = cont->GetPropagator();
  // !!! Watch the '-', apparently different sign convention then for ESD.
  rnrStyle->SetMagField( -0.1*gAlice->Field()->SolenoidField() );

  gEve->AddElement(cont);
  Int_t count = 0;
  Int_t N = stack->GetNtrack();
  Int_t Np = stack->GetNprimary();
  for (Int_t i = 0; i < Np; ++i)
  {
    TParticle* p = stack->Particle(i);
    if (p->GetStatusCode() <= 1)
    {
      if (p->Pt() < min_pt && p->P() < min_p) continue;

      ++count;
      TEveTrack* track = new TEveTrack(p, i, rnrStyle);

      //PH The line below is replaced waiting for a fix in Root
      //PH which permits to use variable siza arguments in CINT
      //PH on some platforms (alphalinuxgcc, solariscc5, etc.)
      //PH    track->SetName(Form("%s [%d]", p->GetName(), i));
      char form[1000];
      sprintf(form,"%s [%d]", p->GetName(), i);
      track->SetName(form);
      track->SetStdTitle();
      Int_t ml = p->GetMother(0);
      if (ml != -1)
      {
        track->SetTitle(Form("%s\nMother label=%d\nMother Pdg=%d",
                             track->GetElementTitle(),
                             ml, stack->Particle(ml)->GetPdgCode()));
      }
      set_track_color(track, pdg_col);

      gEve->AddElement(track, cont);

      if (recurse)
	kine_daughters(track, stack, min_pt, min_p, pdg_col, recurse);
    }
  }

  // set path marks
  AliEveKineTools kt;
  kt.SetDaughterPathMarks(cont, stack, recurse);
  if (use_track_refs && rl->LoadTrackRefs() == 0)
  {
    kt.SetTrackReferences(cont, rl->TreeTR(), recurse);
    rnrStyle->SetEditPathMarks(kTRUE);
  }
  kt.SortPathMarks(cont, recurse);

  //PH  const Text_t* tooltip = Form("min pT=%.2lf, min P=%.2lf), N=%d", min_pt, min_p, count);
  char tooltip[1000];
  sprintf(tooltip,"min pT=%.2lf, min P=%.2lf), N=%d", min_pt, min_p, count);
  cont->SetTitle(tooltip); // Not broadcasted automatically ...

  cont->MakeTracks(recurse);
  gEve->EnableRedraw();
  gEve->Redraw3D();

  return cont;
}

void kine_daughters(TEveTrack* parent,  AliStack* stack,
		    Double_t     min_pt,  Double_t  min_p,
		    Bool_t       pdg_col, Bool_t    recurse)
{
  TParticle *p = stack->Particle(parent->GetLabel());
  if (p->GetNDaughters() > 0)
  {
    TEveTrackPropagator* rs = parent->GetPropagator();
    for (int d=p->GetFirstDaughter(); d>0 && d<=p->GetLastDaughter(); ++d)
    {
      TParticle* dp = stack->Particle(d);
      if (dp->Pt() < min_pt && dp->P() < min_p) continue;

      TEveTrack* dtrack = new TEveTrack(dp, d, rs);
      char form[1000];
      sprintf(form,"%s [%d]", dp->GetName(), d);
      dtrack->SetName(form);
      dtrack->SetStdTitle();
      set_track_color(dtrack, pdg_col);

      gEve->AddElement(dtrack, parent);

      if (recurse)
	kine_daughters(dtrack, stack, min_pt, min_p, pdg_col, recurse);
    }
  }
}

Color_t set_track_color(TEveTrack* t, Bool_t pdg_col)
{
  if (pdg_col)
    t->SetMainColor(get_pdg_color(t->GetPdg()));
  else
    t->SetMainColor(30);
}

Color_t get_pdg_color(Int_t pdg)
{
  // PDG color indices
  static const Color_t DefCol   = 30;
  static const Color_t ECol     = 5;
  static const Color_t MuCol    = 6;
  static const Color_t GammaCol = 7;
  static const Color_t MesCol1  = 3;
  static const Color_t MesCol2  = 38;
  static const Color_t BarCol   = 10;

  Int_t pdga  = TMath::Abs(pdg);
  Color_t col = DefCol;

  // elementary  particles
  if (pdga < 100) {
    switch (pdga) {
      case 11:
	col = ECol; break;
      case 12:
	col = MuCol; break;
      case 22:
	col = GammaCol; break;
    }
  }
  // mesons and barions
  else if (pdga < 100000) {
    Int_t i  = pdga;
    Int_t i0 = i%10; i /= 10;
    Int_t i1 = i%10; i /= 10;
    Int_t i2 = i%10; i /= 10;
    Int_t i3 = i%10; i /= 10;
    Int_t i4 = i%10;
    //printf("pdg(%d) quark indices (%d,%d,%d,%d,%d) \n",pdg, i4,i3,i2, i1, i0);
    // meson
    if ((i3 == 0) && ( i4 < 2)){
      col = MesCol1; // quarks: i1,i2 (spin = i0)
      if(i1 == 3 || i2 == 3)
	col = MesCol2;
    } // barion
    else if ( i2 >= i1 && i3 >= i2 ) {
      col = BarCol; // quarks: i1,i2, i3 (spin = i0))
    }
  }

  return col;
}

/******************************************************************************/

TEveElement*
kine_track(Int_t  label,
	   Bool_t import_mother    = kTRUE, Bool_t import_daughters = kTRUE,
	   Bool_t pdg_col          = kTRUE, Bool_t recurse          = kTRUE,
           TEveElement* cont = 0)

{
  // Create mother and daughters tracks with given label.
  // mother     -> particle with label
  // daughters  -> daughters of label

  if (label < 0) {
    Warning("kine_track", "label not set.");
    return 0;
  }

  AliRunLoader* rl =  AliEveEventManager::AssertRunLoader();
  rl->LoadKinematics();
  AliStack* stack = rl->Stack();
  if (label >= stack->GetNtrack())
  {
    Warning("kine_track", "label out of range.");
    return 0;
  }

  TParticle* p = stack->Particle(label);

  if (import_mother || (import_daughters && p->GetNDaughters()))
  {
    TEveTrack* toptrack = 0;
    TEveTrackList* tracklist = 0;
    TEveTrackPropagator* rs = 0;

    if (cont == 0)
    {
      TEveTrackList* tlist = new TEveTrackList
	(Form("Kinematics of %d", label, p->GetNDaughters()));
      cont = tlist;

      TEveTrackPropagator* rnrStyle = tlist->GetPropagator();
      // !!! Watch the '-', apparently different sign convention then for ESD.
      rnrStyle->SetMagField( -0.1*gAlice->Field()->SolenoidField() );
      char tooltip[1000];
      sprintf(tooltip,"Ndaughters=%d", p->GetNDaughters());
      tlist->SetTitle(tooltip);
      rnrStyle->fMaxOrbs = 2;
      rnrStyle->SetEditPathMarks(kTRUE);

      gEve->AddElement(cont);
      rs = tlist->GetPropagator();
    }
    else
    {
      // check if container is TEveTrackList or TEveTrack (has rnr-style)
      TEveTrack* t = dynamic_cast<TEveTrack*>(cont);
      if (t) {
	rs = t->GetPropagator();
      } else {
        TEveTrackList* l = dynamic_cast<TEveTrackList*>(cont);
        if (l)
	  rs = l->GetPropagator();
        else
	  Error("kine_tracks.C", "TrackRenderStyle not set.");
      }
    }

    if (import_mother)
    {
      TEveTrack* track = new TEveTrack(p, label, rs);
      char form[1000];
      sprintf(form,"%s [%d]", p->GetName(), label);
      track->SetName(form);
      track->SetStdTitle();
      set_track_color(track, pdg_col);

      track->MakeTrack();
      gEve->AddElement(track, cont);
      cont = track;
    }

    if (import_daughters && p->GetNDaughters())
    {
      for (int d=p->GetFirstDaughter(); d>0 && d<=p->GetLastDaughter(); ++d)
      {
	TParticle* dp = stack->Particle(d);
	TEveTrack* track = new TEveTrack(dp, d, rs);
	char form[1000];
	sprintf(form,"%s [%d]", dp->GetName(), d);
	track->SetName(form);
	track->SetStdTitle();
	set_track_color(track, pdg_col);

        track->MakeTrack();
	gEve->AddElement(track, cont);

	if (recurse)
	  kine_daughters(track, stack, 0, 0, pdg_col, recurse);
      }
    }
  }

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