[u/mrichter/AliRoot.git] / FMD / scripts / DrawTrackRefs.C

//____________________________________________________________________
//
// $Id: DrawHits.C 30718 2009-01-22 16:07:40Z cholm $
//
// Script that contains a class to draw hits, using the
// AliFMDInputHits class in the util library. 
//
// It draws the energy loss versus the p/(mq^2).  It can be overlayed
// with the Bethe-Bloc curve to show how the simulation behaves
// relative to the expected. 
//
// Use the script `Compile.C' to compile this class using ACLic. 
//
#include <TH2D.h>
#include <AliFMDHit.h>
#include <AliFMDInput.h>
#include <iostream>
#include <TStyle.h>
#include <TArrayF.h>
#include <TParticle.h>
#include <TCanvas.h>
#include <TGraphErrors.h>
#include <TDatabasePDG.h>
#include <TParticlePDG.h>
#include <TLegend.h>
#include <TArrow.h>
#include <TLatex.h>
#include <TF1.h>
#include <AliStack.h>
#include <AliTrackReference.h>
#include <AliFMDStripIndex.h>
#include <AliFMDGeometry.h>
#include <AliGenEventHeader.h>
#include <AliHeader.h>
#include <THStack.h>

/** @class DrawHits
    @brief Draw hit energy loss
    @code 
    Root> .L Compile.C
    Root> Compile("DrawHits.C")
    Root> DrawHits c
    Root> c.Run();
    @endcode
    @ingroup FMD_script
 */
class DrawTrackRefs : public AliFMDInput
{
private:
  struct Hists 
  {
    TString fPrefix;
    TString fWhere;
    TH2* fAll;
    TH2* fPrimary;
    TH2* fSecondary;
    THStack* fStack;

