Add track plugin to TICL candidate producer

RecoHGCal/TICL/interface/TracksterMomentumPluginBase.h

@@ -12,15 +12,16 @@
 namespace ticl {
   class TracksterMomentumPluginBase {
   public:
-    // TracksterMomentumPluginBase() {}
     TracksterMomentumPluginBase(const edm::ParameterSet&, edm::ConsumesCollector&& iC) {}
     typedef reco::Candidate::LorentzVector LorentzVector;
-    virtual ~TracksterMomentumPluginBase(){};
+    virtual ~TracksterMomentumPluginBase() {}
     virtual LorentzVector calcP4(const ticl::Trackster& trackster) const = 0;
 
+    // Allow access to event so data products can be read once at beginning of event but four-vector calculation can 
+    // be run on single tracksters
     edm::Event& evt() { return *evt_; }
 
-    // To be 
+    // Needs to be called by CMSSW plugins using this plugin before being able to use other methods
     void beginEvent(edm::Event& evt) {
        evt_ = &evt;
        this->beginEvt();

RecoHGCal/TICL/interface/TracksterTrackPluginBase.h

@@ -0,0 +1,38 @@
+#ifndef RecoHGCal_TICL_TracksterTrackPluginBase_H__
+#define RecoHGCal_TICL_TracksterTrackPluginBase_H__
+
+#include "DataFormats/Candidate/interface/Candidate.h"
+#include "DataFormats/VertexReco/interface/Vertex.h"
+#include "DataFormats/CaloRecHit/interface/CaloCluster.h"
+#include "DataFormats/HGCalReco/interface/Trackster.h"
+#include "DataFormats/HGCalReco/interface/TICLCandidate.h"
+#include "FWCore/Framework/interface/ConsumesCollector.h"
+#include "FWCore/Framework/interface/Event.h"
+#include "FWCore/PluginManager/interface/PluginFactory.h"
+
+namespace ticl {
+  class TracksterTrackPluginBase {
+  public:
+    TracksterTrackPluginBase(const edm::ParameterSet&, edm::ConsumesCollector&& iC) {}
+    typedef reco::Candidate::LorentzVector LorentzVector;
+    virtual ~TracksterTrackPluginBase() {}
+    virtual void setTrack(const ticl::Trackster& trackster, ticl::TICLCandidate& ticl_cand) const = 0;
+
+    // Allow access to event so data products can be read once at beginning of event but four-vector calculation can 
+    // be run on single tracksters
+    edm::Event& evt() const { return *evt_; }
+
+    // Needs to be called by CMSSW plugins using this plugin before being able to use other methods
+    void beginEvent(edm::Event& evt) {
+       evt_ = &evt;
+       this->beginEvt();
+     }
+  private:
+    virtual void beginEvt() = 0;
+    edm::Event* evt_;
+  };
+} // namespace
+
+typedef edmplugin::PluginFactory<ticl::TracksterTrackPluginBase*(const edm::ParameterSet&, edm::ConsumesCollector&& iC)> TracksterTrackPluginFactory;
+
+#endif

RecoHGCal/TICL/plugins/PFTICLProducer.cc

@@ -47,14 +47,14 @@ void PFTICLProducer::produce(edm::StreamID, edm::Event& evt, const edm::EventSet
   const auto ticl_candidates = *ticl_cand_h;
 
   auto candidates = std::make_unique<reco::PFCandidateCollection>();
-  // in good particle flow fashion, start from the tracks and go out
+
   for (const auto& ticl_cand : ticl_candidates) {
-    const auto absPdgId = std::abs(ticl_cand.pdgId());
+    const auto abs_pdg_id = std::abs(ticl_cand.pdgId());
     const auto charge = ticl_cand.charge();
     const auto four_mom = ticl_cand.p4();
 
     reco::PFCandidate::ParticleType part_type;
-    switch (absPdgId) {
+    switch (abs_pdg_id) {
       case 11:
         part_type = reco::PFCandidate::e;
         break;

RecoHGCal/TICL/plugins/TICLCandidateFromTrackstersProducer.cc

@@ -17,6 +17,7 @@
 #include "DataFormats/HGCalReco/interface/TICLCandidate.h"
 #include "DataFormats/HGCalReco/interface/Trackster.h"
 #include "RecoHGCal/TICL/interface/TracksterMomentumPluginBase.h"
+#include "RecoHGCal/TICL/interface/TracksterTrackPluginBase.h"
 
 using namespace ticl;
 
@@ -31,6 +32,7 @@ class TICLCandidateFromTrackstersProducer : public edm::stream::EDProducer<> {
 private:
   std::vector<edm::EDGetTokenT<std::vector<Trackster>>> trackster_tokens_;
   std::unique_ptr<TracksterMomentumPluginBase> momentum_algo_;
+  std::unique_ptr<TracksterTrackPluginBase> track_algo_;
 };
 DEFINE_FWK_MODULE(TICLCandidateFromTrackstersProducer);
 
@@ -40,9 +42,10 @@ namespace {
       case 0:
         return 22;
       case 1:
-        return 11;
+        // Pick IDs with positive charge so they can be reset with charge * id downstream 
+        return -11;
       case 2:
-        return 13;
+        return -13;
       case 3:
         return 211;
       case 4:
@@ -59,21 +62,28 @@ TICLCandidateFromTrackstersProducer::TICLCandidateFromTrackstersProducer(const e
   produces<std::vector<TICLCandidate>>();
   auto pset_momentum = ps.getParameter<edm::ParameterSet>("momentumPlugin");
   momentum_algo_ = TracksterMomentumPluginFactory::get()->create(pset_momentum.getParameter<std::string>("plugin"), pset_momentum, consumesCollector());
+
+  auto pset_track = ps.getParameter<edm::ParameterSet>("trackPlugin");
+  track_algo_ = TracksterTrackPluginFactory::get()->create(pset_track.getParameter<std::string>("plugin"), pset_track, consumesCollector());
 }
 
 void TICLCandidateFromTrackstersProducer::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
   edm::ParameterSetDescription desc;
 
-  std::vector<edm::InputTag> source_vector{edm::InputTag("TrackstersMIP"), edm::InputTag("Tracksters")};
+  std::vector<edm::InputTag> source_vector{edm::InputTag("trackstersTrk"), edm::InputTag("trackstersMIP"), edm::InputTag("tracksters")};
   desc.add<std::vector<edm::InputTag>>("tracksterCollections", source_vector);
 
   edm::ParameterSetDescription desc_momentum;
   desc_momentum.add<std::string>("plugin", "TracksterP4FromEnergySum");
   desc_momentum.add<bool>("energyFromRegression", false);
   desc_momentum.add<edm::InputTag>("vertices", edm::InputTag("offlinePrimaryVertices"));
   desc_momentum.add<edm::InputTag>("layerClusters", edm::InputTag("hgcalLayerClusters"));
-
   desc.add<edm::ParameterSetDescription>("momentumPlugin", desc_momentum);
+
+  edm::ParameterSetDescription desc_track;
+  desc_track.add<std::string>("plugin", "TracksterRecoTrackPlugin");
+  desc.add<edm::ParameterSetDescription>("trackPlugin", desc_track);
+
   descriptions.add("ticlCandidateFromTrackstersProducer", desc);
 }
 
@@ -82,6 +92,7 @@ void TICLCandidateFromTrackstersProducer::produce(edm::Event& evt, const edm::Ev
   auto output_mask = std::make_unique<std::vector<float>>();
 
   momentum_algo_->beginEvent(evt);
+  track_algo_->beginEvent(evt);
 
   for (auto& trackster_token : trackster_tokens_) {
     edm::Handle<std::vector<Trackster>> trackster_h;
@@ -94,16 +105,24 @@ void TICLCandidateFromTrackstersProducer::produce(edm::Event& evt, const edm::Ev
       auto max_index = std::distance(id_prob_begin, std::max_element(id_prob_begin,
       trackster.id_probabilities.end()));
       auto pdg_id = pdg_id_from_idx(max_index);
-      int charge = 0;
-      if (pdg_id == 11 || pdg_id == 13 || pdg_id == 211) {
-        charge = 1;
-      }
 
       auto p4 = momentum_algo_->calcP4(trackster);
-      result->emplace_back(reco::LeafCandidate::Charge(charge), p4);
+
+      result->emplace_back(reco::LeafCandidate::Charge(0), p4);
 
       auto& ticl_cand = result->back();
       ticl_cand.setPdgId(pdg_id);
+
+      track_algo_->setTrack(trackster, ticl_cand);
+
+      if (ticl_cand.track_ref().isNonnull()) {
+        auto charge = ticl_cand.track_ref()->charge();
+        ticl_cand.setCharge(charge);
+        ticl_cand.setPdgId(pdg_id*charge);
+      } else if (pdg_id == -11 || pdg_id == -13 || pdg_id == 211) {
+        // FIXME - placeholder for downstream PF code to work, but proper symmetric charge assignment needed
+        ticl_cand.setCharge(1);
+      }
     }
   }
 

RecoHGCal/TICL/plugins/TracksterGeneralTrackPlugin.cc

@@ -0,0 +1,45 @@
+#include "DataFormats/Math/interface/LorentzVector.h"
+#include "FWCore/Framework/interface/ConsumesCollector.h"
+#include "FWCore/ParameterSet/interface/ParameterSetDescription.h"
+#include "RecoHGCal/TICL/interface/TracksterTrackPluginBase.h"
+
+namespace ticl {
+  template<class TrackClass>
+  class TracksterGeneralTrackPlugin final : public TracksterTrackPluginBase {
+  public:
+    typedef edm::Ref<std::vector<TrackClass> > TrackRef;
+    explicit TracksterGeneralTrackPlugin<TrackClass>(const edm::ParameterSet&, edm::ConsumesCollector&& iC);
+    void setTrack(const ticl::Trackster& trackster, ticl::TICLCandidate& ticl_cand) const override;
+  private:
+    void beginEvt() override;
+  };
+
+  template<class TrackClass>
+  TracksterGeneralTrackPlugin<TrackClass>::TracksterGeneralTrackPlugin(const edm::ParameterSet& ps, edm::ConsumesCollector&& ic) : 
+    TracksterTrackPluginBase(ps, std::move(ic)) {// ,
+  }
+
+  template<class TrackClass>
+  void TracksterGeneralTrackPlugin<TrackClass>::beginEvt() {
+  }
+
+  template<class TrackClass>
+  void TracksterGeneralTrackPlugin<TrackClass>::setTrack(const ticl::Trackster& trackster, ticl::TICLCandidate& ticl_cand) const {
+    if (trackster.seedIndex == 0 || !trackster.seedID.isValid()) {
+      return; // leave default empty track ref
+    }
+
+    TrackRef ref(trackster.seedID, trackster.seedIndex, &this->evt().productGetter());
+    ticl_cand.set_track_ref(ref);
+  }
+
+  template class TracksterGeneralTrackPlugin<reco::Track>;
+  typedef TracksterGeneralTrackPlugin<reco::Track> TracksterRecoTrackPlugin;
+
+} // namespace
+
+
+#include "FWCore/Framework/interface/MakerMacros.h"
+DEFINE_EDM_PLUGIN(TracksterTrackPluginFactory,
+                  ticl::TracksterRecoTrackPlugin,
+                  "TracksterRecoTrackPlugin");

RecoHGCal/TICL/src/TracksterTrackPluginBase.cc

@@ -0,0 +1,4 @@
+#include "RecoHGCal/TICL/interface/TracksterTrackPluginBase.h"
+#include "FWCore/Framework/interface/MakerMacros.h"
+
+EDM_REGISTER_PLUGINFACTORY(TracksterTrackPluginFactory, "TracksterTrackPluginFactory");