GenPlusSimParticleProducer.cc

Go to the documentation of this file.

#include "FWCore/Framework/interface/Frameworkfwd.h"
#include "FWCore/Framework/interface/stream/EDProducer.h"
#include "FWCore/Framework/interface/Event.h"
#include "FWCore/Framework/interface/EventSetup.h"
#include "FWCore/Framework/interface/MakerMacros.h"
#include "FWCore/ParameterSet/interface/ParameterSet.h"

#include "DataFormats/HepMCCandidate/interface/GenParticle.h"
#include "DataFormats/HepMCCandidate/interface/GenParticleFwd.h"
#include "SimDataFormats/Track/interface/SimTrackContainer.h"
#include "SimDataFormats/Vertex/interface/SimVertexContainer.h"

#include "CommonTools/Utils/interface/StringCutObjectSelector.h"
#include "SimGeneral/HepPDTRecord/interface/PdtEntry.h"
#include "SimGeneral/HepPDTRecord/interface/ParticleDataTable.h"

#include <ext/algorithm>
#include <memory>

namespace pat {
  class GenPlusSimParticleProducer : public edm::stream::EDProducer<> {
  public:
    explicit GenPlusSimParticleProducer(const edm::ParameterSet &);

  private:
    void produce(edm::Event &, const edm::EventSetup &) override;

    bool firstEvent_;
    edm::EDGetTokenT<edm::SimTrackContainer> simtracksToken_;
    edm::EDGetTokenT<edm::SimVertexContainer> simverticesToken_;
    int setStatus_;
    std::set<int> pdgIds_;        // these are the ones we really use
    std::vector<PdtEntry> pdts_;  // these are needed before we get the EventSetup

    typedef StringCutObjectSelector<reco::GenParticle> StrFilter;
    std::unique_ptr<StrFilter> filter_;

    edm::EDGetTokenT<reco::GenParticleCollection> gensToken_;
    edm::EDGetTokenT<std::vector<int>> genBarcodesToken_;

    edm::ESGetToken<HepPDT::ParticleDataTable, edm::DefaultRecord> tableToken_;

    void addGenParticle(const SimTrack &stMom,
                        const SimTrack &stDau,
                        unsigned int momidx,
                        const edm::SimTrackContainer &simtks,
                        const edm::SimVertexContainer &simvtxs,
                        reco::GenParticleCollection &mergedGens,
                        const reco::GenParticleRefProd ref,
                        std::vector<int> &genBarcodes,
                        bool &barcodesAreSorted) const;
    struct LessById {
      bool operator()(const SimTrack &tk1, const SimTrack &tk2) const { return tk1.trackId() < tk2.trackId(); }
      bool operator()(const SimTrack &tk1, unsigned int id) const { return tk1.trackId() < id; }
      bool operator()(unsigned int id, const SimTrack &tk2) const { return id < tk2.trackId(); }
    };
  };
}  // namespace pat

using namespace std;
using namespace edm;
using namespace reco;
using pat::GenPlusSimParticleProducer;

GenPlusSimParticleProducer::GenPlusSimParticleProducer(const ParameterSet &cfg)
    : firstEvent_(true),
      simtracksToken_(consumes<SimTrackContainer>(cfg.getParameter<InputTag>("src"))),  // source sim tracks & vertices
      simverticesToken_(
          consumes<SimVertexContainer>(cfg.getParameter<InputTag>("src"))),  // source sim tracks & vertices
      setStatus_(cfg.getParameter<int32_t>("setStatus")),                    // set status of GenParticle to this code
      gensToken_(consumes<GenParticleCollection>(
          cfg.getParameter<InputTag>("genParticles"))),  // get the genParticles to add GEANT particles to
      genBarcodesToken_(consumes<std::vector<int>>(
          cfg.getParameter<InputTag>("genParticles")))  // get the genParticles to add GEANT particles to
{
  // Possibly allow a list of particle types
  if (cfg.exists("particleTypes")) {
    pdts_ = cfg.getParameter<vector<PdtEntry>>("particleTypes");
    if (!pdts_.empty()) {
      tableToken_ = esConsumes();
    }
  }

  // Possibly allow a string cut
  if (cfg.existsAs<string>("filter")) {
    string filter = cfg.getParameter<string>("filter");
    if (!filter.empty()) {
      filter_ = std::make_unique<StrFilter>(filter);
    }
  }
  produces<GenParticleCollection>();
  produces<vector<int>>();
}

