PATGenCandsFromSimTracksProducer.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/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 <ext/algorithm>

namespace pat {
  class PATGenCandsFromSimTracksProducer : public edm::stream::EDProducer<> {
  public:
    explicit PATGenCandsFromSimTracksProducer(const edm::ParameterSet&);
    ~PATGenCandsFromSimTracksProducer() override {}
    
  private:
    void produce(edm::Event&, const edm::EventSetup&) override;
    
    bool firstEvent_;
    edm::EDGetTokenT<edm::SimTrackContainer> simTracksToken_;
    edm::EDGetTokenT<edm::SimVertexContainer> simVertexToken_;
    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
    std::set<int>         motherPdgIds_; // these are the ones we really use
    std::vector<PdtEntry> motherPdts_;   // these are needed before we get the EventSetup
    
    typedef StringCutObjectSelector<reco::GenParticle> StrFilter;
    const StrFilter filter_;
    
    bool makeMotherLink_;
    bool writeAncestors_;
    
    edm::EDGetTokenT<reco::GenParticleCollection> genParticlesToken_;
    edm::EDGetTokenT<std::vector<int> > genBarcodesToken_;
    
    struct GlobalContext {
      GlobalContext(const edm::SimTrackContainer &simtks1,
                    const edm::SimVertexContainer &simvtxs1,
                    const edm::Handle<reco::GenParticleCollection> &gens1,
                    const edm::Handle<std::vector<int> >           &genBarcodes1,
                    bool                                            barcodesAreSorted1,
                    reco::GenParticleCollection                     & output1,
                    const edm::RefProd<reco::GenParticleCollection> & refprod1) :
        simtks(simtks1), simvtxs(simvtxs1),
        gens(gens1), genBarcodes(genBarcodes1), barcodesAreSorted(barcodesAreSorted1),
        output(output1), refprod(refprod1), simTksProcessed() {}
      // GEANT info
      const edm::SimTrackContainer &simtks;
      const edm::SimVertexContainer &simvtxs;
      // PYTHIA info
      const edm::Handle<reco::GenParticleCollection> &gens;
      const edm::Handle<std::vector<int> >           &genBarcodes;
      bool                                            barcodesAreSorted;
      // MY OUTPUT info
      reco::GenParticleCollection                     & output;
      const edm::RefProd<reco::GenParticleCollection> & refprod;
      // BOOK-KEEPING
      std::map<unsigned int,int> simTksProcessed; // key = sim track id;
      // val = 0: not processed;
      //       i>0:  (index+1) in my output
      //       i<0: -(index+1) in pythia [NOT USED]
    };
    
    const SimTrack * findGeantMother(const SimTrack &tk, const GlobalContext &g) const ;
    edm::Ref<reco::GenParticleCollection> findRef(const SimTrack &tk, GlobalContext &g) const ;
    
    edm::Ref<reco::GenParticleCollection> generatorRef_(const SimTrack &tk, const GlobalContext &g) const ;
    reco::GenParticle makeGenParticle_(const SimTrack &tk, const edm::Ref<reco::GenParticleCollection> & mother, const GlobalContext &g) 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(); }
    };    
  };
}

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

PATGenCandsFromSimTracksProducer::PATGenCandsFromSimTracksProducer(const ParameterSet& cfg) :
  firstEvent_(true),
  simTracksToken_(consumes<SimTrackContainer>(cfg.getParameter<InputTag>("src"))),            // source sim tracks
  simVertexToken_(consumes<SimVertexContainer>(cfg.getParameter<InputTag>("src"))),            // source sim  vertices
  setStatus_(cfg.getParameter<int32_t>("setStatus")), // set status of GenParticle to this code
  filter_( cfg.existsAs<string>("filter") ? cfg.getParameter<string>("filter") : std::string("1 == 1") ),
  makeMotherLink_(cfg.existsAs<bool>("makeMotherLink") ? cfg.getParameter<bool>("makeMotherLink") : false),
  writeAncestors_(cfg.existsAs<bool>("writeAncestors") ? cfg.getParameter<bool>("writeAncestors") : false),
  genParticlesToken_(mayConsume<GenParticleCollection>(cfg.getParameter<InputTag>("genParticles"))),
  genBarcodesToken_(mayConsume<std::vector<int> >(cfg.getParameter<InputTag>("genParticles")))
{
    // Possibly allow a list of particle types
    if (cfg.exists("particleTypes")) {
        pdts_ = cfg.getParameter<vector<PdtEntry> >("particleTypes");
    }
    if (cfg.exists("motherTypes")) {
        motherPdts_ = cfg.getParameter<vector<PdtEntry> >("motherTypes");
    }
    
    if (writeAncestors_ && !makeMotherLink_) {
        edm::LogWarning("Configuration") << "PATGenCandsFromSimTracksProducer: " <<
            "you have set 'writeAncestors' to 'true' and 'makeMotherLink' to false;" <<
            "GEANT particles with generator level (e.g.PYHIA) mothers won't have mother links.\n";
    }
    produces<GenParticleCollection>();
}

