From 191c322759eddbdd36c5658bf3cae364fd227af3 Mon Sep 17 00:00:00 2001
From: Juraj Smiesko <juraj.smiesko@cern.ch>
Date: Mon, 13 Feb 2023 10:26:07 +0100
Subject: [PATCH] Grouping calo hits by call ID

---
 SimG4Common/include/SimG4Common/Units.h  |  10 ++
 SimG4Components/src/SimG4SaveCalHits.cpp | 124 +++++++++++++++++------
 SimG4Components/src/SimG4SaveCalHits.h   |   7 +-
 3 files changed, 107 insertions(+), 34 deletions(-)

diff --git a/SimG4Common/include/SimG4Common/Units.h b/SimG4Common/include/SimG4Common/Units.h
index 89dc85a..9820d7d 100644
--- a/SimG4Common/include/SimG4Common/Units.h
+++ b/SimG4Common/include/SimG4Common/Units.h
@@ -39,5 +39,15 @@ namespace edm2papas {
 const double length = edmdefault::length / CLHEP::m;
 const double energy = edmdefault::energy / CLHEP::GeV;
 }
+namespace edm2d4h {
+// FIXME: these should be a constexpr, but CLHEP is only const
+const double length = edmdefault::length / CLHEP::cm;
+const double energy = edmdefault::energy / CLHEP::keV;
+}
+namespace d4h2edm {
+// FIXME: these should be a constexpr, but CLHEP is only const
+const double length = CLHEP::cm / edmdefault::length;
+const double energy = CLHEP::keV / edmdefault::energy;
+}
 }
 #endif /* SIMG4COMMON_UNITS_H */
diff --git a/SimG4Components/src/SimG4SaveCalHits.cpp b/SimG4Components/src/SimG4SaveCalHits.cpp
index 26bfd47..87f613f 100644
--- a/SimG4Components/src/SimG4SaveCalHits.cpp
+++ b/SimG4Components/src/SimG4SaveCalHits.cpp
@@ -9,6 +9,7 @@
 #include "G4Event.hh"
 
 // datamodel
+#include "edm4hep/CaloHitContributionCollection.h"
 #include "edm4hep/SimCalorimeterHitCollection.h"
 
 // DD4hep
@@ -21,6 +22,7 @@ DECLARE_COMPONENT(SimG4SaveCalHits)
 SimG4SaveCalHits::SimG4SaveCalHits(const std::string& aType, const std::string& aName, const IInterface* aParent)
     : GaudiTool(aType, aName, aParent), m_geoSvc("GeoSvc", aName), m_eventDataSvc("EventDataSvc", "SimG4SaveCalHits")  {
   declareInterface<ISimG4SaveOutputTool>(this);
+  declareProperty("CaloHitContributions", m_caloHitContribs, "Handle for calo hit contributions");
   declareProperty("CaloHits", m_caloHits, "Handle for calo hits");
   declareProperty("GeoSvc", m_geoSvc);
 }
