Add another test and code for PID

key4hep · Jan 13, 2025 · 910fddd · 910fddd
1 parent 5b0b882
commit 910fddd
Show file tree

Hide file tree

Showing 6 changed files with 65 additions and 17 deletions.
diff --git a/k4MarlinWrapper/k4MarlinWrapper/converters/EDM4hep2Lcio.h b/k4MarlinWrapper/k4MarlinWrapper/converters/EDM4hep2Lcio.h
@@ -65,6 +65,7 @@ class EDM4hep2LcioTool : public AlgTool, virtual public IEDMConverter {
   PodioDataSvc*                   m_podioDataSvc;
   ServiceHandle<IDataProviderSvc> m_eventDataSvc;
   std::vector<std::string>        m_collectionNames;
+  std::map<uint32_t, std::string> m_idToName;
 
   void convertTracks(TrackMap& tracks_vec, const std::string& e4h_coll_name, const std::string& lcio_coll_name,
                      lcio::LCEventImpl* lcio_event);

diff --git a/k4MarlinWrapper/src/components/EDM4hep2Lcio.cpp b/k4MarlinWrapper/src/components/EDM4hep2Lcio.cpp
@@ -20,6 +20,8 @@
 #include <GaudiKernel/StatusCode.h>
 #include "GlobalConvertedObjectsMap.h"
 
+#include "UTIL/PIDHandler.h"
+
 #include "edm4hep/Constants.h"
 
 #include "k4FWCore/DataHandle.h"
@@ -333,7 +335,19 @@ StatusCode EDM4hep2LcioTool::getAvailableCollectionsFromStore() {
     if (pReg->name() == k4FWCore::frameLocation) {
       continue;
     }
+    DataObject* p;
+    sc = m_eventDataSvc->retrieveObject("/Event" + pReg->name(), p);
+    if (sc.isFailure()) {
+        error() << "Could not retrieve object " << pReg->name() << " from the EventStore" << endmsg;
+        return sc;
+    }
+    auto wrapper = dynamic_cast<AnyDataWrapper<std::unique_ptr<podio::CollectionBase>>*>(p);
+    if (!wrapper) {
+      continue;
+    }
+    // Remove the leading /
     m_collectionNames.push_back(pReg->name().substr(1, pReg->name().size() - 1));
+    m_idToName.emplace(wrapper->getData()->getID(), pReg->name());
   }
   return StatusCode::SUCCESS;
 }
@@ -430,7 +444,11 @@ StatusCode EDM4hep2LcioTool::convertCollections(lcio::LCEventImpl* lcio_event) {
     edmEvent          = m_podioDataSvc->getEventFrame();
     m_collectionNames = edmEvent.value().get().getAvailableCollections();
   } else if (m_collectionNames.empty()) {
-    getAvailableCollectionsFromStore();
+    auto sc = getAvailableCollectionsFromStore();
+    if (sc.isFailure()) {
+      warning() << "Could not retrieve available collections from the EventStore" << endmsg;
+      return sc;
+    }
   }
   // Start off with the pre-defined collection name mappings
   auto collsToConvert{m_collNames.value()};
