[WIP] Add the possibility to set edm4hep::utils::ParticleIDMeta via the IMetadataSvc

k4FWCore/CMakeLists.txt

@@ -31,7 +31,7 @@ gaudi_install(PYTHON)
 gaudi_add_library(k4FWCore
 		  SOURCES src/PodioDataSvc.cpp
               src/KeepDropSwitch.cpp
-                  LINK Gaudi::GaudiKernel podio::podioIO ROOT::Core ROOT::RIO ROOT::Tree
+                  LINK Gaudi::GaudiKernel podio::podioIO ROOT::Core ROOT::RIO ROOT::Tree EDM4HEP::utils
                   )
 target_include_directories(k4FWCore PUBLIC
   $<BUILD_INTERFACE:${CMAKE_CURRENT_LIST_DIR}/include>

k4FWCore/include/k4FWCore/IMetadataSvc.h

@@ -21,6 +21,8 @@
 
 #include "GaudiKernel/IInterface.h"
 
+#include "edm4hep/utils/ParticleIDUtils.h"
+
 #include "podio/Frame.h"
 
 class IMetadataSvc : virtual public IInterface {
@@ -31,12 +33,7 @@ class IMetadataSvc : virtual public IInterface {
 
   virtual void setFrame(podio::Frame frame) = 0;
 
-  template <typename T> void put(const std::string& name, const T& obj) {
-    if (!getFrame()) {
-      setFrame(podio::Frame{});
-    }
-    getFrame()->putParameter(name, obj);
-  }
+  template <typename T> void put(const std::string& name, const T& obj) { getFrameForWrite()->putParameter(name, obj); }
 
   template <typename T> std::optional<T> get(const std::string& name) const {
     const auto* frame = getFrame();
@@ -49,6 +46,31 @@ class IMetadataSvc : virtual public IInterface {
 protected:
   virtual podio::Frame*       getFrame()       = 0;
   virtual const podio::Frame* getFrame() const = 0;
+
+private:
+  podio::Frame* getFrameForWrite() {
+    if (!getFrame()) {
+      setFrame(podio::Frame());
+    }
+    return getFrame();
+  }
 };
 
+template <>
+inline void IMetadataSvc::put<edm4hep::utils::ParticleIDMeta>(const std::string&                    collName,
+                                                              const edm4hep::utils::ParticleIDMeta& pidMetaInfo) {
+  edm4hep::utils::PIDHandler::setAlgoInfo(*getFrameForWrite(), collName, pidMetaInfo);
+}
+
+template <>
+inline std::optional<edm4hep::utils::ParticleIDMeta> IMetadataSvc::get<edm4hep::utils::ParticleIDMeta>(
+    const std::string& collName) const {
+  const auto* frame = getFrame();
+  if (!frame) {
+    return std::nullopt;
+  }
+
+  return edm4hep::utils::PIDHandler::getAlgoInfo(*frame, collName);
+}
+
 #endif

test/k4FWCoreTest/CMakeLists.txt

@@ -190,6 +190,8 @@ add_test_with_env(FunctionalProducerRNTuple options/ExampleFunctionalProducerRNT
 add_test_with_env(FunctionalTTreeToRNTuple options/ExampleFunctionalTTreeToRNTuple.py PROPERTIES DEPENDS FunctionalProducer ADD_TO_CHECK_FILES)
 add_test_with_env(GaudiFunctional options/ExampleGaudiFunctional.py PROPERTIES DEPENDS FunctionalProducer ADD_TO_CHECK_FILES)
 
+add_test_with_env(ParticleIDMetadataFramework options/ExampleParticleIDMetadata.py)
+
 
 # The following is done to make the tests work without installing the files in
 # the installation directory. The k4FWCore in the build directory is populated by

test/k4FWCoreTest/options/ExampleParticleIDMetadata.py

@@ -0,0 +1,80 @@
+#!/usr/bin/env python3
+#
+# Copyright (c) 2014-2024 Key4hep-Project.
+#
+# This file is part of Key4hep.
+# See https://key4hep.github.io/key4hep-doc/ for further info.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Example showcasing how to use ParticleID related metadata"""
+
+from Gaudi.Configuration import INFO
+from Configurables import (
+    ExampleParticleIDProducer,
+    ExampleParticleIDConsumer,
+    ExampleFunctionalProducerMultiple,
+    EventDataSvc,
+)
+from k4FWCore import ApplicationMgr, IOSvc
+
+# NOTE: If you are not using the IOSvc (e.g. because you don't need I/O), make
+# sure to add the MetadataSvc to the ExtSvc as that is necessary to store /
+# retrieve the metadata for ParticleIDs
+iosvc = IOSvc()
+iosvc.Output = "example_with_particleids.root"
+iosvc.outputCommands = ["drop *", "keep RecoParticles*"]
+
+reco_producer = ExampleFunctionalProducerMultiple(
+    "RecoProducer", OutputCollectionRecoParticles=["RecoParticles"]
+)
+
+pid_producer1 = ExampleParticleIDProducer(
+    "PIDProducer1",
+    InputCollection=["RecoParticles"],
+    ParticleIDCollection=["RecoParticlesPIDs_1"],
+    PIDAlgoName="PIDAlgo1",
+    PIDParamNames=["single_param"],
+)
+
+pid_producer2 = ExampleParticleIDProducer(
+    "PIDProducer2",
+    InputCollection=["RecoParticles"],
+    ParticleIDCollection=["RecoParticlesPIDs_2"],
+    PIDAlgoName="PIDAlgo2",
+    PIDParamNames=["param_1", "param_2", "param_3"],
+)
+
+pid_consumer = ExampleParticleIDConsumer(
+    "PIDConsumer",
+    RecoParticleCollection=reco_producer.OutputCollectionRecoParticles,
+    # From first producer
+    ParticleIDCollection1=pid_producer1.ParticleIDCollection,
+    PIDAlgoName1=pid_producer1.PIDAlgoName,
+    PIDParamNames1=pid_producer1.PIDParamNames,
+    ParamName1="single_param",
+    # From second producer
+    ParticleIDCollection2=pid_producer2.ParticleIDCollection,
+    PIDAlgoName2=pid_producer2.PIDAlgoName,
+    PIDParamNames2=pid_producer2.PIDParamNames,
+    ParamName2="param_2",
+)
+
+ApplicationMgr(
+    TopAlg=[reco_producer, pid_producer1, pid_producer2, pid_consumer],
+    EvtSel="NONE",
+    EvtMax=10,
+    ExtSvc=[EventDataSvc("EventDataSvc")],
+    OutputLevel=INFO,
+)

test/k4FWCoreTest/src/components/ExampleFunctionalProducerMultiple.cpp

@@ -22,6 +22,7 @@
 #include "k4FWCore/Producer.h"
 
 #include "edm4hep/MCParticleCollection.h"
+#include "edm4hep/ReconstructedParticleCollection.h"
 #include "edm4hep/SimTrackerHitCollection.h"
 #include "edm4hep/TrackCollection.h"
 #include "edm4hep/TrackerHit3DCollection.h"
@@ -33,19 +34,21 @@
 
 using retType =
     std::tuple<podio::UserDataCollection<float>, edm4hep::MCParticleCollection, edm4hep::MCParticleCollection,
-               edm4hep::SimTrackerHitCollection, edm4hep::TrackerHit3DCollection, edm4hep::TrackCollection>;
+               edm4hep::SimTrackerHitCollection, edm4hep::TrackerHit3DCollection, edm4hep::TrackCollection,
+               edm4hep::ReconstructedParticleCollection>;
 
 struct ExampleFunctionalProducerMultiple final : k4FWCore::Producer<retType()> {
   // The pairs in KeyValue can be changed from python and they correspond
   // to the names of the output collections
   ExampleFunctionalProducerMultiple(const std::string& name, ISvcLocator* svcLoc)
-      : Producer(name, svcLoc, {},
-                 {KeyValues("OutputCollectionFloat", {"VectorFloat"}),
-                  KeyValues("OutputCollectionParticles1", {"MCParticles1"}),
-                  KeyValues("OutputCollectionParticles2", {"MCParticles2"}),
-                  KeyValues("OutputCollectionSimTrackerHits", {"SimTrackerHits"}),
-                  KeyValues("OutputCollectionTrackerHits", {"TrackerHits"}),
-                  KeyValues("OutputCollectionTracks", {"Tracks"})}) {}
+      : Producer(
+            name, svcLoc, {},
+            {KeyValues("OutputCollectionFloat", {"VectorFloat"}),
+             KeyValues("OutputCollectionParticles1", {"MCParticles1"}),
+             KeyValues("OutputCollectionParticles2", {"MCParticles2"}),
+             KeyValues("OutputCollectionSimTrackerHits", {"SimTrackerHits"}),
+             KeyValues("OutputCollectionTrackerHits", {"TrackerHits"}), KeyValues("OutputCollectionTracks", {"Tracks"}),
+             KeyValues("OutputCollectionRecoParticles", {"RecoParticles"})}) {}
 
   // This is the function that will be called to produce the data
   retType operator()() const override {
@@ -84,8 +87,14 @@ struct ExampleFunctionalProducerMultiple final : k4FWCore::Producer<retType()> {
     track.addToTrackerHits(trackerHit);
     track.addToTracks(track2);
 
+    auto recos = edm4hep::ReconstructedParticleCollection();
+    for (int i = 1; i < 5; ++i) {
+      auto reco = recos.create();
+      reco.setPDG(i);
+    }
+
     return std::make_tuple(std::move(floatVector), std::move(particles), edm4hep::MCParticleCollection(),
-                           std::move(simTrackerHits), std::move(trackerHits), std::move(tracks));
+                           std::move(simTrackerHits), std::move(trackerHits), std::move(tracks), std::move(recos));
   }
 
 private:

test/k4FWCoreTest/src/components/ExampleParticleIDConsumer.cpp

@@ -0,0 +1,154 @@
+/*
+ * Copyright (c) 2014-2024 Key4hep-Project.
+ *
+ * This file is part of Key4hep.
+ * See https://key4hep.github.io/key4hep-doc/ for further info.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <edm4hep/ReconstructedParticleCollection.h>
+#include "k4FWCore/Consumer.h"
+#include "k4FWCore/MetadataUtils.h"
+
+#include "edm4hep/ParticleIDCollection.h"
+#include "edm4hep/utils/ParticleIDUtils.h"
+
+#include <Gaudi/Property.h>
+
+#include "fmt/format.h"
+#include "fmt/ranges.h"
+
+#include <algorithm>
+#include <stdexcept>
+#include <string>
+
+struct ExampleParticleIDConsumer final
+    : k4FWCore::Consumer<void(const edm4hep::ParticleIDCollection&, const edm4hep::ParticleIDCollection&,
+                              const edm4hep::ReconstructedParticleCollection&)> {
+  ExampleParticleIDConsumer(const std::string& name, ISvcLocator* svcLoc)
+      : Consumer(name, svcLoc,
+                 {KeyValues("ParticleIDCollection1", {"PIDs1"}), KeyValues("ParticleIDCollection2", {"PIDs2"}),
+                  KeyValues("RecoParticleCollection", {"recos"})}) {}
+
+  bool checkAlgoMetadata(const edm4hep::utils::ParticleIDMeta& pidMeta, const Gaudi::Property<std::string>& algoName,
+                         const Gaudi::Property<std::vector<std::string>>& paramNames) const {
+    if (pidMeta.algoName != algoName) {
+      fatal() << fmt::format(
+                     "The PID algorithm name from metadata does not match the expected one from the properties: "
+                     "(expected {}, actual {})",
+                     algoName.value(), pidMeta.algoName)
+              << endmsg;
+      return false;
+    }
+
+    if (!std::ranges::equal(pidMeta.paramNames, paramNames)) {
+      fatal() << fmt::format(
+                     "The PID parameter names retrieved from metadata does not match the expected ones from the "
+                     "properties: (expected {}, actual {})",
+                     paramNames.value(), pidMeta.paramNames)
+              << endmsg;
+      return false;
+    }
+
+    return true;
+  }
+
+  void checkPIDForAlgo(const edm4hep::utils::PIDHandler& pidHandler, const edm4hep::ReconstructedParticle& reco,
+                       const edm4hep::utils::ParticleIDMeta& pidMeta, const int paramIndex) const {
+    auto maybePID = pidHandler.getPID(reco, pidMeta.algoType());
+    if (!maybePID) {
+      throw std::runtime_error(
+          std::format("Could net retrieve the {} PID object for reco particle {}", pidMeta.algoName, reco.id().index));
+    }
+    auto pid      = maybePID.value();
+    auto paramVal = pid.getParameters()[paramIndex];
+
+    // As set in the producer
+    if (paramVal != paramIndex * 0.5f) {
+      throw std::runtime_error(
+          fmt::format("Could not retrieve the correct parameter value for param {} (expected {}, actual {})",
+                      pidMeta.paramNames[paramIndex], paramIndex * 0.5f, paramVal));
+    }
+  }
+
+  StatusCode initialize() final {
+    m_pidMeta1 =
+        k4FWCore::getParameter<edm4hep::utils::ParticleIDMeta>(inputLocations("ParticleIDCollection1")[0]).value();
+
+    m_pidMeta2 =
+        k4FWCore::getParameter<edm4hep::utils::ParticleIDMeta>(inputLocations("ParticleIDCollection2")[0]).value();
+
+    if (!checkAlgoMetadata(m_pidMeta1, m_pidAlgoName1, m_pidParamNames1) ||
+        !checkAlgoMetadata(m_pidMeta2, m_pidAlgoName2, m_pidParamNames2)) {
+      return StatusCode::FAILURE;
+    }
+
+    m_paramIndex1 = edm4hep::utils::getParamIndex(m_pidMeta1, m_paramOfInterest1.value()).value_or(-1);
+    m_paramIndex2 = edm4hep::utils::getParamIndex(m_pidMeta2, m_paramOfInterest2.value()).value_or(-1);
+    if (m_paramIndex1 < 0 || m_paramIndex2 < 0) {
+      error() << fmt::format("Could not get a parameter index for {} (got {}) or {} (got {})",
+                             m_paramOfInterest1.value(), m_paramIndex1, m_paramOfInterest2.value(), m_paramIndex2)
+              << std::endl;
+    }
+
+    return StatusCode::SUCCESS;
+  }
+
+  void operator()(const edm4hep::ParticleIDCollection& pidColl1, const edm4hep::ParticleIDCollection& pidColl2,
+                  const edm4hep::ReconstructedParticleCollection& recos) const {
+    auto pidHandler = edm4hep::utils::PIDHandler::from(pidColl1, pidColl2);
+    pidHandler.addMetaInfos(m_pidMeta1, m_pidMeta2);
+
+    for (const auto r : recos) {
+      auto pids = pidHandler.getPIDs(r);
+      if (pids.size() != 2) {
+        throw std::runtime_error(
+            std::format("Could not get 2 ParticleID objects related to reco particle {}", r.id().index));
+      }
+
+      checkPIDForAlgo(pidHandler, r, m_pidMeta1, m_paramIndex1);
+      checkPIDForAlgo(pidHandler, r, m_pidMeta2, m_paramIndex2);
+    }
+  }
+
+private:
+  edm4hep::utils::ParticleIDMeta m_pidMeta1{};
+  edm4hep::utils::ParticleIDMeta m_pidMeta2{};
+
+  int m_paramIndex1{};
+  int m_paramIndex2{};
+
+  Gaudi::Property<std::string> m_pidAlgoName1{
+      this, "PIDAlgoName1", "fancyPID",
+      "The name of the first ParticleID algorithm that should be used for the metadata"};
+  Gaudi::Property<std::vector<std::string>> m_pidParamNames1{
+      this,
+      "PIDParamNames1",
+      {"p1", "p2", "p3"},
+      "The names of the parameters of the first PID algorithm that will be stored into metadata"};
+  Gaudi::Property<std::string> m_paramOfInterest1{this, "ParamName1", "p1",
+                                                  "The name of the parameter that should be checked"};
+  Gaudi::Property<std::string> m_pidAlgoName2{
+      this, "PIDAlgoName2", "fancyPID",
+      "The name of the second ParticleID algorithm that should be used for the metadata"};
+  Gaudi::Property<std::vector<std::string>> m_pidParamNames2{
+      this,
+      "PIDParamNames2",
+      {"p1", "p2", "p3"},
+      "The names of the parameters of the second PID algorithm that will be stored into metadata"};
+  Gaudi::Property<std::string> m_paramOfInterest2{this, "ParamName2", "p2",
+                                                  "The name of the parameter that should be checked"};
+};
+
+DECLARE_COMPONENT(ExampleParticleIDConsumer);