@@ -60,40 +62,98 @@ StatusCode SimG4SaveCalHits::initialize() {
 StatusCode SimG4SaveCalHits::finalize() { return GaudiTool::finalize(); }
 
 StatusCode SimG4SaveCalHits::saveOutput(const G4Event& aEvent) {
-  G4HCofThisEvent* collections = aEvent.GetHCofThisEvent();
-  G4VHitsCollection* collect;
-  k4::Geant4CaloHit* hit;
-  if (collections != nullptr) {
-    auto edmHits = m_caloHits.createAndPut();
-    for (int iter_coll = 0; iter_coll < collections->GetNumberOfCollections(); iter_coll++) {
-      collect = collections->GetHC(iter_coll);
-      if (std::find(m_readoutNames.begin(), m_readoutNames.end(), collect->GetName()) != m_readoutNames.end()) {
-        // Add CellID encoding string to collection metadata
-        auto lcdd = m_geoSvc->lcdd();
-        auto allReadouts = lcdd->readouts();
-        auto idspec = lcdd->idSpecification(collect->GetName());
-        auto field_str = idspec.fieldDescription();
-        auto& coll_md = m_podioDataSvc->getProvider().getCollectionMetaData( m_caloHits.get()->getID() );
-        coll_md.setValue("CellIDEncodingString", field_str);
-
-        size_t n_hit = collect->GetSize();
-        debug() << "\t" << n_hit << " hits are stored in a collection #" << iter_coll << ": " << collect->GetName()
-                << endmsg;
-        for (size_t iter_hit = 0; iter_hit < n_hit; iter_hit++) {
-          hit = dynamic_cast<k4::Geant4CaloHit*>(collect->GetHit(iter_hit));
-          auto edmHit = edmHits->create();
-          edmHit.setCellID(hit->cellID);
-          //todo
-          //edmHitCore.bits = hit->trackId;
-          edmHit.setEnergy(hit->energyDeposit * sim::g42edm::energy);
-          edmHit.setPosition({
-                       (float) hit->position.x() * (float) sim::g42edm::length,
-                       (float) hit->position.y() * (float) sim::g42edm::length,
-                       (float) hit->position.z() * (float) sim::g42edm::length,
-          });
-        }
+  G4HCofThisEvent* eventCollections = aEvent.GetHCofThisEvent();
+  if (!eventCollections) {
+    debug() << "Event collections not found." << endmsg;
+    return StatusCode::SUCCESS;
+  }
+
+  auto edmHitContribs = m_caloHitContribs.createAndPut();
+  auto edmHits = m_caloHits.createAndPut();
+  for (int iter_coll = 0; iter_coll < eventCollections->GetNumberOfCollections(); iter_coll++) {
+    G4VHitsCollection* collection = eventCollections->GetHC(iter_coll);
+    if (std::find(m_readoutNames.begin(), m_readoutNames.end(), collection->GetName()) == m_readoutNames.end()) {
+      continue;
+    }
+
+    // Add CellID encoding string to collection metadata
+    auto lcdd = m_geoSvc->lcdd();
+    auto allReadouts = lcdd->readouts();
+    auto idspec = lcdd->idSpecification(collection->GetName());
+    auto field_str = idspec.fieldDescription();
+    auto& coll_md = m_podioDataSvc->getProvider().getCollectionMetaData(m_caloHits.get()->getID());
+    coll_md.setValue("CellIDEncodingString", field_str);
+
+    // Lump hit contributions together
+    size_t nHit = collection->GetSize();
+    std::map<uint64_t, std::vector<k4::Geant4CaloHit*>> contribsInCells;
+    for (size_t i = 0; i < nHit; ++i) {
+      auto hit = dynamic_cast<k4::Geant4CaloHit*>(collection->GetHit(i));
+      contribsInCells[hit->cellID].emplace_back(hit);
+    }
+
+    for (const auto& [cellID, contribVec] : contribsInCells) {
+      // Cell energy
+      double energyDeposit = 0;
+      for (const auto* contrib : contribVec) {
+        energyDeposit += contrib->energyDeposit;
+      }
+
+      // Cell position
+      dd4hep::DDSegmentation::CellID volumeID = cellID;
+      auto detElement = lcdd->volumeManager().lookupDetElement(volumeID);
+      auto transformMatrix = detElement.nominal().worldTransformation();
+      auto segmentation = lcdd->readout(collection->GetName()).segmentation().segmentation();
+      const dd4hep::DDSegmentation::Vector3D cellPositionVecLocal = segmentation->position(cellID);
+
+      double cellPositionLocal[] = {cellPositionVecLocal.x(),
+                                    cellPositionVecLocal.y(),
+                                    cellPositionVecLocal.z()};
+      double cellPositionGlobal[3];
+      transformMatrix.LocalToMaster(cellPositionLocal, cellPositionGlobal);
+
+      // Fill the cell hit and contributions
+      auto edmHit = edmHits->create();
+      edmHit.setCellID(cellID);
+      edmHit.setEnergy(energyDeposit * sim::g42edm::energy);
+      edmHit.setPosition({
+        (float) cellPositionGlobal[0] * (float) sim::d4h2edm::length,
+        (float) cellPositionGlobal[1] * (float) sim::d4h2edm::length,
+        (float) cellPositionGlobal[2] * (float) sim::d4h2edm::length,
+      });
+
+      debug() << "Cell ID: " << edmHit.getCellID() << endmsg;
+      debug() << "Cell energy: " << edmHit.getEnergy() << endmsg;
+      debug() << "Cell global position:" << endmsg;
+      debug() << "  x: " << edmHit.getPosition().x << endmsg;
+      debug() << "  y: " << edmHit.getPosition().y << endmsg;
+      debug() << "  z: " << edmHit.getPosition().z << endmsg;
+
+      for (const auto* contrib : contribVec) {
+        auto edmHitContrib = edmHitContribs->create();
+        edmHitContrib.setPDG(contrib->pdgId);
+        edmHitContrib.setEnergy(contrib->energyDeposit * sim::g42edm::energy);
+        edmHitContrib.setTime(contrib->time);
+        edmHitContrib.setStepPosition({
+            (float) contrib->position.x() * (float) sim::g42edm::length,
+            (float) contrib->position.y() * (float) sim::g42edm::length,
+            (float) contrib->position.z() * (float) sim::g42edm::length,
+        });
+        edmHit.addToContributions(edmHitContrib);
+
+        debug() << "Contribution:" << endmsg;
+        debug() << "  PDG ID: " << edmHitContrib.getPDG() << endmsg;
+        debug() << "  energy:  " << edmHitContrib.getEnergy() << endmsg;
+        debug() << "  time:  " << edmHitContrib.getTime() << endmsg;
+        debug() << "  position: " << endmsg;
+        debug() << "    x: " << edmHitContrib.getStepPosition().x << endmsg;
+        debug() << "    y: " << edmHitContrib.getStepPosition().y << endmsg;
+        debug() << "    z: " << edmHitContrib.getStepPosition().z << endmsg;
+        debug() << "  track ID" << contrib->trackId << endmsg;
       }
+
     }
   }
+
   return StatusCode::SUCCESS;
 }
diff --git a/SimG4Components/src/SimG4SaveCalHits.h b/SimG4Components/src/SimG4SaveCalHits.h
index 053d625..39916b1 100644
--- a/SimG4Components/src/SimG4SaveCalHits.h
+++ b/SimG4Components/src/SimG4SaveCalHits.h
@@ -11,7 +11,8 @@ class IGeoSvc;
 
 // datamodel
 namespace edm4hep {
-class SimCalorimeterHitCollection;
+  class CaloHitContributionCollection;
+  class SimCalorimeterHitCollection;
 }
 
 /** @class SimG4SaveCalHits SimG4Components/src/SimG4SaveCalHits.h SimG4SaveCalHits.h
@@ -51,7 +52,9 @@ class SimG4SaveCalHits : public GaudiTool, virtual public ISimG4SaveOutputTool {
   /// Pointer to Podio and Event Data Services
   PodioDataSvc* m_podioDataSvc;
   ServiceHandle<IDataProviderSvc> m_eventDataSvc;
-  /// Handle for calo hits
+  /// Handle for calorimeter hit contributions
+  DataHandle<edm4hep::CaloHitContributionCollection> m_caloHitContribs{"CaloHitContributions", Gaudi::DataHandle::Writer, this};
+  /// Handle for calorimeter hits
   DataHandle<edm4hep::SimCalorimeterHitCollection> m_caloHits{"CaloHits", Gaudi::DataHandle::Writer, this};
   /// Name of the readouts (hits collections) to save
   Gaudi::Property<std::vector<std::string>> m_readoutNames{