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

//____________________________________________________________________
//
// $Id$
//
// Script that contains a class to get the media where a certain track
// was created.   This is used for background studies. 
//
// 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 <TArrayI.h>
#include <AliRun.h>
#include <TNtuple.h>
#include <AliDetector.h>
#include <TFile.h>
#include <TGeoManager.h>
#include <TGeoVolume.h>
#include <TGeoMedium.h>
#include <TParticle.h>
#include <THStack.h>
/** @class Media 
    @brief Media description 
    @ingroup FMD_script
 */
struct Media : public TNamed
{
  TArrayI         fMeds;
  TNtuple*        fCount;
  TH1D*           fHist;
  Media(const char* name, const char* title) 
    : TNamed(name, title), fMeds(0) 
  {
    fHist  = new TH1D(Form("h%s",name), title, 120, -6, 6);
    // fHist->SetFillStyle(3001);
    fCount = new TNtuple(GetName(), GetTitle(), 
			 "E:DET:RNG:SEC:STR:X:Y:Z:EDEP:PDG");
  }
  Media(const char* name, TArrayI* a) 
    : TNamed(name, "Media information"), fMeds(a->fN, a->fArray)
  {
    fHist  = new TH1D(Form("h%s",name), GetTitle(), 120, -6, 6);
    fHist->SetFillStyle(3001);
    fHist->SetFillColor(fMeds.fN > 0 ? fMeds[0] : 0);
    fCount = new TNtuple(GetName(), GetTitle(), 
			 "E:DET:RNG:SEC:STR:X:Y:Z:EDEP:PDG");
    std::cout << "Media list " << name << ":" << std::endl;
    for (Int_t i = 0; i < fMeds.fN; i++) {
      std::cout << " " << fMeds[i] << std::flush;
      if (fMeds[i] == 0 && i > 0) break;
    }
    std::cout << std::endl;
  }
  Bool_t HasMedia(Int_t id) 
  {
    if (fMeds.fN <= 0) return kFALSE;
    for (Int_t i = 0; i < fMeds.fN; i++) {
      if (fMeds[i] == id) return kTRUE;
      if (fMeds[i] == 0 && i > 0) break;
    }
    return kFALSE;
  }
  void Fill(Int_t ev, AliFMDHit* hit) 
  {
    Float_t x[10];
    x[0] = ev;
    x[1] = hit->Detector();
    x[2] = Int_t(hit->Ring());
    x[3] = hit->Sector();
    x[4] = hit->Strip();
    x[5] = hit->X();
    x[6] = hit->Y();
    x[7] = hit->Z();
    x[8] = hit->Edep();
    x[9] = hit->Pdg();
    // return;
    fCount->Fill(x);    
    // r = TMath::Sqrt(hit->X() * hit->X() + hit->Y() * hit->Y());
    // if(r!=0) Float_t wgt=1./(2. * 3.1415*r);
    Double_t r = TMath::Sqrt(hit->X()*hit->X()+hit->Y()*hit->Y());
    Double_t t = TMath::ATan2(r, hit->Z());
    Double_t e = -TMath::Log(TMath::Tan(t/2));
    fHist->Fill(e);
  }
};


//____________________________________________________________________
/** @class GetMedia
    @brief Get media where a particle is produced
    @code 
    aliroot
    Root> gAlice->Init("$ALICE_ROOT/FMD/Config.C");
    Root> .x $ALICE_ROOT/FMD/scriptsWriteMedArrays.C
    Root> gAlice->Run();
    Root> .q
    aliroot
    Root> .L Compile.C
    Root> Compile("GetMedia.C")
    Root> GetMedia c
    Root> c.Run();
    @endcode
    @ingroup FMD_script
 */
