GenPlusSimParticleProducer.cc

/**
  \class    pat::GenPlusSimParticleProducer GenPlusSimParticleProducer.h "PhysicsTools/PatAlgos/interface/GenPlusSimParticleProducer.h"
  \brief    Produces reco::GenParticle from SimTracks

The GenPlusSimParticleProducer produces GenParticles from SimTracks,
so they can be used for MC matching. A copy of the original genParticle list
is made and the genParticles created from the GEANT/FAMOS particles are added
to the list including all ancestors and correct mother/daughter references

Sample useage in cfg.py file:
process.genParticlePlusGEANT = cms.EDProducer("GenPlusSimParticleProducer",
        src           = cms.InputTag("g4SimHits"), # use "fastSimProducer" for FastSim
        setStatus     = cms.int32(8),             # set status = 8 for GEANT GPs
        particleTypes = cms.vstring("pi+"),       # also picks pi- (optional)
        filter        = cms.vstring("pt > 0.0"),  # just for testing
        genParticles   = cms.InputTag("genParticles") # original genParticle list
)

  \author   Jordan Tucker (original module), Keith Ulmer (generalization)
*/

#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_;

    /// Collection of GenParticles I need to make refs to. It must also have its associated vector<int> of barcodes, aligned with them.
    edm::EDGetTokenT<reco::GenParticleCollection> gensToken_;
    edm::EDGetTokenT<std::vector<int>> genBarcodesToken_;

    edm::ESGetToken<HepPDT::ParticleDataTable, edm::DefaultRecord> tableToken_;
    /// Try to link the GEANT particle to the generator particle it came from
    /** Arguments:
   * -- Specific --
   *    gp: GenParticle made from the GEANT particle
   *    st: The GEANT simTrack for which we create a genParticle
   *
   * -- GEANT related --
   *    simtks:  A list of GEANT tracks, sorted by track id
   *    simvtxs: The list of GEANT vertices, in their natural order (skimtks have pointers into this vector!)
   *
   * -- GenParticle related --
   *    gens             : Handle to the GenParticles, to make the ref to
   *    genBarcodes      : Barcodes for each GenParticle, in a vector aligned to the GenParticleCollection.
   *    barcodesAreSorted: true if the barcodes are sorted (which means I can use binary_search on them) */
    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) {
    // Special case for when simTracks need to be added to pythia track filtered by particleType, e.g. a particle with ctau=0.1mm
    if (isimtrk->genpartIndex() != -1 && !pdgIds_.empty() &&
        (pdgIds_.find(std::abs(isimtrk->type())) != pdgIds_.end())) {
      std::vector<int>::const_iterator itIndex;
      if (barcodesAreSorted) {
        itIndex = std::lower_bound(genBarcodes->begin(), genBarcodes->end(), isimtrk->genpartIndex());
      } else {
        itIndex = std::find(genBarcodes->begin(), genBarcodes->end(), isimtrk->genpartIndex());
      }

      if ((itIndex != genBarcodes->end()) && (*itIndex == isimtrk->genpartIndex())) {
        const unsigned int momidx = itIndex - genBarcodes->begin();
        for (SimTrackContainer::const_iterator idau = simtracksSorted->begin(); idau != simtracksSorted->end();
             ++idau) {
          if (idau->noVertex())
            continue;
          const SimVertex &dvtx = (*simvertices)[idau->vertIndex()];
          if (dvtx.noParent())
            continue;
          if (dvtx.parentIndex() == static_cast<int>(isimtrk->trackId())) {
            addGenParticle(*isimtrk,
                           *idau,
                           momidx,
                           *simtracksSorted,
                           *simvertices,
                           *candsPtr,
                           ref,
                           *newGenBarcodes,
                           barcodesAreSorted);
          }
        }
      }
    }

    else {
      // 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);