From 32f0698f7c7b83d2670c4e319b8fda5c1d9f5146 Mon Sep 17 00:00:00 2001 From: tmadlener Date: Wed, 10 Apr 2024 17:07:16 +0200 Subject: [PATCH] Add testing for PIDHandler metadata handling --- test/utils/test_PIDHandler.cpp | 98 ++++++++++++++++++- utils/include/edm4hep/utils/ParticleIDUtils.h | 7 ++ utils/src/ParticleIDUtils.cc | 19 ++++ 3 files changed, 122 insertions(+), 2 deletions(-) diff --git a/test/utils/test_PIDHandler.cpp b/test/utils/test_PIDHandler.cpp index e2331e67e..980790c36 100644 --- a/test/utils/test_PIDHandler.cpp +++ b/test/utils/test_PIDHandler.cpp @@ -3,8 +3,12 @@ #include #include +#include + #include +#include + edm4hep::ReconstructedParticleCollection createRecos() { edm4hep::ReconstructedParticleCollection coll; @@ -23,19 +27,51 @@ edm4hep::ParticleIDCollection createParticleIDs(const edm4hep::ReconstructedPart auto pid = coll.create(); pid.setLikelihood(likelihood); pid.setParticle(reco); + + pid.addToParameters(1.23f); + pid.addToParameters(3.14f); } return coll; } +/// Create an event with some ParticleID collections and a metadata Frame that +/// holds the corresponding metadata +std::tuple createEventAndMetadata() { + auto retTuple = std::make_tuple(podio::Frame{}, podio::Frame{}); + auto& [event, metadata] = retTuple; + + const auto& recos = event.put(createRecos(), "reco_particles"); + auto pidColl1 = createParticleIDs(recos, 1.0f); + auto pidColl2 = createParticleIDs(recos, 2.0f); + + edm4hep::utils::PIDHandler::setAlgoInfo(metadata, pidColl1, "particleIds_1", "pidAlgo_1", 42, + {"first_param", "second_param"}); + edm4hep::utils::PIDHandler::setAlgoInfo(metadata, pidColl2, "particleIds_2", "algo_2", 123, {"1", "2"}); + + event.put(std::move(pidColl1), "particleIds_1"); + event.put(std::move(pidColl2), "particleIds_2"); + + return retTuple; +} + void checkHandlerValidReco(const edm4hep::utils::PIDHandler& handler, const edm4hep::ReconstructedParticle& reco) { const auto pids = handler.getPIDs(reco); REQUIRE(pids.size() == 2); REQUIRE(pids[0].getParticle() == reco); REQUIRE(pids[1].getParticle() == reco); - REQUIRE(pids[0].getLikelihood() == 1.0f); - REQUIRE(pids[1].getLikelihood() == 2.0f); + REQUIRE(pids[0].getParameters()[0] == 1.23f); + REQUIRE(pids[0].getParameters()[1] == 3.14f); + + // Cannot guarantee an order if the handler is constructed from a Frame + const auto llh1 = pids[0].getLikelihood(); + const auto llh2 = pids[1].getLikelihood(); + if (llh1 == 1.0f) { + REQUIRE(llh2 == 2.0f); + } else { + REQUIRE(llh2 == 1.0f); + } } TEST_CASE("PIDHandler basics", "[pid_utils]") { @@ -83,3 +119,61 @@ TEST_CASE("PIDHandler from variadic list of collections", "[pid_utils]") { REQUIRE(pids.empty()); } } + +TEST_CASE("PIDHandler from Frame w/ metadata", "[pid_utils]") { + using namespace edm4hep; + const auto& [event, metadata] = createEventAndMetadata(); + + const auto handler = utils::PIDHandler::from(event, metadata); + + const auto pidAlgo1 = handler.getAlgoType("pidAlgo_1").value(); + const auto pidAlgo2 = handler.getAlgoType("algo_2").value(); + REQUIRE(pidAlgo1 == 42); + REQUIRE(pidAlgo2 == 123); + REQUIRE_FALSE(handler.getAlgoType("non-existant-algo").has_value()); + + // Check that getting a ParticleID object for a reconstructed particle via the + // algorithmType works + const auto& recos = event.get("reco_particles"); + const auto& pidColl1 = event.get("particleIds_1"); + const auto& pidColl2 = event.get("particleIds_2"); + const auto pid1 = handler.getPID(recos[0], pidAlgo1).value(); + REQUIRE(pid1 == pidColl1[0]); + const auto pid2 = handler.getPID(recos[0], pidAlgo2).value(); + REQUIRE(pid2 == pidColl2[0]); + REQUIRE_FALSE(handler.getPID(recos[0], -1).has_value()); // empty optional for non-existant algoType + + // Check that parameter handling works as well + const auto parIndex1 = handler.getParamIndex(pidAlgo1, "first_param").value(); + REQUIRE(parIndex1 == 0); + const auto parIndex2 = handler.getParamIndex(pidAlgo2, "2").value(); + REQUIRE(parIndex2 == 1); + // Valid algo but invalid parameter name + REQUIRE_FALSE(handler.getParamIndex(pidAlgo1, "non-existant-param").has_value()); + // Invalid algorithm, the parameter name is not even checked in this case + REQUIRE_FALSE(handler.getParamIndex(-1, "doesnot matter").has_value()); +} + +TEST_CASE("PIDHandler from Frame w/o metadata", "[pid_utils]") { + using namespace edm4hep; + const auto& [event, _] = createEventAndMetadata(); + + const auto handler = utils::PIDHandler::from(event); + // No metadata available info available in this case + REQUIRE_FALSE(handler.getAlgoType("pidAlgo_1").has_value()); + + // But the rest should still work as expected + const auto& recoColl = event.get("reco_particles"); + + SECTION("Valid PID for reco") { + auto reco = recoColl[0]; + checkHandlerValidReco(handler, reco); + } + + SECTION("Unknown reco") { + const auto reco = edm4hep::ReconstructedParticle(); + const auto pids = handler.getPIDs(reco); + + REQUIRE(pids.empty()); + } +} diff --git a/utils/include/edm4hep/utils/ParticleIDUtils.h b/utils/include/edm4hep/utils/ParticleIDUtils.h index d27b19451..4ed22b160 100644 --- a/utils/include/edm4hep/utils/ParticleIDUtils.h +++ b/utils/include/edm4hep/utils/ParticleIDUtils.h @@ -64,6 +64,13 @@ class PIDHandler { /// Retrieve the algoType for a given algorithm name std::optional getAlgoType(const std::string& algoName) const; + + static void setAlgoInfo(podio::Frame& metadata, edm4hep::ParticleIDCollection& pidColl, const std::string& collName, + const std::string& algoName, const int32_t algoType, + const std::vector& paramNames); + + static void setAlgoInfo(podio::Frame& metadata, const std::string& collName, const std::string& algoName, + const int32_t algoType, const std::vector& paramNames); }; } // namespace edm4hep::utils diff --git a/utils/src/ParticleIDUtils.cc b/utils/src/ParticleIDUtils.cc index ccd0527b8..b040f5648 100644 --- a/utils/src/ParticleIDUtils.cc +++ b/utils/src/ParticleIDUtils.cc @@ -3,6 +3,8 @@ #include "edm4hep/Constants.h" #include "edm4hep/ReconstructedParticle.h" +#include + #include #include #include @@ -90,4 +92,21 @@ std::optional PIDHandler::getAlgoType(const std::string& algoName) cons return std::nullopt; } +void PIDHandler::setAlgoInfo(podio::Frame& metadata, edm4hep::ParticleIDCollection& pidColl, + const std::string& collName, const std::string& algoName, const int32_t algoType, + const std::vector& paramNames) { + for (auto pid : pidColl) { + pid.setAlgorithmType(algoType); + } + + PIDHandler::setAlgoInfo(metadata, collName, algoName, algoType, paramNames); +} + +void PIDHandler::setAlgoInfo(podio::Frame& metadata, const std::string& collName, const std::string& algoName, + const int32_t algoType, const std::vector& paramNames) { + metadata.putParameter(podio::collMetadataParamName(collName, edm4hep::pidAlgoName), algoName); + metadata.putParameter(podio::collMetadataParamName(collName, edm4hep::pidAlgoType), algoType); + metadata.putParameter(podio::collMetadataParamName(collName, edm4hep::pidParameterNames), paramNames); +} + } // namespace edm4hep::utils