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 "ActsExamples/EventData/SimParticle.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<SimBarcode>& 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<SimBarcode>& constituents() const { return m_constituents; }
0109
0110   /// @brief Set the tracks associated to this truth jet
0111   /// @param associatedTracks Vector of tracks associated to this jet
0112   void setAssociatedTracks(
0113       const std::vector<Acts::AnyConstTrackProxy>& associatedTracks) {
0114     m_associatedTracks = associatedTracks;
0115   }
0116
0117   /// @brief Get the tracks associated to this truth jet
0118   /// @return Vector of tracks associated to this jet
0119   const std::vector<Acts::AnyConstTrackProxy>& associatedTracks() const {
0120     return m_associatedTracks;
0121   }
0122
0123  private:
0124   /// @brief  Truth particles as the constituents of the truth jet
0125   std::vector<SimBarcode> m_constituents;
0126   /// @brief Indices of the constituents in the input collection
0127   std::vector<int> m_constituentIndices;
0128   /// @brief The tracks associated to this truth jet
0129   std::vector<Acts::AnyConstTrackProxy> m_associatedTracks;
0130 };
0131
0132 /// @brief Jet constructed from reconstructed tracks
0133 class TrackJet : public Jet {
0134  public:
0135   /// Constructor
0136   /// @param fourMom The jet 4-momentum
0137   /// @param label The jet label
0138   explicit TrackJet(const Acts::Vector4& fourMom, const JetLabel& label)
0139       : Jet(fourMom, label) {}
0140
0141   /// @brief Set the tracks as constituents of this track jet
0142   /// @param constituents Vector of tracks that constitute this jet
0143   void setConstituents(
0144       const std::vector<Acts::AnyConstTrackProxy>& constituents) {
0145     m_constituents = constituents;
0146   }
0147
0148   /// @brief Get the track jet constituents that are tracks
0149   /// @return Vector of tracks that constitute this jet
0150   const std::vector<Acts::AnyConstTrackProxy>& constituents() const {
0151     return m_constituents;
0152   }
0153
0154  private:
0155   /// @brief Tracks as the constituents of the track jet
0156   std::vector<Acts::AnyConstTrackProxy> m_constituents;
0157 };
0158
0159 /// @}
0160 }  // namespace ActsExamples