JetFlavourIdentifier.cc

Go to the documentation of this file.

//
// Translation of BTag MCJetFlavour tool to identify real flavour of a jet
// work with CaloJet objects
// Store Infos by Values in JetFlavour.h
// Author: Attilio
// Date: 05.10.2007
//
//
// \class JetFlavourIdentifier
//
// \brief Interface to pull out the proper flavour identifier from a
//        jet->parton matching collection.
//
// In detail, the definitions are as follows:
//
// Definitions:
// The default behavior is that the "definition" is NULL_DEF,
// so the software will fall back to either "physics" or "algorithmic" definition
// as per the "physDefinition" switch.
//
// If the user specifies "definition", then that definition is taken. However,
// if the requested definition is not defined, the software reverts back to
// either "physics" or "algorithmic" based on the "physDefinition" switch.
// For example, if the user specifies "heaviest" as a flavor ID, and there
// are no bottom, charm, or top quarks in the event that match to the jet,
// then the software will fall back to the "physics" or "algorithmic" definition.
//
// Modifications:
//
//     09.03.2008: Sal Rappoccio.
//                 Added capability for all methods of identification, not just
//                 "physics" or "algorithmic". If the requested method does not exist
//                 (i.e. is unphysical), the "physics" or "algorithmic" definition
//                 is defaulted to.

//=======================================================================

// user include files
#include "FWCore/Framework/interface/global/EDProducer.h"
#include "FWCore/ParameterSet/interface/ParameterSet.h"
#include "FWCore/ParameterSet/interface/ParameterSetfwd.h"
#include "FWCore/Utilities/interface/InputTag.h"

#include "FWCore/Framework/interface/Event.h"
#include "FWCore/Framework/interface/EventSetup.h"
#include "FWCore/Framework/interface/MakerMacros.h"
#include "FWCore/Framework/interface/ESHandle.h"
#include "FWCore/Framework/interface/makeRefToBaseProdFrom.h"
#include "FWCore/MessageLogger/interface/MessageLogger.h"

//#include "SimGeneral/HepPDTRecord/interface/ParticleDataTable.h"

#include "DataFormats/JetReco/interface/Jet.h"
#include "DataFormats/JetReco/interface/JetCollection.h"
#include "DataFormats/JetReco/interface/CaloJetCollection.h"
#include "DataFormats/JetMatching/interface/JetFlavour.h"
#include "DataFormats/JetMatching/interface/JetFlavourMatching.h"
#include "DataFormats/JetMatching/interface/MatchedPartons.h"
#include "DataFormats/JetMatching/interface/JetMatchedPartons.h"

#include "DataFormats/Common/interface/Ref.h"
#include "DataFormats/Candidate/interface/Candidate.h"
#include "DataFormats/Candidate/interface/CandidateFwd.h"

#include "DataFormats/Math/interface/Point3D.h"
#include "DataFormats/Math/interface/LorentzVector.h"

#include <memory>
#include <string>
#include <iostream>
#include <vector>
#include <Math/VectorUtil.h>
#include <TMath.h>

using namespace std;
using namespace reco;
using namespace edm;
using namespace ROOT::Math::VectorUtil;

namespace reco {
  namespace modules {

    //--------------------------------------------------------------------------
    //
    //--------------------------------------------------------------------------
    class JetFlavourIdentifier : public edm::global::EDProducer<> {
    public:
      enum DEFINITION_T { PHYSICS = 0, ALGO, NEAREST_STATUS2, NEAREST_STATUS3, HEAVIEST, N_DEFINITIONS, NULL_DEF };

      JetFlavourIdentifier(const edm::ParameterSet&);
      ~JetFlavourIdentifier() override;

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

      JetFlavour::Leptons findLeptons(const GenParticleRef&) const;
      std::vector<const reco::Candidate*> findCandidates(const reco::Candidate*,
                                                         int,
                                                         math::XYZTLorentzVector const& thePartonLV) const;
      void fillLeptons(const std::vector<const reco::Candidate*>&,
                       JetFlavour::Leptons&,
                       int,
                       int,
                       math::XYZTLorentzVector const& thePartonLV) const;
      static int heaviestFlavour(int);

