AMPTInterface/test/AMPTAnalyzer.cc

0001 // -*- C++ -*-
0002 //
0003 // Package:    AMPTAnalyzer
0004 // Class:      AMPTAnalyzer
0005 //
0006
0007 // system include files
0008 #include <memory>
0009 #include <iostream>
0010 #include <string>
0011 #include <fstream>
0012
0013 // user include files
0014 #include "FWCore/Framework/interface/Frameworkfwd.h"
0015 #include "FWCore/Framework/interface/one/EDAnalyzer.h"
0016
0017 #include "FWCore/Framework/interface/Event.h"
0018 #include "FWCore/Framework/interface/MakerMacros.h"
0019 #include "FWCore/Framework/interface/EventSetup.h"
0020
0021 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0022 #include "FWCore/Utilities/interface/InputTag.h"
0023 #include "FWCore/Framework/interface/ESHandle.h"
0024
0025 #include "FWCore/ServiceRegistry/interface/Service.h"
0026
0027 #include "CommonTools/UtilAlgos/interface/TFileService.h"
0028
0029 #include "SimDataFormats/GeneratorProducts/interface/HepMCProduct.h"
0030 #include "SimDataFormats/CrossingFrame/interface/MixCollection.h"
0031 #include "SimDataFormats/Vertex/interface/SimVertex.h"
0032 #include "SimDataFormats/Vertex/interface/SimVertexContainer.h"
0033
0034 #include "HepMC/GenEvent.h"
0035 #include "HepMC/HeavyIon.h"
0036
0037 #include "SimGeneral/HepPDTRecord/interface/ParticleDataTable.h"
0038
0039 // root include file
0040 #include "TFile.h"
0041 #include "TNtuple.h"
0042
0043 static const int MAXPARTICLES = 5000000;
0044 static const int ETABINS = 3;  // Fix also in branch string
0045
0046 //
0047 // class decleration
0048 //
0049
0050 struct AMPTEvent {
0051   int event;
0052   float b;
0053   float npart;
0054   float ncoll;
0055   float nhard;
0056   float phi0;
0057
0058   int n[ETABINS];
0059   float ptav[ETABINS];
0060
0061   int mult;
0062   float pt[MAXPARTICLES];
0063   float eta[MAXPARTICLES];
0064   float phi[MAXPARTICLES];
0065   int pdg[MAXPARTICLES];
0066   int chg[MAXPARTICLES];
0067
0068   float vx;
0069   float vy;
0070   float vz;
0071   float vr;
0072 };
0073
0074 class AMPTAnalyzer : public edm::one::EDAnalyzer<edm::one::SharedResources> {
0075 public:
0076   explicit AMPTAnalyzer(const edm::ParameterSet&);
0077   ~AMPTAnalyzer() = default;
0078
0079 private:
0080   void beginJob() override;
0081   void analyze(const edm::Event&, const edm::EventSetup&) override;
0082   void endJob() override {}
0083
0084   // ----------member data ---------------------------
0085
0086   std::ofstream out_b;
0087   std::string fBFileName;
0088
0089   std::ofstream out_n;
0090   std::string fNFileName;
0091
0092   std::ofstream out_m;
0093   std::string fMFileName;
0094
0095   TTree* hydjetTree_;
0096   AMPTEvent hev_;
0097
0098   TNtuple* nt;
0099
0100   std::string output;  // Output filename
0101
0102   bool doAnalysis_;
0103   bool printLists_;
0104   bool doCF_;
0105   bool doVertex_;
0106   double etaMax_;
0107   double ptMin_;
0108   edm::InputTag simVerticesTag_;
0109
0110   const edm::ESGetToken<HepPDT::ParticleDataTable, PDTRecord> pdtToken_;
0111   const edm::EDGetTokenT<edm::HepMCProduct> mcToken_;
0112   edm::EDGetTokenT<edm::SimVertexContainer> simVertToken_;
0113   edm::EDGetTokenT<CrossingFrame<edm::HepMCProduct> > cfToken_;
0114 };
0115
0116 //
0117 // constants, enums and typedefs
0118 //
0119
0120 //
0121 // static data member definitions
0122 //
0123
0124 //
0125 // constructors and destructor
0126 //
0127 AMPTAnalyzer::AMPTAnalyzer(const edm::ParameterSet& iConfig)
0128     : pdtToken_(esConsumes<HepPDT::ParticleDataTable, PDTRecord>()),
0129       mcToken_(consumes<edm::HepMCProduct>(
0130           iConfig.getUntrackedParameter<edm::InputTag>("src", edm::InputTag("VtxSmeared")))) {
0131   usesResource(TFileService::kSharedResource);
0132   //now do what ever initialization is needed
0133   fBFileName = iConfig.getUntrackedParameter<std::string>("output_b", "b_values.txt");
0134   fNFileName = iConfig.getUntrackedParameter<std::string>("output_n", "n_values.txt");
0135   fMFileName = iConfig.getUntrackedParameter<std::string>("output_m", "m_values.txt");
0136   doAnalysis_ = iConfig.getUntrackedParameter<bool>("doAnalysis", true);
0137   printLists_ = iConfig.getUntrackedParameter<bool>("printLists", false);
0138   doCF_ = iConfig.getUntrackedParameter<bool>("doMixed", false);
0139   if (doCF_)
0140     cfToken_ = consumes<CrossingFrame<edm::HepMCProduct> >(edm::InputTag("mix", "source"));
0141   doVertex_ = iConfig.getUntrackedParameter<bool>("doVertex", false);
0142   if (doVertex_) {
0143     simVertToken_ = consumes<edm::SimVertexContainer>(iConfig.getParameter<edm::InputTag>("simVerticesTag"));
0144   }
0145   etaMax_ = iConfig.getUntrackedParameter<double>("etaMax", 2);
0146   ptMin_ = iConfig.getUntrackedParameter<double>("ptMin", 0);
0147 }
0148
0149 //
0150 // member functions
0151 //
0152
0153 // ------------ method called to for each event  ------------
0154 void AMPTAnalyzer::analyze(const edm::Event& iEvent, const edm::EventSetup& iSetup) {
0155   const HepPDT::ParticleDataTable* pdt = &iSetup.getData(pdtToken_);
0156
0157   hev_.event = iEvent.id().event();
0158   for (int ieta = 0; ieta < ETABINS; ++ieta) {
0159     hev_.n[ieta] = 0;
0160     hev_.ptav[ieta] = 0;
0161   }
0162   hev_.mult = 0;
0163
0164   double phi0 = 0;
0165   double b = -1;
0166   int npart = -1;
0167   int ncoll = -1;
0168   int nhard = -1;
0169   double vx = -99;
0170   double vy = -99;
0171   double vz = -99;
0172   double vr = -99;
0173   const HepMC::GenEvent* evt;
0174   const HepMC::GenEvent* evt2;
0175
0176   int nmix = -1;
0177   int np = 0;
0178   int sig = -1;
0179   int src = -1;
0180
0181   if (doCF_) {
0182     const edm::Handle<CrossingFrame<edm::HepMCProduct> >& cf = iEvent.getHandle(cfToken_);
0183
0184     MixCollection<edm::HepMCProduct> mix(cf.product());
0185
0186     nmix = mix.size();
0187
0188     edm::LogVerbatim("AMPTAnalysis") << "Mix Collection Size: " << mix;
0189
0190     MixCollection<edm::HepMCProduct>::iterator mbegin = mix.begin();
0191     MixCollection<edm::HepMCProduct>::iterator mend = mix.end();
0192
0193     for (MixCollection<edm::HepMCProduct>::iterator mixit = mbegin; mixit != mend; ++mixit) {
0194       const HepMC::GenEvent* subevt = (*mixit).GetEvent();
0195       int all = subevt->particles_size();
0196       np += all;
0197
0198       /*
0199        HepMC::GenEvent::particle_const_iterator begin = subevt->particles_begin();
0200        HepMC::GenEvent::particle_const_iterator end = subevt->particles_end();
0201        for(HepMC::GenEvent::particle_const_iterator it = begin; it != end; ++it){
0202      if((*it)->status() == 1){
0203            float pdg_id = (*it)->pdg_id();
0204            float eta = (*it)->momentum().eta();
0205            float pt = (*it)->momentum().perp();
0206            const ParticleData * part = pdt->particle(pdg_id );
0207            float charge = part->charge();
0208        }
0209      }
0210        }
0211
0212        */
0213     }
0214   }
0215
0216   const edm::Handle<edm::HepMCProduct>& mc = iEvent.getHandle(mcToken_);
0217   evt = mc->GetEvent();
0218
0219   const edm::Handle<edm::HepMCProduct>& mc2 = iEvent.getHandle(mcToken_);
0220   evt2 = mc2->GetEvent();
0221
0222   const HepMC::HeavyIon* hi = evt->heavy_ion();
0223   if (hi) {
0224     b = hi->impact_parameter();
0225     npart = hi->Npart_proj() + hi->Npart_targ();
0226     ncoll = hi->Ncoll();
0227     nhard = hi->Ncoll_hard();
0228     phi0 = hi->event_plane_angle();
0229
0230     if (printLists_) {
0231       out_b << b << std::endl;
0232       out_n << npart << std::endl;
0233     }
0234   }
0235
0236   src = evt->particles_size();
0237   sig = evt2->particles_size();
0238
0239   HepMC::GenEvent::particle_const_iterator begin = evt->particles_begin();
0240   HepMC::GenEvent::particle_const_iterator end = evt->particles_end();
0241   for (HepMC::GenEvent::particle_const_iterator it = begin; it != end; ++it) {
0242     if ((*it)->status() == 1) {
0243       //if((*it)->status() != 1) edm::LogVerbatim("AMPTAnalysis") << (*it)->status();
0244       int pdg_id = (*it)->pdg_id();
0245       float eta = (*it)->momentum().eta();
0246       float phi = (*it)->momentum().phi();
0247       float pt = (*it)->momentum().perp();
0248       const ParticleData* part = pdt->particle(pdg_id);
0249       int charge = static_cast<int>(part->charge());
0250
0251       hev_.pt[hev_.mult] = pt;
0252       hev_.eta[hev_.mult] = eta;
0253       hev_.phi[hev_.mult] = phi;
0254       hev_.pdg[hev_.mult] = pdg_id;
0255       hev_.chg[hev_.mult] = charge;
0256
0257       eta = fabs(eta);
0258       int etabin = 0;
0259       if (eta > 0.5)
0260         etabin = 1;
0261       if (eta > 1.)
0262         etabin = 2;
0263       if (eta < 2.) {
0264         hev_.ptav[etabin] += pt;
0265         ++(hev_.n[etabin]);
0266       }
0267       ++(hev_.mult);
0268     }
0269   }
0270   //   }
0271
0272   if (doVertex_) {
0273     const edm::Handle<edm::SimVertexContainer>& simVertices = iEvent.getHandle(simVertToken_);
0274
0275     if (!simVertices.isValid())
0276       throw cms::Exception("FatalError") << "No vertices found\n";
0277     int inum = 0;
0278
0279     edm::SimVertexContainer::const_iterator it = simVertices->begin();
0280     SimVertex vertex = (*it);
0281     edm::LogVerbatim("AMPTAnalysis") << " Vertex position " << inum << " " << vertex.position().rho() << " "
0282                                      << vertex.position().z();
0283     vx = vertex.position().x();
0284     vy = vertex.position().y();
0285     vz = vertex.position().z();
0286     vr = vertex.position().rho();
0287   }
0288
0289   for (int i = 0; i < 3; ++i) {
0290     hev_.ptav[i] = hev_.ptav[i] / hev_.n[i];
0291   }
0292
0293   hev_.b = b;
0294   hev_.npart = npart;
0295   hev_.ncoll = ncoll;
0296   hev_.nhard = nhard;
0297   hev_.phi0 = phi0;
0298   hev_.vx = vx;
0299   hev_.vy = vy;
0300   hev_.vz = vz;
0301   hev_.vr = vr;
0302
0303   nt->Fill(nmix, np, src, sig);
0304
0305   hydjetTree_->Fill();
0306 }
0307
0308 // ------------ method called once each job just before starting event loop  ------------
0309 void AMPTAnalyzer::beginJob() {
0310   if (printLists_) {
0311     out_b.open(fBFileName.c_str());
0312     if (out_b.good() == false)
0313       throw cms::Exception("BadFile") << "Can\'t open file " << fBFileName;
0314     out_n.open(fNFileName.c_str());
0315     if (out_n.good() == false)
0316       throw cms::Exception("BadFile") << "Can\'t open file " << fNFileName;
0317     out_m.open(fMFileName.c_str());
0318     if (out_m.good() == false)
0319       throw cms::Exception("BadFile") << "Can\'t open file " << fMFileName;
0320   }
0321
0322   if (doAnalysis_) {
0323     edm::Service<TFileService> f;
0324     nt = f->make<TNtuple>("nt", "Mixing Analysis", "mix:np:src:sig");
0325
0326     hydjetTree_ = f->make<TTree>("hi", "Tree of AMPT Events");
0327     hydjetTree_->Branch("event", &hev_.event, "event/I");
0328     hydjetTree_->Branch("b", &hev_.b, "b/F");
0329     hydjetTree_->Branch("npart", &hev_.npart, "npart/F");
0330     hydjetTree_->Branch("ncoll", &hev_.ncoll, "ncoll/F");
0331     hydjetTree_->Branch("nhard", &hev_.nhard, "nhard/F");
0332     hydjetTree_->Branch("phi0", &hev_.phi0, "phi0/F");
0333
0334     hydjetTree_->Branch("n", hev_.n, "n[3]/I");
0335     hydjetTree_->Branch("ptav", hev_.ptav, "ptav[3]/F");
0336
0337     hydjetTree_->Branch("mult", &hev_.mult, "mult/I");
0338     hydjetTree_->Branch("pt", hev_.pt, "pt[mult]/F");
0339     hydjetTree_->Branch("eta", hev_.eta, "eta[mult]/F");
0340     hydjetTree_->Branch("phi", hev_.phi, "phi[mult]/F");
0341     hydjetTree_->Branch("pdg", hev_.pdg, "pdg[mult]/I");
0342     hydjetTree_->Branch("chg", hev_.chg, "chg[mult]/I");
0343
0344     hydjetTree_->Branch("vx", &hev_.vx, "vx/F");
0345     hydjetTree_->Branch("vy", &hev_.vy, "vy/F");
0346     hydjetTree_->Branch("vz", &hev_.vz, "vz/F");
0347     hydjetTree_->Branch("vr", &hev_.vr, "vr/F");
0348   }
0349 }
0350
0351 //define this as a plug-in
0352 DEFINE_FWK_MODULE(AMPTAnalyzer);