TICL/plugins/TracksterP4FromEnergySumPlugin.cc

0001 // Plugin for getting the four-vector from a Trackster from a simple energy sum and weighted cluster position.
0002 // A simplistic 1/N(tracksters) sharing is applied for hits that belong to multiple tracksters.
0003 // Alternatively takes the energy value from the pre-calculated regressed energy value in the Trackster.
0004
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 TracksterP4FromEnergySum final : public TracksterMomentumPluginBase {
0012   public:
0013     explicit TracksterP4FromEnergySum(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(
0020         const ticl::Trackster& trackster,
0021         const reco::Vertex& vertex,
0022         const std::vector<reco::CaloCluster>& calo_clusters) const;
0023     bool energy_from_regression_;
0024     edm::EDGetTokenT<std::vector<reco::Vertex>> vertex_token_;
0025     edm::EDGetTokenT<std::vector<reco::CaloCluster>> layer_clusters_token_;
0026   };
0027
0028   TracksterP4FromEnergySum::TracksterP4FromEnergySum(const edm::ParameterSet& ps, edm::ConsumesCollector&& ic)
0029       : TracksterMomentumPluginBase(ps, std::move(ic)),
0030         energy_from_regression_(ps.getParameter<bool>("energyFromRegression")),
0031         vertex_token_(ic.consumes<std::vector<reco::Vertex>>(ps.getParameter<edm::InputTag>("vertices"))),
0032         layer_clusters_token_(
0033             ic.consumes<std::vector<reco::CaloCluster>>(ps.getParameter<edm::InputTag>("layerClusters"))) {}
0034
0035   void TracksterP4FromEnergySum::setP4(const std::vector<const Trackster*>& tracksters,
0036                                        std::vector<TICLCandidate>& ticl_cands,
0037                                        edm::Event& event) const {
0038     // Find best vertex
0039     edm::Handle<std::vector<reco::Vertex>> vertex_h;
0040     event.getByToken(vertex_token_, vertex_h);
0041     auto vertex_coll = *vertex_h;
0042     reco::Vertex best_vertex;
0043     if (!vertex_coll.empty()) {
0044       const auto& vertex = vertex_coll[0];
0045       if (vertex.isValid() && !(vertex.isFake())) {
0046         best_vertex = vertex;
0047       }
0048     }
0049
0050     edm::Handle<std::vector<reco::CaloCluster>> layer_clusters_h;
0051     event.getByToken(layer_clusters_token_, layer_clusters_h);
0052
0053     auto size = std::min(tracksters.size(), ticl_cands.size());
0054     for (size_t i = 0; i < size; ++i) {
0055       const auto* trackster = tracksters[i];
0056       auto ret = calcP4(*trackster, best_vertex, *layer_clusters_h);
0057
0058       auto& ticl_cand = ticl_cands[i];
0059       ticl_cand.setP4(std::get<0>(ret));
0060       ticl_cand.setRawEnergy(std::get<1>(ret));
0061     }
0062   }
0063
0064   std::tuple<TracksterMomentumPluginBase::LorentzVector, float> TracksterP4FromEnergySum::calcP4(
0065       const ticl::Trackster& trackster,
0066       const reco::Vertex& vertex,
0067       const std::vector<reco::CaloCluster>& calo_clusters) const {
0068     std::array<double, 3> barycentre{{0., 0., 0.}};
0069     double energy = 0.;
0070     size_t counter = 0;
0071
0072     for (auto idx : trackster.vertices()) {
0073       auto n_vertices = trackster.vertex_multiplicity(counter++);
0074       auto fraction = n_vertices ? 1.f / n_vertices : 1.f;
0075       auto weight = calo_clusters[idx].energy() * fraction;
0076       energy += weight;
0077       barycentre[0] += calo_clusters[idx].x() * weight;
0078       barycentre[1] += calo_clusters[idx].y() * weight;
0079       barycentre[2] += calo_clusters[idx].z() * weight;
0080     }
0081     std::transform(
0082         std::begin(barycentre), std::end(barycentre), std::begin(barycentre), [&energy](double val) -> double {
0083           return energy >= 0. ? val / energy : val;
0084         });
0085
0086     math::XYZVector direction(barycentre[0] - vertex.x(), barycentre[1] - vertex.y(), barycentre[2] - vertex.z());
0087     direction = direction.Unit();
0088     auto raw_energy = energy;
0089     energy = energy_from_regression_ ? trackster.regressed_energy() : raw_energy;
0090     direction *= energy;
0091
0092     math::XYZTLorentzVector cartesian(direction.X(), direction.Y(), direction.Z(), energy);
0093     // Convert px, py, pz, E vector to CMS standard pt/eta/phi/m vector
0094     TracksterP4FromEnergySum::LorentzVector p4(cartesian);
0095     return std::tuple(p4, raw_energy);
0096   }
0097 }  // namespace ticl
0098
0099 #include "FWCore/Framework/interface/MakerMacros.h"
0100 DEFINE_EDM_PLUGIN(TracksterMomentumPluginFactory, ticl::TracksterP4FromEnergySum, "TracksterP4FromEnergySum");