Add the possibility to set edm4hep::utils::ParticleIDMeta via the IMetadataSvc

CMakeLists.txt

@@ -30,6 +30,7 @@ find_package(ROOT COMPONENTS RIO Tree REQUIRED)
 find_package(Gaudi REQUIRED)
 find_package(podio 1.0.1 REQUIRED)
 find_package(EDM4HEP 0.99 REQUIRED)
+find_package(fmt REQUIRED)
 
 include(cmake/Key4hepConfig.cmake)
 

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

@@ -195,6 +195,8 @@ add_test_with_env(FunctionalProducerRNTuple options/ExampleFunctionalProducerRNT
 add_test_with_env(FunctionalTTreeToRNTuple options/ExampleFunctionalTTreeToRNTuple.py PROPERTIES FIXTURES_REQUIRED ProducerFile ADD_TO_CHECK_FILES)
 add_test_with_env(GaudiFunctional options/ExampleGaudiFunctional.py PROPERTIES FIXTURES_REQUIRED ProducerFile 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
@@ -204,3 +206,10 @@ add_test_with_env(GaudiFunctional options/ExampleGaudiFunctional.py PROPERTIES F
 add_custom_command(TARGET k4FWCoreTestPlugins POST_BUILD
                    COMMAND ${CMAKE_COMMAND} -E copy_directory
                    ${PROJECT_SOURCE_DIR}/python/k4FWCore/ ${PROJECT_BINARY_DIR}/k4FWCore/genConfDir/k4FWCore)
+
+
+add_executable(check_ParticleIDOutputs src/check_ParticleIDOutputs.cpp)
+target_link_libraries(check_ParticleIDOutputs PRIVATE podio::podioIO EDM4HEP::edm4hep EDM4HEP::utils fmt::fmt)
+add_test(NAME check_ParticleIDOutputs COMMAND check_ParticleIDOutputs example_with_particleids.root)
+set_test_env(check_ParticleIDOutputs)
+set_tests_properties(check_ParticleIDOutputs PROPERTIES DEPENDS ParticleIDMetadataFramework)

test/k4FWCoreTest/options/ExampleFunctionalConsumerRuntimeCollectionsMultiple.py

@@ -37,6 +37,7 @@
     OutputCollectionSimTrackerHits=["SimTrackerHits0"],
     OutputCollectionTrackerHits=["TrackerHits0"],
     OutputCollectionTracks=["Tracks0"],
+    OutputCollectionRecoParticles=["Recos0"],
     ExampleInt=5,
 )
 producer1 = ExampleFunctionalProducerMultiple(
@@ -47,6 +48,7 @@
     OutputCollectionSimTrackerHits=["SimTrackerHits1"],
     OutputCollectionTrackerHits=["TrackerHits1"],
     OutputCollectionTracks=["Tracks1"],
+    OutputCollectionRecoParticles=["Recos1"],
     ExampleInt=5,
 )
 producer2 = ExampleFunctionalProducerMultiple(
@@ -57,6 +59,7 @@
     OutputCollectionSimTrackerHits=["SimTrackerHits2"],
     OutputCollectionTrackerHits=["TrackerHits2"],
     OutputCollectionTracks=["Tracks2"],
+    OutputCollectionRecoParticles=["Recos2"],
     ExampleInt=5,
 )
 

test/k4FWCoreTest/options/ExampleFunctionalTransformerRuntimeCollectionsMultiple.py

@@ -38,6 +38,7 @@
     OutputCollectionSimTrackerHits=["SimTrackerHits0"],
     OutputCollectionTrackerHits=["TrackerHits0"],
     OutputCollectionTracks=["Tracks0"],
+    OutputCollectionRecoParticles=["Recos0"],
     ExampleInt=5,
 )
 producer1 = ExampleFunctionalProducerMultiple(
@@ -48,6 +49,7 @@
     OutputCollectionSimTrackerHits=["SimTrackerHits1"],
     OutputCollectionTrackerHits=["TrackerHits1"],
     OutputCollectionTracks=["Tracks1"],
+    OutputCollectionRecoParticles=["Recos1"],
     ExampleInt=5,
 )
 producer2 = ExampleFunctionalProducerMultiple(
@@ -58,6 +60,7 @@
     OutputCollectionSimTrackerHits=["SimTrackerHits2"],
     OutputCollectionTrackerHits=["TrackerHits2"],
     OutputCollectionTracks=["Tracks2"],
+    OutputCollectionRecoParticles=["Recos2"],
     ExampleInt=5,
 )
 

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/options/runFunctionalMix.py

@@ -90,6 +90,7 @@
     OutputCollectionSimTrackerHits=["FunctionalSimTrackerHits"],
     OutputCollectionTrackerHits=["FunctionalTrackerHits"],
     OutputCollectionTracks=["FunctionalTracks"],
+    OutputCollectionRecoParticles=["FunctionalRecos"],
     ExampleInt=5,
 )
 

test/k4FWCoreTest/scripts/CheckOutputFiles.py

@@ -83,11 +83,19 @@ def check_metadata(filename, expected_metadata):
         "Tracks",
         "Counter",
         "NewMCParticles",
+        "RecoParticles",
     ],
 )
 check_collections(
     "functional_transformer_multiple_output_commands.root",
-    ["VectorFloat", "MCParticles1", "MCParticles2", "SimTrackerHits", "TrackerHits"],
+    [
+        "VectorFloat",
+        "MCParticles1",
+        "MCParticles2",
+        "SimTrackerHits",
+        "TrackerHits",
+        "RecoParticles",
+    ],
 )
 check_collections("/tmp/a/b/c/functional_producer.root", ["MCParticles"])
 check_collections(
@@ -104,6 +112,7 @@ def check_metadata(filename, expected_metadata):
         "TrackerHits",
         "Tracks",
         "NewMCParticles",
+        "RecoParticles",
     ],
 )
 
