JetFlavourClustering.cc

Go to the documentation of this file.

// -*- C++ -*-
//
// Package:    JetFlavourClustering
// Class:      JetFlavourClustering
//
//
// Original Author:  Dinko Ferencek
//         Created:  Wed Nov  6 00:49:55 CET 2013
//
//

// system include files
#include <memory>

// user include files
#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/JetReco/interface/Jet.h"
#include "DataFormats/JetReco/interface/JetCollection.h"
#include "SimDataFormats/JetMatching/interface/JetFlavourInfo.h"
#include "SimDataFormats/JetMatching/interface/JetFlavourInfoMatching.h"
#include "DataFormats/HepMCCandidate/interface/GenParticle.h"
#include "DataFormats/HepMCCandidate/interface/GenParticleFwd.h"
#include "DataFormats/Math/interface/deltaR.h"
#include "PhysicsTools/JetMCUtils/interface/CandMCTag.h"

#include "fastjet/JetDefinition.hh"
#include "fastjet/ClusterSequence.hh"
#include "fastjet/Selector.hh"
#include "fastjet/PseudoJet.hh"

//
// constants, enums and typedefs
//
typedef std::shared_ptr<fastjet::ClusterSequence> ClusterSequencePtr;
typedef std::shared_ptr<fastjet::JetDefinition> JetDefPtr;

//
// class declaration
//
class GhostInfo : public fastjet::PseudoJet::UserInfoBase {
public:
  GhostInfo(const bool& isHadron,
            const bool& isbHadron,
            const bool& isParton,
            const bool& isLepton,
            const reco::GenParticleRef& particleRef)
      : m_particleRef(particleRef) {
    m_type = 0;
    if (isHadron)
      m_type |= (1 << 0);
    if (isbHadron)
      m_type |= (1 << 1);
    if (isParton)
      m_type |= (1 << 2);
    if (isLepton)
      m_type |= (1 << 3);
  }

  const bool isHadron() const { return (m_type & (1 << 0)); }
  const bool isbHadron() const { return (m_type & (1 << 1)); }
  const bool isParton() const { return (m_type & (1 << 2)); }
  const bool isLepton() const { return (m_type & (1 << 3)); }
  const reco::GenParticleRef& particleRef() const { return m_particleRef; }

protected:
  const reco::GenParticleRef m_particleRef;
  int m_type;
};

class JetFlavourClustering : public edm::stream::EDProducer<> {
public:
  explicit JetFlavourClustering(const edm::ParameterSet&);
  ~JetFlavourClustering() override;

  static void fillDescriptions(edm::ConfigurationDescriptions& descriptions);

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

  void insertGhosts(const edm::Handle<reco::GenParticleRefVector>& particles,
                    const double ghostRescaling,
                    const bool isHadron,
                    const bool isbHadron,
                    const bool isParton,
                    const bool isLepton,
                    std::vector<fastjet::PseudoJet>& constituents);

  void matchReclusteredJets(const edm::Handle<edm::View<reco::Jet> >& jets,
                            const std::vector<fastjet::PseudoJet>& matchedJets,
                            std::vector<int>& matchedIndices);
  void matchGroomedJets(const edm::Handle<edm::View<reco::Jet> >& jets,
                        const edm::Handle<edm::View<reco::Jet> >& matchedJets,
                        std::vector<int>& matchedIndices);
  void matchSubjets(const std::vector<int>& groomedIndices,
                    const edm::Handle<edm::View<reco::Jet> >& groomedJets,
                    const edm::Handle<edm::View<reco::Jet> >& subjets,
                    std::vector<std::vector<int> >& matchedIndices);

  void setFlavours(const reco::GenParticleRefVector& clusteredbHadrons,
                   const reco::GenParticleRefVector& clusteredcHadrons,
                   const reco::GenParticleRefVector& clusteredPartons,
                   int& hadronFlavour,
                   int& partonFlavour);

  void assignToSubjets(const reco::GenParticleRefVector& clusteredParticles,
                       const edm::Handle<edm::View<reco::Jet> >& subjets,
                       const std::vector<int>& subjetIndices,
                       std::vector<reco::GenParticleRefVector>& assignedParticles);

