ActsExamples/EventData/Jets.hpp

0001 // This file is part of the ACTS project.
0002 //
0003 // Copyright (C) 2016 CERN for the benefit of the ACTS project
0004 //
0005 // This Source Code Form is subject to the terms of the Mozilla Public
0006 // License, v. 2.0. If a copy of the MPL was not distributed with this
0007 // file, You can obtain one at https://mozilla.org/MPL/2.0/.
0008
0009 #pragma once
0010
0011 #include "Acts/Definitions/Algebra.hpp"
0012 #include "Acts/EventData/AnyTrackProxy.hpp"
0013 #include "ActsFatras/EventData/Barcode.hpp"
0014
0015 #include <vector>
0016
0017 namespace ActsExamples {
0018
0019 /// Jet labeling enum
0020
0021 enum class JetLabel {
0022   Unknown = 0,
0023   BJet = 5,
0024   CJet = 4,
0025   LightJet = 1,
0026 };
0027
0028 /// Output stream operator for JetLabel
0029 /// @param os The output stream
0030 /// @param label The jet label to output
0031 /// @return The output stream
0032 inline std::ostream& operator<<(std::ostream& os, const JetLabel& label) {
0033   switch (label) {
0034     case JetLabel::Unknown:
0035       os << "Unknown";
0036       break;
0037     case JetLabel::BJet:
0038       os << "BJet";
0039       break;
0040     case JetLabel::CJet:
0041       os << "CJet";
0042       break;
0043     case JetLabel::LightJet:
0044       os << "LightJet";
0045       break;
0046     default:
0047       os << "InvalidJetLabel";
0048       break;
0049   }
0050   return os;
0051 }
0052
0053 /// Common class for jets
0054 class Jet {
0055  public:
0056   /// Constructor
0057   /// @param fourMom The jet 4-momentum
0058   /// @param label The jet label
0059   Jet(const Acts::Vector4& fourMom, const JetLabel& label)
0060       : m_fourMomentum{fourMom}, m_jetLabel{label} {}
0061
0062   /// @brief Get the jet 4-momentum
0063   /// @return the jet 4-momentum as an Acts::Vector4
0064   Acts::Vector4 fourMomentum() const { return m_fourMomentum; }
0065
0066   /// @brief Get the jet label
0067   /// @return the jet label as JetLabel enum
0068   JetLabel jetLabel() const { return m_jetLabel; }
0069
0070   /// @brief Set the jet label
0071   /// @param label the jet label as JetLabel enum
0072   void setJetLabel(const JetLabel& label) { m_jetLabel = label; }
0073
0074  private:
0075   Acts::Vector4 m_fourMomentum{Acts::Vector4::Zero()};
0076   JetLabel m_jetLabel{JetLabel::Unknown};
0077   /// @brief Print the jet information
0078   friend std::ostream& operator<<(std::ostream& os, const Jet& jet) {
0079     os << "Jet 4-momentum: " << jet.fourMomentum().transpose() << std::endl;
0080     os << "Jet label: " << jet.jetLabel() << std::endl;
0081     return os;
0082   }
0083 };
0084
0085 /// @brief Jet constructed from truth particles
0086 class TruthJet : public Jet {
0087  public:
0088   /// Constructor
0089   /// @param fourMom The jet 4-momentum
0090   /// @param label The jet label
0091   TruthJet(const Acts::Vector4& fourMom, const JetLabel& label)
0092       : Jet(fourMom, label) {}
0093
0094   /// @brief Set the truth particles as constituents of this truth jet from its barcode
0095   /// @param constituents Vector of barcodes for constituent particles
0096   void setConstituents(const std::vector<ActsFatras::Barcode>& constituents) {
0097     m_constituents = constituents;
0098   }
0099
0100   /// @brief Set the indices of the truth particles that are constituents of this truth jet
0101   /// @param indices
0102   void setConstituentIndices(const std::vector<int>& indices) {
0103     m_constituentIndices = indices;
0104   }
0105
0106   /// @brief Get the truth particles that are truth jet constituents
0107   /// @return Vector of barcodes for constituent particles
0108   const std::vector<ActsFatras::Barcode>& constituents() const {
0109     return m_constituents;
0110   }
0111
0112   /// @brief Set the tracks associated to this truth jet
0113   /// @param associatedTracks Vector of tracks associated to this jet
0114   void setAssociatedTracks(
0115       const std::vector<Acts::AnyConstTrackProxy>& associatedTracks) {
0116     m_associatedTracks = associatedTracks;
0117   }
0118
0119   /// @brief Get the tracks associated to this truth jet
0120   /// @return Vector of tracks associated to this jet
0121   const std::vector<Acts::AnyConstTrackProxy>& associatedTracks() const {
0122     return m_associatedTracks;
0123   }
0124
0125  private:
0126   /// @brief  Truth particles as the constituents of the truth jet
0127   std::vector<ActsFatras::Barcode> m_constituents;
0128   /// @brief Indices of the constituents in the input collection
0129   std::vector<int> m_constituentIndices;
0130   /// @brief The tracks associated to this truth jet
0131   std::vector<Acts::AnyConstTrackProxy> m_associatedTracks;
0132 };
0133
0134 /// @brief Jet constructed from reconstructed tracks
0135 class TrackJet : public Jet {
0136  public:
0137   /// Constructor
0138   /// @param fourMom The jet 4-momentum
0139   /// @param label The jet label
0140   explicit TrackJet(const Acts::Vector4& fourMom, const JetLabel& label)
0141       : Jet(fourMom, label) {}
0142
0143   /// @brief Set the tracks as constituents of this track jet
0144   /// @param constituents Vector of tracks that constitute this jet
0145   void setConstituents(
0146       const std::vector<Acts::AnyConstTrackProxy>& constituents) {
0147     m_constituents = constituents;
0148   }
0149
0150   /// @brief Get the track jet constituents that are tracks
0151   /// @return Vector of tracks that constitute this jet
0152   const std::vector<Acts::AnyConstTrackProxy>& constituents() const {
0153     return m_constituents;
0154   }
0155
0156  private:
0157   /// @brief Tracks as the constituents of the track jet
0158   std::vector<Acts::AnyConstTrackProxy> m_constituents;
0159 };
0160
0161 /// @}
0162 }  // namespace ActsExamples