Almost final version of scanning tools.

[u/mrichter/AliRoot.git] / EVE / alice-macros / esd_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.                                                 *
 **************************************************************************/

#if !defined(__CINT__) || defined(__MAKECINT__)

#include "TString.h"
#include "TMath.h"
#include "TGListTree.h"
#include "TEveVSDStructs.h"
#include "TEveManager.h"
#include "TEveTrackPropagator.h"

#include "AliESDEvent.h"
#include "AliESDtrackCuts.h"
#include "AliESDtrack.h"
#include "AliExternalTrackParam.h"

#include "EVE/EveBase/AliEveTrack.h"
#include "EVE/EveBase/AliEveMagField.h"
#include "EVE/EveBase/AliEveEventManager.h"

#endif


// Use inner-tpc track params when its refit failed.
Bool_t g_esd_tracks_use_ip_on_failed_its_refit = kFALSE;

// Use magnetic-field as retrieved from GRP.
Bool_t g_esd_tracks_true_field = kTRUE;

// Use Runge-Kutta track stepper.
Bool_t g_esd_tracks_rk_stepper = kFALSE;


//==============================================================================

void esd_track_propagator_setup(TEveTrackPropagator* trkProp,
				Float_t magF, Float_t maxR)
{
  if (g_esd_tracks_true_field)
  {
    trkProp->SetMagFieldObj(new AliEveMagField);
  }
  else
  {
    trkProp->SetMagField(magF);
  }
  if (g_esd_tracks_rk_stepper)
  {
    trkProp->SetStepper(TEveTrackPropagator::kRungeKutta);
  }
  trkProp->SetMaxR(maxR);
}

//==============================================================================

TString esd_track_title(AliESDtrack* t)
{
  TString s;

  Int_t label = t->GetLabel(), index = t->GetID();
  TString idx(index == kMinInt ? "<undef>" : Form("%d", index));
  TString lbl(label == kMinInt ? "<undef>" : Form("%d", label));

  Double_t p[3], v[3];
  t->GetXYZ(v);
  t->GetPxPyPz(p);
  Double_t pt    = t->Pt();
  Double_t ptsig = TMath::Sqrt(t->GetSigma1Pt2());
  Double_t ptsq  = pt*pt;
  Double_t ptm   = pt / (1.0 + pt*ptsig);
  Double_t ptM   = pt / (1.0 - pt*ptsig);

  s = Form("Index=%s, Label=%s\nChg=%d, Pdg=%d\n"
	   "pT = %.3f + %.3f - %.3f [%.3f]\n"
           "P  = (%.3f, %.3f, %.3f)\n"
           "V  = (%.3f, %.3f, %.3f)\n",
	   idx.Data(), lbl.Data(), t->Charge(), 0,
	   pt, ptM - pt, pt - ptm, ptsig*ptsq,
           p[0], p[1], p[2],
           v[0], v[1], v[2]);

  Int_t   o;
  s += "Det (in,out,refit,pid):\n";
  o  = AliESDtrack::kITSin;
  s += Form("ITS (%d,%d,%d,%d)  ",  t->IsOn(o), t->IsOn(o<<1), t->IsOn(o<<2), t->IsOn(o<<3));
  o  = AliESDtrack::kTPCin;
  s += Form("TPC(%d,%d,%d,%d)\n",   t->IsOn(o), t->IsOn(o<<1), t->IsOn(o<<2), t->IsOn(o<<3));
  o  = AliESDtrack::kTRDin;
  s += Form("TRD(%d,%d,%d,%d) ",    t->IsOn(o), t->IsOn(o<<1), t->IsOn(o<<2), t->IsOn(o<<3));
  o  = AliESDtrack::kTOFin;
  s += Form("TOF(%d,%d,%d,%d)\n",   t->IsOn(o), t->IsOn(o<<1), t->IsOn(o<<2), t->IsOn(o<<3));
  o  = AliESDtrack::kHMPIDout;
  s += Form("HMPID(out=%d,pid=%d)\n", t->IsOn(o), t->IsOn(o<<1));
  s += Form("ESD pid=%d", t->IsOn(AliESDtrack::kESDpid));

  if (t->IsOn(AliESDtrack::kESDpid))
  {
    Double_t pid[5];
    t->GetESDpid(pid);
    s += Form("\n[%.2f %.2f %.2f %.2f %.2f]", pid[0], pid[1], pid[2], pid[3], pid[4]);
  }

  return s;
}

