Saving generator protons from pile-up for standard mixing and premixing

Configuration/StandardSequences/python/Digi_cff.py

@@ -28,11 +28,12 @@
 # add updating the GEN information by default
 from Configuration.StandardSequences.Generator_cff import *
 from GeneratorInterface.Core.generatorSmeared_cfi import *
+from PhysicsTools.HepMCCandAlgos.genPUProtons_cfi import *
 
 doAllDigi = cms.Sequence(generatorSmeared*calDigi+muonDigi)
-pdigi = cms.Sequence(generatorSmeared*fixGenInfo*cms.SequencePlaceholder("randomEngineStateProducer")*cms.SequencePlaceholder("mix")*doAllDigi*addPileupInfo)
+pdigi = cms.Sequence(generatorSmeared*fixGenInfo*cms.SequencePlaceholder("randomEngineStateProducer")*cms.SequencePlaceholder("mix")*doAllDigi*addPileupInfo*genPUProtons)
 pdigi_valid = cms.Sequence(pdigi)
-pdigi_nogen=cms.Sequence(generatorSmeared*cms.SequencePlaceholder("randomEngineStateProducer")*cms.SequencePlaceholder("mix")*doAllDigi*addPileupInfo)
+pdigi_nogen=cms.Sequence(generatorSmeared*cms.SequencePlaceholder("randomEngineStateProducer")*cms.SequencePlaceholder("mix")*doAllDigi*addPileupInfo*genPUProtons)
 pdigi_valid_nogen=cms.Sequence(pdigi_nogen)
 
 from GeneratorInterface.HiGenCommon.HeavyIon_cff import *

PhysicsTools/HepMCCandAlgos/plugins/GenPUProtonProducer.cc

