diff --git a/k4FWCore/components/CollectionMerger.cpp b/k4FWCore/components/CollectionMerger.cpp new file mode 100644 index 00000000..cbc314f6 --- /dev/null +++ b/k4FWCore/components/CollectionMerger.cpp @@ -0,0 +1,143 @@ +/* + * 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. + */ + +/** @class CollectionMerger + * + * Merges collections of the same type into a single collection. + * The output collection is created as a subset collection if the Copy property is set to false (default). + * The collections are merged by copying the elements of the input collections into the output collection. + * + * In both cases, the relations in the new objects point to the original objects. This means that collections + * having relations to objects in different collections will need these collections to be present to fully + * work. If the collections are not present (for example, after dropping them with the output commands) + * and this relations are used, that will almost certainly crash. An example of an usage of this algorithm + * is to merge several collections to give to an algorithm so that it can use all of them at the same time. + * + * The collections to be merged are specified in the InputCollections property, which is a list of collection names. + * The output collection is specified in the OutputCollection property. + */ + +#include "edm4hep/edm4hep.h" + +#include "k4FWCore/Transformer.h" + +#include +#include +#include +#include + +struct CollectionMerger final : k4FWCore::Transformer( + const std::vector*>&)> { + CollectionMerger(const std::string& name, ISvcLocator* svcLoc) + : Transformer(name, svcLoc, {KeyValues("InputCollections", {"MCParticles"})}, + {KeyValues("OutputCollection", {"NewMCParticles"})}) { + if (System::cmdLineArgs()[0].find("genconf") != std::string::npos) { + return; + } + m_map["edm4hep::MCParticleCollection"] = &CollectionMerger::mergeCollections; + m_map["edm4hep::SimTrackerHitCollection"] = &CollectionMerger::mergeCollections; + m_map["edm4hep::CaloHitContributionCollection"] = + &CollectionMerger::mergeCollections; + m_map["edm4hep::SimCalorimeterHitCollection"] = + &CollectionMerger::mergeCollections; + m_map["edm4hep::RawCalorimeterHitCollection"] = + &CollectionMerger::mergeCollections; + m_map["edm4hep::CalorimeterHitCollection"] = &CollectionMerger::mergeCollections; + m_map["edm4hep::ParticleIDCollection"] = &CollectionMerger::mergeCollections; + m_map["edm4hep::ClusterCollection"] = &CollectionMerger::mergeCollections; + m_map["edm4hep::TrackerHit3DCollection"] = &CollectionMerger::mergeCollections; + m_map["edm4hep::TrackerHitPlaneCollection"] = + &CollectionMerger::mergeCollections; + m_map["edm4hep::RawTimeSeriesCollection"] = &CollectionMerger::mergeCollections; + m_map["edm4hep::TrackCollection"] = &CollectionMerger::mergeCollections; + m_map["edm4hep::VertexCollection"] = &CollectionMerger::mergeCollections; + m_map["edm4hep::ReconstructedParticleCollection"] = + &CollectionMerger::mergeCollections; + m_map["edm4hep::MCRecoParticleAssociationCollection"] = + &CollectionMerger::mergeCollections; + m_map["edm4hep::MCRecoCaloAssociationCollection"] = + &CollectionMerger::mergeCollections; + m_map["edm4hep::MCRecoTrackerAssociationCollection"] = + &CollectionMerger::mergeCollections; + m_map["edm4hep::MCRecoCaloParticleAssociationCollection"] = + &CollectionMerger::mergeCollections; + m_map["edm4hep::MCRecoClusterParticleAssociationCollection"] = + &CollectionMerger::mergeCollections; + m_map["edm4hep::MCRecoTrackParticleAssociationCollection"] = + &CollectionMerger::mergeCollections; + m_map["edm4hep::RecoParticleVertexAssociationCollection"] = + &CollectionMerger::mergeCollections; + m_map["edm4hep::TimeSeriesCollection"] = &CollectionMerger::mergeCollections; + m_map["edm4hep::RecDqdxCollection"] = &CollectionMerger::mergeCollections; + m_map["edm4hep::GeneratorEventParametersCollection"] = + &CollectionMerger::mergeCollections; + m_map["edm4hep::GeneratorPdfInfoCollection"] = + &CollectionMerger::mergeCollections; + } + + std::shared_ptr operator()( + const std::vector*>& input) const override { + std::shared_ptr ret; + debug() << "Merging " << input.size() << " collections" << endmsg; + std::string_view type = ""; + for (const auto& coll : input) { + debug() << "Merging collection of type " << (*coll)->getTypeName() << " with " << (*coll)->size() << " elements" + << endmsg; + if (type.empty()) { + type = (*coll)->getTypeName(); + } else if (type != (*coll)->getTypeName()) { + throw std::runtime_error("Different collection types are not supported"); + return ret; + } + (this->*m_map.at((*coll)->getTypeName()))(*coll, ret); + } + return ret; + } + +private: + using MergeType = void (CollectionMerger::*)(const std::shared_ptr&, + std::shared_ptr&) const; + std::map m_map; + Gaudi::Property m_copy{this, "Copy", false, + "Copy the elements of the collections instead of creating a subset collection"}; + + template + void mergeCollections(const std::shared_ptr& source, + std::shared_ptr& ret) const { + if (!ret) { + ret.reset(new T()); + if (!m_copy) { + ret->setSubsetCollection(); + } + } + const auto ptr = std::static_pointer_cast(ret); + const auto sourceColl = std::static_pointer_cast(source); + if (m_copy) { + for (const auto& elem : *sourceColl) { + ptr->push_back(elem.clone()); + } + } else { + for (const auto& elem : *sourceColl) { + ptr->push_back(elem); + } + } + } +}; + +DECLARE_COMPONENT(CollectionMerger) diff --git a/k4FWCore/include/k4FWCore/FunctionalUtils.h b/k4FWCore/include/k4FWCore/FunctionalUtils.h index 58268021..36e81af5 100644 --- a/k4FWCore/include/k4FWCore/FunctionalUtils.h +++ b/k4FWCore/include/k4FWCore/FunctionalUtils.h @@ -42,7 +42,14 @@ namespace k4FWCore { namespace details { // This function will be used to modify std::shared_ptr to the actual collection type - template const auto& maybeTransformToEDM4hep(const P& arg) { return arg; } + template auto maybeTransformToEDM4hep(P& arg) { return arg; } + + template + requires std::same_as>> + auto maybeTransformToEDM4hep(P& arg) { + return arg; + } + template requires std::is_base_of_v const auto& maybeTransformToEDM4hep(P* arg) { @@ -69,7 +76,8 @@ namespace k4FWCore { template struct isVectorLike : std::false_type {}; template - requires std::is_base_of_v> + requires std::is_base_of_v> || + std::is_same_v, std::remove_cvref_t> struct isVectorLike> : std::true_type {}; template @@ -96,6 +104,12 @@ namespace k4FWCore { return std::shared_ptr(std::make_shared(std::move(arg))); } + template + requires std::is_same_v> + auto convertToSharedPtr(T&& arg) { + return std::move(arg); + } + template struct filter_evtcontext_tt { static_assert(!std::disjunction_v...>, "EventContext can only appear as first argument"); @@ -127,11 +141,14 @@ namespace k4FWCore { std::remove_pointer_t>::value_type>; auto inputMap = std::vector(); for (auto& handle : std::get(handles)) { - auto in = get(handle, thisClass, Gaudi::Hive::currentContext()); - inputMap.push_back(static_cast(in.get())); + if constexpr (std::is_same_v>) { + inputMap.push_back(&get(handle, thisClass, Gaudi::Hive::currentContext())); + } else { + auto in = get(handle, thisClass, Gaudi::Hive::currentContext()); + inputMap.push_back(static_cast(in.get())); + } } std::get(inputTuple) = std::move(inputMap); - } else { try { auto in = get(std::get(handles)[0], thisClass, Gaudi::Hive::currentContext()); diff --git a/k4FWCore/include/k4FWCore/Transformer.h b/k4FWCore/include/k4FWCore/Transformer.h index ec803cc1..323c0246 100644 --- a/k4FWCore/include/k4FWCore/Transformer.h +++ b/k4FWCore/include/k4FWCore/Transformer.h @@ -43,11 +43,11 @@ namespace k4FWCore { : Gaudi::Functional::details::DataHandleMixin, std::tuple<>, Traits_> { using Gaudi::Functional::details::DataHandleMixin, std::tuple<>, Traits_>::DataHandleMixin; - static_assert( - ((std::is_base_of_v || isVectorLike_v)&&...), - "Transformer and Producer input types must be EDM4hep collections or vectors of collection pointers"); - static_assert((std::is_base_of_v || isVectorLike_v), - "Transformer and Producer output types must be EDM4hep collections or vectors of collections"); + static_assert(((std::is_base_of_v || isVectorLike_v)&&...), + "Transformer and Producer input types must be EDM4hep collections or maps to collections"); + static_assert((std::is_base_of_v || isVectorLike_v || + std::is_same_v, Out>), + "Transformer and Producer output types must be EDM4hep collections or maps to collections"); template using InputHandle_t = Gaudi::Functional::details::InputHandle_t>; diff --git a/test/k4FWCoreTest/CMakeLists.txt b/test/k4FWCoreTest/CMakeLists.txt index 857fc92c..33c6581a 100644 --- a/test/k4FWCoreTest/CMakeLists.txt +++ b/test/k4FWCoreTest/CMakeLists.txt @@ -130,9 +130,11 @@ add_test_with_env(FunctionalTransformerRuntimeCollections options/ExampleFunctio add_test_with_env(FunctionalTransformerRuntimeEmpty options/ExampleFunctionalTransformerRuntimeEmpty.py) add_test_with_env(FunctionalTransformerRuntimeCollectionsMultiple options/ExampleFunctionalTransformerRuntimeCollectionsMultiple.py) add_test_with_env(FunctionalTransformerHist options/ExampleFunctionalTransformerHist.py) +add_test_with_env(FunctionalCollectionMerger options/ExampleFunctionalCollectionMerger.py) add_test(NAME FunctionalCheckFiles COMMAND python3 ${CMAKE_CURRENT_LIST_DIR}/options/CheckOutputFiles.py) -set_tests_properties(FunctionalCheckFiles PROPERTIES DEPENDS "FunctionalFile;FunctionalMTFile;FunctionalMultipleFile;FunctionalOutputCommands;FunctionalProducerAbsolutePath;FunctionalTransformerRuntimeEmpty;FunctionalMix;FunctionalMixIOSvc;FunctionalTransformerHist") +set_tests_properties(FunctionalCheckFiles PROPERTIES DEPENDS "FunctionalFile;FunctionalMTFile;FunctionalMultipleFile;FunctionalOutputCommands;FunctionalProducerAbsolutePath;FunctionalTransformerRuntimeEmpty;FunctionalMix;FunctionalMixIOSvc;FunctionalTransformerHist;FunctionalCollectionMerger") + # Do this after checking the files not to overwrite them add_test_with_env(FunctionalFile_toolong options/ExampleFunctionalFile.py -n 999 PROPERTIES DEPENDS FunctionalCheckFiles PASS_REGULAR_EXPRESSION "Application Manager Terminated successfully with a user requested ScheduledStop") diff --git a/test/k4FWCoreTest/options/CheckOutputFiles.py b/test/k4FWCoreTest/options/CheckOutputFiles.py index e439154c..6968a4f6 100644 --- a/test/k4FWCoreTest/options/CheckOutputFiles.py +++ b/test/k4FWCoreTest/options/CheckOutputFiles.py @@ -129,3 +129,14 @@ def check_collections(filename, names): raise RuntimeError( "Directory structure does not match expected for functional_transformer_hist.root" ) + +check_collections( + "functional_merged_collections.root", + ["MCParticles1", "MCParticles2", "MCParticles3", "NewMCParticles", "SimTrackerHits"], +) + +podio_reader = podio.root_io.Reader("functional_merged_collections.root") +frames = podio_reader.get("events") +ev = frames[0] +if len(ev.get("NewMCParticles")) != 4: + raise RuntimeError(f"Expected 4 NewMCParticles but got {len(ev.get('NewMCParticles'))}") diff --git a/test/k4FWCoreTest/options/ExampleFunctionalCollectionMerger.py b/test/k4FWCoreTest/options/ExampleFunctionalCollectionMerger.py new file mode 100644 index 00000000..2c6b8a73 --- /dev/null +++ b/test/k4FWCoreTest/options/ExampleFunctionalCollectionMerger.py @@ -0,0 +1,66 @@ +# +# 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. +# + +# This is an example reading from a file and using a consumer with several inputs +# to check that the contents of the file are the expected ones + +from Gaudi.Configuration import INFO, DEBUG +from Configurables import CollectionMerger +from Configurables import EventDataSvc +from k4FWCore import ApplicationMgr, IOSvc + +from Configurables import ExampleFunctionalProducer + +svc = IOSvc("IOSvc") +svc.input = "functional_producer_multiple.root" +svc.output = "functional_merged_collections.root" +svc.outputCommands = [ + "drop *", + "keep MCParticles1", + "keep MCParticles2", + "keep MCParticles3", + "keep NewMCParticles", + "keep SimTrackerHits", +] + + +particle_producer = ExampleFunctionalProducer( + OutputCollection=["MCParticles3"], +) + + +merger = CollectionMerger( + "CollectionMerger", + # List of collections to concatenate + InputCollections=["MCParticles2", "MCParticles1", "MCParticles3"], + # Name of the single output collection + OutputCollection=["NewMCParticles"], + OutputLevel=DEBUG, +) + +# If we want to copy instead of creating a subset collection +# merger.Copy = True + +mgr = ApplicationMgr( + TopAlg=[particle_producer, merger], + EvtSel="NONE", + EvtMax=-1, + ExtSvc=[EventDataSvc("EventDataSvc")], + OutputLevel=INFO, +)