edm4hep/utils/ParticleIDUtils.h

0001 #ifndef EDM4HEP_UTILS_PARTICLEIDUTILS_H
0002 #define EDM4HEP_UTILS_PARTICLEIDUTILS_H
0003
0004 #include <edm4hep/ParticleIDCollection.h>
0005 #include <edm4hep/ReconstructedParticle.h>
0006
0007 #include <podio/Frame.h>
0008
0009 #include <map>
0010 #include <optional>
0011 #include <string>
0012 #include <vector>
0013
0014 namespace edm4hep::utils {
0015
0016 /// A simple struct bundling relevant metadata for a ParticleID collection
0017 struct ParticleIDMeta {
0018   ParticleIDMeta(const std::string& algName, int32_t algType, const std::vector<std::string>& parNames);
0019   ParticleIDMeta(const std::string& algName, const std::vector<std::string>& parNames);
0020
0021   ~ParticleIDMeta() = default;
0022   ParticleIDMeta() = default;
0023   ParticleIDMeta(const ParticleIDMeta&) = default;
0024   ParticleIDMeta& operator=(const ParticleIDMeta&) = default;
0025   ParticleIDMeta(ParticleIDMeta&&) = default;
0026   ParticleIDMeta& operator=(ParticleIDMeta&&) = default;
0027
0028   std::string algoName{};                ///< The name of the algorithm
0029   std::vector<std::string> paramNames{}; ///< The names of the parameters
0030
0031   int32_t algoType() const {
0032     return m_algoType;
0033   }
0034
0035 private:
0036   int32_t m_algoType{0}; ///< The (user defined) algorithm type
0037 };
0038
0039 /// Get the index of the parameter in the passed ParticleID meta info
0040 std::optional<int> getParamIndex(const ParticleIDMeta& pidMetaInfo, const std::string& param);
0041
0042 /// Utility class to invert the ParticleID to ReconstructedParticle relation
0043 /// See [this page](@ref md_doc_2_p_i_d_handler) for example usage and more information.
0044 class PIDHandler {
0045
0046   using RecoPidMapT = std::multimap<edm4hep::ReconstructedParticle, edm4hep::ParticleID>;
0047
0048   RecoPidMapT m_recoPidMap{}; ///< The internal map from recos to pids
0049
0050   std::map<std::string, int> m_algoTypes{}; ///< Maps algo names to algo types
0051
0052   /// Maps algo types to the full particle id meta information
0053   std::map<int, edm4hep::utils::ParticleIDMeta> m_algoPidMeta{};
0054
0055 public:
0056   PIDHandler() = default;
0057   ~PIDHandler() = default;
0058   PIDHandler(const PIDHandler&) = default;
0059   PIDHandler& operator=(const PIDHandler&) = default;
0060   PIDHandler(PIDHandler&&) = default;
0061   PIDHandler& operator=(PIDHandler&&) = default;
0062
0063   /// Construct a PIDHandler from an arbitrary number of ParticleIDCollections
0064   template <typename... PIDColls>
0065   static PIDHandler from(const ParticleIDCollection& coll, const PIDColls&... pidColls) {
0066     static_assert((std::is_same_v<PIDColls, edm4hep::ParticleIDCollection> && ...),
0067                   "PIDHandler can only be constructed from ParticleIDCollections");
0068     PIDHandler handler{};
0069     handler.addColl(coll);
0070     (handler.addColl(pidColls), ...);
0071     return handler;
0072   }
0073
0074   /// Create a PIDHandler from a Frame potentially also populating some metadata information.
0075   static PIDHandler from(const podio::Frame& event, const podio::Frame& metadata = {});
0076
0077   /// Add the information from one ParticleIDCollection to the handler
0078   void addColl(const edm4hep::ParticleIDCollection& coll);
0079
0080   /// Add the information from one ParticleIDCollection to the handler together
0081   /// with its meta data
0082   void addColl(const edm4hep::ParticleIDCollection& coll, const edm4hep::utils::ParticleIDMeta& pidInfo);
0083
0084   /// Add meta information for a collection
0085   void addMetaInfo(const edm4hep::utils::ParticleIDMeta& pidInfo);
0086
0087   /// Retrieve all ParticleIDs that are related to the passed
0088   /// ReconstructedParticle
0089   std::vector<edm4hep::ParticleID> getPIDs(const edm4hep::ReconstructedParticle& reco) const;
0090
0091   /// Retrieve the ParticleID for a given algorithm type
0092   std::optional<edm4hep::ParticleID> getPID(const edm4hep::ReconstructedParticle& reco, int algoType) const;
0093
0094   /// Retrieve the index in the parameters for a given parameter name and
0095   /// algoType
0096   std::optional<int> getParamIndex(int32_t algoType, const std::string& paramName) const;
0097
0098   /// Retrieve the algoType for a given algorithm name
0099   std::optional<int32_t> getAlgoType(const std::string& algoName) const;
0100
0101   /// Set the metadata information for the passed collection in the metadata Frame.
0102   ///
0103   /// This also sets the algorithmType of all elements in the collection to the
0104   /// one that is found in the meta information.
0105   static void setAlgoInfo(podio::Frame& metadata, edm4hep::ParticleIDCollection& pidcoll, const std::string& collname,
0106                           const edm4hep::utils::ParticleIDMeta& pidMetaInfo);
0107
0108   /// Set the metadata information for a given collection name in the metadata Frame.
0109   ///
0110   /// @note It is user responsibility to ensure that the meta information that
0111   /// is passed here and the one that is present in the collection with the
0112   /// given name is consistent
0113   static void setAlgoInfo(podio::Frame& metadata, const std::string& collname,
0114                           const edm4hep::utils::ParticleIDMeta& pidMetaInfo);
0115
0116   /// Get the ParticleID meta information for a given collection name from the metadata Frame.
0117   static std::optional<edm4hep::utils::ParticleIDMeta> getAlgoInfo(const podio::Frame& metadata,
0118                                                                    const std::string& collName);
0119 };
0120 } // namespace edm4hep::utils
0121
0122 #endif // EDM4HEP_UTILS_PARTICLEIDUTILS_H