class GetMedia : public AliFMDInput
{
private:
  TString    fModList;
  TObjArray  fMedia;
  Media*     fOther;
  Media*     fAll;
  Media*     fPrim;
  Int_t      fEv;
  THStack*   fStack;
  TFile*     fOutput;
public:
  //__________________________________________________________________
  GetMedia(const char* modlist="FMD:ITS:BODY:ABSO:T0:PIPE", 
	   const char* output="media.root") 
    :  fModList(modlist)
  { 
    AddLoad(kGeometry);
    AddLoad(kHits);
    AddLoad(kKinematics);

    fOutput = TFile::Open(output, "RECREATE");
    fStack  = new THStack("summed", "Summed hits");
    fOther  = AddMedia("other", "Unknown media");
    fAll    = AddMedia("All", "All media");
    fPrim   = AddMedia("Primary", "Primary particles");
    fEv     = 0;

    
    fAll->fHist->SetFillColor(0);
    fAll->fHist->SetFillStyle(3004);
    fPrim->fHist->SetFillColor(1);
    fPrim->fHist->SetFillStyle(3001);
    fStack->Add(fPrim->fHist);
  }
  //__________________________________________________________________
  Media* AddMedia(const char* name, const char* title, TArrayI* a=0)
  {
    Media* media = 0;
    if (!a) media = new Media(name, title);
    else    media = new Media(name, a);
    if (a)  {
      fMedia.Add(media);
      fStack->Add(media->fHist);
    }
    return media;
  }
  //__________________________________________________________________
  Media* FindMedia(Int_t med) 
  {
    for (Int_t i = 0; i < fMedia.GetEntries(); i++) {
      Media* media = static_cast<Media*>(fMedia.At(i));
      if (!media) continue;
      if (media->HasMedia(med)) return media;
    }
    return 0;
  }
  //__________________________________________________________________
  Bool_t Init()
  {
    Bool_t ret = AliFMDInputHits::Init();
    if (!gGeoManager) {
      Error("Init", "No geometry defined - make sure you have that");
      return kFALSE;
    }
    if (!fRun) {
      Error("Init", "gAlice not defined");
      return kFALSE;
    }
    Int_t now = 0;
    Int_t colon;
    TFile* file = TFile::Open("medid.root", "READ");
    if (!file) {
      Error("Init", "couldn't open medid.root");
      return kFALSE;
    }
    fOutput->cd();
    while ((colon = fModList.Index(":", now)) >= 0) {
      TString d(fModList(now, colon-now));
      now = colon+1;
      TArrayI* a = 0;
      file->GetObject(d.Data(), a);
      if (!a) {
	Warning("Init", "No medium array for %s", d.Data());
	continue;
      }
      AddMedia(d.Data(),0, a);
    }
    if (now < fModList.Length()) {
      colon = fModList.Length();
      TString d(fModList(now, colon-now));
      TArrayI* a = 0;
      file->GetObject(d.Data(), a);
      if (!a) 
	Warning("Init", "No medium array for %s", d.Data());
      else 
	AddMedia(d.Data(),0, a);
    }
    file->Close();
    if (fMedia.GetEntries() <= 0) return kFALSE;
    return ret;
  }  
  //__________________________________________________________________
  /** Begining of event
      @param ev Event number
      @return @c false on error */
  Bool_t Begin(Int_t ev) 
  {
    fEv = ev;
    return AliFMDInputHits::Begin(ev);
  }
  //__________________________________________________________________
  Bool_t ProcessHit(AliFMDHit* hit, TParticle* track) 
  {
    if (!hit || !track) {
      std::cout << "No hit or track (hit: " << hit << ", track: " 
		<< track << ")" << std::endl;
      return kFALSE;
    }
    fAll->Fill(fEv, hit);
    if (track->IsPrimary()) { 
      fPrim->Fill(fEv, hit);
      return kTRUE;
    }
      
    // Get production vertex 
    Double_t vx = track->Vx();
    Double_t vy = track->Vy();
    Double_t vz = track->Vz();
    // Get node 
    TGeoNode* prodNode = gGeoManager->FindNode(vx,vy,vz);
    if (!prodNode) { 
      Warning("ProcessHit", "didn't find a node for production vertex");
      return kTRUE;
    }
    //Info("ProcessHit", "Production vertex in node %s", prodNode->GetName());

    // Get volume 
    TGeoVolume* prodVol = prodNode->GetVolume();
    if (!prodVol) { 
      Warning("ProcessHit", "didn't find a volume for production vertex");
      return kTRUE;
    }
    //Info("ProcessHit", "Production vertex in volume %s",prodVol->GetName());
    
    // Med medium 
    TGeoMedium* prodMed = prodVol->GetMedium();
    if (!prodMed) { 
      Warning("ProcessHit", "didn't find a medium for production vertex");
      return kTRUE;
    }
    // Info("ProcessHit", "Production vertex in medium %s %d", 
    //      prodMed->GetName(), prodMed->GetId());
    
    
    Media* media = FindMedia(prodMed->GetId());
    if (!media) { 
      Warning("ProcessHit", "Media not found %s (%d)", 
	      prodMed->GetName(), prodMed->GetId());
      media = fOther; 
    }
    // Info("ProcessHit", "Adding to %s", media->GetName());
    media->Fill(fEv, hit);
    return kTRUE;
  }
  //__________________________________________________________________
  Bool_t Finish()
  {
    fOutput->cd();
    fStack->Write();
    fStack->Draw();
    fOutput->Write();
    fOutput->Close();
    fOutput = 0;
    fMedia.Delete();
    return kTRUE;
  }
};