  // ----------member data ---------------------------
  const edm::EDGetTokenT<edm::View<reco::Jet> > jetsToken_;           // Input jet collection
  edm::EDGetTokenT<edm::View<reco::Jet> > groomedJetsToken_;          // Input groomed jet collection
  edm::EDGetTokenT<edm::View<reco::Jet> > subjetsToken_;              // Input subjet collection
  const edm::EDGetTokenT<reco::GenParticleRefVector> bHadronsToken_;  // Input b hadron collection
  const edm::EDGetTokenT<reco::GenParticleRefVector> cHadronsToken_;  // Input c hadron collection
  const edm::EDGetTokenT<reco::GenParticleRefVector> partonsToken_;   // Input parton collection
  edm::EDGetTokenT<reco::GenParticleRefVector> leptonsToken_;         // Input lepton collection

  const std::string jetAlgorithm_;
  const double rParam_;
  const double jetPtMin_;
  const double ghostRescaling_;
  const double relPtTolerance_;
  const bool hadronFlavourHasPriority_;
  const bool useSubjets_;
  const bool useLeptons_;

  ClusterSequencePtr fjClusterSeq_;
  JetDefPtr fjJetDefinition_;
};

//
// static data member definitions
//

//
// constructors and destructor
//
JetFlavourClustering::JetFlavourClustering(const edm::ParameterSet& iConfig)
    :

      jetsToken_(consumes<edm::View<reco::Jet> >(iConfig.getParameter<edm::InputTag>("jets"))),
      bHadronsToken_(consumes<reco::GenParticleRefVector>(iConfig.getParameter<edm::InputTag>("bHadrons"))),
      cHadronsToken_(consumes<reco::GenParticleRefVector>(iConfig.getParameter<edm::InputTag>("cHadrons"))),
      partonsToken_(consumes<reco::GenParticleRefVector>(iConfig.getParameter<edm::InputTag>("partons"))),
      jetAlgorithm_(iConfig.getParameter<std::string>("jetAlgorithm")),
      rParam_(iConfig.getParameter<double>("rParam")),
      jetPtMin_(
          0.),  // hardcoded to 0. since we simply want to recluster all input jets which already had some PtMin applied
      ghostRescaling_(iConfig.exists("ghostRescaling") ? iConfig.getParameter<double>("ghostRescaling") : 1e-18),
      relPtTolerance_(
          iConfig.exists("relPtTolerance")
              ? iConfig.getParameter<double>("relPtTolerance")
              : 1e-03),  // 0.1% relative difference in Pt should be sufficient to detect possible misconfigurations
      hadronFlavourHasPriority_(iConfig.getParameter<bool>("hadronFlavourHasPriority")),
      useSubjets_(iConfig.exists("groomedJets") && iConfig.exists("subjets")),
      useLeptons_(iConfig.exists("leptons"))

{
  // register your products
  produces<reco::JetFlavourInfoMatchingCollection>();
  if (useSubjets_)
    produces<reco::JetFlavourInfoMatchingCollection>("SubJets");

  // set jet algorithm
  if (jetAlgorithm_ == "Kt")
    fjJetDefinition_ = std::make_shared<fastjet::JetDefinition>(fastjet::kt_algorithm, rParam_);
  else if (jetAlgorithm_ == "CambridgeAachen")
    fjJetDefinition_ = std::make_shared<fastjet::JetDefinition>(fastjet::cambridge_algorithm, rParam_);
  else if (jetAlgorithm_ == "AntiKt")
    fjJetDefinition_ = std::make_shared<fastjet::JetDefinition>(fastjet::antikt_algorithm, rParam_);
  else
    throw cms::Exception("InvalidJetAlgorithm") << "Jet clustering algorithm is invalid: " << jetAlgorithm_
                                                << ", use CambridgeAachen | Kt | AntiKt" << std::endl;

  if (useSubjets_) {
    groomedJetsToken_ = consumes<edm::View<reco::Jet> >(iConfig.getParameter<edm::InputTag>("groomedJets"));
    subjetsToken_ = consumes<edm::View<reco::Jet> >(iConfig.getParameter<edm::InputTag>("subjets"));
  }
  if (useLeptons_) {
    leptonsToken_ = consumes<reco::GenParticleRefVector>(iConfig.getParameter<edm::InputTag>("leptons"));
  }
}