      EDGetTokenT<JetMatchedPartonsCollection> sourceByReferToken_;
      bool physDefinition;
      bool leptonInfo_;
      DEFINITION_T definition;
    };
  }  // namespace modules
}  // namespace reco
using reco::modules::JetFlavourIdentifier;

//=========================================================================

JetFlavourIdentifier::JetFlavourIdentifier(const edm::ParameterSet& iConfig) {
  produces<JetFlavourMatchingCollection>();
  sourceByReferToken_ = consumes<JetMatchedPartonsCollection>(iConfig.getParameter<InputTag>("srcByReference"));
  physDefinition = iConfig.getParameter<bool>("physicsDefinition");
  leptonInfo_ = iConfig.exists("leptonInfo") ? iConfig.getParameter<bool>("leptonInfo") : false;
  // If we have a definition of which parton to identify, use it,
  // otherwise we default to the "old" behavior of either "physics" or "algorithmic".
  // Furthermore, if the specified definition is not sensible for the given jet,
  // then the "physDefinition" switch is used to identify the flavour of the jet.
  if (iConfig.exists("definition")) {
    definition = static_cast<DEFINITION_T>(iConfig.getParameter<int>("definition"));
  } else {
    definition = NULL_DEF;
  }
}

//=========================================================================

JetFlavourIdentifier::~JetFlavourIdentifier() {}

// ------------ method called to produce the data  ------------