//==============================================================================

void esd_track_add_param(AliEveTrack* track, const AliExternalTrackParam* tp)
{
  // Add additional track parameters as a path-mark to track.

  if (tp == 0)
    return;

  Double_t pbuf[3], vbuf[3];
  tp->GetXYZ(vbuf);
  tp->GetPxPyPz(pbuf);

  TEvePathMark pm(TEvePathMark::kReference);
  pm.fV.Set(vbuf);
  pm.fP.Set(pbuf);
  track->AddPathMark(pm);
}

//==============================================================================

AliEveTrack* esd_make_track(AliESDtrack *at, TEveTrackList* cont)
{
  // Make a standard track representation and put it into given container.

  // Choose which parameters to use a track's starting point.
  // If gkFixFailedITSExtr is TRUE (FALSE by default) and
  // if ITS refit failed, take track parameters at inner TPC radius.

  const AliExternalTrackParam* tp = at;

  Bool_t innerTaken = kFALSE;
  if ( ! at->IsOn(AliESDtrack::kITSrefit) && g_esd_tracks_use_ip_on_failed_its_refit)
  {
    tp = at->GetInnerParam();
    innerTaken = kTRUE;
  }

  AliEveTrack* track = new AliEveTrack(at, cont->GetPropagator());
  track->SetAttLineAttMarker(cont);
  track->SetName(Form("AliEveTrack %d", at->GetID()));
  track->SetElementTitle(esd_track_title(at));
  track->SetSourceObject(at);

  // Add inner/outer track parameters as path-marks.
  if (at->IsOn(AliESDtrack::kTPCrefit))
  {
    if ( ! innerTaken)
    {
      esd_track_add_param(track, at->GetInnerParam());
    }
    esd_track_add_param(track, at->GetOuterParam());
  }

  return track;
}


//==============================================================================
// esd_tracks()
//==============================================================================

TEveTrackList* esd_tracks()
{
  AliESDEvent* esd = AliEveEventManager::AssertESD();

  TEveTrackList* cont = new TEveTrackList("ESD Tracks");
  cont->SetMainColor(6);

  esd_track_propagator_setup(cont->GetPropagator(),
			     0.1*esd->GetMagneticField(), 520);

  gEve->AddElement(cont);

  Int_t count = 0;
  for (Int_t n = 0; n < esd->GetNumberOfTracks(); ++n)
  {
    ++count;
    AliEveTrack* track = esd_make_track(esd->GetTrack(n), cont);

    cont->AddElement(track);
  }
  cont->SetTitle(Form("N=%d", count));
  cont->MakeTracks();

  gEve->Redraw3D();

  return cont;
}

TEveElementList* esd_tracks_MI()
{
  AliESDEvent* esd = AliEveEventManager::AssertESD();

  TEveElementList* cont = new TEveElementList("ESD Tracks MI");
  gEve->AddElement(cont);

  Int_t count = 0;
  for (Int_t n = 0; n < esd->GetNumberOfTracks(); ++n)
  {
    ++count;
    AliESDtrack* at = esd->GetTrack(n);
    TEveLine* l = new TEveLine; 
    l->SetLineColor(5);
    at->FillPolymarker(l, esd->GetMagneticField(), 0, 250, 5);
    l->SetName(Form("ESDTrackMI %d", at->GetID()));
    l->SetElementTitle(esd_track_title(at));
    l->SetSourceObject(at);

    cont->AddElement(l);
  }
  cont->SetTitle(Form("N=%d", count));

  gEve->Redraw3D();

  return cont;
}


//==============================================================================
// esd_tracks_from_array()
//==============================================================================

