7bb1143819d7d2a121d2d54cdde36683823e0da1
[u/mrichter/AliRoot.git] / TAmpt / macros / createAmptMC.C
1 #if !defined(__CINT__) || defined(__MAKECINT__)
2 #include <Riostream.h>
3 #include <TFile.h>
4 #include <TTree.h>
5 #include <TMath.h>
6 #include <TParticle.h>
7 #include <TPDGCode.h>
8 #include <TDatabasePDG.h>
9 #include <TRandom3.h>
10 #include <TChain.h>
11 #include "AliRun.h"
12 #include "AliRunLoader.h"
13 #include "AliHeader.h"
14 #include "AliStack.h"
15 #include "AliPDG.h"
16 #include "AliGenAmpt.h"
17 #include "TAmpt.h"
18 #endif
19
20 class MyHeader
21 {
22 public:
23   MyHeader() : fNATT(0), fEATT(0),fJATT(0),fNT(0),fNP(0),fN0(0),fN01(0),fN10(0),fN11(0),fBB(0),fRP(0) {;}
24   virtual ~MyHeader() {;}
25   Int_t       fNATT;
26   Float_t     fEATT;
27   Int_t       fJATT;
28   Int_t       fNT;
29   Int_t       fNP;
30   Int_t       fN0;
31   Int_t       fN01;
32   Int_t       fN10;
33   Int_t       fN11;
34   Float_t     fBB;
35   Float_t     fRP;
36   ClassDef(MyHeader,1) // My header class
37 };
38
39 class MyNuc : public TObject
40 {
41 public:
42   MyNuc(Int_t pid=0, Int_t st=0, Int_t type=0, Double_t x=0, Double_t y=0, Double_t z=0) :
43     TObject(), fPid(pid), fSt(st), fType(type), fX(x), fY(y), fZ(z) {;}
44   Double_t    X()    const { return fX; }
45   Double_t    Y()    const { return fY; }
46   Double_t    Z()    const { return fZ; }
47   Int_t       Pid()  const { return fPid; }
48   Int_t       St()   const { return fSt; }
49   Int_t       Type() const { return fType; }
50 protected:
51   Int_t       fPid;
52   Int_t       fSt;
53   Int_t       fType;
54   Double32_t  fX;
55   Double32_t  fY;
56   Double32_t  fZ;
57   ClassDef(MyNuc,1) // My nucleon class in cartesian coordinates
58 };
59
60 class MyPart : public TObject
61 {
62 public:
63   MyPart(Int_t pid=0, Int_t st=0, Int_t type=0, Double_t pt=0, Double_t eta=0, Double_t phi=0) :
64     TObject(), fPid(pid), fSt(st), fType(type), fPt(pt), fEta(eta), fPhi(phi) {;}
65   Double_t    Px()   const { return fPt*TMath::Cos(fPhi);  }
66   Double_t    Py()   const { return fPt*TMath::Sin(fPhi);  }
67   Double_t    Pz()   const { return fPt*TMath::SinH(fEta); }
68   Double_t    Pt()   const { return fPt;  }
69   Double_t    Eta()  const { return fEta; }
70   Double_t    Phi()  const { return fPhi; }
71   Int_t       Pid()  const { return fPid; }
72   Int_t       St()   const { return fSt; }
73   Int_t       Type() const { return fType; }
74 protected:
75   Int_t       fPid;
76   Int_t       fSt;
77   Int_t       fType;
78   Double32_t  fPt;
79   Double32_t  fEta;
80   Double32_t  fPhi;
81   ClassDef(MyPart,1) // My particle class in cylindrical coordinates
82 };
83
84 void createAmptMC(Int_t nEvents,
85                   const char *outFileName = "amptout.root");
86
87 void anaAmptMC(Int_t nEvents,
88                const char *inFileNames = "/eliza6/alice/loizides/ampt/run_1/ampt*.root");
89
90 //-----------------------------------------------------------------------------------------------------
91
92 void createAmptMC(Int_t nEvents,
93                   const char *outFileName) 
94 {
95   TClass::GetClass("MyNuc")->IgnoreTObjectStreamer();
96   TClass::GetClass("MyPart")->IgnoreTObjectStreamer();
97
98   AliPDG::AddParticlesToPdgDataBase();
99   TDatabasePDG::Instance();
100
101   // Run loader
102   TFolder *folder = new TFolder("myfolder","myfolder");
103   AliRunLoader* rl = new AliRunLoader(folder);
104   rl->MakeHeader();
105   rl->MakeStack();
106   AliStack* stack = rl->Stack();
107   //AliHeader* rheader = rl->GetHeader();
108
109   AliGenAmpt *genHi = new AliGenAmpt(-1);
110   genHi->SetStack(stack);
111   genHi->SetEnergyCMS(2760);
112   genHi->SetReferenceFrame("CMS");
113   genHi->SetProjectile("A", 208, 82);
114   genHi->SetTarget    ("A", 208, 82);
115   genHi->SetPtHardMin (2);
116   genHi->SetImpactParameterRange(0.,30);
117   genHi->SetJetQuenching(0); // enable jet quenching
118   genHi->SetShadowing(1);    // enable shadowing
119   genHi->SetDecaysOff(1);    // neutral pion and heavy particle decays switched off
120   genHi->SetSpectators(1);   // track spectators 
121   genHi->SetIsoft(1);        // standard AMPT
122   genHi->SetNtMax(3);        // time bins
123   genHi->SetXmu(3.2264);     // parton xsection
124   genHi->Init();
125
126   TClonesArray *inucs = new TClonesArray("TParticle",10000);
127   TClonesArray *parts = new TClonesArray("MyPart",10000);
128   TClonesArray *nucs  = new TClonesArray("MyNuc",500);
129   MyHeader *myheader = new MyHeader;
130
131   TFile *outFile = TFile::Open(outFileName, "RECREATE");
132   outFile->SetCompressionLevel(5);
133   TDirectory::TContext context(outFile);
134
135   TTree *tree = new TTree("ampt", "AMPT tree");
136   tree->Branch("parts",&parts, 256*1024, 99);
137   tree->Branch("nucs",&nucs, 64*1024, 99);
138   tree->Branch("info",&myheader, 32*1024, 99);
139
140   // create events and fill them
141   for (Int_t iEvent = 0; iEvent < nEvents; ++iEvent) {
142
143     cout << "Event " << iEvent+1 << "/" << nEvents << endl;;
144     stack->Reset();
145     genHi->Generate();
146     parts->Clear();
147     const TClonesArray *iarr = genHi->GetParticles();
148     if (iarr) {
149       for(Int_t i=0;i<iarr->GetEntriesFast();++i) {
150         TParticle *p = (TParticle*)iarr->At(i);
151         new((*parts)[i]) MyPart(p->GetPdgCode(),
152                                 p->GetStatusCode(),
153                                 p->GetUniqueID(),
154                                 p->Pt(),
155                                 p->Eta(),
156                                 p->Phi());
157       }
158     }
159     TAmpt *ampt = genHi->Ampt();
160     if (ampt) {
161       ampt->ImportNucleons(inucs);
162       nucs->Clear();
163       for(Int_t i=0;i<inucs->GetEntriesFast();++i) {
164         TParticle *p = (TParticle*)inucs->At(i);
165         new((*nucs)[i]) MyNuc(p->GetPdgCode(),
166                               p->GetStatusCode(),
167                               p->GetUniqueID(),
168                               p->Vx(),
169                               p->Vy(),
170                               p->Vz());
171       }
172
173       myheader->fNATT=ampt->GetNATT();
174       myheader->fEATT=ampt->GetEATT();
175       myheader->fJATT=ampt->GetJATT();
176       myheader->fNT=ampt->GetNT();
177       myheader->fNP=ampt->GetNP();
178       myheader->fN0=ampt->GetN0();
179       myheader->fN01=ampt->GetN01();
180       myheader->fN10=ampt->GetN10();
181       myheader->fN11=ampt->GetN11();
182       myheader->fBB=ampt->GetBB();
183       myheader->fRP=ampt->GetPhi();
184     }
185     tree->Fill();
186   } // end of event loop
187   tree->Write();
188   outFile->Close();
189 }
190
191 //-----------------------------------------------------------------------------------------------------
192
193 void anaAmptMC(Int_t nEvents,
194                const char *inFileNames)
195 {
196
197   TChain *c = new TChain("ampt");
198   c->Add(inFileNames);
199
200
201   TClonesArray *parts = 0;
202   TClonesArray *nucs  = 0;
203   MyHeader *info = 0;
204
205   c->SetBranchAddress("info",&info);
206   c->SetBranchAddress("nucs",&nucs);
207   c->SetBranchAddress("parts",&parts);
208
209   Int_t nRead = nEvents;
210   if (nRead<0)
211     nRead = c->GetEntries();
212   else if (0 && (nRead>c->GetEntries()))
213     nRead = c->GetEntries();
214
215    for (Int_t ev=0;ev<nRead;++ev) {
216       c->GetEntry(ev);
217
218       Int_t fAN=0, fBN=0, fAStat[250], fBStat[250];
219       Double_t fXA[250], fXB[250], fYA[250], fYB[250];
220       for (Int_t k=0;k<nucs->GetEntries();++k) {
221         MyNuc *n = (MyNuc*)nucs->At(k);
222         if (n->Type()>0) {
223           fAStat[fAN] =0;
224           fXA[fAN] = n->X();
225           fYA[fAN] = n->Y();
226           ++fAN;
227         } else {
228           fBStat[fBN] =0;
229           fXB[fBN] = n->X();
230           fYB[fBN] = n->Y();
231           ++fBN;
232         }
233       }
234
235       // Glauber calculation
236       Double_t fXSect = 65; //mbarn
237       Double_t d2 = (Double_t)fXSect/(TMath::Pi()*10); // in fm^2
238
239       // For each of the A nucleons in nucleus B
240       for (Int_t i = 0; i<fBN; i++) {
241         for (Int_t j = 0 ; j < fAN ;j++) {
242           Double_t dx = fXB[i]-fXA[j];
243           Double_t dy = fYB[i]-fYA[j];
244           Double_t dij = dx*dx+dy*dy;
245           if (dij < d2) {
246             fBStat[i]++;
247             fAStat[j]++;
248           }
249         }
250       }
251       // Calculate npart
252       Int_t npart=0;
253       for (Int_t i = 0; i<fAN; i++) {
254         if (fAStat[i]>0) {
255           npart++;
256         }
257       }
258       for (Int_t i = 0; i<fBN; i++) {
259         if (fBStat[i]>0) {
260           npart++;
261         }
262       }
263       cout << ev << " : "  << info->fBB << " np " << npart << " vs " << info->fNP+info->fNT << endl;
264    }
265 }