void GenPlusSimParticleProducer::addGenParticle(const SimTrack &stMom,
                                                const SimTrack &stDau,
                                                unsigned int momidx,
                                                const SimTrackContainer &simtracksSorted,
                                                const SimVertexContainer &simvertices,
                                                reco::GenParticleCollection &mergedGens,
                                                const GenParticleRefProd ref,
                                                std::vector<int> &genBarcodes,
                                                bool &barcodesAreSorted) const {
  // Make the genParticle for stDau and add it to the new collection and update the parent-child relationship
  // Make up a GenParticleCandidate from the GEANT track info.
  int charge = static_cast<int>(stDau.charge());
  const Particle::LorentzVector &p4 = stDau.momentum();
  Particle::Point vtx;  // = (0,0,0) by default
  if (!stDau.noVertex())
    vtx = simvertices[stDau.vertIndex()].position();
  GenParticle genp(charge, p4, vtx, stDau.type(), setStatus_, true);

  // Maybe apply filter on the particle
  if (filter_.get() != nullptr) {
    if (!(*filter_)(genp))
      return;
  }

  reco::GenParticleRef parentRef(ref, momidx);
  genp.addMother(parentRef);
  mergedGens.push_back(genp);
  // get the index for the daughter just added
  unsigned int dauidx = mergedGens.size() - 1;

  // update add daughter relationship
  reco::GenParticle &cand = mergedGens[momidx];
  cand.addDaughter(GenParticleRef(ref, dauidx));

  //look for simtrack daughters of stDau to see if we need to recur further down the chain

  for (SimTrackContainer::const_iterator isimtrk = simtracksSorted.begin(); isimtrk != simtracksSorted.end();
       ++isimtrk) {
    if (!isimtrk->noVertex()) {
      // Pick the vertex (isimtrk.vertIndex() is really an index)
      const SimVertex &vtx = simvertices[isimtrk->vertIndex()];

      // Check if the vertex has a parent track (otherwise, we're lost)
      if (!vtx.noParent()) {
        // Now note that vtx.parentIndex() is NOT an index, it's a track id, so I have to search for it
        unsigned int idx = vtx.parentIndex();
        SimTrackContainer::const_iterator it =
            std::lower_bound(simtracksSorted.begin(), simtracksSorted.end(), idx, LessById());
        if ((it != simtracksSorted.end()) && (it->trackId() == idx)) {
          if (it->trackId() == stDau.trackId()) {
            //need the genparticle index of stDau which is dauidx
            addGenParticle(
                stDau, *isimtrk, dauidx, simtracksSorted, simvertices, mergedGens, ref, genBarcodes, barcodesAreSorted);
          }
        }
      }
    }
  }
}

