Add an option to copy CellID metadata when overlaying

key4hep · Oct 17, 2024 · 6b4e0ae · 6b4e0ae
1 parent f93b00e
commit 6b4e0ae
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 0 deletions.
diff --git a/k4Reco/Overlay/components/OverlayTiming.cpp b/k4Reco/Overlay/components/OverlayTiming.cpp
@@ -21,8 +21,11 @@
 
 #include "podio/Frame.h"
 
+#include "edm4hep/Constants.h"
 #include "edm4hep/MutableCaloHitContribution.h"
 
+#include "k4FWCore/MetadataUtils.h"
+
 #include <TMath.h>
 
 #include <limits>
@@ -378,4 +381,24 @@ retType OverlayTiming::operator()(const edm4hep::EventHeaderCollection&
                          std::move(ocaloHitContribs));
 }
 
+StatusCode OverlayTiming::finalize() {
+  if (m_copyCellIDMetadata) {
+    for (auto& [input, output] :
+         {std::make_pair(inputLocations(SIMTRACKERHIT_INDEX_POSITION), outputLocations("OutputSimTrackerHits")),
+          std::make_pair(inputLocations(SIMCALOHIT_INDEX_POSITION), outputLocations("OutputSimCalorimeterHits"))}) {
+      for (size_t i = 0; i < input.size(); ++i) {
+        auto name  = input[i];
+        auto value = k4FWCore::getParameter<std::string>(name + "__" + edm4hep::labels::CellIDEncoding, this);
+        if (value.has_value()) {
+          k4FWCore::putParameter(output[i] + "__" + edm4hep::labels::CellIDEncoding, value.value(), this);
+        } else {
+          warning() << "No metadata found for " << name << endmsg;
+        }
+      }
+    }
+  }
+
+  return Gaudi::Algorithm::finalize();
+}
+
 DECLARE_COMPONENT(OverlayTiming)
diff --git a/k4Reco/Overlay/components/OverlayTiming.h b/k4Reco/Overlay/components/OverlayTiming.h
@@ -94,6 +94,7 @@ struct OverlayTiming : public k4FWCore::MultiTransformer<retType(
   template <typename T> void overlayCollection(std::string collName, const podio::CollectionBase& inColl);
 
   virtual StatusCode initialize() final;
+  virtual StatusCode finalize() final;
 
   retType virtual operator()(
       const edm4hep::EventHeaderCollection& headers, const edm4hep::MCParticleCollection& mcParticles,
@@ -140,6 +141,9 @@ struct OverlayTiming : public k4FWCore::MultiTransformer<retType(
       this, "TimeWindows", std::map<std::string, std::vector<float>>(), "Time windows for the different collections"};
   Gaudi::Property<bool> m_allowReusingBackgroundFiles{
       this, "AllowReusingBackgroundFiles", false, "If true the same background file can be used for the same event"};
+  Gaudi::Property<bool> m_copyCellIDMetadata{this, "CopyCellIDMetadata", false,
+                                             "If metadata is found in the signal file, copy it to the output file, "
+                                             "replacing the old names with the new names"};
 
   // Gaudi::Property<int> m_maxCachedFrames{
   //   this, "MaxCachedFrames", 0, "Maximum number of frames cached from background files"};

diff --git a/k4Reco/Overlay/options/runOverlayTiming.py b/k4Reco/Overlay/options/runOverlayTiming.py
@@ -57,6 +57,7 @@
       ["background_group2_1.root"],
 ]
 overlay.TimeWindows = {"MCParticle": [0, 23.5], "VertexBarrelCollection": [0, 23.5], "VertexEndcapCollection": [0, 23.5], "HCalRingCollection": [0, 23.5]}
+overlay.CopyCellIDMetadata = True
 
 ApplicationMgr(TopAlg=[overlay],
                EvtSel="NONE",