diff --git a/README.md b/README.md index 83563ecd3..e2ef22b4a 100644 --- a/README.md +++ b/README.md @@ -23,6 +23,7 @@ A generic event data model for future HEP collider experiments. | | | | |-|-|-| + | [EventHeader](https://github.com/key4hep/EDM4hep/blob/main/edm4hep.yaml#L246) | [MCParticle](https://github.com/key4hep/EDM4hep/blob/main/edm4hep.yaml#L258) | [SimTrackerHit](https://github.com/key4hep/EDM4hep/blob/main/edm4hep.yaml#L326) | | [CaloHitContribution](https://github.com/key4hep/EDM4hep/blob/main/edm4hep.yaml#L368) | [SimCalorimeterHit](https://github.com/key4hep/EDM4hep/blob/main/edm4hep.yaml#L380) | [RawCalorimeterHit](https://github.com/key4hep/EDM4hep/blob/main/edm4hep.yaml#L392) | | [CalorimeterHit](https://github.com/key4hep/EDM4hep/blob/main/edm4hep.yaml#L401) | [ParticleID](https://github.com/key4hep/EDM4hep/blob/main/edm4hep.yaml#L413) | [Cluster](https://github.com/key4hep/EDM4hep/blob/main/edm4hep.yaml#L426) | diff --git a/doc/PIDHandler.md b/doc/PIDHandler.md new file mode 100644 index 000000000..76474ab0c --- /dev/null +++ b/doc/PIDHandler.md @@ -0,0 +1,239 @@ +# PIDHandler introduction and usage + +This page contains some bigger picture introduction for the utilities that are +available to work with `ParticleID`s and its related metadata. It also contains +examples for the most common usage patterns. + +## PIDHandler basics +The `PIDHandler` can be use to work with `ParticleIDCollection`s. The main +features are +- the retrieval meta information for a `ParticleIDCollection`, making it + possible to e.g. retrieve the indices of parameters from parameter names +- the possibility to invert the relation to `ReconstructedParticle`s, which + allows one to, e.g. get all `ParticleID`s that point to a specific + `ReconstructedParticle`. + +### When (not) to use the PIDHandler +NOTE: Depending on how you use the `PIDHandler` it might incur an unnecessary +performance overhead. The main purpose is to use it in cases, where the relation +between `ParticleID`s and `ReconstructedParticle`s is not trivial, e.g. +- When the two collections do **not** run in parallel +- When not all `ReconstructedParticle`s have a `ParticleID` attached +- When you want to have access to several `ParticleID`s for a given + `ReonstructedParticle` and there is no trivial relation between them. + +In case your collections run in parallel it will be much quicker to just loop +over them in parallel. + +In case you just want look at one `ParticleID` collection, but need to look at +some properties of the `ReconstructedParticle`, simply use the existing relation +to do that. + + +## ParticleIDMeta basics +`ParticleIDMeta` is a simple struct that bundles all ParticleID meta information +for one collection together. Whenever metadata is involved for ParticleIDs this +will be the thing to use. + +## General considerations + +There is no strictly enforced coupling between the meta information that is set +for a collection and the contents of that collection. In order for everything to +work as expected the following assumptions have to be met +- The collection name that is passed to the `setAlgoInfo` methods to set meta + information has to match the collection name used for putting collections into + events. +- All elements of a `ParticleIDCollection` have the same value set for + `algorithmType`. Additionally, it is usually assumed that this is a unique + value. +- The `ParticleIDMeta::paramNames` is assumed to match the parameters as they + are set in the `ParticleID` elements. + +Additionally there are a few usage considerations to be kept in mind +- The `PIDHandler` can be used without any meta data about the contained + `ParticleID`s. However, in this case it will only be useful for inverting + relations and for getting `ParticleID`s that point to a given + `ReconstructedParticle`. +- There is no guarantee that the meta information that is set in a `PIDHandler` + is consistent with the elements it knows about. Some constructors attempt to + construct a consistent handler, but it is always possible to set meta + information via `setMetaInfo`, which might be completely unrelated to the + actual contents of a handler. + +Given that there is no enforced consistency the interface and the utility +functionality makes heavy use of +[`std::optional`](https://en.cppreference.com/w/cpp/utility/optional) in its +return values / types. This makes it possible to check whether the desired +information was actually present, without having to reserve some special values +to indicate that fact. + +## Usage examples + +**The following examples assume the following includes** (which will not be visible in the examples) + +```cpp +#include +#include + +#include +``` + +**The following examples additionally assume, that there is a `metadata` and an +`event` Frame present.** The `metadata` Frame can be obtained like this from +standard EDM4hep files. + +```cpp +#include +#include + +podio::Frame getMetadataFrame(const std::string& filename) { + podio::ROOTReader reader{}; + reader.openFile(filename); + return podio::Frame(reader.readNextEntry(podio::Category::Metadata)); +} +``` + +Alternatively it can be created as an empty Frame and then just written using +the `podio::Category::Metadata` category name. + +**Finally most of the examples assume that the desired values were found and +simply get the `value` from the returned `std::optional` directly.** This means +that most of the times you will see lines like this + +```cpp +const auto value = pidHandler.getAlgoType("someAlgo").value(); +``` + +This will throw an exception if some information is not available. Check if the +optional has a value when actually using these utilities. + +### Creating a PIDHandler for a complete event + +If you want to work with a `PIDHandler` that has a somewhat global view of all +`ParticleID`s that are present you can simply construct one from an event Frame + +```cpp +const auto pidHandler = edm4hep::utils::PIDHandler::from(event); +``` + +You can also construct a `pidHandler` that populates some meta information +internally by also passing in the `metadata` Frame to this constructor +```cpp +const auto pidHandler = edm4hep::utils::PIDHandler::from(event, metadata); +``` + +### Creating a PIDHandler from an arbitrary number of collections + +If you simply want to use a `PIDHandler` to invert the relations for a given set +of `ParticleIDCollection`s you can do that via + +```cpp +const auto& tofPID = event.get("ToFPID"); +const auto& dNdXPID = event.get("dNdXPID"); +const auto& mvaPID = event.get("mvaPID"); + +const auto pidHandler = edm4hep::utils::PIDHandler::from(tofPID, dNdXPID, mvaPID); +``` + +**This handler will not have any access to metadata!** If you need metadata as +well, you can add that after the fact using the `addMetaInfo` method. + +### Using a PIDHandler to get the ParticleIDs for a reco particle + +If you have constructed a `PIDHandler` you can use it to get all or specific +related `ParticleID` objects of a `ReconstructedParticle` + +```cpp +const auto pidHandler = /* constructed as above somehow */; + +const auto recos = event.get("recos"); +const auto reco = recos[0]; + +// Get all related ParticleIDs +const auto pids = pidHandler.getPIDs(reco); + +// Only get a specific ParticleID +const auto tofAlgoType = pidHandler.getAlgoType("ToFPID").value(); +const auto tofPID = pidHandler.getPID(reco, tofAlgoType).value(); +``` + +### Retrieving meta information for a collection + +If you simply want to get the meta information for a collection (called "ToFPID" +in this example) do + +```cpp +const auto algoInfo = edm4hep::utils::PIDHandler::getAlgoInfo(metadata, "ToFPID").value(); +``` + + +### Using meta information of a collection + +If you have retrieved the meta information you can directly use that to learn +something about the `ParticleID`s it describes + +```cpp +const auto& tofPIDs = event.get("ToFPID"); +const auto algoInfo = edm4hep::utils::PIDHandler::getAlgoInfo(metadata, "ToFPID").value(); +const auto timeIndex = edm4hep::utils::getParamIndex(algoInfo, "time").value(); + +for (const auto tof : tofPIDs) { + const auto time = tof.getParameters(timeIndex); + // ... +} +``` + +If you have a `PIDHandler` with enough meta information set it is also possible +to retrieve the `timeIndex` from that + +```cpp +const auto pidHandler = /* construct somehow see above */ +const auto tofAlgoType = pidHandler.getAlgoType("ToFPID").value(); +const auto timeIndex = pidHandler.getParamIndex("time").value(); + +// ... as above +``` + + +### Setting meta information for a collection + +If you have a `ParticleIDCollection` and want to persist some meta information +for it the following snippet shows you how to do that + +```cpp +// Create ParticleID meta information for an algorithm +const auto tofMeta = edm4hep::utils::ParticleIDMeta{"TimeOfFlight", 42, {"time", "time_error"}}; + +auto tofPIDs = edm4hep::ParticleIDCollection{}; +// ... fill collection + +edm4hep::utils::PIDHandler::setAlgoInfo(metadata, tofPIDs, "ToFPID", tofMeta); +event.put(std::move(tofPIDs), "ToFPID"); +``` + +This will put the necessary metadata into the `metadata` Frame and also take +care of setting the `algorithmType` field of all elements of the `tofPIDs` +collection to `42` (the value it is set to in the `tofMeta` object). + +**Things to note:** +- It is important to use the same names for putting the collection into the + `event` Frame and the one that is passed to `PIDHandler::setAlgoInfo`! + +### Setting meta information using a collection name + +If you want to set the meta information for a `ParticleIDCollection` that you no +longer have mutable access to, this can be done via + +```cpp +// Create ParticleID meta information for an algorithm +const auto tofMeta = edm4hep::utils::ParticleIDMeta{"TimeOfFlight", 42, {"time", "time_error"}}; + +edm4hep::utils::PIDHandler::setAlgoInfo(metadata, "ToFPID", tofMeta); +``` + +**Things to note:** +- In this case it is user responsibility to set the `algoType` of the + `ParticleIDMeta` to the same value as the one that is used for the elements of + the collection. +- It is important to use the same names for putting the collection into the + `event` Frame and the one that is passed to `PIDHandler::setAlgoInfo`! diff --git a/edm4hep.yaml b/edm4hep.yaml index 2f51180c4..9f3544c3e 100644 --- a/edm4hep.yaml +++ b/edm4hep.yaml @@ -420,6 +420,8 @@ datatypes: - float likelihood // likelihood of this hypothesis - in a user defined normalization VectorMembers: - float parameters // parameters associated with this hypothesis + OneToOneRelations: + - edm4hep::ReconstructedParticle particle // the particle from which this PID has been computed #------ Cluster @@ -586,12 +588,10 @@ datatypes: - edm4hep::CovMatrix4f covMatrix // covariance matrix of the reconstructed particle 4vector OneToOneRelations: - edm4hep::Vertex startVertex // start vertex associated to this particle - - edm4hep::ParticleID particleIDUsed // particle Id used for the kinematics of this particle OneToManyRelations: - edm4hep::Cluster clusters // clusters that have been used for this particle - edm4hep::Track tracks // tracks that have been used for this particle - edm4hep::ReconstructedParticle particles // reconstructed particles that have been combined to this particle - - edm4hep::ParticleID particleIDs // particle Ids (not sorted by their likelihood) ExtraCode: includes: "#include " declaration: " diff --git a/include/edm4hep/Constants.h b/include/edm4hep/Constants.h index 64ac846a3..d75959942 100644 --- a/include/edm4hep/Constants.h +++ b/include/edm4hep/Constants.h @@ -43,6 +43,13 @@ enum class TrackParams : DimType { d0 = 0, phi, omega, z0, tanLambda, time }; /// Enum for accessing the covariance matrix in the TrackerPulse enum class TrackerPulseDims : DimType { charge = 0, time }; + +/// The collection parameter name for accessing the names of the parameters for +/// a ParticleID collection +static constexpr const char* pidParameterNames = "ParameterNames"; +static constexpr const char* pidAlgoName = "AlgoName"; +static constexpr const char* pidAlgoType = "AlgoType"; + } // namespace edm4hep #endif // EDM4HEP_CONSTANTS_H diff --git a/test/utils/CMakeLists.txt b/test/utils/CMakeLists.txt index b481312b1..fb84ecda3 100644 --- a/test/utils/CMakeLists.txt +++ b/test/utils/CMakeLists.txt @@ -16,7 +16,8 @@ endif() include(Catch) add_executable(unittests_edm4hep - test_kinematics.cpp test_vector_utils.cpp test_covmatrix_utils.cpp) + test_kinematics.cpp test_vector_utils.cpp test_covmatrix_utils.cpp test_PIDHandler.cpp) + target_link_libraries(unittests_edm4hep edm4hep EDM4HEP::utils Catch2::Catch2 Catch2::Catch2WithMain) option(SKIP_CATCH_DISCOVERY "Skip the Catch2 test discovery" OFF) diff --git a/test/utils/test_PIDHandler.cpp b/test/utils/test_PIDHandler.cpp new file mode 100644 index 000000000..1b018a08e --- /dev/null +++ b/test/utils/test_PIDHandler.cpp @@ -0,0 +1,219 @@ +#include + +#include +#include + +#include + +#include + +#include + +edm4hep::ReconstructedParticleCollection createRecos() { + edm4hep::ReconstructedParticleCollection coll; + + for (size_t i = 1; i < 10; ++i) { + auto reco = coll.create(); + reco.setPDG(i); + } + + return coll; +} + +edm4hep::ParticleIDCollection createParticleIDs(const edm4hep::ReconstructedParticleCollection& recos, + float likelihood) { + edm4hep::ParticleIDCollection coll; + for (const auto reco : recos) { + 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].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]") { + using namespace edm4hep; + + const auto recoColl = createRecos(); + const auto pidColl1 = createParticleIDs(recoColl, 1.0f); + const auto pidColl2 = createParticleIDs(recoColl, 2.0f); + + auto handler = utils::PIDHandler(); + handler.addColl(pidColl1); + handler.addColl(pidColl2); + + 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()); + } +} + +TEST_CASE("PIDHandler from variadic list of collections", "[pid_utils]") { + using namespace edm4hep; + + const auto recoColl = createRecos(); + const auto pidColl1 = createParticleIDs(recoColl, 1.0f); + const auto pidColl2 = createParticleIDs(recoColl, 2.0f); + + const auto handler = utils::PIDHandler::from(pidColl1, pidColl2); + + SECTION("Valid PID for reco") { + auto reco = recoColl[3]; + checkHandlerValidReco(handler, reco); + } + + SECTION("Unkown reco") { + const auto reco = edm4hep::ReconstructedParticle(); + const auto pids = handler.getPIDs(reco); + + REQUIRE(pids.empty()); + } +} + +TEST_CASE("PIDHandler w/ addMetaInfo", "[pid_utils]") { + using namespace edm4hep; + auto handler = utils::PIDHandler(); + + const auto recoColl = createRecos(); + auto pidColl1 = createParticleIDs(recoColl, 1.0f); + for (auto pid : pidColl1) { + pid.setAlgorithmType(42); + } + const auto pidInfo1 = utils::ParticleIDMeta{"fancyAlgo", 42, {"p1", "p2"}}; + + handler.addColl(pidColl1, pidInfo1); + + REQUIRE(handler.getAlgoType("fancyAlgo").value_or(0) == 42); + REQUIRE(handler.getParamIndex(42, "p2").value_or(-1) == 1); + REQUIRE(handler.getPID(recoColl[0], 42).value() == pidColl1[0]); + + // Technically, we can even just add meta data without having a corresponding + // ParticleID collection to match + handler.addMetaInfo(utils::ParticleIDMeta{"anotherAlgo", 123, {}}); + REQUIRE(handler.getAlgoType("anotherAlgo").value() == 123); + + // Expected exceptions also get thrown + REQUIRE_THROWS_AS(handler.addMetaInfo(utils::ParticleIDMeta{"anotherAlgo", 321, {"param"}}), std::runtime_error); + // No information about this meta data can be obtained + REQUIRE_FALSE(handler.getParamIndex(321, "param").has_value()); + + REQUIRE_THROWS_AS(handler.addMetaInfo(utils::ParticleIDMeta{"newAlgo", 42, {"PARAM"}}), std::runtime_error); + // Existing meta info is unchanged + REQUIRE_FALSE(handler.getParamIndex(42, "PARAM").has_value()); + REQUIRE(handler.getParamIndex(42, "p2").value_or(-1) == 1); +} + +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()); + + const auto pidInfo = utils::PIDHandler::getAlgoInfo(metadata, "particleIds_1").value(); + REQUIRE(pidInfo.algoName == "pidAlgo_1"); + REQUIRE(pidInfo.algoType == 42); + REQUIRE(pidInfo.paramNames.size() == 2); + REQUIRE(pidInfo.paramNames[0] == "first_param"); + REQUIRE(pidInfo.paramNames[1] == "second_param"); +} + +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/CMakeLists.txt b/utils/CMakeLists.txt index 8194ae0a0..821938008 100644 --- a/utils/CMakeLists.txt +++ b/utils/CMakeLists.txt @@ -8,10 +8,17 @@ target_include_directories(kinematics target_link_libraries(kinematics PUBLIC INTERFACE ROOT::Core) target_compile_features(kinematics INTERFACE cxx_std_17) -add_library(utils INTERFACE) +set(utils_sources + src/ParticleIDUtils.cc +) + +add_library(utils SHARED ${utils_sources}) set_target_properties(utils PROPERTIES OUTPUT_NAME "edm4hepUtils") add_library(EDM4HEP::utils ALIAS utils) -target_link_libraries(utils INTERFACE kinematics) +target_include_directories(utils + PUBLIC $ + $) +target_link_libraries(utils PUBLIC EDM4HEP::edm4hep kinematics) set(sources src/dataframe.cc) diff --git a/utils/include/edm4hep/utils/ParticleIDUtils.h b/utils/include/edm4hep/utils/ParticleIDUtils.h new file mode 100644 index 000000000..0878c3a09 --- /dev/null +++ b/utils/include/edm4hep/utils/ParticleIDUtils.h @@ -0,0 +1,107 @@ +#ifndef EDM4HEP_UTILS_PARTICLEIDUTILS_H +#define EDM4HEP_UTILS_PARTICLEIDUTILS_H + +#include +#include + +#include + +#include +#include +#include +#include +#include + +namespace edm4hep::utils { + +/// A simple struct bundling relevant metadata for a ParticleID collection +struct ParticleIDMeta { + std::string algoName{}; ///< The name of the algorithm + int32_t algoType{0}; ///< The (user defined) algorithm type + std::vector paramNames{}; ///< The names of the parameters +}; + +/// Get the index of the parameter in the passed ParticleID meta info +std::optional getParamIndex(const ParticleIDMeta& pidMetaInfo, const std::string& param); + +/// Utility class to invert the ParticleID to ReconstructedParticle relation +/// See [this page](@ref md_doc_2_p_i_d_handler) for example usage and more information. +class PIDHandler { + + using RecoPidMapT = std::multimap; + + RecoPidMapT m_recoPidMap{}; ///< The internal map from recos to pids + + std::map m_algoTypes{}; ///< Maps algo names to algo types + + /// Maps algo types to the full particle id meta information + std::map m_algoPidMeta{}; + +public: + PIDHandler() = default; + ~PIDHandler() = default; + PIDHandler(const PIDHandler&) = default; + PIDHandler& operator=(const PIDHandler&) = default; + PIDHandler(PIDHandler&&) = default; + PIDHandler& operator=(PIDHandler&&) = default; + + /// Construct a PIDHandler from an arbitrary number of ParticleIDCollections + template + static PIDHandler from(const ParticleIDCollection& coll, const PIDColls&... pidColls) { + static_assert((std::is_same_v && ...), + "PIDHandler can only be constructed from ParticleIDCollections"); + PIDHandler handler{}; + handler.addColl(coll); + (handler.addColl(pidColls), ...); + return handler; + } + + /// Create a PIDHandler from a Frame potentially also populating some metadata information. + static PIDHandler from(const podio::Frame& event, const podio::Frame& metadata = {}); + + /// Add the information from one ParticleIDCollection to the handler + void addColl(const edm4hep::ParticleIDCollection& coll); + + /// Add the information from one ParticleIDCollection to the handler together + /// with its meta data + void addColl(const edm4hep::ParticleIDCollection& coll, const edm4hep::utils::ParticleIDMeta& pidInfo); + + /// Add meta information for a collection + void addMetaInfo(const edm4hep::utils::ParticleIDMeta& pidInfo); + + /// Retrieve all ParticleIDs that are related to the passed + /// ReconstructedParticle + std::vector getPIDs(const edm4hep::ReconstructedParticle& reco) const; + + /// Retrieve the ParticleID for a given algorithm type + std::optional getPID(const edm4hep::ReconstructedParticle& reco, int algoType) const; + + /// Retrieve the index in the parameters for a given parameter name and + /// algoType + std::optional getParamIndex(int32_t algoType, const std::string& paramName) const; + + /// Retrieve the algoType for a given algorithm name + std::optional getAlgoType(const std::string& algoName) const; + + /// Set the metadata information for the passed collection in the metadata Frame. + /// + /// This also sets the algorithmType of all elements in the collection to the + /// one that is found in the meta information. + static void setAlgoInfo(podio::Frame& metadata, edm4hep::ParticleIDCollection& pidcoll, const std::string& collname, + const edm4hep::utils::ParticleIDMeta& pidMetaInfo); + + /// Set the metadata information for a given collection name in the metadata Frame. + /// + /// @note It is user responsibility to ensure that the meta information that + /// is passed here and the one that is present in the collection with the + /// given name is consisent + static void setAlgoInfo(podio::Frame& metadata, const std::string& collname, + const edm4hep::utils::ParticleIDMeta& pidMetaInfo); + + /// Get the ParticleID meta information for a given collection name from the metadata Frame. + static std::optional getAlgoInfo(const podio::Frame& metadata, + const std::string& collName); +}; +} // namespace edm4hep::utils + +#endif // EDM4HEP_UTILS_PARTICLEIDUTILS_H diff --git a/utils/src/ParticleIDUtils.cc b/utils/src/ParticleIDUtils.cc new file mode 100644 index 000000000..43c1c61ab --- /dev/null +++ b/utils/src/ParticleIDUtils.cc @@ -0,0 +1,140 @@ +#include + +#include "edm4hep/Constants.h" + +#include + +#include +#include +#include +#include + +namespace edm4hep::utils { + +std::optional getParamIndex(const ParticleIDMeta& pidMetaInfo, const std::string& param) { + const auto nameIt = std::find(pidMetaInfo.paramNames.begin(), pidMetaInfo.paramNames.end(), param); + if (nameIt != pidMetaInfo.paramNames.end()) { + return std::distance(pidMetaInfo.paramNames.begin(), nameIt); + } + return std::nullopt; +} + +void PIDHandler::addColl(const edm4hep::ParticleIDCollection& coll) { + for (const auto pid : coll) { + m_recoPidMap.emplace(pid.getParticle(), pid); + } +} + +void PIDHandler::addColl(const edm4hep::ParticleIDCollection& coll, const edm4hep::utils::ParticleIDMeta& pidInfo) { + addColl(coll); + addMetaInfo(pidInfo); +} + +void PIDHandler::addMetaInfo(const edm4hep::utils::ParticleIDMeta& pidInfo) { + const auto [algoIt, inserted] = m_algoTypes.emplace(pidInfo.algoName, pidInfo.algoType); + if (!inserted) { + throw std::runtime_error("Cannot have duplicate algorithm names (" + pidInfo.algoName + " already exists)"); + } + + const auto [__, metaInserted] = m_algoPidMeta.emplace(pidInfo.algoType, pidInfo); + if (!metaInserted) { + if (inserted) { + m_algoTypes.erase(algoIt); + } + throw std::runtime_error("Cannot have duplicate algorithm types (" + std::to_string(pidInfo.algoType) + + " already exists)"); + } +} + +PIDHandler PIDHandler::from(const podio::Frame& event, const podio::Frame& metadata) { + PIDHandler handler{}; + for (const auto& name : event.getAvailableCollections()) { + const auto* coll = event.get(name); + if (const auto pidColl = dynamic_cast(coll)) { + handler.addColl(*pidColl); + + auto maybeMetaInfo = PIDHandler::getAlgoInfo(metadata, name); + if (maybeMetaInfo) { + handler.addMetaInfo(std::move(maybeMetaInfo.value())); + } + } + } + + return handler; +} + +std::vector PIDHandler::getPIDs(const edm4hep::ReconstructedParticle& reco) const { + std::vector pids; + const auto& [begin, end] = m_recoPidMap.equal_range(reco); + pids.reserve(std::distance(begin, end)); + + for (auto it = begin; it != end; ++it) { + pids.emplace_back(it->second); + } + + return pids; +} + +std::optional PIDHandler::getPID(const edm4hep::ReconstructedParticle& reco, + int32_t algoType) const { + const auto& [begin, end] = m_recoPidMap.equal_range(reco); + for (auto it = begin; it != end; ++it) { + if (it->second.getAlgorithmType() == algoType) { + return it->second; + } + } + + return std::nullopt; +} + +std::optional PIDHandler::getParamIndex(int32_t algoType, const std::string& paramName) const { + if (const auto it = m_algoPidMeta.find(algoType); it != m_algoPidMeta.end()) { + return edm4hep::utils::getParamIndex(it->second, paramName); + } + + return std::nullopt; +} + +std::optional PIDHandler::getAlgoType(const std::string& algoName) const { + if (const auto it = m_algoTypes.find(algoName); it != m_algoTypes.end()) { + return it->second; + } + + return std::nullopt; +} + +void PIDHandler::setAlgoInfo(podio::Frame& metadata, edm4hep::ParticleIDCollection& pidColl, + const std::string& collName, const edm4hep::utils::ParticleIDMeta& pidMetaInfo) { + for (auto pid : pidColl) { + pid.setAlgorithmType(pidMetaInfo.algoType); + } + + PIDHandler::setAlgoInfo(metadata, collName, pidMetaInfo); +} + +void PIDHandler::setAlgoInfo(podio::Frame& metadata, const std::string& collName, + const edm4hep::utils::ParticleIDMeta& pidMetaInfo) { + metadata.putParameter(podio::collMetadataParamName(collName, edm4hep::pidAlgoName), pidMetaInfo.algoName); + metadata.putParameter(podio::collMetadataParamName(collName, edm4hep::pidAlgoType), pidMetaInfo.algoType); + metadata.putParameter(podio::collMetadataParamName(collName, edm4hep::pidParameterNames), pidMetaInfo.paramNames); +} + +std::optional PIDHandler::getAlgoInfo(const podio::Frame& metadata, + const std::string& collName) { + ParticleIDMeta pidInfo{}; + + pidInfo.algoName = metadata.getParameter(podio::collMetadataParamName(collName, edm4hep::pidAlgoName)); + // Use the algoName as proxy to see whether we could actually get the + // information from the metadata + if (pidInfo.algoName.empty()) { + return std::nullopt; + } + + pidInfo.algoType = metadata.getParameter(podio::collMetadataParamName(collName, edm4hep::pidAlgoType)); + pidInfo.paramNames = metadata.getParameter>( + podio::collMetadataParamName(collName, edm4hep::pidParameterNames)); + + return pidInfo; +} + +} // namespace edm4hep::utils