Make retrieving param indices free function

key4hep · Apr 19, 2024 · 432629e · 432629e
1 parent 672e570
commit 432629e
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 10 deletions.
diff --git a/utils/include/edm4hep/utils/ParticleIDUtils.h b/utils/include/edm4hep/utils/ParticleIDUtils.h
@@ -21,6 +21,9 @@ struct ParticleIDMeta {
   std::vector<std::string> paramNames{}; ///< The names of the parameters
 };
 
+/// Get the index of the parameter in the passed ParticleID meta info
+std::optional<int> getParamIndex(const ParticleIDMeta& pidMetaInfo, const std::string& param);
+
 /// Utility class to invert the ParticleID to ReconstructedParticle relation
 class PIDHandler {
 
@@ -30,8 +33,8 @@ class PIDHandler {
 
   std::map<std::string, int> m_algoTypes{}; ///< Maps algo names to algo types
 
-  /// Maps algo types to the parameter names for each algorithm
-  std::map<int, std::vector<std::string>> m_algoParamNames{};
+  /// Maps algo types to the full particle id meta information
+  std::map<int, edm4hep::utils::ParticleIDMeta> m_algoPidMeta{};
 
 public:
   PIDHandler() = default;

diff --git a/utils/src/ParticleIDUtils.cc b/utils/src/ParticleIDUtils.cc
@@ -11,6 +11,14 @@
 
 namespace edm4hep::utils {
 
+std::optional<int> getParamIndex(const ParticleIDMeta& pidMetaInfo, const std::string& param) {
+  const auto nameIt = std::find(pidMetaInfo.paramNames.begin(), pidMetaInfo.paramNames.end(), param);
+  if (nameIt != pidMetaInfo.paramNames.end()) {
+    return std::distance(pidMetaInfo.paramNames.begin(), nameIt);
+  }
+  return std::nullopt;
+}
+
 void PIDHandler::addColl(const edm4hep::ParticleIDCollection& coll) {
   for (const auto pid : coll) {
     m_recoPidMap.emplace(pid.getParticle(), pid);
@@ -28,8 +36,8 @@ void PIDHandler::addMetaInfo(const edm4hep::utils::ParticleIDMeta& pidInfo) {
     throw std::runtime_error("Cannot have duplicate algorithm names (" + pidInfo.algoName + " already exists)");
   }
 
-  const auto [__, parInserted] = m_algoParamNames.emplace(pidInfo.algoType, pidInfo.paramNames);
-  if (!parInserted) {
+  const auto [__, metaInserted] = m_algoPidMeta.emplace(pidInfo.algoType, pidInfo);
+  if (!metaInserted) {
     if (inserted) {
       m_algoTypes.erase(algoIt);
     }
@@ -79,12 +87,8 @@ std::optional<edm4hep::ParticleID> PIDHandler::getPID(const edm4hep::Reconstruct
 }
 
 std::optional<int> PIDHandler::getParamIndex(int32_t algoType, const std::string& paramName) const {
-  if (const auto it = m_algoParamNames.find(algoType); it != m_algoParamNames.end()) {
-    const auto& names = it->second;
-    const auto nameIt = std::find(names.begin(), names.end(), paramName);
-    if (nameIt != names.end()) {
-      return std::distance(names.begin(), nameIt);
-    }
+  if (const auto it = m_algoPidMeta.find(algoType); it != m_algoPidMeta.end()) {
+    return edm4hep::utils::getParamIndex(it->second, paramName);
   }
 
   return std::nullopt;