void JetFlavourIdentifier::produce(StreamID, Event& iEvent, const EventSetup& iEs) const {
  // Get the JetMatchedPartons
  Handle<JetMatchedPartonsCollection> theTagByRef;
  iEvent.getByToken(sourceByReferToken_, theTagByRef);

  // Create a JetFlavourMatchingCollection
  JetFlavourMatchingCollection* jfmc;
  if (!theTagByRef->empty()) {
    RefToBase<Jet> jj = theTagByRef->begin()->first;
    jfmc = new JetFlavourMatchingCollection(edm::makeRefToBaseProdFrom(jj, iEvent));
  } else {
    jfmc = new JetFlavourMatchingCollection();
  }
  std::unique_ptr<reco::JetFlavourMatchingCollection> jetFlavMatching(jfmc);

  // Loop over the matched partons and see which match.
  for (JetMatchedPartonsCollection::const_iterator j = theTagByRef->begin(); j != theTagByRef->end(); j++) {
    // Consider this match.
    const MatchedPartons aMatch = (*j).second;

    // This will hold the 4-vector, vertex, flavour and the leptonian decays (0: no lepton, xyz: x leptons in layer 2, y in layer 1 and z in layer 0) of the requested object.
    math::XYZTLorentzVector thePartonLorentzVector(0, 0, 0, 0);
    math::XYZPoint thePartonVertex(0, 0, 0);
    int thePartonFlavour = 0;
    JetFlavour::Leptons theLeptons;

    // get the partons based on which definition to use.
    switch (definition) {
      case PHYSICS: {
        const GenParticleRef& aPartPhy = aMatch.physicsDefinitionParton();
        if (aPartPhy.isNonnull()) {
          thePartonLorentzVector = aPartPhy.get()->p4();
          thePartonVertex = aPartPhy.get()->vertex();
          thePartonFlavour = aPartPhy.get()->pdgId();
          if (leptonInfo_)
            theLeptons = findLeptons(aPartPhy);
        }
        break;
      }
      case ALGO: {
        const GenParticleRef& aPartAlg = aMatch.algoDefinitionParton();
        if (aPartAlg.isNonnull()) {
          thePartonLorentzVector = aPartAlg.get()->p4();
          thePartonVertex = aPartAlg.get()->vertex();
          thePartonFlavour = aPartAlg.get()->pdgId();
          if (leptonInfo_)
            theLeptons = findLeptons(aPartAlg);
        }
        break;
      }
      case NEAREST_STATUS2: {
        const GenParticleRef& aPartN2 = aMatch.nearest_status2();
        if (aPartN2.isNonnull()) {
          thePartonLorentzVector = aPartN2.get()->p4();
          thePartonVertex = aPartN2.get()->vertex();
          thePartonFlavour = aPartN2.get()->pdgId();
          if (leptonInfo_)
            theLeptons = findLeptons(aPartN2);
        }
        break;
      }
      case NEAREST_STATUS3: {
        const GenParticleRef& aPartN3 = aMatch.nearest_status3();
        if (aPartN3.isNonnull()) {
          thePartonLorentzVector = aPartN3.get()->p4();
          thePartonVertex = aPartN3.get()->vertex();
          thePartonFlavour = aPartN3.get()->pdgId();
          if (leptonInfo_)
            theLeptons = findLeptons(aPartN3);
        }
        break;
      }
      case HEAVIEST: {
        const GenParticleRef aPartHeaviest = aMatch.heaviest();
        if (aPartHeaviest.isNonnull()) {
          thePartonLorentzVector = aPartHeaviest.get()->p4();
          thePartonVertex = aPartHeaviest.get()->vertex();
          thePartonFlavour = aPartHeaviest.get()->pdgId();
          if (leptonInfo_)
            theLeptons = findLeptons(aPartHeaviest);
        }
        break;
      }
      // Default case is backwards-compatible
      default: {
        if (physDefinition) {
          const GenParticleRef& aPartPhy = aMatch.physicsDefinitionParton();
          if (aPartPhy.isNonnull()) {
            thePartonLorentzVector = aPartPhy.get()->p4();
            thePartonVertex = aPartPhy.get()->vertex();
            thePartonFlavour = aPartPhy.get()->pdgId();
            if (leptonInfo_)
              theLeptons = findLeptons(aPartPhy);
          }
        } else {
          const GenParticleRef& aPartAlg = aMatch.algoDefinitionParton();
          if (aPartAlg.isNonnull()) {
            thePartonLorentzVector = aPartAlg.get()->p4();
            thePartonVertex = aPartAlg.get()->vertex();
            thePartonFlavour = aPartAlg.get()->pdgId();
            if (leptonInfo_)
              theLeptons = findLeptons(aPartAlg);
          }
        }
      } break;
    }  // end switch on definition

    // Now make sure we have a match. If the user specified "heaviest", for instance,
    // and there is no b- or c-quarks, then fall back to the "physDefinition" switch.

    if (thePartonFlavour == 0) {
      if (physDefinition) {
        const GenParticleRef& aPartPhy = aMatch.physicsDefinitionParton();
        if (aPartPhy.isNonnull()) {
          thePartonLorentzVector = aPartPhy.get()->p4();
          thePartonVertex = aPartPhy.get()->vertex();
          thePartonFlavour = aPartPhy.get()->pdgId();
          if (leptonInfo_)
            theLeptons = findLeptons(aPartPhy);
        }
      } else {
        const GenParticleRef& aPartAlg = aMatch.algoDefinitionParton();
        if (aPartAlg.isNonnull()) {
          thePartonLorentzVector = aPartAlg.get()->p4();
          thePartonVertex = aPartAlg.get()->vertex();
          thePartonFlavour = aPartAlg.get()->pdgId();
          if (leptonInfo_)
            theLeptons = findLeptons(aPartAlg);
        }
      }
    }

    /*
     std::cout << "Leptons of " <<thePartonFlavour << " Jet: " << std::endl;
     std::cout << "  electrons: " <<theLeptons.electron << std::endl;
     std::cout << "  muons    : " <<theLeptons.muon << std::endl;
     std::cout << "  tau      : " <<theLeptons.tau << std::endl;
*/
    // Add the jet->flavour match to the map.
    (*jetFlavMatching)[(*j).first] = JetFlavour(thePartonLorentzVector, thePartonVertex, thePartonFlavour, theLeptons);
  }  // end loop over jets

  // Put the object into the event.
  iEvent.put(std::move(jetFlavMatching));
}

JetFlavour::Leptons JetFlavourIdentifier::findLeptons(const GenParticleRef& parton) const {
  JetFlavour::Leptons theLeptons;

  auto const& thePartonLV = parton->p4();

  const reco::Candidate* mcstring = parton->daughter(0);
  int partonFlavour = std::abs(parton->pdgId());
  //  std::cout << "parton DeltaR: " << DeltaR(thePartonLV, parton->p4()) << std::endl;

  std::vector<const reco::Candidate*> candidates = findCandidates(mcstring, partonFlavour, parton->p4());
  //  std::cout << "Candidates are:" << std::endl;
  //  for(unsigned int j = 0; j < candidates.size(); j++) std::cout << "   --> " << candidates[j]->pdgId() << std::endl;

  fillLeptons(candidates, theLeptons, 1, partonFlavour, thePartonLV);

  return theLeptons;
}