    //__________________________________________________________________
    Hists(const char* prefix, const char* where)
      : fPrefix(prefix), fWhere(where), 
	fAll(0), fPrimary(0), fSecondary(0), fStack(0)
    {
      fAll       = MakeHist(Form("%sAll", prefix), 
			    Form("All tracks in %s", where));
      fPrimary   = MakeHist(Form("%sPrimary", prefix), 
			    Form("All primaries in %s", where));
      fSecondary = MakeHist(Form("%sSecondary", prefix), 
			    Form("All secondaries in %s", where));
    }
     //__________________________________________________________________
    TH2* MakeHist(const char* name, const char* title) 
    {
      TH2D* ret = new TH2D(name, title, 200, -4, 6, 40, 0, 2*TMath::Pi());
      ret->SetDrawOption("colz");
      return ret;
    }
    //__________________________________________________________________
    void Draw(TVirtualPad* pad, Bool_t stack=false, Option_t* option="colz") 
    {
      if (stack) { 
	pad->cd();
	pad->Divide(1,2, 0, 0);
	pad->cd(1);
	fStack = new THStack(Form("%sDisplay", fPrefix.Data()), 
			     Form("2nd over 1st particles in %s", 
				  fWhere.Data()));
	TH1D* prim = 0;
	TH1D* sec = 0;
	fStack->Add((prim = fPrimary->ProjectionX()));
	fStack->Add((sec  = fSecondary->ProjectionX()));
	prim->SetFillColor(kGray);
	sec->SetFillColor(kRed);
	sec->SetFillStyle(3001);
	fStack->Draw();

	pad->cd(2);
	TH1* ratio = new TH1D(*sec);
	ratio->SetName(Form("%sRatio", fPrefix.Data()));
	ratio->SetTitle(Form("2nd over 1st particles in %s", fWhere.Data()));
	ratio->Divide(prim);
	ratio->SetFillColor(kRed);
	ratio->SetFillStyle(3001);
	ratio->Draw();

	return;
      }
      pad->Divide(1,3,0,0);
      TVirtualPad* p1 = pad->cd(1);
      p1->SetLogz();
      fAll->Draw(option);
      p1 = pad->cd(2);
      p1->SetLogz();
      fPrimary->Draw(option);
      p1 = pad->cd(3);
      p1->SetLogz();
      fSecondary->Draw(option);
      pad->Modified();
      pad->Update();
      pad->cd();
    }
  };
  Hists fTotal;
  Hists fFMD1I;
  Hists fFMD2I;
  Hists fFMD2O;
  Hists fFMD3I;
  Hists fFMD3O;
public:
  //__________________________________________________________________
  DrawTrackRefs() 
    : AliFMDInput("galice.root"), 
      fTotal("total", "ALICE"), 
      fFMD1I("fmd1I", "FMD1i"),
      fFMD2I("fmd2I", "FMD2i"),
      fFMD2O("fmd2O", "FMD2o"),
      fFMD3I("fmd3I", "FMD3i"),
      fFMD3O("fmd3O", "FMD3o")
  { 
    AddLoad(kKinematics);
    AddLoad(kTrackRefs);
    AddLoad(kGeometry);
    AddLoad(kHeader);
  }
  //__________________________________________________________________
  Bool_t ProcessTrackRef(AliTrackReference* trackRef, TParticle* p)
  {
    if (trackRef->DetectorId() != AliTrackReference::kFMD) return kTRUE;

    if (!p->GetPDG() || p->GetPDG()->Charge() == 0) return kTRUE;

    // Double_t eta = p->Eta();
    // Double_t phi = p->Phi();

    UShort_t d   = 0;
    Char_t   r   = '\0';
    UShort_t s   = 0;
    UShort_t t   = 0;
    AliFMDStripIndex::Unpack(trackRef->UserId(),d,r,s,t);
    
    Double_t x, y, z;
    AliFMDGeometry::Instance()->Detector2XYZ(d,r,s,t,x,y,z);

    AliGenEventHeader* genHeader = fHeader->GenEventHeader();
    TArrayF v;
    genHeader->PrimaryVertex(v);
    z -= v.fArray[2];
    
    Double_t eta, phi, theta, rr;
    AliFMDGeometry::XYZ2REtaPhiTheta(x, y, z, rr, eta, phi, theta);
    if (phi < 0)               phi += 2*TMath::Pi();
    if (phi > 2 * TMath::Pi()) phi -= 2*TMath::Pi();

    Hists* hists = 0;
    switch (d) { 
    case 1: hists = &fFMD1I; break;
    case 2: hists = (r == 'I' || r == 'i') ? &fFMD2I : &fFMD2O; break;
    case 3: hists = (r == 'I' || r == 'i') ? &fFMD3I : &fFMD3O; break;
    }
    if (!hists) return kTRUE;

    hists->fAll->Fill(eta, phi);
    
    if (fStack->IsPhysicalPrimary(trackRef->GetTrack())) 
      hists->fPrimary->Fill(eta,phi); 
    else
      hists->fSecondary->Fill(eta, phi);

    return kTRUE;
  }
  //__________________________________________________________________
  Bool_t ProcessParticle(Int_t id, TParticle* p)
  {
    if (!p) return kTRUE;

    if (!p->GetPDG() || p->GetPDG()->Charge() == 0) return kTRUE;
    // std::cout << p->GetPDG()->GetName() << std::endl;

    Double_t eta = p->Eta();
    Double_t phi = p->Phi();

    fTotal.fAll->Fill(eta, phi);
    
    if (fStack->IsPhysicalPrimary(id)) fTotal.fPrimary->Fill(eta,phi); 
    else                               fTotal.fSecondary->Fill(eta, phi);

    return kTRUE;
  }