@@ -0,0 +1,255 @@
+/* \class GenPUProtonProducer
+ *
+ * Modification of GenParticleProducer.
+ * Saves final state protons from HepMC events in Crossing Frame, in the generator-particle format.
+ *
+ * Note: Use the option USER_CXXFLAGS=-DEDM_ML_DEBUG with SCRAM in order to enable debug messages.
+ *
+ *    March  9, 2017   : Initial version.
+ *    March 14, 2017   : Updated debug messages.
+ *     July 27, 2017   : Removed extra loop initially inherited from GenParticleProducer. 
+ *   August 17, 2017   : Replaced std::auto_ptr with std::unique_ptr. 
+ * September 6, 2017   : Updated module to edm::global::EDProducer with ConvertParticle as RunCache following GenParticleProducer. 
+ *
+ */
+
+#include "SimDataFormats/GeneratorProducts/interface/HepMCProduct.h"
+#include "SimGeneral/HepPDTRecord/interface/ParticleDataTable.h"
+#include "DataFormats/HepMCCandidate/interface/GenParticle.h"
+#include "FWCore/Utilities/interface/EDMException.h"
+
+#include <vector>
+#include <string>
+#include <unordered_map>
+
+namespace {
+
+  class ConvertParticle {
+    public:
+     static constexpr int PDGCacheMax = 32768;
+     static constexpr double mmToCm = 0.1;
+
+     ConvertParticle() :
+        abortOnUnknownPDGCode_( true ),
+        initialized_( false ),
+        chargeP_( PDGCacheMax, 0 ), chargeM_( PDGCacheMax, 0 ) {};
+
+     ConvertParticle(bool abortOnUnknownPDGCode) :
+        abortOnUnknownPDGCode_( abortOnUnknownPDGCode ),
+        initialized_( false ),
+        chargeP_( PDGCacheMax, 0 ), chargeM_( PDGCacheMax, 0 ) {};
+
+     ~ConvertParticle() {};
+
+     bool initialized() const { return initialized_; }
+
+     void init(HepPDT::ParticleDataTable const& pdt) {
+	if( !initialized_ ) {
+	   for( HepPDT::ParticleDataTable::const_iterator p = pdt.begin(); p != pdt.end(); ++p ) {
+	      HepPDT::ParticleID const& id = p->first;
+	      int pdgId = id.pid(), apdgId = std::abs( pdgId );
+	      int q3 = id.threeCharge();
+	      if ( apdgId < PDGCacheMax && pdgId > 0 ) {
+		 chargeP_[ apdgId ] = q3;
+		 chargeM_[ apdgId ] = -q3;
+	      } else if ( apdgId < PDGCacheMax ) {
+		 chargeP_[ apdgId ] = -q3;
+		 chargeM_[ apdgId ] = q3;
+	      } else {
+	         chargeMap_.emplace( pdgId, q3);
+	         chargeMap_.emplace( -pdgId, -q3);
+	      }
+	   }
+	   initialized_ = true;
+	}
+     }
+
+     bool operator() (reco::GenParticle& cand, HepMC::GenParticle const* part) const {
+	reco::Candidate::LorentzVector p4( part->momentum() );
+	int pdgId = part->pdg_id();
+	cand.setThreeCharge( chargeTimesThree( pdgId ) );
+	cand.setPdgId( pdgId );
+	cand.setStatus( part->status() );
+	cand.setP4( p4 );
+	cand.setCollisionId(0);
+	HepMC::GenVertex const* v = part->production_vertex();
+	if ( v != nullptr ) {
+	   HepMC::ThreeVector vtx = v->point3d();
+	   reco::Candidate::Point vertex( vtx.x() * mmToCm, vtx.y() * mmToCm, vtx.z() * mmToCm );
+	   cand.setVertex( vertex );
+	} else {
+	   cand.setVertex( reco::Candidate::Point( 0, 0, 0 ) );
+	}
+	return true;
+     }
+
+    private:
+     bool abortOnUnknownPDGCode_;
+     bool initialized_; 
+     std::vector<int> chargeP_, chargeM_;
+     std::unordered_map<int, int> chargeMap_;
+
+     int chargeTimesThree( int id ) const {
+	if( std::abs( id ) < PDGCacheMax )
+	   return id > 0 ? chargeP_[ id ] : chargeM_[ - id ];
+
+	auto f = chargeMap_.find( id );
+	if ( f == chargeMap_.end() )  {
+	   if ( abortOnUnknownPDGCode_ )
+	      throw edm::Exception( edm::errors::LogicError ) << "invalid PDG id: " << id << std::endl;
+	   else
+	      return HepPDT::ParticleID(id).threeCharge();
+	}
+	return f->second;
+     }
+
+  }; 
+
+  class SelectProton { 
+     public:
+        bool operator() ( HepMC::GenParticle const* part, double minPz ) const {
+           bool selection = ( ( !part->end_vertex() && part->status() == 1 ) && 
+                              ( part->pdg_id() == 2212 ) && ( TMath::Abs( part->momentum().pz() ) >= minPz ) );
+           return selection;
+        }
+  };
+
+} // Anonymous namespace
+
+#include "FWCore/Framework/interface/global/EDProducer.h"
+#include "FWCore/Utilities/interface/InputTag.h"
+#include "DataFormats/Candidate/interface/CandidateFwd.h"
+#include "DataFormats/HepMCCandidate/interface/GenParticleFwd.h"
+
+#include "SimDataFormats/CrossingFrame/interface/CrossingFrame.h"
+#include "SimDataFormats/GeneratorProducts/interface/HepMCProduct.h"
+#include "PhysicsTools/HepMCCandAlgos/interface/MCTruthHelper.h"
+#include "SimGeneral/HepPDTRecord/interface/ParticleDataTable.h"
+#include "DataFormats/HepMCCandidate/interface/GenParticle.h"
+
+namespace edm { class ParameterSet; }
+
+class GenPUProtonProducer : public edm::global::EDProducer<edm::RunCache<ConvertParticle> > {
+ public:
+  GenPUProtonProducer( const edm::ParameterSet & );
+  ~GenPUProtonProducer() override;
+
+  void produce( edm::StreamID, edm::Event& e, const edm::EventSetup&) const override;
+  std::shared_ptr<ConvertParticle> globalBeginRun(const edm::Run&, const edm::EventSetup&) const override;
+  void globalEndRun(edm::Run const&, edm::EventSetup const&) const override {};
+
+ private:
+  edm::EDGetTokenT<CrossingFrame<edm::HepMCProduct> > mixToken_;
+
+  bool abortOnUnknownPDGCode_;
+  std::vector<int> bunchList_;
+  double minPz_; 
+  SelectProton select_;
+
+};
+
+#include "FWCore/ParameterSet/interface/ParameterSet.h"
+#include "DataFormats/Common/interface/Handle.h"
+#include "FWCore/Framework/interface/ESHandle.h"
+#include "FWCore/Framework/interface/Event.h"
+#include "FWCore/Framework/interface/Run.h"
+#include "FWCore/Framework/interface/EventSetup.h"
+#include "FWCore/Utilities/interface/EDMException.h"
+
+#include "DataFormats/HepMCCandidate/interface/GenParticle.h"
+#include "SimDataFormats/CrossingFrame/interface/MixCollection.h"
+#include "SimGeneral/HepPDTRecord/interface/ParticleDataTable.h"
+
+#include <algorithm>
+
+using namespace edm;
+using namespace reco;
+using namespace std;
+using namespace HepMC;
+
+GenPUProtonProducer::GenPUProtonProducer( const ParameterSet & cfg ) :
+  abortOnUnknownPDGCode_( cfg.getUntrackedParameter<bool>( "abortOnUnknownPDGCode", true ) ),
+  bunchList_( cfg.getParameter<vector<int> >( "bunchCrossingList" ) ),
+  minPz_( cfg.getParameter<double>( "minPz" ) ) 
+{
+  produces<GenParticleCollection>();
+  mixToken_ = consumes<CrossingFrame<HepMCProduct> >( InputTag(cfg.getParameter<std::string>( "mix" ),"generatorSmeared") );
+}
+
+GenPUProtonProducer::~GenPUProtonProducer() { }
+
+std::shared_ptr<ConvertParticle> GenPUProtonProducer::globalBeginRun(const Run&, const EventSetup& es) const {
+  ESHandle<HepPDT::ParticleDataTable> pdt;
+  es.getData( pdt );
+  auto convert_ptr = std::make_shared<ConvertParticle>( abortOnUnknownPDGCode_ );
+  if( !convert_ptr->initialized() ) convert_ptr->init( *pdt );
+
+  return convert_ptr;
+}
+
+void GenPUProtonProducer::produce( StreamID, Event& evt, const EventSetup& es ) const {
+
+   size_t totalSize = 0;
+   size_t npiles = 1;
+
+   Handle<CrossingFrame<HepMCProduct> > cf;
+   evt.getByToken(mixToken_,cf);
+   std::unique_ptr<MixCollection<HepMCProduct> > cfhepmcprod( new MixCollection<HepMCProduct>( cf.product() ) );
+   npiles = cfhepmcprod->size();
+
+   LogDebug("GenPUProtonProducer") << " Number of pile-up events : " << npiles << endl;
+
+   for(size_t icf = 0; icf < npiles; ++icf){
+      LogDebug("GenPUProtonProducer") << "CF " << icf << " size : " << cfhepmcprod->getObject(icf).GetEvent()->particles_size() << endl;
+      totalSize += cfhepmcprod->getObject(icf).GetEvent()->particles_size();
+   }
+   LogDebug("GenPUProtonProducer") << "Total size : " << totalSize << endl;
+
+   // Initialise containers
+   auto candsPtr = std::make_unique<GenParticleCollection>();
+   GenParticleCollection& cands = *candsPtr;
+
+   // Loop over pile-up events
+   ConvertParticle const& convertParticle_ = *runCache( evt.getRun().index() );
+
+   MixCollection<HepMCProduct>::MixItr mixHepMC_itr;
+   unsigned int total_number_of_protons = 0;
+   size_t idx_mix = 0;
+   // Fill collection
+   for( mixHepMC_itr = cfhepmcprod->begin() ; mixHepMC_itr != cfhepmcprod->end() ; ++mixHepMC_itr, ++idx_mix ){
+      int bunch = mixHepMC_itr.bunch();
+
+      if( find( bunchList_.begin(), bunchList_.end(), bunch ) != bunchList_.end() ){
+
+	 auto event = (*mixHepMC_itr).GetEvent();
+
+	 size_t num_particles = event->particles_size();
+
+	 // Fill output collection
+	 unsigned int number_of_protons = 0;
+	 for( auto p = event->particles_begin() ; p != event->particles_end() ; ++p ) {
+	    HepMC::GenParticle const* part = *p;  
+	    if( select_(part, minPz_) ) {
+	       reco::GenParticle cand;
+	       convertParticle_( cand, part );
+               ++number_of_protons;
+               cands.push_back( cand );
+	    } 
+	 }
+	 LogDebug("GenPUProtonProducer") << "Idx : " << idx_mix << " Bunch : " << bunch 
+                                         << " Number of particles : " << num_particles 
+                                         << " Number of protons : " << number_of_protons << endl;
+
+         total_number_of_protons += number_of_protons;
+      }
+   }
+   LogDebug("GenPUProtonProducer") << "Total number of protons : " << total_number_of_protons << endl;
+   LogDebug("GenPUProtonProducer") << "Output collection size : " << cands.size() << endl;
+
+   evt.put( std::move(candsPtr) );
+}
+
+#include "FWCore/Framework/interface/MakerMacros.h"
+
+DEFINE_FWK_MODULE( GenPUProtonProducer );
+

