TICL/plugins/TracksterP4FromTrackAndPCA.cc

0001 // Plugin for getting the four-vector of a Trackster from the track, if
0002 // present, or from PCA decomposition.
0003
0004 #include "FWCore/Framework/interface/Event.h"
0005 #include "DataFormats/Math/interface/LorentzVector.h"
0006 #include "FWCore/Framework/interface/ConsumesCollector.h"
0007 #include "FWCore/ParameterSet/interface/ParameterSetDescription.h"
0008 #include "RecoHGCal/TICL/interface/TracksterMomentumPluginBase.h"
0009
0010 namespace ticl {
0011   class TracksterP4FromTrackAndPCA final : public TracksterMomentumPluginBase {
0012   public:
0013     explicit TracksterP4FromTrackAndPCA(const edm::ParameterSet&, edm::ConsumesCollector&& iC);
0014     void setP4(const std::vector<const Trackster*>& tracksters,
0015                std::vector<TICLCandidate>& ticl_cands,
0016                edm::Event& event) const override;
0017
0018   private:
0019     std::tuple<TracksterMomentumPluginBase::LorentzVector, float> calcP4(const ticl::Trackster& trackster) const;
0020     bool energy_from_regression_;
0021     edm::EDGetTokenT<reco::TrackCollection> tracks_token_;
0022   };
0023
0024   TracksterP4FromTrackAndPCA::TracksterP4FromTrackAndPCA(const edm::ParameterSet& ps, edm::ConsumesCollector&& ic)
0025       : TracksterMomentumPluginBase(ps, std::move(ic)),
0026         energy_from_regression_(ps.getParameter<bool>("energyFromRegression")),
0027         tracks_token_(ic.consumes<reco::TrackCollection>(ps.getParameter<edm::InputTag>("tracks"))) {}
0028
0029   void TracksterP4FromTrackAndPCA::setP4(const std::vector<const Trackster*>& tracksters,
0030                                          std::vector<TICLCandidate>& ticl_cands,
0031                                          edm::Event& event) const {
0032     edm::Handle<reco::TrackCollection> tracks_h;
0033     event.getByToken(tracks_token_, tracks_h);
0034     edm::ProductID trkId = tracks_h.id();
0035     const reco::TrackCollection& trackCollection = *tracks_h.product();
0036
0037     auto size = std::min(tracksters.size(), ticl_cands.size());
0038     for (size_t i = 0; i < size; ++i) {
0039       const auto* trackster = tracksters[i];
0040       // If there's a track, use it.
0041       if (trackster->seedIndex() != -1) {
0042         assert(trackster->seedID() == trkId);
0043         auto const& tkRef = trackCollection[trackster->seedIndex()];
0044         auto const& three_mom = tkRef.momentum();
0045         constexpr double mpion2 = 0.13957 * 0.13957;
0046         double energy = std::sqrt(tkRef.momentum().mag2() + mpion2);
0047         math::XYZTLorentzVector trk_p4(three_mom.x(), three_mom.y(), three_mom.z(), energy);
0048         auto& ticl_cand = ticl_cands[i];
0049         ticl_cand.setP4(trk_p4);
0050         ticl_cand.setRawEnergy(energy);
0051       } else {
0052         auto direction = trackster->eigenvectors(0).Unit();
0053         auto energy = energy_from_regression_ ? trackster->regressed_energy() : trackster->raw_energy();
0054         direction *= energy;
0055         math::XYZTLorentzVector cartesian(direction.X(), direction.Y(), direction.Z(), energy);
0056         auto& ticl_cand = ticl_cands[i];
0057         ticl_cand.setP4(cartesian);
0058         ticl_cand.setRawEnergy(energy);
0059       }
0060     }
0061   }
0062 }  // namespace ticl
0063
0064 #include "FWCore/Framework/interface/MakerMacros.h"
0065 DEFINE_EDM_PLUGIN(TracksterMomentumPluginFactory, ticl::TracksterP4FromTrackAndPCA, "TracksterP4FromTrackAndPCA");