  //__________________________________________________________________
  Bool_t Finish()
  {
    gStyle->SetPalette(1);
    // gStyle->SetOptTitle(0);
    gStyle->SetTitleBorderSize(1);
    gStyle->SetTitleFillColor(0);
    gStyle->SetCanvasColor(0);
    gStyle->SetCanvasColor(0);
    gStyle->SetCanvasBorderSize(0);
    gStyle->SetPadColor(0);
    gStyle->SetPadBorderSize(0);

    TCanvas* c2d = new TCanvas("c2d", "2D plots", 1200, 700);
    c2d->Divide(6);
    fTotal.Draw(c2d->cd(1));   c2d->cd();
    fFMD1I.Draw(c2d->cd(6));   c2d->cd();
    fFMD2I.Draw(c2d->cd(5));   c2d->cd();
    fFMD2O.Draw(c2d->cd(4));   c2d->cd();
    fFMD3I.Draw(c2d->cd(2));   c2d->cd();
    fFMD3O.Draw(c2d->cd(3));   c2d->cd();

    c2d->Modified();
    c2d->Update();
    c2d->cd();
    c2d->SaveAs("kine_etaphi.png");

    TCanvas* c1d = new TCanvas("c1d", "1D plots", 1200, 700);
    c1d->Divide(6);
    fTotal.Draw(c1d->cd(1), kTRUE);   c1d->cd();
    fFMD1I.Draw(c1d->cd(6), kTRUE);   c1d->cd();
    fFMD2I.Draw(c1d->cd(5), kTRUE);   c1d->cd();
    fFMD2O.Draw(c1d->cd(4), kTRUE);   c1d->cd();
    fFMD3I.Draw(c1d->cd(2), kTRUE);   c1d->cd();
    fFMD3O.Draw(c1d->cd(3), kTRUE);   c1d->cd();

    c1d->Modified();
    c1d->Update();
    c1d->cd();
    c1d->SaveAs("kine_eta.png");


    return kTRUE;
  }
  
  ClassDef(DrawTrackRefs,0);
};