JetFlavourClustering::~JetFlavourClustering() {
  // do anything here that needs to be done at desctruction time
  // (e.g. close files, deallocate resources etc.)
}

//
// member functions
//

// ------------ method called to produce the data  ------------
void JetFlavourClustering::produce(edm::Event& iEvent, const edm::EventSetup& iSetup) {
  edm::Handle<edm::View<reco::Jet> > jets;
  iEvent.getByToken(jetsToken_, jets);

  edm::Handle<edm::View<reco::Jet> > groomedJets;
  edm::Handle<edm::View<reco::Jet> > subjets;
  if (useSubjets_) {
    iEvent.getByToken(groomedJetsToken_, groomedJets);
    iEvent.getByToken(subjetsToken_, subjets);
  }

  edm::Handle<reco::GenParticleRefVector> bHadrons;
  iEvent.getByToken(bHadronsToken_, bHadrons);

  edm::Handle<reco::GenParticleRefVector> cHadrons;
  iEvent.getByToken(cHadronsToken_, cHadrons);

  edm::Handle<reco::GenParticleRefVector> partons;
  iEvent.getByToken(partonsToken_, partons);

  edm::Handle<reco::GenParticleRefVector> leptons;
  if (useLeptons_)
    iEvent.getByToken(leptonsToken_, leptons);

  auto jetFlavourInfos = std::make_unique<reco::JetFlavourInfoMatchingCollection>(reco::JetRefBaseProd(jets));
  std::unique_ptr<reco::JetFlavourInfoMatchingCollection> subjetFlavourInfos;
  if (useSubjets_)
    subjetFlavourInfos = std::make_unique<reco::JetFlavourInfoMatchingCollection>(reco::JetRefBaseProd(subjets));

  // vector of constituents for reclustering jets and "ghosts"
  std::vector<fastjet::PseudoJet> fjInputs;
  unsigned int reserve = jets->size() * 128 + bHadrons->size() + cHadrons->size() + partons->size();
  if (useLeptons_)
    reserve += leptons->size();
  fjInputs.reserve(reserve);
  // loop over all input jets and collect all their constituents
  for (edm::View<reco::Jet>::const_iterator it = jets->begin(); it != jets->end(); ++it) {
    std::vector<edm::Ptr<reco::Candidate> > constituents = it->getJetConstituents();
    std::vector<edm::Ptr<reco::Candidate> >::const_iterator m;
    for (m = constituents.begin(); m != constituents.end(); ++m) {
      reco::CandidatePtr constit = *m;
      if (!constit.isNonnull() || !constit.isAvailable()) {
        edm::LogError("MissingJetConstituent") << "Jet constituent required for jet reclustering is missing. "
                                                  "Reclustered jets are not guaranteed to reproduce the original jets!";
        continue;
      }
      if (constit->pt() == 0) {
        edm::LogWarning("NullTransverseMomentum") << "dropping input candidate with pt=0";
        continue;
      }
      fjInputs.push_back(fastjet::PseudoJet(constit->px(), constit->py(), constit->pz(), constit->energy()));
    }
  }
  // insert "ghost" b hadrons in the vector of constituents
  insertGhosts(bHadrons, ghostRescaling_, true, true, false, false, fjInputs);
  // insert "ghost" c hadrons in the vector of constituents
  insertGhosts(cHadrons, ghostRescaling_, true, false, false, false, fjInputs);
  // insert "ghost" partons in the vector of constituents
  insertGhosts(partons, ghostRescaling_, false, false, true, false, fjInputs);
  // if used, insert "ghost" leptons in the vector of constituents
  if (useLeptons_)
    insertGhosts(leptons, ghostRescaling_, false, false, false, true, fjInputs);

  // define jet clustering sequence
  fjClusterSeq_ = std::make_shared<fastjet::ClusterSequence>(fjInputs, *fjJetDefinition_);
  // recluster jet constituents and inserted "ghosts"
  std::vector<fastjet::PseudoJet> inclusiveJets = fastjet::sorted_by_pt(fjClusterSeq_->inclusive_jets(jetPtMin_));

  if (inclusiveJets.size() < jets->size())
    edm::LogError("TooFewReclusteredJets")
        << "There are fewer reclustered (" << inclusiveJets.size() << ") than original jets (" << jets->size()
        << "). Please check that the jet algorithm and jet size match those used for the original jet collection.";

  // match reclustered and original jets
  std::vector<int> reclusteredIndices;
  matchReclusteredJets(jets, inclusiveJets, reclusteredIndices);

  // match groomed and original jets
  std::vector<int> groomedIndices;
  if (useSubjets_) {
    if (groomedJets->size() > jets->size())
      edm::LogError("TooManyGroomedJets")
          << "There are more groomed (" << groomedJets->size() << ") than original jets (" << jets->size()
          << "). Please check that the two jet collections belong to each other.";

    matchGroomedJets(jets, groomedJets, groomedIndices);
  }

  // match subjets and original jets
  std::vector<std::vector<int> > subjetIndices;
  if (useSubjets_) {
    matchSubjets(groomedIndices, groomedJets, subjets, subjetIndices);
  }

  // determine jet flavour
  for (size_t i = 0; i < jets->size(); ++i) {
    reco::GenParticleRefVector clusteredbHadrons;
    reco::GenParticleRefVector clusteredcHadrons;
    reco::GenParticleRefVector clusteredPartons;
    reco::GenParticleRefVector clusteredLeptons;

    // if matching reclustered to original jets failed
    if (reclusteredIndices.at(i) < 0) {
      // set an empty JetFlavourInfo for this jet
      (*jetFlavourInfos)[jets->refAt(i)] =
          reco::JetFlavourInfo(clusteredbHadrons, clusteredcHadrons, clusteredPartons, clusteredLeptons, 0, 0);
    } else if (jets->at(i).pt() == 0) {
      edm::LogWarning("NullTransverseMomentum")
          << "The original jet " << i << " has Pt=0. This is not expected so the jet will be skipped.";

      // set an empty JetFlavourInfo for this jet
      (*jetFlavourInfos)[jets->refAt(i)] =
          reco::JetFlavourInfo(clusteredbHadrons, clusteredcHadrons, clusteredPartons, clusteredLeptons, 0, 0);

      // if subjets are used
      if (useSubjets_ && !subjetIndices.at(i).empty()) {
        // loop over subjets
        for (size_t sj = 0; sj < subjetIndices.at(i).size(); ++sj) {
          // set an empty JetFlavourInfo for this subjet
          (*subjetFlavourInfos)[subjets->refAt(subjetIndices.at(i).at(sj))] =
              reco::JetFlavourInfo(reco::GenParticleRefVector(),
                                   reco::GenParticleRefVector(),
                                   reco::GenParticleRefVector(),
                                   reco::GenParticleRefVector(),
                                   0,
                                   0);
        }
      }
    } else {
      // since the "ghosts" are extremely soft, the configuration and ordering of the reclustered and original jets should in principle stay the same
      if ((std::abs(inclusiveJets.at(reclusteredIndices.at(i)).pt() - jets->at(i).pt()) / jets->at(i).pt()) >
          relPtTolerance_) {
        if (jets->at(i).pt() < 10.)  // special handling for low-Pt jets (Pt<10 GeV)
          edm::LogWarning("JetPtMismatchAtLowPt")
              << "The reclustered and original jet " << i << " have different Pt's ("
              << inclusiveJets.at(reclusteredIndices.at(i)).pt() << " vs " << jets->at(i).pt()
              << " GeV, respectively).\n"
              << "Please check that the jet algorithm and jet size match those used for the original jet collection "
                 "and also make sure the original jets are uncorrected. In addition, make sure you are not using "
                 "CaloJets which are presently not supported.\n"
              << "Since the mismatch is at low Pt (Pt<10 GeV), it is ignored and only a warning is issued.\n"
              << "\nIn extremely rare instances the mismatch could be caused by a difference in the machine precision "
                 "in which case make sure the original jet collection is produced and reclustering is performed in the "
                 "same job.";
        else
          edm::LogError("JetPtMismatch")
              << "The reclustered and original jet " << i << " have different Pt's ("
              << inclusiveJets.at(reclusteredIndices.at(i)).pt() << " vs " << jets->at(i).pt()
              << " GeV, respectively).\n"
              << "Please check that the jet algorithm and jet size match those used for the original jet collection "
                 "and also make sure the original jets are uncorrected. In addition, make sure you are not using "
                 "CaloJets which are presently not supported.\n"
              << "\nIn extremely rare instances the mismatch could be caused by a difference in the machine precision "
                 "in which case make sure the original jet collection is produced and reclustering is performed in the "
                 "same job.";
      }

      // get jet constituents (sorted by Pt)
      std::vector<fastjet::PseudoJet> constituents =
          fastjet::sorted_by_pt(inclusiveJets.at(reclusteredIndices.at(i)).constituents());

      // loop over jet constituents and try to find "ghosts"
      for (std::vector<fastjet::PseudoJet>::const_iterator it = constituents.begin(); it != constituents.end(); ++it) {
        if (!it->has_user_info())
          continue;  // skip if not a "ghost"

        // "ghost" hadron
        if (it->user_info<GhostInfo>().isHadron()) {
          // "ghost" b hadron
          if (it->user_info<GhostInfo>().isbHadron())
            clusteredbHadrons.push_back(it->user_info<GhostInfo>().particleRef());
          // "ghost" c hadron
          else
            clusteredcHadrons.push_back(it->user_info<GhostInfo>().particleRef());
        }
        // "ghost" parton
        else if (it->user_info<GhostInfo>().isParton())
          clusteredPartons.push_back(it->user_info<GhostInfo>().particleRef());
        // "ghost" lepton
        else if (it->user_info<GhostInfo>().isLepton())
          clusteredLeptons.push_back(it->user_info<GhostInfo>().particleRef());
      }

      int hadronFlavour = 0;  // default hadron flavour set to 0 (= undefined)
      int partonFlavour = 0;  // default parton flavour set to 0 (= undefined)

      // set hadron- and parton-based flavours
      setFlavours(clusteredbHadrons, clusteredcHadrons, clusteredPartons, hadronFlavour, partonFlavour);

      // set the JetFlavourInfo for this jet
      (*jetFlavourInfos)[jets->refAt(i)] = reco::JetFlavourInfo(
          clusteredbHadrons, clusteredcHadrons, clusteredPartons, clusteredLeptons, hadronFlavour, partonFlavour);
    }

    // if subjets are used, determine their flavour
    if (useSubjets_) {
      if (subjetIndices.at(i).empty())
        continue;  // continue if the original jet does not have subjets assigned

      // define vectors of GenParticleRefVectors for hadrons and partons assigned to different subjets
      std::vector<reco::GenParticleRefVector> assignedbHadrons(subjetIndices.at(i).size(),
                                                               reco::GenParticleRefVector());
      std::vector<reco::GenParticleRefVector> assignedcHadrons(subjetIndices.at(i).size(),
                                                               reco::GenParticleRefVector());
      std::vector<reco::GenParticleRefVector> assignedPartons(subjetIndices.at(i).size(), reco::GenParticleRefVector());
      std::vector<reco::GenParticleRefVector> assignedLeptons(subjetIndices.at(i).size(), reco::GenParticleRefVector());

      // loop over clustered b hadrons and assign them to different subjets based on smallest dR
      assignToSubjets(clusteredbHadrons, subjets, subjetIndices.at(i), assignedbHadrons);
      // loop over clustered c hadrons and assign them to different subjets based on smallest dR
      assignToSubjets(clusteredcHadrons, subjets, subjetIndices.at(i), assignedcHadrons);
      // loop over clustered partons and assign them to different subjets based on smallest dR
      assignToSubjets(clusteredPartons, subjets, subjetIndices.at(i), assignedPartons);
      // if used, loop over clustered leptons and assign them to different subjets based on smallest dR
      if (useLeptons_)
        assignToSubjets(clusteredLeptons, subjets, subjetIndices.at(i), assignedLeptons);

      // loop over subjets and determine their flavour
      for (size_t sj = 0; sj < subjetIndices.at(i).size(); ++sj) {
        int subjetHadronFlavour = 0;  // default hadron flavour set to 0 (= undefined)
        int subjetPartonFlavour = 0;  // default parton flavour set to 0 (= undefined)

        // set hadron- and parton-based flavours
        setFlavours(assignedbHadrons.at(sj),
                    assignedcHadrons.at(sj),
                    assignedPartons.at(sj),
                    subjetHadronFlavour,
                    subjetPartonFlavour);

        // set the JetFlavourInfo for this subjet
        (*subjetFlavourInfos)[subjets->refAt(subjetIndices.at(i).at(sj))] =
            reco::JetFlavourInfo(assignedbHadrons.at(sj),
                                 assignedcHadrons.at(sj),
                                 assignedPartons.at(sj),
                                 assignedLeptons.at(sj),
                                 subjetHadronFlavour,
                                 subjetPartonFlavour);
      }
    }
  }

  //deallocate only at the end of the event processing
  fjClusterSeq_.reset();

  // put jet flavour infos in the event
  iEvent.put(std::move(jetFlavourInfos));
  // put subjet flavour infos in the event
  if (useSubjets_)
    iEvent.put(std::move(subjetFlavourInfos), "SubJets");
}