PhysicsTools/HepMCCandAlgos/python/genPUProtons_cfi.py

@@ -0,0 +1,7 @@
+import FWCore.ParameterSet.Config as cms
+
+genPUProtons = cms.EDProducer("GenPUProtonProducer",
+	mix = cms.string("mix"),
+	bunchCrossingList = cms.vint32(0),
+	minPz = cms.double( 2400. )
+	)

PhysicsTools/PatAlgos/python/slimming/MicroEventContent_cff.py

@@ -90,6 +90,7 @@
         'keep GenLumiInfoHeader_generator_*_*',
         'keep GenLumiInfoProduct_*_*_*',
         'keep GenEventInfoProduct_generator_*_*',
+        'keep recoGenParticles_genPUProtons_*_*', 
         # RUN
         'keep LHERunInfoProduct_*_*_*',
         'keep GenRunInfoProduct_*_*_*',

SimGeneral/Configuration/python/SimGeneral_EventContent_cff.py

@@ -12,17 +12,20 @@
     outputCommands = cms.untracked.vstring('keep CrossingFramePlaybackInfoNew_*_*_*',
                                            'keep PileupSummaryInfos_*_*_*',
                                            'keep int6stdbitsetstdpairs_*_AffectedAPVList_*',
-                                           'keep int_*_bunchSpacing_*')
+                                           'keep int_*_bunchSpacing_*',
+                                           'keep *_genPUProtons_*_*') 
 )
 #RECO content
 SimGeneralRECO = cms.PSet(
     outputCommands = cms.untracked.vstring('keep PileupSummaryInfos_*_*_*',
-                                           'keep int_*_bunchSpacing_*')
+                                           'keep int_*_bunchSpacing_*',
+                                           'keep *_genPUProtons_*_*') 
 )
 #AOD content
 SimGeneralAOD = cms.PSet(
     outputCommands = cms.untracked.vstring('keep PileupSummaryInfos_*_*_*',
-                                           'keep int_*_bunchSpacing_*')
+                                           'keep int_*_bunchSpacing_*',
+                                           'keep *_genPUProtons_*_*') 
 )
 
 # mods for HGCAL