const SimTrack *
PATGenCandsFromSimTracksProducer::findGeantMother(const SimTrack &tk, const GlobalContext &g) const {
   assert(tk.genpartIndex() == -1); // MUST NOT be called with a PYTHIA track
   if (!tk.noVertex()) {
       const SimVertex &vtx = g.simvtxs[tk.vertIndex()];
       if (!vtx.noParent()) {
           unsigned int idx = vtx.parentIndex();
           SimTrackContainer::const_iterator it = std::lower_bound(g.simtks.begin(), g.simtks.end(), idx, LessById());
           if ((it != g.simtks.end()) && (it->trackId() == idx)) {
                return &*it;
           }
       }
   }
   return nullptr;
}

edm::Ref<reco::GenParticleCollection>
PATGenCandsFromSimTracksProducer::findRef(const SimTrack &tk, GlobalContext &g) const {
    if (tk.genpartIndex() != -1) return makeMotherLink_ ? generatorRef_(tk, g) : edm::Ref<reco::GenParticleCollection>();
    const SimTrack * simMother = findGeantMother(tk, g);

    edm::Ref<reco::GenParticleCollection> motherRef;
    if (simMother != nullptr) motherRef = findRef(*simMother,g);

    if (writeAncestors_) {
        // If writing ancestors, I need to serialize myself, and then to return a ref to me
        // But first check if I've already been serialized
        std::map<unsigned int,int>::const_iterator it = g.simTksProcessed.find(tk.trackId());
        if (it != g.simTksProcessed.end()) {
            // just return a ref to it
            assert(it->second > 0);
            return edm::Ref<reco::GenParticleCollection>(g.refprod, (it->second) - 1);
        } else {
            // make genParticle, save, update the map, and return ref to myself
            reco::GenParticle p = makeGenParticle_(tk, motherRef, g);
            g.output.push_back(p);
            g.simTksProcessed[tk.trackId()] = g.output.size();
            return edm::Ref<reco::GenParticleCollection>(g.refprod, g.output.size()-1 );
        }
    } else {
        // Otherwise, I just return a ref to my mum
        return motherRef;
    }
}

edm::Ref<reco::GenParticleCollection>
PATGenCandsFromSimTracksProducer::generatorRef_(const SimTrack &st, const GlobalContext &g) const {
    assert(st.genpartIndex() != -1);
    // Note that st.genpartIndex() is the barcode, not the index within GenParticleCollection, so I have to search the particle
    std::vector<int>::const_iterator it;
    if (g.barcodesAreSorted) {
        it = std::lower_bound(g.genBarcodes->begin(), g.genBarcodes->end(), st.genpartIndex());
    } else {
        it = std::find(       g.genBarcodes->begin(), g.genBarcodes->end(), st.genpartIndex());
    }

    // Check that I found something
    // I need to check '*it == st.genpartIndex()' because lower_bound just finds the right spot for an item in a sorted list, not the item
    if ((it != g.genBarcodes->end()) && (*it == st.genpartIndex())) {
        return reco::GenParticleRef(g.gens, it - g.genBarcodes->begin());
    } else {
        return reco::GenParticleRef();
    }
}

reco::GenParticle
PATGenCandsFromSimTracksProducer::makeGenParticle_(const SimTrack &tk, const edm::Ref<reco::GenParticleCollection> & mother, const GlobalContext &g) const {
    // Make up a GenParticleCandidate from the GEANT track info.
    int charge = static_cast<int>(tk.charge());
    const Particle::LorentzVector& p4 = tk.momentum();
    Particle::Point vtx; // = (0,0,0) by default
    if (!tk.noVertex()) vtx = g.simvtxs[tk.vertIndex()].position();
    GenParticle gp(charge, p4, vtx, tk.type(), setStatus_, true);
    if (mother.isNonnull()) gp.addMother(mother);
    return gp;
}