// ------------ method that inserts "ghost" particles in the vector of jet constituents ------------
void JetFlavourClustering::insertGhosts(const edm::Handle<reco::GenParticleRefVector>& particles,
                                        const double ghostRescaling,
                                        const bool isHadron,
                                        const bool isbHadron,
                                        const bool isParton,
                                        const bool isLepton,
                                        std::vector<fastjet::PseudoJet>& constituents) {
  // insert "ghost" particles in the vector of jet constituents
  for (reco::GenParticleRefVector::const_iterator it = particles->begin(); it != particles->end(); ++it) {
    if ((*it)->pt() == 0) {
      edm::LogInfo("NullTransverseMomentum") << "dropping input ghost candidate with pt=0";
      continue;
    }
    fastjet::PseudoJet p((*it)->px(), (*it)->py(), (*it)->pz(), (*it)->energy());
    p *= ghostRescaling;  // rescale particle momentum
    p.set_user_info(new GhostInfo(isHadron, isbHadron, isParton, isLepton, *it));
    constituents.push_back(p);
  }
}

// ------------ method that matches reclustered and original jets based on minimum dR ------------
void JetFlavourClustering::matchReclusteredJets(const edm::Handle<edm::View<reco::Jet> >& jets,
                                                const std::vector<fastjet::PseudoJet>& reclusteredJets,
                                                std::vector<int>& matchedIndices) {
  std::vector<bool> matchedLocks(reclusteredJets.size(), false);

  for (size_t j = 0; j < jets->size(); ++j) {
    double matchedDR2 = 1e9;
    int matchedIdx = -1;

    for (size_t rj = 0; rj < reclusteredJets.size(); ++rj) {
      if (matchedLocks.at(rj))
        continue;  // skip jets that have already been matched

      double tempDR2 = reco::deltaR2(jets->at(j).rapidity(),
                                     jets->at(j).phi(),
                                     reclusteredJets.at(rj).rapidity(),
                                     reclusteredJets.at(rj).phi_std());
      if (tempDR2 < matchedDR2) {
        matchedDR2 = tempDR2;
        matchedIdx = rj;
      }
    }

    if (matchedIdx >= 0) {
      if (matchedDR2 > rParam_ * rParam_) {
        edm::LogError("JetMatchingFailed") << "Matched reclustered jet " << matchedIdx << " and original jet " << j
                                           << " are separated by dR=" << sqrt(matchedDR2)
                                           << " which is greater than the jet size R=" << rParam_ << ".\n"
                                           << "This is not expected so please check that the jet algorithm and jet "
                                              "size match those used for the original jet collection.";
      } else
        matchedLocks.at(matchedIdx) = true;
    } else
      edm::LogError("JetMatchingFailed") << "Matching reclustered to original jets failed. Please check that the jet "
                                            "algorithm and jet size match those used for the original jet collection.";

    matchedIndices.push_back(matchedIdx);
  }
}