//____________________________________________________________________
//
// EOF
//
Commit	Line	Data
2040edda	1	//____________________________________________________________________
	2	//
	3	// $Id: DrawHits.C 30718 2009-01-22 16:07:40Z cholm $
	4	//
	5	// Script that contains a class to draw hits, using the
	6	// AliFMDInputHits class in the util library.
	7	//
	8	// It draws the energy loss versus the p/(mq^2). It can be overlayed
	9	// with the Bethe-Bloc curve to show how the simulation behaves
	10	// relative to the expected.
	11	//
	12	// Use the script `Compile.C' to compile this class using ACLic.
	13	//
	14	#include <TH2D.h>
	15	#include <AliFMDHit.h>
	16	#include <AliFMDInput.h>
	17	#include <iostream>
	18	#include <TStyle.h>
	19	#include <TArrayF.h>
	20	#include <TParticle.h>
	21	#include <TCanvas.h>
	22	#include <TGraphErrors.h>
	23	#include <TDatabasePDG.h>
	24	#include <TParticlePDG.h>
	25	#include <TLegend.h>
	26	#include <TArrow.h>
	27	#include <TLatex.h>
	28	#include <TF1.h>
	29	#include <AliStack.h>
	30	#include <AliTrackReference.h>
	31	#include <AliFMDStripIndex.h>
	32	#include <AliFMDGeometry.h>
	33	#include <AliGenEventHeader.h>
	34	#include <AliHeader.h>
	35	#include <THStack.h>
	36
	37	/** @class DrawHits
	38	@brief Draw hit energy loss
	39	@code
	40	Root> .L Compile.C
	41	Root> Compile("DrawHits.C")
	42	Root> DrawHits c
	43	Root> c.Run();
	44	@endcode
	45	@ingroup FMD_script
	46	*/
	47	class DrawTrackRefs : public AliFMDInput
	48	{
	49	private:
	50	struct Hists
	51	{
	52	TString fPrefix;
	53	TString fWhere;
	54	TH2* fAll;
	55	TH2* fPrimary;
	56	TH2* fSecondary;
	57	THStack* fStack;
	58
	59	//__________________________________________________________________
	60	Hists(const char* prefix, const char* where)
	61	: fPrefix(prefix), fWhere(where),
	62	fAll(0), fPrimary(0), fSecondary(0), fStack(0)
	63	{
	64	fAll = MakeHist(Form("%sAll", prefix),
65	Form("All tracks in %s", where));
66	fPrimary = MakeHist(Form("%sPrimary", prefix),
67	Form("All primaries in %s", where));
68	fSecondary = MakeHist(Form("%sSecondary", prefix),
69	Form("All secondaries in %s", where));
70	}
71	//__________________________________________________________________
72	TH2* MakeHist(const char* name, const char* title)
73	{
74	TH2D* ret = new TH2D(name, title, 200, -4, 6, 40, 0, 2*TMath::Pi());
75	ret->SetDrawOption("colz");
76	return ret;
77	}
78	//__________________________________________________________________
79	void Draw(TVirtualPad* pad, Bool_t stack=false, Option_t* option="colz")
80	{
81	if (stack) {
82	pad->cd();
83	pad->Divide(1,2, 0, 0);
84	pad->cd(1);
85	fStack = new THStack(Form("%sDisplay", fPrefix.Data()),
86	Form("2nd over 1st particles in %s",
87	fWhere.Data()));
88	TH1D* prim = 0;
89	TH1D* sec = 0;
90	fStack->Add((prim = fPrimary->ProjectionX()));
91	fStack->Add((sec = fSecondary->ProjectionX()));
92	prim->SetFillColor(kGray);
93	sec->SetFillColor(kRed);
94	sec->SetFillStyle(3001);
95	fStack->Draw();
96
97	pad->cd(2);
98	TH1* ratio = new TH1D(*sec);
99	ratio->SetName(Form("%sRatio", fPrefix.Data()));
100	ratio->SetTitle(Form("2nd over 1st particles in %s", fWhere.Data()));
101	ratio->Divide(prim);
102	ratio->SetFillColor(kRed);
103	ratio->SetFillStyle(3001);
104	ratio->Draw();
105
106	return;
107	}
108	pad->Divide(1,3,0,0);
109	TVirtualPad* p1 = pad->cd(1);
110	p1->SetLogz();
111	fAll->Draw(option);
112	p1 = pad->cd(2);
113	p1->SetLogz();
114	fPrimary->Draw(option);
115	p1 = pad->cd(3);
116	p1->SetLogz();
117	fSecondary->Draw(option);
118	pad->Modified();
119	pad->Update();
120	pad->cd();
121	}
122	};
123	Hists fTotal;
124	Hists fFMD1I;
125	Hists fFMD2I;
126	Hists fFMD2O;
127	Hists fFMD3I;
128	Hists fFMD3O;
129	public:
130	//__________________________________________________________________
131	DrawTrackRefs()
132	: AliFMDInput("galice.root"),
133	fTotal("total", "ALICE"),
134	fFMD1I("fmd1I", "FMD1i"),
135	fFMD2I("fmd2I", "FMD2i"),
136	fFMD2O("fmd2O", "FMD2o"),
137	fFMD3I("fmd3I", "FMD3i"),
138	fFMD3O("fmd3O", "FMD3o")
139	{
140	AddLoad(kKinematics);
141	AddLoad(kTrackRefs);
142	AddLoad(kGeometry);
143	AddLoad(kHeader);
144	}
145	//__________________________________________________________________
146	Bool_t ProcessTrackRef(AliTrackReference* trackRef, TParticle* p)
147	{
148	if (trackRef->DetectorId() != AliTrackReference::kFMD) return kTRUE;
149
150	if (!p->GetPDG() \|\| p->GetPDG()->Charge() == 0) return kTRUE;
151
152	// Double_t eta = p->Eta();
153	// Double_t phi = p->Phi();
154
155	UShort_t d = 0;
156	Char_t r = '\0';
157	UShort_t s = 0;
158	UShort_t t = 0;
159	AliFMDStripIndex::Unpack(trackRef->UserId(),d,r,s,t);
160
161	Double_t x, y, z;
162	AliFMDGeometry::Instance()->Detector2XYZ(d,r,s,t,x,y,z);
163
164	AliGenEventHeader* genHeader = fHeader->GenEventHeader();
165	TArrayF v;
166	genHeader->PrimaryVertex(v);
167	z -= v.fArray[2];
168
169	Double_t eta, phi, theta, rr;
170	AliFMDGeometry::XYZ2REtaPhiTheta(x, y, z, rr, eta, phi, theta);
171	if (phi < 0) phi += 2*TMath::Pi();
172	if (phi > 2 * TMath::Pi()) phi -= 2*TMath::Pi();
173
174	Hists* hists = 0;
175	switch (d) {
176	case 1: hists = &fFMD1I; break;
177	case 2: hists = (r == 'I' \|\| r == 'i') ? &fFMD2I : &fFMD2O; break;
178	case 3: hists = (r == 'I' \|\| r == 'i') ? &fFMD3I : &fFMD3O; break;
179	}
180	if (!hists) return kTRUE;
181
182	hists->fAll->Fill(eta, phi);
183
184	if (fStack->IsPhysicalPrimary(trackRef->GetTrack()))
185	hists->fPrimary->Fill(eta,phi);
186	else
187	hists->fSecondary->Fill(eta, phi);
188
189	return kTRUE;
190	}
191	//__________________________________________________________________
192	Bool_t ProcessParticle(Int_t id, TParticle* p)
193	{
194	if (!p) return kTRUE;
195
196	if (!p->GetPDG() \|\| p->GetPDG()->Charge() == 0) return kTRUE;
197	// std::cout << p->GetPDG()->GetName() << std::endl;
198
199	Double_t eta = p->Eta();
200	Double_t phi = p->Phi();
201
202	fTotal.fAll->Fill(eta, phi);
203
204	if (fStack->IsPhysicalPrimary(id)) fTotal.fPrimary->Fill(eta,phi);
205	else fTotal.fSecondary->Fill(eta, phi);
206
207	return kTRUE;
208	}
209
210	//__________________________________________________________________
211	Bool_t Finish()
212	{
213	gStyle->SetPalette(1);
214	// gStyle->SetOptTitle(0);
215	gStyle->SetTitleBorderSize(1);
216	gStyle->SetTitleFillColor(0);
217	gStyle->SetCanvasColor(0);
218	gStyle->SetCanvasColor(0);
219	gStyle->SetCanvasBorderSize(0);
220	gStyle->SetPadColor(0);
221	gStyle->SetPadBorderSize(0);
222
223	TCanvas* c2d = new TCanvas("c2d", "2D plots", 1200, 700);
224	c2d->Divide(6);
225	fTotal.Draw(c2d->cd(1)); c2d->cd();
226	fFMD1I.Draw(c2d->cd(6)); c2d->cd();
227	fFMD2I.Draw(c2d->cd(5)); c2d->cd();
228	fFMD2O.Draw(c2d->cd(4)); c2d->cd();
229	fFMD3I.Draw(c2d->cd(2)); c2d->cd();
230	fFMD3O.Draw(c2d->cd(3)); c2d->cd();
231
232	c2d->Modified();
233	c2d->Update();
234	c2d->cd();
235	c2d->SaveAs("kine_etaphi.png");
236
237	TCanvas* c1d = new TCanvas("c1d", "1D plots", 1200, 700);
238	c1d->Divide(6);
239	fTotal.Draw(c1d->cd(1), kTRUE); c1d->cd();
240	fFMD1I.Draw(c1d->cd(6), kTRUE); c1d->cd();
241	fFMD2I.Draw(c1d->cd(5), kTRUE); c1d->cd();
242	fFMD2O.Draw(c1d->cd(4), kTRUE); c1d->cd();
243	fFMD3I.Draw(c1d->cd(2), kTRUE); c1d->cd();
244	fFMD3O.Draw(c1d->cd(3), kTRUE); c1d->cd();
245
246	c1d->Modified();
247	c1d->Update();
248	c1d->cd();
249	c1d->SaveAs("kine_eta.png");
250
251
252	return kTRUE;
253	}
254
255	ClassDef(DrawTrackRefs,0);
256	};
257
258	//____________________________________________________________________
259	//
260	// EOF
261	//