MixingModule/plugins/MixEvtVtxGenerator.cc

0001 #ifndef HI_MixEvtVtxGenerator_H
0002 #define HI_MixEvtVtxGenerator_H
0003 /*
0004 */
0005 #include "FWCore/PluginManager/interface/ModuleDef.h"
0006 #include "FWCore/Framework/interface/MakerMacros.h"
0007
0008 #include "FWCore/Framework/interface/stream/EDProducer.h"
0009 #include "FWCore/Utilities/interface/InputTag.h"
0010 #include "FWCore/Framework/interface/Event.h"
0011 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0012 #include "FWCore/ServiceRegistry/interface/Service.h"
0013 #include "FWCore/Utilities/interface/Exception.h"
0014 #include "FWCore/Utilities/interface/EDMException.h"
0015
0016 #include "SimDataFormats/CrossingFrame/interface/MixCollection.h"
0017 #include "SimDataFormats/GeneratorProducts/interface/HepMCProduct.h"
0018 #include "DataFormats/Provenance/interface/Provenance.h"
0019
0020 #include "DataFormats/VertexReco/interface/Vertex.h"
0021 #include "DataFormats/VertexReco/interface/VertexFwd.h"
0022
0023 #include "TMatrixD.h"
0024 #include <iostream>
0025
0026 using namespace edm;
0027 using namespace std;
0028
0029 namespace HepMC {
0030   class FourVector;
0031 }
0032
0033 class MixEvtVtxGenerator : public edm::stream::EDProducer<> {
0034 public:
0035   // ctor & dtor
0036   explicit MixEvtVtxGenerator(const edm::ParameterSet&);
0037   ~MixEvtVtxGenerator() override;
0038
0039   void produce(edm::Event&, const edm::EventSetup&) override;
0040
0041   virtual HepMC::FourVector* getVertex(edm::Event&);
0042   virtual HepMC::FourVector* getRecVertex(edm::Event&);
0043
0044 protected:
0045   HepMC::FourVector* fVertex;
0046   TMatrixD* boost_;
0047
0048 private:
0049   edm::EDGetTokenT<reco::VertexCollection> hiLabel;
0050   edm::EDGetTokenT<HepMCProduct> signalLabel;
0051   edm::EDGetTokenT<CrossingFrame<HepMCProduct> > cfLabel;
0052
0053   bool useRecVertex;
0054   std::vector<double> vtxOffset;
0055   bool useCF_;
0056 };
0057
0058 MixEvtVtxGenerator::MixEvtVtxGenerator(const ParameterSet& pset)
0059     : fVertex(nullptr),
0060       boost_(nullptr),
0061       useRecVertex(pset.exists("useRecVertex") ? pset.getParameter<bool>("useRecVertex") : false)
0062
0063 {
0064   produces<edm::HepMCProduct>();
0065   vtxOffset.resize(3);
0066   if (pset.exists("vtxOffset"))
0067     vtxOffset = pset.getParameter<std::vector<double> >("vtxOffset");
0068
0069   if (useRecVertex)
0070     useCF_ = false;
0071   else {
0072     useCF_ = pset.getUntrackedParameter<bool>("useCF", false);
0073     cfLabel = consumes<CrossingFrame<HepMCProduct> >(pset.getParameter<edm::InputTag>("mixLabel"));
0074   }
0075   signalLabel = consumes<HepMCProduct>(pset.getParameter<edm::InputTag>("signalLabel"));
0076 }
0077
0078 MixEvtVtxGenerator::~MixEvtVtxGenerator() {
0079   delete fVertex;
0080   if (boost_ != nullptr)
0081     delete boost_;
0082   // no need since now it's done in HepMCProduct
0083   // delete fEvt ;
0084 }
0085
0086 HepMC::FourVector* MixEvtVtxGenerator::getVertex(Event& evt) {
0087   HepMC::GenVertex* genvtx = nullptr;
0088   const HepMC::GenEvent* inev = nullptr;
0089
0090   if (useCF_) {
0091     Handle<CrossingFrame<HepMCProduct> > cf;
0092     evt.getByToken(cfLabel, cf);
0093     MixCollection<HepMCProduct> mix(cf.product());
0094     if (mix.size() < 2) {
0095       throw cms::Exception("MatchVtx") << "Mixing has " << mix.size() << " sub-events, should have been at least 2"
0096                                        << endl;
0097     }
0098     const HepMCProduct& bkg = mix.getObject(1);
0099     if (!(bkg.isVtxGenApplied())) {
0100       throw cms::Exception("MatchVtx") << "Input background does not have smeared vertex!" << endl;
0101     } else {
0102       inev = bkg.GetEvent();
0103     }
0104   } else {
0105     Handle<HepMCProduct> input;
0106     evt.getByToken(signalLabel, input);
0107     inev = input->GetEvent();
0108   }
0109
0110   genvtx = inev->signal_process_vertex();
0111   if (!genvtx) {
0112     HepMC::GenEvent::particle_const_iterator pt = inev->particles_begin();
0113     HepMC::GenEvent::particle_const_iterator ptend = inev->particles_end();
0114     while (!genvtx || (genvtx->particles_in_size() == 1 && pt != ptend)) {
0115       if (pt == ptend)
0116         cout << "End reached, No Gen Vertex!" << endl;
0117       genvtx = (*pt)->production_vertex();
0118       ++pt;
0119     }
0120   }
0121   double aX, aY, aZ, aT;
0122
0123   aX = genvtx->position().x();
0124   aY = genvtx->position().y();
0125   aZ = genvtx->position().z();
0126   aT = genvtx->position().t();
0127
0128   if (!fVertex) {
0129     fVertex = new HepMC::FourVector();
0130   }
0131   LogInfo("MatchVtx") << " setting vertex "
0132                       << " aX " << aX << " aY " << aY << " aZ " << aZ << " aT " << aT << endl;
0133   fVertex->set(aX, aY, aZ, aT);
0134
0135   return fVertex;
0136 }
0137
0138 HepMC::FourVector* MixEvtVtxGenerator::getRecVertex(Event& evt) {
0139   Handle<reco::VertexCollection> input;
0140   evt.getByToken(hiLabel, input);
0141
0142   double aX, aY, aZ;
0143
0144   aX = input->begin()->position().x() + vtxOffset[0];
0145   aY = input->begin()->position().y() + vtxOffset[1];
0146   aZ = input->begin()->position().z() + vtxOffset[2];
0147
0148   if (!fVertex)
0149     fVertex = new HepMC::FourVector();
0150   fVertex->set(10.0 * aX, 10.0 * aY, 10.0 * aZ, 0.0);  // HepMC positions in mm (RECO in cm)
0151
0152   return fVertex;
0153 }
0154
0155 void MixEvtVtxGenerator::produce(Event& evt, const EventSetup&) {
0156   Handle<HepMCProduct> HepUnsmearedMCEvt;
0157
0158   evt.getByToken(signalLabel, HepUnsmearedMCEvt);
0159
0160   // generate new vertex & apply the shift
0161   //
0162   if (HepUnsmearedMCEvt->isVtxGenApplied())
0163     throw cms::Exception("MatchVtx")
0164         << "Signal HepMCProduct is not compatible for embedding - it's vertex is already smeared." << std::endl;
0165   // Copy the HepMC::GenEvent
0166   HepMC::GenEvent* genevt = new HepMC::GenEvent(*HepUnsmearedMCEvt->GetEvent());
0167   std::unique_ptr<edm::HepMCProduct> HepMCEvt(new edm::HepMCProduct(genevt));
0168   // generate new vertex & apply the shift
0169   //
0170   HepMCEvt->applyVtxGen(useRecVertex ? getRecVertex(evt) : getVertex(evt));
0171
0172   //   HepMCEvt->boostToLab( GetInvLorentzBoost(), "vertex" );
0173   //   HepMCEvt->boostToLab( GetInvLorentzBoost(), "momentum" );
0174
0175   evt.put(std::move(HepMCEvt));
0176
0177   return;
0178 }
0179
0180 DEFINE_FWK_MODULE(MixEvtVtxGenerator);
0181
0182 #endif