// ------------ method that matches groomed and original jets based on minimum dR ------------
void JetFlavourClustering::matchGroomedJets(const edm::Handle<edm::View<reco::Jet> >& jets,
                                            const edm::Handle<edm::View<reco::Jet> >& groomedJets,
                                            std::vector<int>& matchedIndices) {
  std::vector<bool> jetLocks(jets->size(), false);
  std::vector<int> jetIndices;

  for (size_t gj = 0; gj < groomedJets->size(); ++gj) {
    double matchedDR2 = 1e9;
    int matchedIdx = -1;

    if (groomedJets->at(gj).pt() > 0.)  // skip pathological cases of groomed jets with Pt=0
    {
      for (size_t j = 0; j < jets->size(); ++j) {
        if (jetLocks.at(j))
          continue;  // skip jets that have already been matched

        double tempDR2 = reco::deltaR2(
            jets->at(j).rapidity(), jets->at(j).phi(), groomedJets->at(gj).rapidity(), groomedJets->at(gj).phi());
        if (tempDR2 < matchedDR2) {
          matchedDR2 = tempDR2;
          matchedIdx = j;
        }
      }
    }

    if (matchedIdx >= 0) {
      if (matchedDR2 > rParam_ * rParam_) {
        edm::LogWarning("MatchedJetsFarApart")
            << "Matched groomed jet " << gj << " and original jet " << matchedIdx
            << " are separated by dR=" << sqrt(matchedDR2) << " which is greater than the jet size R=" << rParam_
            << ".\n"
            << "This is not expected so the matching of these two jets has been discarded. Please check that the two "
               "jet collections belong to each other.";
        matchedIdx = -1;
      } else
        jetLocks.at(matchedIdx) = true;
    }
    jetIndices.push_back(matchedIdx);
  }

  for (size_t j = 0; j < jets->size(); ++j) {
    std::vector<int>::iterator matchedIndex = std::find(jetIndices.begin(), jetIndices.end(), j);

    matchedIndices.push_back(matchedIndex != jetIndices.end() ? std::distance(jetIndices.begin(), matchedIndex) : -1);
  }
}

