diff --git a/k4FWCore/CMakeLists.txt b/k4FWCore/CMakeLists.txt index e9b8bc03..6797bb6a 100644 --- a/k4FWCore/CMakeLists.txt +++ b/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 $ diff --git a/k4FWCore/include/k4FWCore/IMetadataSvc.h b/k4FWCore/include/k4FWCore/IMetadataSvc.h index 65748d6e..897581eb 100644 --- a/k4FWCore/include/k4FWCore/IMetadataSvc.h +++ b/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 void put(const std::string& name, const T& obj) { - if (!getFrame()) { - setFrame(podio::Frame{}); - } - getFrame()->putParameter(name, obj); - } + template void put(const std::string& name, const T& obj) { getFrameForWrite()->putParameter(name, obj); } template std::optional 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(const std::string& collName, + const edm4hep::utils::ParticleIDMeta& pidMetaInfo) { + edm4hep::utils::PIDHandler::setAlgoInfo(*getFrameForWrite(), collName, pidMetaInfo); +} + +template <> +inline std::optional IMetadataSvc::get( + const std::string& collName) const { + const auto* frame = getFrame(); + if (!frame) { + return std::nullopt; + } + + return edm4hep::utils::PIDHandler::getAlgoInfo(*frame, collName); +} + #endif diff --git a/test/k4FWCoreTest/CMakeLists.txt b/test/k4FWCoreTest/CMakeLists.txt index 0e7291e4..f956b698 100644 --- a/test/k4FWCoreTest/CMakeLists.txt +++ b/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 diff --git a/test/k4FWCoreTest/options/ExampleParticleIDMetadata.py b/test/k4FWCoreTest/options/ExampleParticleIDMetadata.py new file mode 100644 index 00000000..a7a22ffb --- /dev/null +++ b/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, +) diff --git a/test/k4FWCoreTest/src/components/ExampleFunctionalProducerMultiple.cpp b/test/k4FWCoreTest/src/components/ExampleFunctionalProducerMultiple.cpp index 96a6dcb2..a56ec7bc 100644 --- a/test/k4FWCoreTest/src/components/ExampleFunctionalProducerMultiple.cpp +++ b/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, edm4hep::MCParticleCollection, edm4hep::MCParticleCollection, - edm4hep::SimTrackerHitCollection, edm4hep::TrackerHit3DCollection, edm4hep::TrackCollection>; + edm4hep::SimTrackerHitCollection, edm4hep::TrackerHit3DCollection, edm4hep::TrackCollection, + edm4hep::ReconstructedParticleCollection>; struct ExampleFunctionalProducerMultiple final : k4FWCore::Producer { // 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 { 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: diff --git a/test/k4FWCoreTest/src/components/ExampleParticleIDConsumer.cpp b/test/k4FWCoreTest/src/components/ExampleParticleIDConsumer.cpp new file mode 100644 index 00000000..712741f8 --- /dev/null +++ b/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 +#include "k4FWCore/Consumer.h" +#include "k4FWCore/MetadataUtils.h" + +#include "edm4hep/ParticleIDCollection.h" +#include "edm4hep/utils/ParticleIDUtils.h" + +#include + +#include "fmt/format.h" +#include "fmt/ranges.h" + +#include +#include +#include + +struct ExampleParticleIDConsumer final + : k4FWCore::Consumer { + 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& algoName, + const Gaudi::Property>& 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(inputLocations("ParticleIDCollection1")[0]).value(); + + m_pidMeta2 = + k4FWCore::getParameter(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 m_pidAlgoName1{ + this, "PIDAlgoName1", "fancyPID", + "The name of the first ParticleID algorithm that should be used for the metadata"}; + Gaudi::Property> 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 m_paramOfInterest1{this, "ParamName1", "p1", + "The name of the parameter that should be checked"}; + Gaudi::Property m_pidAlgoName2{ + this, "PIDAlgoName2", "fancyPID", + "The name of the second ParticleID algorithm that should be used for the metadata"}; + Gaudi::Property> 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 m_paramOfInterest2{this, "ParamName2", "p2", + "The name of the parameter that should be checked"}; +}; + +DECLARE_COMPONENT(ExampleParticleIDConsumer); diff --git a/test/k4FWCoreTest/src/components/ExampleParticleIDProducer.cpp b/test/k4FWCoreTest/src/components/ExampleParticleIDProducer.cpp new file mode 100644 index 00000000..fc102349 --- /dev/null +++ b/test/k4FWCoreTest/src/components/ExampleParticleIDProducer.cpp @@ -0,0 +1,71 @@ +/* + * 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 "k4FWCore/MetadataUtils.h" +#include "k4FWCore/Transformer.h" + +#include "edm4hep/ParticleIDCollection.h" +#include "edm4hep/ReconstructedParticleCollection.h" +#include "edm4hep/utils/ParticleIDUtils.h" + +#include "Gaudi/Property.h" + +#include + +struct ExampleParticleIDProducer final + : k4FWCore::Transformer { + ExampleParticleIDProducer(const std::string& name, ISvcLocator* svcLoc) + : Transformer(name, svcLoc, {KeyValues("InputCollection", {"RecoParticles"})}, + KeyValues("ParticleIDCollection", {"reco_PIDs"})) {} + + StatusCode initialize() final { + m_pidMeta = {m_pidAlgoName, m_pidParamNames}; + std::string collname = outputLocations("ParticleIDCollection")[0]; + k4FWCore::putParameter(collname, m_pidMeta); + return StatusCode::SUCCESS; + } + + edm4hep::ParticleIDCollection operator()(const edm4hep::ReconstructedParticleCollection& recos) const { + auto pidColl = edm4hep::ParticleIDCollection{}; + for (const auto r : recos) { + auto pid = pidColl.create(); + pid.setAlgorithmType(m_pidMeta.algoType()); + pid.setPDG(r.getPDG() - 10); + pid.setParticle(r); + for (size_t i = 0; i < m_pidMeta.paramNames.size(); ++i) { + pid.addToParameters(i * 0.5f); + } + } + + return pidColl; + } + +private: + Gaudi::Property m_pidAlgoName{ + this, "PIDAlgoName", "fancyPID", "The name of the ParticleID algorithm that should be used for the metadata"}; + Gaudi::Property> m_pidParamNames{ + this, + "PIDParamNames", + {"p1", "p2", "p3"}, + "The names of the parameters of the PID algorithm that will be stored into metadata"}; + + edm4hep::utils::ParticleIDMeta m_pidMeta{}; +}; + +DECLARE_COMPONENT(ExampleParticleIDProducer);