SimGeneral/DataMixingModule/plugins/BuildFile.xml

@@ -17,6 +17,7 @@
 <use   name="DataFormats/SiPixelDigi"/>
 <use   name="DataFormats/SiStripDigi"/>
 <use   name="DataFormats/TrackReco"/>
+<use   name="DataFormats/HepMCCandidate"/>
 <use   name="FWCore/Framework"/>
 <use   name="FWCore/MessageLogger"/>
 <use   name="FWCore/ParameterSet"/>

SimGeneral/DataMixingModule/plugins/DataMixingModule.cc

@@ -24,6 +24,7 @@
 //
 //
 #include "SimDataFormats/CrossingFrame/interface/CrossingFramePlaybackInfoNew.h"
+#include "DataFormats/HepMCCandidate/interface/GenParticle.h"
 #include "DataMixingModule.h"
 #include "SimGeneral/MixingModule/interface/PileUpEventPrincipal.h"
 
@@ -71,9 +72,9 @@ namespace edm
 
     // Check to see if we are working in Full or Fast Simulation
 
-    MergeTrackerDigis_ = (ps.getParameter<std::string>("TrackerMergeType")).compare("Digis") == 0;
-    MergeEMDigis_ = (ps.getParameter<std::string>("EcalMergeType")).compare("Digis") == 0;
-    MergeHcalDigis_ = (ps.getParameter<std::string>("HcalMergeType")).compare("Digis") == 0;
+    MergeTrackerDigis_ = (ps.getParameter<std::string>("TrackerMergeType")) == "Digis";
+    MergeEMDigis_ = (ps.getParameter<std::string>("EcalMergeType")) == "Digis";
+    MergeHcalDigis_ = (ps.getParameter<std::string>("HcalMergeType")) == "Digis";
     if(MergeHcalDigis_) MergeHcalDigisProd_ = (ps.getParameter<std::string>("HcalDigiMerge")=="FullProd");
 
     addMCDigiNoise_ = false;
@@ -235,6 +236,12 @@ namespace edm
       produces< int >("bunchSpacing");
       produces<CrossingFramePlaybackInfoNew>();
 
+      std::vector<edm::InputTag> GenPUProtonsInputTags;
+      GenPUProtonsInputTags = ps.getParameter<std::vector<edm::InputTag> >("GenPUProtonsInputTags");
+      for(std::vector<edm::InputTag>::const_iterator it_InputTag = GenPUProtonsInputTags.begin(); 
+                                                     it_InputTag != GenPUProtonsInputTags.end(); ++it_InputTag) 
+         produces< std::vector<reco::GenParticle> >( it_InputTag->label() );
+
       PUWorker_ = new DataMixingPileupCopy(ps, consumesCollector());
     }