//____________________________________________________________________
//
// EOF
//
Commit	Line	Data
bf000c32	1	//____________________________________________________________________
088f8e79	2	//
	3	// $Id$
	4	//
	5	// Script that contains a class to get the media where a certain track
	6	// was created. This is used for background studies.
	7	//
	8	// Use the script `Compile.C' to compile this class using ACLic.
	9	//
	10	#include <TH2D.h>
	11	#include <AliFMDHit.h>
	12	#include <AliFMDInput.h>
	13	#include <iostream>
	14	#include <TStyle.h>
	15	#include <TArrayF.h>
	16	#include <TArrayI.h>
	17	#include <AliRun.h>
	18	#include <TNtuple.h>
55f0ce5b	19	#include <AliDetector.h>
	20	#include <TFile.h>
	21	#include <TGeoManager.h>
	22	#include <TGeoVolume.h>
	23	#include <TGeoMedium.h>
	24	#include <TParticle.h>
	25	#include <THStack.h>
9b48326f	26	/** @class Media
	27	@brief Media description
	28	@ingroup FMD_script
	29	*/
088f8e79	30	struct Media : public TNamed
088f8e79	31	{
55f0ce5b	32	TArrayI fMeds;
088f8e79	33	TNtuple* fCount;
55f0ce5b	34	TH1D* fHist;
088f8e79	35	Media(const char* name, const char* title)
	36	: TNamed(name, title), fMeds(0)
	37	{
55f0ce5b	38	fHist = new TH1D(Form("h%s",name), title, 120, -6, 6);
	39	// fHist->SetFillStyle(3001);
	40	fCount = new TNtuple(GetName(), GetTitle(),
	41	"E:DET:RNG:SEC:STR:X:Y:Z:EDEP:PDG");
088f8e79	42	}
55f0ce5b	43	Media(const char* name, TArrayI* a)
55f0ce5b	44	: TNamed(name, "Media information"), fMeds(a->fN, a->fArray)
088f8e79	45	{
55f0ce5b	46	fHist = new TH1D(Form("h%s",name), GetTitle(), 120, -6, 6);
	47	fHist->SetFillStyle(3001);
	48	fHist->SetFillColor(fMeds.fN > 0 ? fMeds[0] : 0);
	49	fCount = new TNtuple(GetName(), GetTitle(),
	50	"E:DET:RNG:SEC:STR:X:Y:Z:EDEP:PDG");
	51	std::cout << "Media list " << name << ":" << std::endl;
	52	for (Int_t i = 0; i < fMeds.fN; i++) {
	53	std::cout << " " << fMeds[i] << std::flush;
	54	if (fMeds[i] == 0 && i > 0) break;
088f8e79	55	}
55f0ce5b	56	std::cout << std::endl;
	57	}
	58	Bool_t HasMedia(Int_t id)
	59	{
	60	if (fMeds.fN <= 0) return kFALSE;
	61	for (Int_t i = 0; i < fMeds.fN; i++) {
	62	if (fMeds[i] == id) return kTRUE;
	63	if (fMeds[i] == 0 && i > 0) break;
088f8e79	64	}
55f0ce5b	65	return kFALSE;
	66	}
	67	void Fill(Int_t ev, AliFMDHit* hit)
	68	{
	69	Float_t x[10];
	70	x[0] = ev;
	71	x[1] = hit->Detector();
	72	x[2] = Int_t(hit->Ring());
	73	x[3] = hit->Sector();
	74	x[4] = hit->Strip();
	75	x[5] = hit->X();
	76	x[6] = hit->Y();
	77	x[7] = hit->Z();
	78	x[8] = hit->Edep();
	79	x[9] = hit->Pdg();
	80	// return;
	81	fCount->Fill(x);
	82	// r = TMath::Sqrt(hit->X() * hit->X() + hit->Y() * hit->Y());
	83	// if(r!=0) Float_t wgt=1./(2. * 3.1415*r);
	84	Double_t r = TMath::Sqrt(hit->X()hit->X()+hit->Y()hit->Y());
	85	Double_t t = TMath::ATan2(r, hit->Z());
	86	Double_t e = -TMath::Log(TMath::Tan(t/2));
	87	fHist->Fill(e);
088f8e79	88	}
	89	};
	90
	91
