GenPUProtonProducer.cc

Go to the documentation of this file.

/* \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 "FWCore/MessageLogger/interface/MessageLogger.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 "FWCore/Utilities/interface/ESGetToken.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 "SimGeneral/HepPDTRecord/interface/ParticleDataTable.h"
#include "DataFormats/HepMCCandidate/interface/GenParticle.h"
#include "SimGeneral/HepPDTRecord/interface/ParticleDataTable.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_;
  edm::ESGetToken<HepPDT::ParticleDataTable, edm::DefaultRecord> pdtToken_;
  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 <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"));
  pdtToken_ = esConsumes<HepPDT::ParticleDataTable, edm::DefaultRecord, edm::Transition::BeginRun>();
}

GenPUProtonProducer::~GenPUProtonProducer() {}

std::shared_ptr<ConvertParticle> GenPUProtonProducer::globalBeginRun(const Run&, const EventSetup& es) const {
  ESHandle<HepPDT::ParticleDataTable> pdt = es.getHandle(pdtToken_);
  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);