TEveTrackList* esd_tracks_from_array(TCollection* col, AliESDEvent* esd=0)
{
  // Retrieves AliESDTrack's from collection.
  // See example usage with AliAnalysisTrackCuts in the next function.

  if (esd == 0) esd = AliEveEventManager::AssertESD();

  TEveTrackList* cont = new TEveTrackList("ESD Tracks");
  cont->SetMainColor(6);

  esd_track_propagator_setup(cont->GetPropagator(),
			     0.1*esd->GetMagneticField(), 520);

  gEve->AddElement(cont);

  Int_t    count = 0;
  TIter    next(col);
  TObject *obj;
  while ((obj = next()) != 0)
  {
    if (obj->IsA()->InheritsFrom("AliESDtrack") == kFALSE)
    {
      Warning("esd_tracks_from_array", "Object '%s', '%s' is not an AliESDtrack.",
	      obj->GetName(), obj->GetTitle());
      continue;
    }

    ++count;
    AliESDtrack* at = reinterpret_cast<AliESDtrack*>(obj);

    AliEveTrack* track = esd_make_track(at, cont);
    cont->AddElement(track);
  }
  cont->SetTitle(Form("N=%d", count));
  cont->MakeTracks();

  gEve->Redraw3D();

  return cont;
}

void esd_tracks_alianalcuts_demo()
{
  AliESDEvent* esd = AliEveEventManager::AssertESD();

  AliESDtrackCuts atc;
  atc.SetPtRange(0.1, 5);
  atc.SetRapRange(-1, 1);

  esd_tracks_from_array(atc.GetAcceptedTracks(esd), esd);
}


//==============================================================================
// esd_tracks_by_category
//==============================================================================

Float_t get_sigma_to_vertex(AliESDtrack* esdTrack)
{
  // Taken from: PWG0/esdTrackCuts/AliESDtrackCuts.cxx
  // Float_t AliESDtrackCuts::GetSigmaToVertex(AliESDtrack* esdTrack)

  Float_t b[2];
  Float_t bRes[2];
  Float_t bCov[3];
  esdTrack->GetImpactParameters(b,bCov);
  if (bCov[0] <= 0 || bCov[2] <= 0)
  {
    printf("Estimated b resolution lower or equal zero!\n");
    bCov[0] = bCov[2] = 0;
  }
  bRes[0] = TMath::Sqrt(bCov[0]);
  bRes[1] = TMath::Sqrt(bCov[2]);

  // -----------------------------------
  // How to get to a n-sigma cut?
  //
  // The accumulated statistics from 0 to d is
  //
  // ->  Erf(d/Sqrt(2)) for a 1-dim gauss (d = n_sigma)
  // ->  1 - Exp(-d**2) for a 2-dim gauss (d*d = dx*dx + dy*dy != n_sigma)
  //
  // It means that for a 2-dim gauss: n_sigma(d) = Sqrt(2)*ErfInv(1 - Exp((-x**2)/2)
  // Can this be expressed in a different way?

  if (bRes[0] == 0 || bRes[1] == 0)
    return -1;

  Float_t d = TMath::Sqrt(TMath::Power(b[0]/bRes[0],2) + TMath::Power(b[1]/bRes[1],2));

  // stupid rounding problem screws up everything:
  // if d is too big, TMath::Exp(...) gets 0, and TMath::ErfInverse(1) that should be infinite, gets 0 :(
  if (TMath::Exp(-d * d / 2) < 1e-10)
    return 1000;

  d = TMath::ErfInverse(1 - TMath::Exp(-d * d / 2)) * TMath::Sqrt(2);
  return d;
}

TEveElementList* g_esd_tracks_by_category_container = 0;

