From 3725a3f68965c3b968e2b2ebf74aae8dc5040b8a Mon Sep 17 00:00:00 2001 From: jmcarcell Date: Wed, 20 Sep 2023 13:19:11 +0200 Subject: [PATCH] Add a working example of reading several collections at the same time --- k4FWCore/components/PodioInput.cpp | 163 +++++++++++++++++- k4FWCore/include/k4FWCore/PodioDataSvc.h | 51 +++++- test/k4FWCoreTest/CMakeLists.txt | 8 + ...unExampleFunctionalConsumerSeveralColls.py | 26 +++ .../ExampleFunctionalConsumerSeveralColls.cpp | 48 ++++++ 5 files changed, 285 insertions(+), 11 deletions(-) create mode 100644 test/k4FWCoreTest/options/runExampleFunctionalConsumerSeveralColls.py create mode 100644 test/k4FWCoreTest/src/components/ExampleFunctionalConsumerSeveralColls.cpp diff --git a/k4FWCore/components/PodioInput.cpp b/k4FWCore/components/PodioInput.cpp index a2d8dab5..9045729f 100644 --- a/k4FWCore/components/PodioInput.cpp +++ b/k4FWCore/components/PodioInput.cpp @@ -209,6 +209,154 @@ void PodioInput::fillReaders() { [&](std::string_view collName) { maybeRead>(collName); }; + m_readers["std::vector"] = + [&](std::string_view collName) { + maybeRead>(collName); + }; + m_readers["std::vector"] = + [&](std::string_view collName) { + maybeRead>(collName); + }; + m_readers["std::vector"] = + [&](std::string_view collName) { + maybeRead>(collName); + }; + m_readers["std::vector"] = + [&](std::string_view collName) { + maybeRead>(collName); + }; + m_readers["std::vector"] = + [&](std::string_view collName) { + maybeRead>(collName); + }; + m_readers["std::vector"] = + [&](std::string_view collName) { + maybeRead>(collName); + }; + m_readers["std::vector"] = + [&](std::string_view collName) { + maybeRead>(collName); + }; + m_readers["std::vector"] = + [&](std::string_view collName) { + maybeRead>(collName); + }; + m_readers["std::vector"] = + [&](std::string_view collName) { + maybeRead>(collName); + }; + m_readers["std::vector"] = + [&](std::string_view collName) { + maybeRead>(collName); + }; + m_readers["std::vector"] = + [&](std::string_view collName) { + maybeRead>(collName); + }; + m_readers["std::vector"] = + [&](std::string_view collName) { + maybeRead>(collName); + }; + m_readers["std::vector"] = + [&](std::string_view collName) { + maybeRead>(collName); + }; + m_readers["std::vector"] = + [&](std::string_view collName) { + maybeRead>(collName); + }; + m_readers["std::vector"] = + [&](std::string_view collName) { + maybeRead>(collName); + }; + m_readers["std::vector"] = + [&](std::string_view collName) { + maybeRead>(collName); + }; + m_readers["std::vector"] = + [&](std::string_view collName) { + maybeRead>(collName); + }; + m_readers["std::vector"] = + [&](std::string_view collName) { + maybeRead>(collName); + }; + m_readers["std::vector"] = + [&](std::string_view collName) { + maybeRead>(collName); + }; + m_readers["std::vector"] = + [&](std::string_view collName) { + maybeRead>(collName); + }; + m_readers["std::vector"] = + [&](std::string_view collName) { + maybeRead>(collName); + }; + m_readers["std::vector"] = + [&](std::string_view collName) { + maybeRead>(collName); + }; + m_readers["std::vector"] = + [&](std::string_view collName) { + maybeRead>(collName); + }; + m_readers["std::vector"] = + [&](std::string_view collName) { + maybeRead>(collName); + }; + m_readers["std::vector"] = + [&](std::string_view collName) { + maybeRead>(collName); + }; + m_readers["std::vector"] = + [&](std::string_view collName) { + maybeRead>(collName); + }; + m_readers["std::vector>"] = + [&](std::string_view collName) { + maybeRead*>>(collName); + }; + m_readers["std::vector>"] = + [&](std::string_view collName) { + maybeRead*>>(collName); + }; + m_readers["std::vector>"] = + [&](std::string_view collName) { + maybeRead*>>(collName); + }; + m_readers["std::vector>"] = + [&](std::string_view collName) { + maybeRead*>>(collName); + }; + m_readers["std::vector>"] = + [&](std::string_view collName) { + maybeRead*>>(collName); + }; + m_readers["std::vector>"] = + [&](std::string_view collName) { + maybeRead*>>(collName); + }; + m_readers["std::vector>"] = + [&](std::string_view collName) { + maybeRead*>>(collName); + }; + m_readers["std::vector>"] = + [&](std::string_view collName) { + maybeRead*>>(collName); + }; + m_readers["std::vector>"] = + [&](std::string_view collName) { + maybeRead*>>(collName); + }; + m_readers["std::vector>"] = + [&](std::string_view collName) { + maybeRead*>>(collName); + }; + m_readers["std::vector>"] = + [&](std::string_view collName) { + maybeRead*>>(collName); + }; } PodioInput::PodioInput(const std::string& name, ISvcLocator* svcLoc) : Consumer(name, svcLoc) { @@ -222,13 +370,24 @@ PodioInput::PodioInput(const std::string& name, ISvcLocator* svcLoc) : Consumer( void PodioInput::operator()() const { for (auto& collName : m_collectionNames) { - debug() << "Registering collection to read " << collName << endmsg; - auto type = m_podioDataSvc->getCollectionType(collName); + info() << "Registering collection to read " << collName << endmsg; + + // We use the space as a separator for when reading multiple collections + std::string type; + std::string name; + if (collName.find(" ") != std::string::npos) { + auto first = collName.substr(0, collName.find(" ")); + type = "std::vector<" + std::string(m_podioDataSvc->getCollectionType(first)) + ">"; + } + else { + type = m_podioDataSvc->getCollectionType(collName); + } if (m_readers.find(type) != m_readers.end()) { m_readers[type](collName); } else { maybeRead(collName); } + } // Tell data service that we are done with requested collections diff --git a/k4FWCore/include/k4FWCore/PodioDataSvc.h b/k4FWCore/include/k4FWCore/PodioDataSvc.h index 76d92a1f..4f699558 100644 --- a/k4FWCore/include/k4FWCore/PodioDataSvc.h +++ b/k4FWCore/include/k4FWCore/PodioDataSvc.h @@ -68,17 +68,50 @@ class PodioDataSvc : public DataSvc { const std::string_view getCollectionType(const std::string& collName); + template + struct is_vector + { + static constexpr bool value = false; + }; + + template class C, typename U> + struct is_vector> + { + static constexpr bool value = + std::is_same,std::vector>::value; + }; + template - StatusCode readCollection(const std::string& collName) { - const T* collection(nullptr); - collection = static_cast(m_eventframe.get(collName)); - if (collection == nullptr) { - error() << "Collection " << collName << " does not exist." << endmsg; + std::enable_if_t::value,StatusCode> + readCollection(const std::string& collName) { + const T* collection = static_cast(m_eventframe.get(collName)); + if (!collection) { + error() << "Collection " << collName << " does not exist." << endmsg; + } + auto wrapper = new DataWrapper; + wrapper->setData(collection); + m_podio_datawrappers.push_back(wrapper); + return DataSvc::registerObject("/Event", "/" + collName, wrapper); } - auto wrapper = new DataWrapper; - wrapper->setData(collection); - m_podio_datawrappers.push_back(wrapper); - return DataSvc::registerObject("/Event", "/" + collName, wrapper); + + template + std::enable_if_t::value,StatusCode> + readCollection(const std::string& collName) { + std::istringstream iss(collName); + std::string token; + auto vec = new std::vector(); + // Assume collName is a space-separated list of collection names + while (iss >> token) { + auto collection = dynamic_cast(const_cast(m_eventframe.get(token))); + if (!collection) { + error() << "Collection " << token << " does not exist." << endmsg; + } + vec->push_back(collection); + } + auto wrapper = new DataWrapper; + wrapper->setData(vec); + m_podio_datawrappers.push_back(wrapper); + return DataSvc::registerObject("/Event", "/" + collName, wrapper); } const podio::Frame& getEventFrame() const { return m_eventframe; } diff --git a/test/k4FWCoreTest/CMakeLists.txt b/test/k4FWCoreTest/CMakeLists.txt index ec33ee94..7711d96b 100644 --- a/test/k4FWCoreTest/CMakeLists.txt +++ b/test/k4FWCoreTest/CMakeLists.txt @@ -38,6 +38,7 @@ set(k4fwcoretest_plugin_sources src/components/ExampleFunctionalProducerMultiple.cpp src/components/ExampleFunctionalConsumerMultiple.cpp src/components/ExampleFunctionalTransformerMultiple.cpp + src/components/ExampleFunctionalConsumerSeveralColls.cpp ) gaudi_add_module(k4FWCoreTestPlugins @@ -243,3 +244,10 @@ add_test(NAME ExampleFunctionalTransformerMultiple set_test_env(ExampleFunctionalTransformerMultiple) set_tests_properties(ExampleFunctionalTransformerMultiple PROPERTIES DEPENDS "ExampleFunctionalProducerMultiple;ExampleFunctionalConsumerMultiple") + +add_test(NAME ExampleFunctionalConsumerSeveralColls + WORKING_DIRECTORY ${CMAKE_CURRENT_LIST_DIR} + COMMAND ${K4RUN} options/runExampleFunctionalConsumerSeveralColls.py) +set_test_env(ExampleFunctionalConsumerSeveralColls) +set_tests_properties(ExampleFunctionalConsumerSeveralColls PROPERTIES + DEPENDS "ExampleFunctionalProducerMultiple") diff --git a/test/k4FWCoreTest/options/runExampleFunctionalConsumerSeveralColls.py b/test/k4FWCoreTest/options/runExampleFunctionalConsumerSeveralColls.py new file mode 100644 index 00000000..0464f0f9 --- /dev/null +++ b/test/k4FWCoreTest/options/runExampleFunctionalConsumerSeveralColls.py @@ -0,0 +1,26 @@ +from Gaudi.Configuration import INFO +from Configurables import ExampleFunctionalConsumerSeveralColls +from Configurables import ApplicationMgr +from Configurables import k4DataSvc +from Configurables import PodioInput + +podioevent = k4DataSvc("EventDataSvc") +podioevent.input = "output_k4test_exampledata_producer_multiple.root" + +inp = PodioInput() +# We pass a space-separated list of collections to PodioInput to make sure +# they are pushed to the TES +inp.collections = [ + "MCParticles1 MCParticles2", +] + +consumer = ExampleFunctionalConsumerSeveralColls("ExampleFunctionalConsumerSeveralColls", + InputCollection="MCParticles1 MCParticles2", + ) + +ApplicationMgr(TopAlg=[inp, consumer], + EvtSel="NONE", + EvtMax=10, + ExtSvc=[podioevent], + OutputLevel=INFO, + ) diff --git a/test/k4FWCoreTest/src/components/ExampleFunctionalConsumerSeveralColls.cpp b/test/k4FWCoreTest/src/components/ExampleFunctionalConsumerSeveralColls.cpp new file mode 100644 index 00000000..c9ea4bb0 --- /dev/null +++ b/test/k4FWCoreTest/src/components/ExampleFunctionalConsumerSeveralColls.cpp @@ -0,0 +1,48 @@ +#include "Gaudi/Property.h" +#include "GaudiAlg/Consumer.h" + +#include "edm4hep/MCParticleCollection.h" + +// Define BaseClass_t +#include "k4FWCore/BaseClass.h" + +#include + +// Which type of collection we are reading +// When reading multiple collections +using colltype = std::vector; + +struct ExampleFunctionalConsumerSeveralColls final : Gaudi::Functional::Consumer { + // The pair in KeyValue can be changed from python and it corresponds + // to the name of the input collection + ExampleFunctionalConsumerSeveralColls(const std::string& name, ISvcLocator* svcLoc) + : Consumer(name, svcLoc, + KeyValue("InputCollection", "MCParticles1 MCParticles 2") + ) {} + + // This is the function that will be called to transform the data + // Note that the function has to be const, as well as the collections + // we get from the input + void operator()(const colltype& input) const override { + int i = 0; + for (auto ptr : input) { + if (i == 0) { + for (const auto& particle : *ptr) { + if ((particle.getPDG() != 1 + i) || (particle.getGeneratorStatus() != 2 + i) || + (particle.getSimulatorStatus() != 3 + i) || (particle.getCharge() != 4 + i) || + (particle.getTime() != 5 + i) || (particle.getMass() != 6 + i)) { + fatal() << "Wrong data in MCParticle collection"; + } + i++; + } + } + else { + if (ptr->size() != 0) { + fatal() << "Wrong data in MCParticle collection"; + } + } + } + } +}; + +DECLARE_COMPONENT(ExampleFunctionalConsumerSeveralColls)