bf000c32	92	//____________________________________________________________________
9b48326f	93	/** @class GetMedia
	94	@brief Get media where a particle is produced
	95	@code
55f0ce5b	96	aliroot
	97	Root> gAlice->Init("$ALICE_ROOT/FMD/Config.C");
	98	Root> .x $ALICE_ROOT/FMD/scriptsWriteMedArrays.C
	99	Root> gAlice->Run();
	100	Root> .q
	101	aliroot
9b48326f	102	Root> .L Compile.C
	103	Root> Compile("GetMedia.C")
	104	Root> GetMedia c
	105	Root> c.Run();
	106	@endcode
	107	@ingroup FMD_script
	108	*/
8ec606c2	109	class GetMedia : public AliFMDInput
088f8e79	110	{
	111	private:
	112	TString fModList;
	113	TObjArray fMedia;
	114	Media* fOther;
	115	Media* fAll;
55f0ce5b	116	Media* fPrim;
088f8e79	117	Int_t fEv;
55f0ce5b	118	THStack* fStack;
088f8e79	119	TFile* fOutput;
088f8e79	120	public:
bf000c32	121	//__________________________________________________________________
ababa197	122	GetMedia(const char* modlist="FMD:ITS:BODY:ABSO:T0:PIPE",
088f8e79	123	const char* output="media.root")
55f0ce5b	124	: fModList(modlist)
088f8e79	125	{
55f0ce5b	126	AddLoad(kGeometry);
8ec606c2	127	AddLoad(kHits);
55f0ce5b	128	AddLoad(kKinematics);
55f0ce5b	129
088f8e79	130	fOutput = TFile::Open(output, "RECREATE");
55f0ce5b	131	fStack = new THStack("summed", "Summed hits");
	132	fOther = AddMedia("other", "Unknown media");
	133	fAll = AddMedia("All", "All media");
	134	fPrim = AddMedia("Primary", "Primary particles");
088f8e79	135	fEv = 0;
55f0ce5b	136
	137
	138	fAll->fHist->SetFillColor(0);
	139	fAll->fHist->SetFillStyle(3004);
	140	fPrim->fHist->SetFillColor(1);
	141	fPrim->fHist->SetFillStyle(3001);
	142	fStack->Add(fPrim->fHist);
088f8e79	143	}
bf000c32	144	//__________________________________________________________________
55f0ce5b	145	Media* AddMedia(const char* name, const char* title, TArrayI* a=0)
088f8e79	146	{
088f8e79	147	Media* media = 0;
55f0ce5b	148	if (!a) media = new Media(name, title);
	149	else media = new Media(name, a);
	150	if (a) {
	151	fMedia.Add(media);
	152	fStack->Add(media->fHist);
	153	}
	154	return media;
	155	}
	156	//__________________________________________________________________
	157	Media* FindMedia(Int_t med)
	158	{
	159	for (Int_t i = 0; i < fMedia.GetEntries(); i++) {
	160	Media* media = static_cast<Media*>(fMedia.At(i));
	161	if (!media) continue;
	162	if (media->HasMedia(med)) return media;
088f8e79	163	}
	164	return 0;
	165	}
bf000c32	166	//__________________________________________________________________
088f8e79	167	Bool_t Init()
088f8e79	168	{
55f0ce5b	169	Bool_t ret = AliFMDInputHits::Init();
088f8e79	170	if (!gGeoManager) {
55f0ce5b	171	Error("Init", "No geometry defined - make sure you have that");
088f8e79	172	return kFALSE;
088f8e79	173	}
55f0ce5b	174	if (!fRun) {
55f0ce5b	175	Error("Init", "gAlice not defined");
088f8e79	176	return kFALSE;
	177	}
	178	Int_t now = 0;
	179	Int_t colon;
55f0ce5b	180	TFile* file = TFile::Open("medid.root", "READ");
	181	if (!file) {
	182	Error("Init", "couldn't open medid.root");
	183	return kFALSE;
	184	}
	185	fOutput->cd();
088f8e79	186	while ((colon = fModList.Index(":", now)) >= 0) {
55f0ce5b	187	TString d(fModList(now, colon-now));
088f8e79	188	now = colon+1;
55f0ce5b	189	TArrayI* a = 0;
	190	file->GetObject(d.Data(), a);
	191	if (!a) {
	192	Warning("Init", "No medium array for %s", d.Data());
	193	continue;
	194	}
	195	AddMedia(d.Data(),0, a);
088f8e79	196	}
55f0ce5b	197	if (now < fModList.Length()) {
	198	colon = fModList.Length();
	199	TString d(fModList(now, colon-now));
	200	TArrayI* a = 0;
	201	file->GetObject(d.Data(), a);
	202	if (!a)
	203	Warning("Init", "No medium array for %s", d.Data());
	204	else
	205	AddMedia(d.Data(),0, a);
	206	}
	207	file->Close();
088f8e79	208	if (fMedia.GetEntries() <= 0) return kFALSE;
55f0ce5b	209	return ret;
088f8e79	210	}
bf000c32	211	//__________________________________________________________________
9b48326f	212	/** Begining of event
	213	@param ev Event number
	214	@return @c false on error */
088f8e79	215	Bool_t Begin(Int_t ev)
	216	{
	217	fEv = ev;
	218	return AliFMDInputHits::Begin(ev);
	219	}
bf000c32	220	//__________________________________________________________________
088f8e79	221	Bool_t ProcessHit(AliFMDHit* hit, TParticle* track)
	222	{
	223	if (!hit \|\| !track) {
	224	std::cout << "No hit or track (hit: " << hit << ", track: "
	225	<< track << ")" << std::endl;
	226	return kFALSE;
	227	}
55f0ce5b	228	fAll->Fill(fEv, hit);
	229	if (track->IsPrimary()) {
	230	fPrim->Fill(fEv, hit);
	231	return kTRUE;
	232	}
	233
088f8e79	234	// Get production vertex
	235	Double_t vx = track->Vx();
	236	Double_t vy = track->Vy();
	237	Double_t vz = track->Vz();
088f8e79	238	// Get node
088f8e79	239	TGeoNode* prodNode = gGeoManager->FindNode(vx,vy,vz);
55f0ce5b	240	if (!prodNode) {
	241	Warning("ProcessHit", "didn't find a node for production vertex");
	242	return kTRUE;
	243	}
	244	//Info("ProcessHit", "Production vertex in node %s", prodNode->GetName());
	245
088f8e79	246	// Get volume
088f8e79	247	TGeoVolume* prodVol = prodNode->GetVolume();
55f0ce5b	248	if (!prodVol) {
	249	Warning("ProcessHit", "didn't find a volume for production vertex");
	250	return kTRUE;
	251	}
	252	//Info("ProcessHit", "Production vertex in volume %s",prodVol->GetName());
088f8e79	253
	254	// Med medium
	255	TGeoMedium* prodMed = prodVol->GetMedium();
55f0ce5b	256	if (!prodMed) {
	257	Warning("ProcessHit", "didn't find a medium for production vertex");
	258	return kTRUE;
	259	}
	260	// Info("ProcessHit", "Production vertex in medium %s %d",
	261	// prodMed->GetName(), prodMed->GetId());
088f8e79	262
088f8e79	263
55f0ce5b	264	Media* media = FindMedia(prodMed->GetId());
	265	if (!media) {
	266	Warning("ProcessHit", "Media not found %s (%d)",
	267	prodMed->GetName(), prodMed->GetId());
	268	media = fOther;
	269	}
	270	// Info("ProcessHit", "Adding to %s", media->GetName());
	271	media->Fill(fEv, hit);
088f8e79	272	return kTRUE;
088f8e79	273	}
bf000c32	274	//__________________________________________________________________
088f8e79	275	Bool_t Finish()
088f8e79	276	{
55f0ce5b	277	fOutput->cd();
	278	fStack->Write();
	279	fStack->Draw();
088f8e79	280	fOutput->Write();
088f8e79	281	fOutput->Close();
55f0ce5b	282	fOutput = 0;
55f0ce5b	283	fMedia.Delete();
088f8e79	284	return kTRUE;
	285	}
	286	};
	287
	288	//____________________________________________________________________
	289	//
	290	// EOF
	291	//