TEveElementList* esd_tracks_by_category()
{
  // Import ESD tracks, separate them into several containers
  // according to primary-vertex cut and ITS&TPC refit status.

  AliESDEvent* esd = AliEveEventManager::AssertESD();

  TEveElementList* cont = new TEveElementList("ESD Tracks by category");
  gEve->AddElement(cont);

  g_esd_tracks_by_category_container = cont;

  const Int_t   nCont = 7;
  const Float_t maxR  = 520;
  const Float_t magF  = 0.1*esd->GetMagneticField();

  TEveTrackList *tl[nCont];
  Int_t          tc[nCont];
  Int_t          count = 0;

  tl[0] = new TEveTrackList("Sigma < 3");
  tc[0] = 0;
  esd_track_propagator_setup(tl[0]->GetPropagator(), magF, maxR);
  tl[0]->SetMainColor(3);
  cont->AddElement(tl[0]);

  tl[1] = new TEveTrackList("3 < Sigma < 5");
  tc[1] = 0;
  esd_track_propagator_setup(tl[1]->GetPropagator(), magF, maxR);
  tl[1]->SetMainColor(7);
  cont->AddElement(tl[1]);

  tl[2] = new TEveTrackList("5 < Sigma");
  tc[2] = 0;
  esd_track_propagator_setup(tl[2]->GetPropagator(), magF, maxR);
  tl[2]->SetMainColor(46);
  cont->AddElement(tl[2]);

  tl[3] = new TEveTrackList("no ITS refit; Sigma < 5");
  tc[3] = 0;
  esd_track_propagator_setup(tl[3]->GetPropagator(), magF, maxR);
  tl[3]->SetMainColor(41);
  cont->AddElement(tl[3]);

  tl[4] = new TEveTrackList("no ITS refit; Sigma > 5");
  tc[4] = 0;
  esd_track_propagator_setup(tl[4]->GetPropagator(), magF, maxR);
  tl[4]->SetMainColor(48);
  cont->AddElement(tl[4]);

  tl[5] = new TEveTrackList("no TPC refit");
  tc[5] = 0;
  esd_track_propagator_setup(tl[5]->GetPropagator(), magF, maxR);
  tl[5]->SetMainColor(kRed);
  cont->AddElement(tl[5]);

  tl[6] = new TEveTrackList("ITS stand-alone");
  tc[6] = 0;
  esd_track_propagator_setup(tl[6]->GetPropagator(), magF, maxR);
  tl[6]->SetMainColor(kMagenta - 9);
  cont->AddElement(tl[6]);

  for (Int_t n = 0; n < esd->GetNumberOfTracks(); ++n)
  {
    AliESDtrack* at = esd->GetTrack(n);

    Float_t s  = get_sigma_to_vertex(at);
    Int_t   ti;
    if      (s <  3) ti = 0;
    else if (s <= 5) ti = 1;
    else             ti = 2;

    if (at->IsOn(AliESDtrack::kITSin) && ! at->IsOn(AliESDtrack::kTPCin))
    {
      ti = 6;
    }
    else if (at->IsOn(AliESDtrack::kTPCin) && ! at->IsOn(AliESDtrack::kTPCrefit))
    {
      ti = 5;
    }
    else if ( ! at->IsOn(AliESDtrack::kITSrefit))
    {
      ti = (ti == 2) ? 4 : 3;
    }

    TEveTrackList* tlist = tl[ti];
    ++tc[ti];
    ++count;

    AliEveTrack* track = esd_make_track(at, tlist);
    track->SetName(Form("AliEveTrack idx=%d, sigma=%5.3f", at->GetID(), s));
    tlist->AddElement(track);
  }

  for (Int_t ti = 0; ti < nCont; ++ti)
  {
    TEveTrackList* tlist = tl[ti];
    tlist->SetName(Form("%s [%d]", tlist->GetName(), tlist->NumChildren()));
    tlist->SetTitle(Form("N tracks=%d", tc[ti]));

    tlist->MakeTracks();

    Bool_t good_cont = ti <= 1;
    if (AliEveTrackCounter::IsActive())
    {
      AliEveTrackCounter::fgInstance->RegisterTracks(tlist, good_cont);
    }
    else
    {
      if ( ! good_cont)
        tlist->SetLineStyle(6);
    }
  }
  cont->SetTitle(Form("N all tracks = %d", count));
  // ??? The following does not always work:
  cont->FindListTreeItem(gEve->GetListTree())->SetOpen(kTRUE);

  gEve->Redraw3D();

  return cont;
}