// ------------ method that matches subjets and original jets ------------
void JetFlavourClustering::matchSubjets(const std::vector<int>& groomedIndices,
                                        const edm::Handle<edm::View<reco::Jet> >& groomedJets,
                                        const edm::Handle<edm::View<reco::Jet> >& subjets,
                                        std::vector<std::vector<int> >& matchedIndices) {
  for (size_t g = 0; g < groomedIndices.size(); ++g) {
    std::vector<int> subjetIndices;

    if (groomedIndices.at(g) >= 0) {
      for (size_t s = 0; s < groomedJets->at(groomedIndices.at(g)).numberOfDaughters(); ++s) {
        const edm::Ptr<reco::Candidate>& subjet = groomedJets->at(groomedIndices.at(g)).daughterPtr(s);

        for (size_t sj = 0; sj < subjets->size(); ++sj) {
          if (subjet == edm::Ptr<reco::Candidate>(subjets->ptrAt(sj))) {
            subjetIndices.push_back(sj);
            break;
          }
        }
      }

      if (subjetIndices.empty())
        edm::LogError("SubjetMatchingFailed") << "Matching subjets to original jets failed. Please check that the "
                                                 "groomed jet and subjet collections belong to each other.";

      matchedIndices.push_back(subjetIndices);
    } else
      matchedIndices.push_back(subjetIndices);
  }
}