@@ -468,7 +486,6 @@ StatusCode EDM4hep2LcioTool::convertCollections(lcio::LCEventImpl* lcio_event) {
   debug() << "Event: " << lcio_event->getEventNumber() << " Run: " << lcio_event->getRunNumber() << endmsg;
 
   EDM4hep2LCIOConv::sortParticleIDs(pidCollections);
-  // TODO: collections produced by algorithms will never be found in the frame for IOSvc
   if (!m_podioDataSvc) {
     DataObject* p;
     StatusCode  code = m_eventDataSvc->retrieveObject("/Event" + k4FWCore::frameLocation, p);
@@ -481,10 +498,25 @@ StatusCode EDM4hep2LcioTool::convertCollections(lcio::LCEventImpl* lcio_event) {
     }
   }
   for (const auto& pidCollMeta : pidCollections) {
-    const auto algoId = attachParticleIDMetaData(lcio_event, edmEvent.value(), pidCollMeta);
+    auto algoId = attachParticleIDMetaData(lcio_event, edmEvent.value(), pidCollMeta);
     if (!algoId.has_value()) {
-      warning() << "Could not determine algorithm type for ParticleID collection " << pidCollMeta.name
-                << " for setting consistent metadata" << endmsg;
+      // Now go over the collections that have been produced in a functional algorithm (if any)
+      bool found = false;
+      if (!m_podioDataSvc) {
+        const auto id = (*pidCollMeta.coll)[0].getParticle().id().collectionID;
+        if (auto it = m_idToName.find(id); it != m_idToName.end()) {
+          auto name = it->second;
+          if (pidCollMeta.metadata.has_value()) {
+            UTIL::PIDHandler pidHandler(lcio_event->getCollection(name));
+            algoId = pidHandler.addAlgorithm(pidCollMeta.metadata.value().algoName, pidCollMeta.metadata.value().paramNames);
+            found  = true;
+          }
+        }
+      }
+      if (!found) {
+        warning() << "Could not determine algorithm type for ParticleID collection " << pidCollMeta.name
+                  << " for setting consistent metadata" << endmsg;
+      }
     }
     convertParticleIDs(collection_pairs.particleIDs, pidCollMeta.name, algoId.value_or(-1));
   }

diff --git a/k4MarlinWrapper/src/components/Lcio2EDM4hep.cpp b/k4MarlinWrapper/src/components/Lcio2EDM4hep.cpp
@@ -67,7 +67,7 @@ bool Lcio2EDM4hepTool::collectionExist(const std::string& collection_name) {
     collections = m_podioDataSvc->getEventFrame().getAvailableCollections();
   }
   for (const auto& name : collections) {
-    if (collection_name == name) {
+    if (name == collection_name) {
       debug() << "Collection named " << name << " already registered, skipping conversion." << endmsg;
       return true;
     }

diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt
@@ -94,6 +94,10 @@ ExternalData_Add_Test( marlinwrapper_tests
   NAME link_conversion_edm4hep_to_lcio
   COMMAND ${K4RUN} ${CMAKE_CURRENT_SOURCE_DIR}/gaudi_opts/test_link_conversion_edm4hep.py --inputfile DATA{${PROJECT_SOURCE_DIR}/test/input_files/ttbar_20240223_edm4hep.root}
 )
+ExternalData_Add_Test( marlinwrapper_tests
+  NAME link_conversion_edm4hep_to_lcio_iosvc
+  COMMAND ${K4RUN} ${CMAKE_CURRENT_SOURCE_DIR}/gaudi_opts/test_link_conversion_edm4hep.py --iosvc --inputfile DATA{${PROJECT_SOURCE_DIR}/test/input_files/ttbar_20240223_edm4hep.root}
+)
 
 add_test( event_header_conversion bash -c "k4run ${CMAKE_CURRENT_SOURCE_DIR}/gaudi_opts/createEventHeader.py && anajob test.slcio | grep 'EVENT: 42'" )
 

diff --git a/test/gaudi_opts/test_global_converter_maps.py b/test/gaudi_opts/test_global_converter_maps.py
@@ -33,10 +33,10 @@
     EventDataSvc,
 )
 
-from k4FWCore.parseArgs import parser
-
 from k4FWCore import ApplicationMgr, IOSvc
 
+from k4FWCore.parseArgs import parser
+
 parser.add_argument("--iosvc", action="store_true", default=False, help="Use IOSvc instead of PodioDataSvc")
 
 args = parser.parse_known_args()[0]

diff --git a/test/gaudi_opts/test_link_conversion_edm4hep.py b/test/gaudi_opts/test_link_conversion_edm4hep.py
@@ -23,24 +23,31 @@
 from Configurables import (
     PodioInput,
     k4DataSvc,
-    ApplicationMgr,
     PseudoRecoAlgorithm,
     TrivialMCRecoLinker,
     MarlinProcessorWrapper,
     EDM4hep2LcioTool,
+    EventDataSvc,
 )
 
+from k4FWCore import ApplicationMgr, IOSvc
+
 from k4FWCore.parseArgs import parser
 
 parser.add_argument("--inputfile", help="Input file")
-my_args = parser.parse_known_args()[0]
-
-evtsvc = k4DataSvc("EventDataSvc")
-evtsvc.input = my_args.inputfile
+parser.add_argument("--iosvc", action="store_true", default=False, help="Use IOSvc instead of PodioDataSvc")
+args = parser.parse_known_args()[0]
 
-podioInput = PodioInput("InputReader")
-podioInput.collections = ["MCParticles"]
-podioInput.OutputLevel = INFO
+if args.iosvc:
+    evtsvc = EventDataSvc("EventDataSvc")
+    iosvc = IOSvc()
+    iosvc.Input = args.inputfile
+else:
+    evtsvc = k4DataSvc("EventDataSvc")
+    evtsvc.input = args.inputfile
+    podioInput = PodioInput("InputReader")
+    podioInput.collections = ["MCParticles"]
+    podioInput.OutputLevel = INFO
 
 
 PseudoRecoAlg = PseudoRecoAlgorithm(
@@ -74,9 +81,13 @@
 mcLinkConverter.OutputLevel = DEBUG
 MarlinMCLinkChecker.EDM4hep2LcioTool = mcLinkConverter
 
+algList = [PseudoRecoAlg, MCRecoLinker, MarlinMCLinkChecker]
+
+if not args.iosvc:
+    algList =[podioInput] + algList
 
 ApplicationMgr(
-    TopAlg=[podioInput, PseudoRecoAlg, MCRecoLinker, MarlinMCLinkChecker],
+    TopAlg=algList,
     ExtSvc=[evtsvc],
     EvtMax=-1,
     EvtSel="NONE",