Introduce and adapt parameters ala arXiv:1101.2231v1
[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->SetXmu(3.2264);     // parton xsection
123   genHi->SetNtMax(3);        // time bins
124   if (0) { //RHIC settings
125     genHi->SetAlpha(0.47140452);
126     genHi->SetStringFrag(2.2,0.5);
127   }
128   genHi->Init();
129
130   TClonesArray *inucs = new TClonesArray("TParticle",10000);
131   TClonesArray *parts = new TClonesArray("MyPart",10000);
132   TClonesArray *nucs  = new TClonesArray("MyNuc",500);
133   MyHeader *myheader = new MyHeader;
134
135   TFile *outFile = TFile::Open(outFileName, "RECREATE");
136   outFile->SetCompressionLevel(5);
137   TDirectory::TContext context(outFile);
138
139   TTree *tree = new TTree("ampt", "AMPT tree");
140   tree->Branch("parts",&parts, 256*1024, 99);
141   tree->Branch("nucs",&nucs, 64*1024, 99);
142   tree->Branch("info",&myheader, 32*1024, 99);
143
144   // create events and fill them
145   for (Int_t iEvent = 0; iEvent < nEvents; ++iEvent) {
146
147     cout << "Event " << iEvent+1 << "/" << nEvents << endl;;
148     stack->Reset();
149     genHi->Generate();
150     parts->Clear();
151     const TClonesArray *iarr = genHi->GetParticles();
152     if (iarr) {
153       for(Int_t i=0;i<iarr->GetEntriesFast();++i) {
154         TParticle *p = (TParticle*)iarr->At(i);
155         new((*parts)[i]) MyPart(p->GetPdgCode(),
156                                 p->GetStatusCode(),
157                                 p->GetUniqueID(),
158                                 p->Pt(),
159                                 p->Eta(),
160                                 p->Phi());
161       }
162     }
163     TAmpt *ampt = genHi->Ampt();
164     if (ampt) {
165       ampt->ImportNucleons(inucs);
166       nucs->Clear();
167       for(Int_t i=0;i<inucs->GetEntriesFast();++i) {
168         TParticle *p = (TParticle*)inucs->At(i);
169         new((*nucs)[i]) MyNuc(p->GetPdgCode(),
170                               p->GetStatusCode(),
171                               p->GetUniqueID(),
172                               p->Vx(),
173                               p->Vy(),
174                               p->Vz());
175       }
176
177       myheader->fNATT=ampt->GetNATT();
178       myheader->fEATT=ampt->GetEATT();
179       myheader->fJATT=ampt->GetJATT();
180       myheader->fNT=ampt->GetNT();
181       myheader->fNP=ampt->GetNP();
182       myheader->fN0=ampt->GetN0();
183       myheader->fN01=ampt->GetN01();
184       myheader->fN10=ampt->GetN10();
185       myheader->fN11=ampt->GetN11();
186       myheader->fBB=ampt->GetBB();
187       myheader->fRP=ampt->GetPhi();
188     }
189     tree->Fill();
190   } // end of event loop
191   tree->Write();
192   outFile->Close();
193 }
194
195 //-----------------------------------------------------------------------------------------------------
196
197 void anaAmptMC(Int_t nEvents,
198                const char *inFileNames)
199 {
200
201   TChain *c = new TChain("ampt");
202   c->Add(inFileNames);
203
204
205   TClonesArray *parts = 0;
206   TClonesArray *nucs  = 0;
207   MyHeader *info = 0;
208
209   c->SetBranchAddress("info",&info);
210   c->SetBranchAddress("nucs",&nucs);
211   c->SetBranchAddress("parts",&parts);
212
213   Int_t nRead = nEvents;
214   if (nRead<0)
215     nRead = c->GetEntries();
216   else if (0 && (nRead>c->GetEntries()))
217     nRead = c->GetEntries();
218
219    for (Int_t ev=0;ev<nRead;++ev) {
220       c->GetEntry(ev);
221
222       Int_t fAN=0, fBN=0, fAStat[250], fBStat[250];
223       Double_t fXA[250], fXB[250], fYA[250], fYB[250];
224       for (Int_t k=0;k<nucs->GetEntries();++k) {
225         MyNuc *n = (MyNuc*)nucs->At(k);
226         if (n->Type()>0) {
227           fAStat[fAN] =0;
228           fXA[fAN] = n->X();
229           fYA[fAN] = n->Y();
230           ++fAN;
231         } else {
232           fBStat[fBN] =0;
233           fXB[fBN] = n->X();
234           fYB[fBN] = n->Y();
235           ++fBN;
236         }
237       }
238
239       // Glauber calculation
240       Double_t fXSect = 65; //mbarn
241       Double_t d2 = (Double_t)fXSect/(TMath::Pi()*10); // in fm^2
242
243       // For each of the A nucleons in nucleus B
244       for (Int_t i = 0; i<fBN; i++) {
245         for (Int_t j = 0 ; j < fAN ;j++) {
246           Double_t dx = fXB[i]-fXA[j];
247           Double_t dy = fYB[i]-fYA[j];
248           Double_t dij = dx*dx+dy*dy;
249           if (dij < d2) {
250             fBStat[i]++;
251             fAStat[j]++;
252           }
253         }
254       }
255       // Calculate npart
256       Int_t npart=0;
257       for (Int_t i = 0; i<fAN; i++) {
258         if (fAStat[i]>0) {
259           npart++;
260         }
261       }
262       for (Int_t i = 0; i<fBN; i++) {
263         if (fBStat[i]>0) {
264           npart++;
265         }
266       }
267       cout << ev << " : "  << info->fBB << " np " << npart << " vs " << info->fNP+info->fNT << endl;
268    }
269 }