std::vector<const reco::Candidate*> JetFlavourIdentifier::findCandidates(
    const reco::Candidate* cand, int partonFlavour, math::XYZTLorentzVector const& thePartonLV) const {
  std::vector<const reco::Candidate*> cands;
  if (!cand)
    return cands;

  for (unsigned int i = 0; i < cand->numberOfDaughters(); i++) {
    /*
    std::cout << "DeltaR - " << partonFlavour << " ";
    if (DeltaR(thePartonLV, cand->daughter(i)->p4()) > 0.7) std::cout << "(";
    std::cout << cand->daughter(i)->pdgId() << ": " << DeltaR(thePartonLV, cand->daughter(i)->p4());
    if (DeltaR(thePartonLV, cand->daughter(i)->p4()) > 0.7) std::cout << ")";
    std::cout << std::endl;
*/
    if (DeltaR(thePartonLV, cand->daughter(i)->p4()) < 0.7) {
      int pdgId = std::abs(cand->daughter(i)->pdgId());
      int flavour = heaviestFlavour(pdgId);
      if (flavour == partonFlavour || (flavour >= 10 && partonFlavour >= 10)) {
        //        std::cout << "<------- " << std::endl;
        std::vector<const reco::Candidate*> newcands = findCandidates(cand->daughter(i), partonFlavour, thePartonLV);
        //        std::cout << " ------->" << std::endl;
        std::copy(newcands.begin(), newcands.end(), std::back_inserter(cands));
      }
      if (partonFlavour >= 10)
        cands.push_back(cand->daughter(i));
    }
  }

  if (cands.empty() && std::abs(cand->pdgId()) > 110 &&
      !(partonFlavour >= 4 && partonFlavour < 10 && heaviestFlavour(cand->pdgId()) < 4))
    cands.push_back(cand);

  return cands;
}

void JetFlavourIdentifier::fillLeptons(const std::vector<const reco::Candidate*>& cands,
                                       JetFlavour::Leptons& leptons,
                                       int rank,
                                       int flavour,
                                       math::XYZTLorentzVector const& thePartonLV) const {
  for (unsigned int j = 0; j < cands.size(); j++) {
    for (unsigned int i = 0; i < cands[j]->numberOfDaughters(); i++) {
      int pdgId = std::abs(cands[j]->daughter(i)->pdgId());

      //      for(int z = 1; z <= rank; z *= 10) std::cout << " ------ ";
      //      std::cout << pdgId << std::endl;

      if (pdgId == 12)
        leptons.electron += rank;
      else if (pdgId == 14)
        leptons.muon += rank;
      else if (pdgId == 16)
        leptons.tau += rank;
      else {
        int heaviest = heaviestFlavour(pdgId);
        int heaviest_ = heaviest < 10 ? heaviest : 0;
        if (!heaviest || (flavour < 4 ? (heaviest_ < 4) : (heaviest >= 4))) {
          std::vector<const reco::Candidate*> newcands = findCandidates(cands[j]->daughter(i), heaviest, thePartonLV);
          if (pdgId <= 110)
            newcands.push_back(cands[j]->daughter(i));
          fillLeptons(newcands, leptons, rank * 10, std::max(heaviest_, flavour), thePartonLV);
        }
      }
    }
  }
}

int JetFlavourIdentifier::heaviestFlavour(int pdgId) {
  int flavour = 0;

  pdgId = std::abs(pdgId) % 100000;
  if (pdgId > 110) {
    while (pdgId % 10 > 0 && pdgId % 10 < 6) {
      pdgId /= 10;
      if (pdgId % 10 > flavour)
        flavour = pdgId % 10;
    }
  } else
    flavour = pdgId;

  return flavour;
}

//define this as a plug-in
DEFINE_FWK_MODULE(JetFlavourIdentifier);