TICL/plugins/TICLCandidateFromTrackstersProducer.cc

0001 // user include files
0002 #include <algorithm>
0003 #include <vector>
0004
0005 #include "FWCore/Framework/interface/ESHandle.h"
0006 #include "FWCore/Framework/interface/Event.h"
0007 #include "FWCore/Framework/interface/Frameworkfwd.h"
0008 #include "FWCore/Framework/interface/MakerMacros.h"
0009 #include "FWCore/Framework/interface/stream/EDProducer.h"
0010 #include "FWCore/Framework/interface/ConsumesCollector.h"
0011 #include "FWCore/MessageLogger/interface/MessageLogger.h"
0012 #include "FWCore/ParameterSet/interface/ConfigurationDescriptions.h"
0013 #include "FWCore/ParameterSet/interface/ParameterSetDescription.h"
0014
0015 #include "FWCore/Utilities/interface/transform.h"
0016
0017 #include "DataFormats/HGCalReco/interface/TICLCandidate.h"
0018 #include "DataFormats/HGCalReco/interface/Trackster.h"
0019 #include "RecoHGCal/TICL/interface/TracksterMomentumPluginBase.h"
0020 #include "RecoHGCal/TICL/interface/TracksterTrackPluginBase.h"
0021
0022 using namespace ticl;
0023
0024 class TICLCandidateFromTrackstersProducer : public edm::stream::EDProducer<> {
0025 public:
0026   TICLCandidateFromTrackstersProducer(const edm::ParameterSet&);
0027   ~TICLCandidateFromTrackstersProducer() override {}
0028   static void fillDescriptions(edm::ConfigurationDescriptions& descriptions);
0029
0030   void produce(edm::Event&, const edm::EventSetup&) override;
0031
0032 private:
0033   std::vector<edm::EDGetTokenT<std::vector<Trackster>>> trackster_tokens_;
0034   std::unique_ptr<TracksterMomentumPluginBase> momentum_algo_;
0035   std::unique_ptr<TracksterTrackPluginBase> track_algo_;
0036   float min_particle_prob_;
0037 };
0038 DEFINE_FWK_MODULE(TICLCandidateFromTrackstersProducer);
0039
0040 namespace {
0041   int pdg_id_from_particle_type(ticl::Trackster::ParticleType type) {
0042     switch (type) {
0043       case ticl::Trackster::ParticleType::photon:
0044         return 22;
0045       case ticl::Trackster::ParticleType::electron:
0046         // Pick IDs with positive charge so they can be reset with charge * id downstream
0047         return -11;
0048       case ticl::Trackster::ParticleType::muon:
0049         return -13;
0050       case ticl::Trackster::ParticleType::neutral_pion:
0051         return 111;
0052       case ticl::Trackster::ParticleType::charged_hadron:
0053         return 211;
0054       case ticl::Trackster::ParticleType::neutral_hadron:
0055         return 130;
0056       default:
0057         return 0;
0058     }
0059     return 0;
0060   }
0061 }  // namespace
0062
0063 TICLCandidateFromTrackstersProducer::TICLCandidateFromTrackstersProducer(const edm::ParameterSet& ps)
0064     : min_particle_prob_(ps.getParameter<double>("minParticleProbability")) {
0065   trackster_tokens_ =
0066       edm::vector_transform(ps.getParameter<std::vector<edm::InputTag>>("tracksterCollections"),
0067                             [this](edm::InputTag const& tag) { return consumes<std::vector<Trackster>>(tag); });
0068   produces<std::vector<TICLCandidate>>();
0069   auto pset_momentum = ps.getParameter<edm::ParameterSet>("momentumPlugin");
0070   momentum_algo_ = TracksterMomentumPluginFactory::get()->create(
0071       pset_momentum.getParameter<std::string>("plugin"), pset_momentum, consumesCollector());
0072
0073   auto pset_track = ps.getParameter<edm::ParameterSet>("trackPlugin");
0074   track_algo_ = TracksterTrackPluginFactory::get()->create(
0075       pset_track.getParameter<std::string>("plugin"), pset_track, consumesCollector());
0076 }
0077
0078 void TICLCandidateFromTrackstersProducer::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
0079   edm::ParameterSetDescription desc;
0080
0081   std::vector<edm::InputTag> source_vector{edm::InputTag("trackstersTrk"),
0082                                            edm::InputTag("trackstersMIP"),
0083                                            edm::InputTag("trackstersEM"),
0084                                            edm::InputTag("trackstersHAD")};
0085   desc.add<std::vector<edm::InputTag>>("tracksterCollections", source_vector);
0086   desc.add<double>("minParticleProbability", 0.);
0087
0088   edm::ParameterSetDescription desc_momentum;
0089   desc_momentum.add<std::string>("plugin", "TracksterP4FromEnergySum");
0090   desc_momentum.add<bool>("energyFromRegression", false);
0091   desc_momentum.add<edm::InputTag>("vertices", edm::InputTag("offlinePrimaryVertices"));
0092   desc_momentum.add<edm::InputTag>("layerClusters", edm::InputTag("hgcalMergeLayerClusters"));
0093   desc_momentum.add<edm::InputTag>("tracks", edm::InputTag("generalTracks"));
0094   desc.add<edm::ParameterSetDescription>("momentumPlugin", desc_momentum);
0095
0096   edm::ParameterSetDescription desc_track;
0097   desc_track.add<std::string>("plugin", "TracksterRecoTrackPlugin");
0098   desc.add<edm::ParameterSetDescription>("trackPlugin", desc_track);
0099
0100   descriptions.add("ticlCandidateFromTrackstersProducer", desc);
0101 }
0102
0103 void TICLCandidateFromTrackstersProducer::produce(edm::Event& evt, const edm::EventSetup& es) {
0104   auto result = std::make_unique<std::vector<TICLCandidate>>();
0105
0106   std::vector<const Trackster*> trackster_ptrs;
0107
0108   // adds one TICLCandidate for each trackster that has a minimum particle probability
0109   for (auto& trackster_token : trackster_tokens_) {
0110     edm::Handle<std::vector<Trackster>> trackster_h;
0111     evt.getByToken(trackster_token, trackster_h);
0112     for (size_t i_trackster = 0; i_trackster < trackster_h->size(); ++i_trackster) {
0113       auto const& trackster = trackster_h->at(i_trackster);
0114       auto id_prob_begin = std::begin(trackster.id_probabilities());
0115       auto max_id_prob_it = std::max_element(id_prob_begin, std::end(trackster.id_probabilities()));
0116       float max_id_prob = *max_id_prob_it;
0117       if (max_id_prob < min_particle_prob_)
0118         continue;
0119
0120       trackster_ptrs.push_back(&trackster);
0121       auto& ticl_cand = result->emplace_back(edm::Ptr<ticl::Trackster>(trackster_h, i_trackster));
0122       auto max_index = std::distance(id_prob_begin, max_id_prob_it);
0123       auto pdg_id = pdg_id_from_particle_type(static_cast<ticl::Trackster::ParticleType>(max_index));
0124       ticl_cand.setPdgId(pdg_id);
0125     }
0126   }
0127
0128   momentum_algo_->setP4(trackster_ptrs, *result, evt);
0129   track_algo_->setTrack(trackster_ptrs, *result, evt);
0130
0131   // charge assignment
0132   for (size_t i = 0; i < result->size(); ++i) {
0133     auto& ticl_cand = result->at(i);
0134     auto pdg_id = ticl_cand.pdgId();
0135     if (pdg_id == -11 || pdg_id == -13 || pdg_id == 211) {
0136       if (ticl_cand.trackPtr().isNonnull()) {
0137         auto charge = ticl_cand.trackPtr()->charge();
0138         ticl_cand.setCharge(charge);
0139         ticl_cand.setPdgId(pdg_id * charge);
0140         // If this is a special 0 energy trackster that has been injected
0141         // starting from a seed that has not been linked to any shower,
0142         // re-assign the energy and momentum to be the ones of the track linked
0143         // to the original seeding region.
0144         if (ticl_cand.rawEnergy() == 0.) {
0145           auto const& three_mom = ticl_cand.trackPtr()->momentum();
0146           constexpr double mpion2 = 0.13957 * 0.13957;
0147           double energy = std::sqrt(three_mom.mag2() + mpion2);
0148           math::XYZTLorentzVector trk_p4(three_mom.x(), three_mom.y(), three_mom.z(), energy);
0149           ticl_cand.setP4(trk_p4);
0150           ticl_cand.setRawEnergy(energy);
0151         }
0152       } else {
0153         // Demote them to be neutral, since there's not track associated.
0154         ticl_cand.setCharge(0);
0155         if (pdg_id == -11) {
0156           ticl_cand.setPdgId(22);
0157         } else {
0158           ticl_cand.setPdgId(130);
0159         }
0160       }
0161     }
0162   }
0163
0164   evt.put(std::move(result));
0165 }