From 537b91da9e331086f2b19e977ea27cdafb0e50ac Mon Sep 17 00:00:00 2001 From: Thomas Madlener Date: Thu, 19 Dec 2024 20:35:38 +0100 Subject: [PATCH] Add a producer and a consumer for PID with metadata --- .../components/ExampleParticleIDConsumer.cpp | 154 ++++++++++++++++++ .../components/ExampleParticleIDProducer.cpp | 71 ++++++++ 2 files changed, 225 insertions(+) create mode 100644 test/k4FWCoreTest/src/components/ExampleParticleIDConsumer.cpp create mode 100644 test/k4FWCoreTest/src/components/ExampleParticleIDProducer.cpp 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);