@@ -116,13 +125,15 @@ def check_metadata(filename, expected_metadata):
     "SimTrackerHits",
     "TrackerHits",
     "Tracks",
+    "RecoParticles",
     # Produced by functional
     "FunctionalVectorFloat",
     "FunctionalMCParticles",
     "FunctionalMCParticles2",
     "FunctionalSimTrackerHits",
     "FunctionalTrackerHits",
     "FunctionalTracks",
+    "FunctionalRecos",
     # Produced by an old algorithm
     "OldAlgorithmMCParticles",
     "OldAlgorithmSimTrackerHits",

test/k4FWCoreTest/src/check_ParticleIDOutputs.cpp

@@ -0,0 +1,118 @@
+/*
+ * 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 <edm4hep/utils/ParticleIDUtils.h>
+
+#include <podio/Reader.h>
+
+#include <fmt/format.h>
+#include <fmt/ranges.h>
+
+#include <algorithm>
+
+bool checkPIDForAlgo(const edm4hep::utils::PIDHandler& pidHandler, const edm4hep::ReconstructedParticle& reco,
+                     const edm4hep::utils::ParticleIDMeta& pidMeta, const int paramIndex) {
+  auto maybePID = pidHandler.getPID(reco, pidMeta.algoType());
+  if (!maybePID) {
+    fmt::print("Could not retrieve the {} PID object for reco particle {}", pidMeta.algoName, reco.id().index);
+    return false;
+  }
+  auto pid      = maybePID.value();
+  auto paramVal = pid.getParameters()[paramIndex];
+
+  // As set in the producer
+  if (paramVal != paramIndex * 0.5f) {
+    fmt::print("Could not retrieve the correct parameter value for param {} (expected {}, actual {})",
+               pidMeta.paramNames[paramIndex], paramIndex * 0.5f, paramVal);
+    return false;
+  }
+
+  return true;
+}
+
+bool checkAlgoMetadata(const edm4hep::utils::ParticleIDMeta& pidMeta, const std::string& algoName,
+                       const std::vector<std::string>& paramNames) {
+  if (pidMeta.algoName != algoName) {
+    fmt::print(
+        "The PID algorithm name from metadata does not match the expected one from the properties: (expected {}, "
+        "actual {})\n",
+        algoName, pidMeta.algoName);
+    return false;
+  }
+
+  if (!std::ranges::equal(pidMeta.paramNames, paramNames)) {
+    fmt::print(
+        "The PID parameter names retrieved from metadata does not match the expected ones from the properties: "
+        "(expected {}, actual {})\n",
+        paramNames, pidMeta.paramNames);
+    return false;
+  }
+
+  return true;
+}
+
+// Test configuration (this needs to match the settings in
+// options/ExampleParticleIDMetadata.py)
+constexpr auto                 pidCollectionName1 = "RecoParticlesPIDs_1";
+constexpr auto                 pidCollectionName2 = "RecoParticlesPIDs_2";
+constexpr auto                 pidAlgo1           = "PIDAlgo1";
+constexpr auto                 pidAlgo2           = "PIDAlgo2";
+constexpr auto                 pidParam1          = "single_param";
+constexpr auto                 pidParam2          = "param_2";
+const std::vector<std::string> paramNames1        = {"single_param"};
+const std::vector<std::string> paramNames2        = {"param_1", "param_2", "param_3"};
+
+int main(int, char* argv[]) {
+  auto       reader   = podio::makeReader(argv[1]);
+  const auto metadata = reader.readFrame(podio::Category::Metadata, 0);
+
+  const auto pidMeta1 = edm4hep::utils::PIDHandler::getAlgoInfo(metadata, pidCollectionName1).value();
+  const auto pidMeta2 = edm4hep::utils::PIDHandler::getAlgoInfo(metadata, pidCollectionName2).value();
+
+  if (!checkAlgoMetadata(pidMeta1, pidAlgo1, paramNames1) || !checkAlgoMetadata(pidMeta2, pidAlgo2, paramNames2)) {
+    return 1;
+  }
+
+  const auto paramIndex1 = edm4hep::utils::getParamIndex(pidMeta1, pidParam1).value_or(-1);
+  const auto paramIndex2 = edm4hep::utils::getParamIndex(pidMeta2, pidParam2).value_or(-1);
+  if (paramIndex1 == -1 || paramIndex2 == -1) {
+    fmt::print("Could not get a parameter index for '{}' (got {}) or '{}' (got {})\n", pidParam1, paramIndex1,
+               pidParam2, paramIndex2);
+  }
+
+  const auto event      = reader.readEvent(0);
+  const auto pidHandler = edm4hep::utils::PIDHandler::from(event, metadata);
+
+  const auto& recos = event.get<edm4hep::ReconstructedParticleCollection>("RecoParticles");
+  for (const auto r : recos) {
+    auto pids = pidHandler.getPIDs(r);
+    if (pids.size() != 2) {
+      fmt::print("Failed to retrieve the two expected ParticlID objects related to reco particle {}", r.id().index);
+      return 1;
+    }
+
+    if (!checkPIDForAlgo(pidHandler, r, pidMeta1, paramIndex1) ||
+        !checkPIDForAlgo(pidHandler, r, pidMeta2, paramIndex2)) {
+      return 1;
+    }
+  }
+
+  return 0;
+}