// ------------ method that sets hadron- and parton-based flavours ------------
void JetFlavourClustering::setFlavours(const reco::GenParticleRefVector& clusteredbHadrons,
                                       const reco::GenParticleRefVector& clusteredcHadrons,
                                       const reco::GenParticleRefVector& clusteredPartons,
                                       int& hadronFlavour,
                                       int& partonFlavour) {
  reco::GenParticleRef hardestParton;
  reco::GenParticleRef hardestLightParton;
  reco::GenParticleRef flavourParton;

  // loop over clustered partons (already sorted by Pt)
  for (reco::GenParticleRefVector::const_iterator it = clusteredPartons.begin(); it != clusteredPartons.end(); ++it) {
    // hardest parton
    if (hardestParton.isNull())
      hardestParton = (*it);
    // hardest light-flavour parton
    if (hardestLightParton.isNull()) {
      if (CandMCTagUtils::isLightParton(*(*it)))
        hardestLightParton = (*it);
    }
    // c flavour
    if (flavourParton.isNull() && (std::abs((*it)->pdgId()) == 4))
      flavourParton = (*it);
    // b flavour gets priority
    if (std::abs((*it)->pdgId()) == 5) {
      if (flavourParton.isNull())
        flavourParton = (*it);
      else if (std::abs(flavourParton->pdgId()) != 5)
        flavourParton = (*it);
    }
  }

  // set hadron-based flavour
  if (!clusteredbHadrons.empty())
    hadronFlavour = 5;
  else if (!clusteredcHadrons.empty() && clusteredbHadrons.empty())
    hadronFlavour = 4;
  // set parton-based flavour
  if (flavourParton.isNull()) {
    if (hardestParton.isNonnull())
      partonFlavour = hardestParton->pdgId();
  } else
    partonFlavour = flavourParton->pdgId();

  // if enabled, check for conflicts between hadron- and parton-based flavours and give priority to the hadron-based flavour
  if (hadronFlavourHasPriority_) {
    if (hadronFlavour == 0 && (std::abs(partonFlavour) == 4 || std::abs(partonFlavour) == 5))
      partonFlavour = (hardestLightParton.isNonnull() ? hardestLightParton->pdgId() : 0);
    else if (hadronFlavour != 0 && std::abs(partonFlavour) != hadronFlavour)
      partonFlavour = hadronFlavour;
  }
}

