IsolatedTrackCleaner.cc

Go to the documentation of this file.

#include "FWCore/Framework/interface/stream/EDProducer.h"
#include "FWCore/Framework/interface/Event.h"
#include "FWCore/MessageLogger/interface/MessageLogger.h"
#include "FWCore/ParameterSet/interface/ParameterSet.h"
#include "DataFormats/Candidate/interface/Candidate.h"
#include "DataFormats/PatCandidates/interface/IsolatedTrack.h"
#include "DataFormats/Common/interface/RefToPtr.h"
#include "CommonTools/Utils/interface/StringCutObjectSelector.h"

#include <vector>


class IsolatedTrackCleaner : public edm::stream::EDProducer<> {
    public:
        IsolatedTrackCleaner( edm::ParameterSet const & params ) :
            tracks_(consumes<std::vector<pat::IsolatedTrack>>(params.getParameter<edm::InputTag>("tracks"))),
            cut_(params.getParameter<std::string>("cut"))
        {
            for (const edm::InputTag & tag : params.getParameter<std::vector<edm::InputTag>>("finalLeptons")) {
                leptons_.push_back(consumes<reco::CandidateView>(tag));
            }
            produces<std::vector<pat::IsolatedTrack>>();
        }

        ~IsolatedTrackCleaner() override {}

        void produce(edm::Event& iEvent, const edm::EventSetup& iSetup) override {
            auto out  = std::make_unique<std::vector<pat::IsolatedTrack>>();

            std::vector<reco::CandidatePtr> leptonPfCands;
            edm::Handle<reco::CandidateView> leptons;
            for (const auto & token : leptons_) {
                iEvent.getByToken(token, leptons);
                for (const auto & lep : *leptons) {
                    for (unsigned int i = 0, n = lep.numberOfSourceCandidatePtrs(); i < n; ++i) {
                        auto ptr = lep.sourceCandidatePtr(i);
                        if (ptr.isNonnull()) leptonPfCands.push_back(ptr);
                    }
                }
            }
            std::sort(leptonPfCands.begin(), leptonPfCands.end());

            edm::Handle<std::vector<pat::IsolatedTrack>> tracks;
            iEvent.getByToken(tracks_, tracks);
            for (const auto & track : *tracks) {
                if (!cut_(track)) continue; 
                if (track.packedCandRef().isNonnull()) {
                    reco::CandidatePtr pfCand(edm::refToPtr(track.packedCandRef()));
                    if (std::binary_search(leptonPfCands.begin(), leptonPfCands.end(), pfCand)) {
                        continue;
                    }
                }
                out->push_back(track);
            }

            iEvent.put(std::move(out));
        }

    protected:
        edm::EDGetTokenT<std::vector<pat::IsolatedTrack>> tracks_;
        StringCutObjectSelector<pat::IsolatedTrack> cut_;
        std::vector<edm::EDGetTokenT<reco::CandidateView>> leptons_;
};

#include "FWCore/Framework/interface/MakerMacros.h"
DEFINE_FWK_MODULE(IsolatedTrackCleaner);