HiGenCommon/plugins/MixBoostEvtVtxGenerator.cc

0001
0002 /*
0003 ________________________________________________________________________
0004
0005  MixBoostEvtVtxGenerator
0006
0007  Smear vertex according to the Beta function on the transverse plane
0008  and a Gaussian on the z axis. It allows the beam to have a crossing
0009  angle (slopes dxdz and dydz).
0010
0011  Based on GaussEvtVtxGenerator
0012  implemented by Francisco Yumiceva (yumiceva@fnal.gov)
0013
0014  FERMILAB
0015  2006
0016 ________________________________________________________________________
0017 */
0018
0019 //lingshan: add beta for z-axis boost
0020
0021 //#include "IOMC/EventVertexGenerators/interface/BetafuncEvtVtxGenerator.h"
0022
0023 #include "FWCore/Framework/interface/Event.h"
0024 #include "FWCore/Utilities/interface/Exception.h"
0025 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0026
0027 #include "FWCore/Framework/interface/one/EDProducer.h"
0028 #include "FWCore/Utilities/interface/InputTag.h"
0029
0030 #include "SimDataFormats/CrossingFrame/interface/MixCollection.h"
0031 #include "SimDataFormats/GeneratorProducts/interface/HepMCProduct.h"
0032 #include "DataFormats/VertexReco/interface/Vertex.h"
0033 #include "DataFormats/VertexReco/interface/VertexFwd.h"
0034
0035 #include <CLHEP/Units/SystemOfUnits.h>
0036 #include <CLHEP/Units/GlobalPhysicalConstants.h>
0037 //#include "CLHEP/Vector/ThreeVector.h"
0038 #include "HepMC/SimpleVector.h"
0039 #include "TMatrixD.h"
0040
0041 #include <iostream>
0042
0043 using namespace edm;
0044 using namespace std;
0045 using namespace CLHEP;
0046
0047 class RandGaussQ;
0048 class FourVector;
0049
0050 class MixBoostEvtVtxGenerator : public edm::one::EDProducer<> {
0051 public:
0052   MixBoostEvtVtxGenerator(const edm::ParameterSet& p);
0053   /** Copy constructor */
0054   MixBoostEvtVtxGenerator(const MixBoostEvtVtxGenerator& p) = delete;
0055   /** Copy assignment operator */
0056   MixBoostEvtVtxGenerator& operator=(const MixBoostEvtVtxGenerator& rhs) = delete;
0057   ~MixBoostEvtVtxGenerator() override;
0058
0059   /// return a new event vertex
0060   void produce(edm::Event&, const edm::EventSetup&) override;
0061   virtual TMatrixD* GetInvLorentzBoost();
0062   virtual HepMC::FourVector* getVertex(edm::Event&);
0063   virtual HepMC::FourVector* getRecVertex(edm::Event&);
0064
0065   /// set resolution in Z in cm
0066   void sigmaZ(double s = 1.0);
0067
0068   /// set mean in X in cm
0069   void X0(double m = 0) { fX0 = m; }
0070   /// set mean in Y in cm
0071   void Y0(double m = 0) { fY0 = m; }
0072   /// set mean in Z in cm
0073   void Z0(double m = 0) { fZ0 = m; }
0074
0075   /// set half crossing angle
0076   void Phi(double m = 0) { phi_ = m; }
0077   /// angle between crossing plane and horizontal plane
0078   void Alpha(double m = 0) { alpha_ = m; }
0079   void Beta(double m = 0) { beta_ = m; }
0080
0081   /// set beta_star
0082   void betastar(double m = 0) { fbetastar = m; }
0083   /// emittance (no the normalized)
0084   void emittance(double m = 0) { femittance = m; }
0085
0086   /// beta function
0087   double BetaFunction(double z, double z0);
0088
0089 private:
0090   double alpha_, phi_;
0091   //TMatrixD boost_;
0092   double beta_;
0093   double fX0, fY0, fZ0;
0094   double fSigmaZ;
0095   //double fdxdz, fdydz;
0096   double fbetastar, femittance;
0097   double falpha;
0098
0099   HepMC::FourVector* fVertex;
0100   TMatrixD* boost_;
0101   double fTimeOffset;
0102
0103   const edm::EDGetTokenT<reco::VertexCollection> vtxLabel;
0104   const edm::EDGetTokenT<HepMCProduct> signalLabel;
0105   const edm::EDGetTokenT<CrossingFrame<HepMCProduct> > mixLabel;
0106   const bool useRecVertex;
0107   std::vector<double> vtxOffset;
0108 };
0109
0110 MixBoostEvtVtxGenerator::MixBoostEvtVtxGenerator(const edm::ParameterSet& pset)
0111     : fVertex(nullptr),
0112       boost_(nullptr),
0113       fTimeOffset(0),
0114       vtxLabel(mayConsume<reco::VertexCollection>(pset.getParameter<edm::InputTag>("vtxLabel"))),
0115       signalLabel(consumes<HepMCProduct>(pset.getParameter<edm::InputTag>("signalLabel"))),
0116       mixLabel(consumes<CrossingFrame<HepMCProduct> >(pset.getParameter<edm::InputTag>("mixLabel"))),
0117       useRecVertex(pset.exists("useRecVertex") ? pset.getParameter<bool>("useRecVertex") : false) {
0118   beta_ = pset.getParameter<double>("Beta");
0119   alpha_ = 0;
0120   phi_ = 0;
0121   if (pset.exists("Alpha")) {
0122     alpha_ = pset.getParameter<double>("Alpha") * radian;
0123     phi_ = pset.getParameter<double>("Phi") * radian;
0124   }
0125
0126   vtxOffset.resize(3);
0127   if (pset.exists("vtxOffset"))
0128     vtxOffset = pset.getParameter<std::vector<double> >("vtxOffset");
0129
0130   produces<edm::HepMCProduct>();
0131 }
0132
0133 MixBoostEvtVtxGenerator::~MixBoostEvtVtxGenerator() {
0134   if (fVertex != nullptr)
0135     delete fVertex;
0136   if (boost_ != nullptr)
0137     delete boost_;
0138 }
0139
0140 double MixBoostEvtVtxGenerator::BetaFunction(double z, double z0) {
0141   return sqrt(femittance * (fbetastar + (((z - z0) * (z - z0)) / fbetastar)));
0142 }
0143
0144 void MixBoostEvtVtxGenerator::sigmaZ(double s) {
0145   if (s >= 0) {
0146     fSigmaZ = s;
0147   } else {
0148     throw cms::Exception("LogicError") << "Error in MixBoostEvtVtxGenerator::sigmaZ: "
0149                                        << "Illegal resolution in Z (negative)";
0150   }
0151 }
0152
0153 TMatrixD* MixBoostEvtVtxGenerator::GetInvLorentzBoost() {
0154   //alpha_ = 0;
0155   //phi_ = 142.e-6;
0156   //    if (boost_ != 0 ) return boost_;
0157
0158   //boost_.ResizeTo(4,4);
0159   //boost_ = new TMatrixD(4,4);
0160   TMatrixD tmpboost(4, 4);
0161   TMatrixD tmpboostZ(4, 4);
0162   TMatrixD tmpboostXYZ(4, 4);
0163
0164   //if ( (alpha_ == 0) && (phi_==0) ) { boost_->Zero(); return boost_; }
0165
0166   // Lorentz boost to frame where the collision is head-on
0167   // phi is the half crossing angle in the plane ZS
0168   // alpha is the angle to the S axis from the X axis in the XY plane
0169
0170   tmpboost(0, 0) = 1. / cos(phi_);
0171   tmpboost(0, 1) = -cos(alpha_) * sin(phi_);
0172   tmpboost(0, 2) = -tan(phi_) * sin(phi_);
0173   tmpboost(0, 3) = -sin(alpha_) * sin(phi_);
0174   tmpboost(1, 0) = -cos(alpha_) * tan(phi_);
0175   tmpboost(1, 1) = 1.;
0176   tmpboost(1, 2) = cos(alpha_) * tan(phi_);
0177   tmpboost(1, 3) = 0.;
0178   tmpboost(2, 0) = 0.;
0179   tmpboost(2, 1) = -cos(alpha_) * sin(phi_);
0180   tmpboost(2, 2) = cos(phi_);
0181   tmpboost(2, 3) = -sin(alpha_) * sin(phi_);
0182   tmpboost(3, 0) = -sin(alpha_) * tan(phi_);
0183   tmpboost(3, 1) = 0.;
0184   tmpboost(3, 2) = sin(alpha_) * tan(phi_);
0185   tmpboost(3, 3) = 1.;
0186   //cout<<"beta "<<beta_;
0187   double gama = 1.0 / sqrt(1 - beta_ * beta_);
0188   tmpboostZ(0, 0) = gama;
0189   tmpboostZ(0, 1) = 0.;
0190   tmpboostZ(0, 2) = -1.0 * beta_ * gama;
0191   tmpboostZ(0, 3) = 0.;
0192   tmpboostZ(1, 0) = 0.;
0193   tmpboostZ(1, 1) = 1.;
0194   tmpboostZ(1, 2) = 0.;
0195   tmpboostZ(1, 3) = 0.;
0196   tmpboostZ(2, 0) = -1.0 * beta_ * gama;
0197   tmpboostZ(2, 1) = 0.;
0198   tmpboostZ(2, 2) = gama;
0199   tmpboostZ(2, 3) = 0.;
0200   tmpboostZ(3, 0) = 0.;
0201   tmpboostZ(3, 1) = 0.;
0202   tmpboostZ(3, 2) = 0.;
0203   tmpboostZ(3, 3) = 1.;
0204
0205   tmpboostXYZ = tmpboostZ * tmpboost;
0206   tmpboostXYZ.Invert();
0207
0208   //cout<<"Boosting with beta : "<<beta_<<endl;
0209
0210   boost_ = new TMatrixD(tmpboostXYZ);
0211   boost_->Print();
0212
0213   return boost_;
0214 }
0215
0216 HepMC::FourVector* MixBoostEvtVtxGenerator::getVertex(Event& evt) {
0217   const HepMC::GenEvent* inev = nullptr;
0218
0219   const edm::Handle<CrossingFrame<HepMCProduct> >& cf = evt.getHandle(mixLabel);
0220   MixCollection<HepMCProduct> mix(cf.product());
0221
0222   const HepMCProduct& bkg = mix.getObject(1);
0223   if (!(bkg.isVtxGenApplied())) {
0224     throw cms::Exception("MatchVtx") << "Input background does not have smeared vertex!" << endl;
0225   } else {
0226     inev = bkg.GetEvent();
0227   }
0228
0229   HepMC::GenVertex* genvtx = inev->signal_process_vertex();
0230   if (!genvtx) {
0231     cout << "No Signal Process Vertex!" << endl;
0232     HepMC::GenEvent::particle_const_iterator pt = inev->particles_begin();
0233     HepMC::GenEvent::particle_const_iterator ptend = inev->particles_end();
0234     while (!genvtx || (genvtx->particles_in_size() == 1 && pt != ptend)) {
0235       if (!genvtx)
0236         cout << "No Gen Vertex!" << endl;
0237       if (pt == ptend)
0238         cout << "End reached!" << endl;
0239       genvtx = (*pt)->production_vertex();
0240       ++pt;
0241     }
0242   }
0243   double aX, aY, aZ, aT;
0244
0245   aX = genvtx->position().x();
0246   aY = genvtx->position().y();
0247   aZ = genvtx->position().z();
0248   aT = genvtx->position().t();
0249
0250   if (!fVertex)
0251     fVertex = new HepMC::FourVector();
0252   fVertex->set(aX, aY, aZ, aT);
0253
0254   return fVertex;
0255 }
0256
0257 HepMC::FourVector* MixBoostEvtVtxGenerator::getRecVertex(Event& evt) {
0258   const edm::Handle<reco::VertexCollection>& input = evt.getHandle(vtxLabel);
0259
0260   double aX, aY, aZ;
0261
0262   aX = input->begin()->position().x() + vtxOffset[0];
0263   aY = input->begin()->position().y() + vtxOffset[1];
0264   aZ = input->begin()->position().z() + vtxOffset[2];
0265
0266   if (!fVertex)
0267     fVertex = new HepMC::FourVector();
0268   fVertex->set(10.0 * aX, 10.0 * aY, 10.0 * aZ, 0.0);  // HepMC positions in mm (RECO in cm)
0269
0270   return fVertex;
0271 }
0272
0273 void MixBoostEvtVtxGenerator::produce(Event& evt, const EventSetup&) {
0274   const edm::Handle<HepMCProduct>& HepUnsmearedMCEvt = evt.getHandle(signalLabel);
0275
0276   // Copy the HepMC::GenEvent
0277   HepMC::GenEvent* genevt = new HepMC::GenEvent(*HepUnsmearedMCEvt->GetEvent());
0278   std::unique_ptr<edm::HepMCProduct> HepMCEvt(new edm::HepMCProduct(genevt));
0279   // generate new vertex & apply the shift
0280   //
0281
0282   HepMCEvt->boostToLab(GetInvLorentzBoost(), "vertex");
0283   HepMCEvt->boostToLab(GetInvLorentzBoost(), "momentum");
0284
0285   HepMCEvt->applyVtxGen(useRecVertex ? getRecVertex(evt) : getVertex(evt));
0286
0287   evt.put(std::move(HepMCEvt));
0288   return;
0289 }
0290
0291 #include "FWCore/Framework/interface/MakerMacros.h"
0292 DEFINE_FWK_MODULE(MixBoostEvtVtxGenerator);