// ------------ method that assigns clustered particles to subjets ------------
void JetFlavourClustering::assignToSubjets(const reco::GenParticleRefVector& clusteredParticles,
                                           const edm::Handle<edm::View<reco::Jet> >& subjets,
                                           const std::vector<int>& subjetIndices,
                                           std::vector<reco::GenParticleRefVector>& assignedParticles) {
  // loop over clustered particles and assign them to different subjets based on smallest dR
  for (reco::GenParticleRefVector::const_iterator it = clusteredParticles.begin(); it != clusteredParticles.end();
       ++it) {
    std::vector<double> dR2toSubjets;

    for (size_t sj = 0; sj < subjetIndices.size(); ++sj)
      dR2toSubjets.push_back(reco::deltaR2((*it)->rapidity(),
                                           (*it)->phi(),
                                           subjets->at(subjetIndices.at(sj)).rapidity(),
                                           subjets->at(subjetIndices.at(sj)).phi()));

    // find the closest subjet
    int closestSubjetIdx =
        std::distance(dR2toSubjets.begin(), std::min_element(dR2toSubjets.begin(), dR2toSubjets.end()));

    assignedParticles.at(closestSubjetIdx).push_back(*it);
  }
}

// ------------ method fills 'descriptions' with the allowed parameters for the module  ------------
void JetFlavourClustering::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
  //The following says we do not know what parameters are allowed so do no validation
  // Please change this to state exactly what you do use, even if it is no parameters
  edm::ParameterSetDescription desc;
  desc.setUnknown();
  descriptions.addDefault(desc);
}

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