void GenPlusSimParticleProducer::produce(Event &event, const EventSetup &iSetup) {
  if (firstEvent_) {
    if (!pdts_.empty()) {
      auto const &pdt = iSetup.getData(tableToken_);
      pdgIds_.clear();
      for (vector<PdtEntry>::iterator itp = pdts_.begin(), edp = pdts_.end(); itp != edp; ++itp) {
        itp->setup(pdt);  // decode string->pdgId and vice-versa
        pdgIds_.insert(std::abs(itp->pdgId()));
      }
      pdts_.clear();
    }
    firstEvent_ = false;
  }

  // Simulated tracks (i.e. GEANT particles).
  Handle<SimTrackContainer> simtracks;
  event.getByToken(simtracksToken_, simtracks);

  // Need to check that SimTrackContainer is sorted; otherwise, copy and sort :-(
  std::unique_ptr<SimTrackContainer> simtracksTmp;
  const SimTrackContainer *simtracksSorted = &*simtracks;
  if (!__gnu_cxx::is_sorted(simtracks->begin(), simtracks->end(), LessById())) {
    simtracksTmp = std::make_unique<SimTrackContainer>(*simtracks);
    std::sort(simtracksTmp->begin(), simtracksTmp->end(), LessById());
    simtracksSorted = &*simtracksTmp;
  }

  // Get the associated vertices
  Handle<SimVertexContainer> simvertices;
  event.getByToken(simverticesToken_, simvertices);

  // Get the GenParticles and barcodes, if needed to set mother and daughter links
  Handle<GenParticleCollection> gens;
  Handle<std::vector<int>> genBarcodes;
  bool barcodesAreSorted = true;
  event.getByToken(gensToken_, gens);
  event.getByToken(genBarcodesToken_, genBarcodes);
  if (gens->size() != genBarcodes->size())
    throw cms::Exception("Corrupt data") << "Barcodes not of the same size as GenParticles!\n";

  // make the output collection
  auto candsPtr = std::make_unique<GenParticleCollection>();
  GenParticleCollection &cands = *candsPtr;

  const GenParticleRefProd ref = event.getRefBeforePut<GenParticleCollection>();

  // add the original genParticles to the merged output list
  for (size_t i = 0; i < gens->size(); ++i) {
    reco::GenParticle cand((*gens)[i]);
    cands.push_back(cand);
  }

  // make new barcodes vector and fill it with the original list
  auto newGenBarcodes = std::make_unique<vector<int>>();
  for (unsigned int i = 0; i < genBarcodes->size(); ++i) {
    newGenBarcodes->push_back((*genBarcodes)[i]);
  }
  barcodesAreSorted = __gnu_cxx::is_sorted(newGenBarcodes->begin(), newGenBarcodes->end());

  for (size_t i = 0; i < cands.size(); ++i) {
    reco::GenParticle &cand = cands[i];
    size_t nDaus = cand.numberOfDaughters();
    GenParticleRefVector daus = cand.daughterRefVector();
    cand.resetDaughters(ref.id());
    for (size_t d = 0; d < nDaus; ++d) {
      cand.addDaughter(GenParticleRef(ref, daus[d].key()));
    }

    size_t nMoms = cand.numberOfMothers();
    GenParticleRefVector moms = cand.motherRefVector();
    cand.resetMothers(ref.id());
    for (size_t m = 0; m < nMoms; ++m) {
      cand.addMother(GenParticleRef(ref, moms[m].key()));
    }
  }

  for (SimTrackContainer::const_iterator isimtrk = simtracks->begin(); isimtrk != simtracks->end(); ++isimtrk) {
    // Skip PYTHIA tracks.
    if (isimtrk->genpartIndex() != -1)
      continue;

    // Maybe apply the PdgId filter
    if (!pdgIds_.empty()) {  // if we have a filter on pdg ids
      if (pdgIds_.find(std::abs(isimtrk->type())) == pdgIds_.end())
        continue;
    }

    // find simtrack that has a genParticle match to its parent
    // Look at the production vertex. If there is no vertex, I can do nothing...
    if (!isimtrk->noVertex()) {
      // Pick the vertex (isimtrk.vertIndex() is really an index)
      const SimVertex &vtx = (*simvertices)[isimtrk->vertIndex()];

      // Check if the vertex has a parent track (otherwise, we're lost)
      if (!vtx.noParent()) {
        // Now note that vtx.parentIndex() is NOT an index, it's a track id, so I have to search for it
        unsigned int idx = vtx.parentIndex();
        SimTrackContainer::const_iterator it =
            std::lower_bound(simtracksSorted->begin(), simtracksSorted->end(), idx, LessById());
        if ((it != simtracksSorted->end()) && (it->trackId() == idx)) {  //it is the parent sim track
          if (it->genpartIndex() != -1) {
            std::vector<int>::const_iterator itIndex;
            if (barcodesAreSorted) {
              itIndex = std::lower_bound(genBarcodes->begin(), genBarcodes->end(), it->genpartIndex());
            } else {
              itIndex = std::find(genBarcodes->begin(), genBarcodes->end(), it->genpartIndex());
            }

            // Check that I found something
            // I need to check '*itIndex == it->genpartIndex()' because lower_bound just finds the right spot for an item in a sorted list, not the item
            if ((itIndex != genBarcodes->end()) && (*itIndex == it->genpartIndex())) {
              // Ok, I'll make the genParticle for st and add it to the new collection updating the map and parent-child relationship
              // pass the mother and daughter sim tracks and the mother genParticle to method to create the daughter genParticle and recur
              unsigned int momidx = itIndex - genBarcodes->begin();
              addGenParticle(*it,
                             *isimtrk,
                             momidx,
                             *simtracksSorted,
                             *simvertices,
                             *candsPtr,
                             ref,
                             *newGenBarcodes,
                             barcodesAreSorted);
            }
          }
        }
      }
    }
  }

  event.put(std::move(candsPtr));
  event.put(std::move(newGenBarcodes));
}

DEFINE_FWK_MODULE(GenPlusSimParticleProducer);