void PATGenCandsFromSimTracksProducer::produce(Event& event,
                                               const EventSetup& iSetup) {

  if (firstEvent_){
    if (!pdts_.empty()) {
      pdgIds_.clear();
      for (vector<PdtEntry>::iterator itp = pdts_.begin(), edp = pdts_.end(); itp != edp; ++itp) {
    itp->setup(iSetup); // decode string->pdgId and vice-versa
    pdgIds_.insert(std::abs(itp->pdgId()));
      }
      pdts_.clear();
    }
    if (!motherPdts_.empty()) {
      motherPdgIds_.clear();
      for (vector<PdtEntry>::iterator itp = motherPdts_.begin(), edp = motherPdts_.end(); itp != edp; ++itp) {
    itp->setup(iSetup); // decode string->pdgId and vice-versa
    motherPdgIds_.insert(std::abs(itp->pdgId()));
      }
      motherPdts_.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 (makeMotherLink_ || writeAncestors_) {
      if (!__gnu_cxx::is_sorted(simtracks->begin(), simtracks->end(), LessById())) {
          simtracksTmp.reset(new SimTrackContainer(*simtracks));
          std::sort(simtracksTmp->begin(), simtracksTmp->end(), LessById());
          simtracksSorted = &* simtracksTmp;
      }
  }

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

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


  // make the output collection
  auto cands = std::make_unique<GenParticleCollection>();
  edm::RefProd<GenParticleCollection> refprod = event.getRefBeforePut<GenParticleCollection>();

  GlobalContext globals(*simtracksSorted, *simvertices, gens, genBarcodes, barcodesAreSorted, *cands, refprod);

  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;
      }

      GenParticle genp = makeGenParticle_(*isimtrk, Ref<GenParticleCollection>(), globals);

      // Maybe apply filter on the particle
      if (!(filter_(genp))) continue;


      if (!motherPdgIds_.empty()) {
           const SimTrack *motherSimTk = findGeantMother(*isimtrk, globals);
           if (motherSimTk == nullptr) continue;
           if (motherPdgIds_.find(std::abs(motherSimTk->type())) == motherPdgIds_.end()) continue;
      }

      if (makeMotherLink_ || writeAncestors_) {
          Ref<GenParticleCollection> motherRef;
          const SimTrack * mother = findGeantMother(*isimtrk, globals);
          if (mother != nullptr) motherRef = findRef(*mother, globals);
          if (motherRef.isNonnull()) genp.addMother(motherRef);
      }

      cands->push_back(genp);
  }

  // Write to the Event, and get back a handle (which can be useful for debugging)
  edm::OrphanHandle<reco::GenParticleCollection> orphans = event.put(std::move(cands));

#ifdef DEBUG_PATGenCandsFromSimTracksProducer
  std::cout << "Produced a list of " << orphans->size() << " genParticles." << std::endl;
  for (GenParticleCollection::const_iterator it = orphans->begin(), ed = orphans->end(); it != ed; ++it) {
      std::cout << "    ";
      std::cout << "GenParticle #" << (it - orphans->begin()) << ": pdgId " << it->pdgId()
                << ", pt = " << it->pt() << ", eta = " << it->eta() << ", phi = " << it->phi()
                << ", rho = " << it->vertex().Rho() << ", z = " << it->vertex().Z() << std::endl;
      edm::Ref<GenParticleCollection> mom = it->motherRef();
      size_t depth = 2;
      while (mom.isNonnull()) {
          if (mom.id() == orphans.id()) {
              // I need to re-make the ref because they are not working until this module returns.
              mom = edm::Ref<GenParticleCollection>(orphans, mom.key());
          }
          for (size_t i = 0; i < depth; ++i) std::cout << "    ";
          std::cout << "GenParticleRef [" << mom.id() << "/" << mom.key() << "]: pdgId " << mom->pdgId() << ", status = " << mom->status()
                    << ", pt = " << mom->pt() << ", eta = " << mom->eta() << ", phi = " << mom->phi()
                    << ", rho = " << mom->vertex().Rho() << ", z = " << mom->vertex().Z() << std::endl;
          if (mom.id() != orphans.id()) break;
          if ((mom->motherRef().id() == mom.id()) && (mom->motherRef().key() == mom.key())) {
              throw cms::Exception("Corrupt Data") << "A particle is it's own mother.\n";
          }
          mom = mom->motherRef();
          depth++;
      }
  }
  std::cout << std::endl;
#endif

}

DEFINE_FWK_MODULE(PATGenCandsFromSimTracksProducer);