WIP: add examples to read/write edm4hep.

Examples/CMakeLists.txt

@@ -4,12 +4,14 @@ gaudi_subdir(Examples v0r0)
 find_package(podio REQUIRED)
 find_package(plcio REQUIRED)
 find_package(LCIO REQUIRED)
+find_package(EDM4HEP REQUIRED)
 
 set(Examples_srcs
     src/HelloWorld/*.cpp
     src/FirstSvc/*.cpp
     src/SecondAlg/*.cpp
     src/PlcioTest/*.cpp
+    src/Edm4hepTest/*.cpp
 )
 
 # Headers and Libraries
@@ -19,7 +21,8 @@ gaudi_install_headers(Examples)
 # Modules
 gaudi_add_module(Examples ${Examples_srcs}
     INCLUDE_DIRS GaudiKernel FWCore ${plcio_INCLUDE_DIRS} ${podio_INCLUDE_DIRS} ${LCIO_INCLUDE_DIRS}
-    LINK_LIBRARIES GaudiKernel FWCore ${podio_LIBRARIES} ${LCIO_LIBRARIES} $ENV{PLCIO}/lib/libplcio.so
+    LINK_LIBRARIES GaudiKernel FWCore ${podio_LIBRARIES} ${LCIO_LIBRARIES} $ENV{PLCIO}/lib/libplcio.so 
+      EDM4HEP::edm4hep EDM4HEP::edm4hepDict
 )
 
 # Unit tests
@@ -35,6 +38,12 @@ gaudi_add_test(PlcioWriteAlg
 gaudi_add_test(PlcioReadAlg
                FRAMEWORK options/plcio_read.py)
 
+gaudi_add_test(Edm4hepWriteAlg
+               FRAMEWORK options/edm4hep_write.py)
+
+gaudi_add_test(Edm4hepReadAlg
+               FRAMEWORK options/edm4hep_read.py)
+
 gaudi_add_test(LCIOReadAlg
                FRAMEWORK options/LCIO_read.py)
 

Examples/options/edm4hep_read.py

+#!/usr/bin/env python
+
+from Gaudi.Configuration import *
+
+from Configurables import CEPCDataSvc
+dsvc = CEPCDataSvc("EventDataSvc", input="test.root")
+
+from Configurables import Edm4hepReadAlg
+alg = Edm4hepReadAlg("Edm4hepReadAlg")
+alg.HeaderCol.Path = "EventHeader"
+alg.InputCol.Path = "MCParticle"
+
+from Configurables import PodioInput
+podioinput = PodioInput("PodioReader", collections=[
+    "EventHeader",
+    "MCParticle"
+    ])
+
+# ApplicationMgr
+from Configurables import ApplicationMgr
+ApplicationMgr( TopAlg = [podioinput, alg],
+                EvtSel = 'NONE',
+                EvtMax = 10,
+                ExtSvc = [dsvc],
+                OutputLevel=DEBUG
+)

Examples/options/edm4hep_write.py

+#!/usr/bin/env python
+
+from Gaudi.Configuration import *
+
+from Configurables import CEPCDataSvc
+dsvc = CEPCDataSvc("EventDataSvc")
+
+from Configurables import Edm4hepWriteAlg
+alg = Edm4hepWriteAlg("Edm4hepWriteAlg")
+alg.HeaderCol.Path = "EventHeader"
+alg.OutputCol.Path = "MCParticle"
+
+from Configurables import PodioOutput
+out = PodioOutput("out")
+out.filename = "test.root"
+out.outputCommands = ["keep *"]
+
+# ApplicationMgr
+from Configurables import ApplicationMgr
+ApplicationMgr( TopAlg = [alg, out],
+                EvtSel = 'NONE',
+                EvtMax = 10,
+                ExtSvc=[dsvc],
+                OutputLevel=DEBUG
+)

Examples/src/Edm4hepTest/Edm4hepReadAlg.cpp

+#include "Edm4hepReadAlg.h"
+#include "edm4hep/EventHeaderCollection.h"
+#include "edm4hep/MCParticleCollection.h"
+
+DECLARE_COMPONENT(Edm4hepReadAlg)
+
+Edm4hepReadAlg::Edm4hepReadAlg(const std::string& name, ISvcLocator* svcLoc)
+    : GaudiAlgorithm(name, svcLoc)
+{
+    declareProperty("HeaderCol", m_headerCol);
+    declareProperty("InputCol", m_mcParCol, "MCParticle collection (input)");
+}
+
+StatusCode Edm4hepReadAlg::initialize()
+{
+    debug() << "begin initialize Edm4hepReadAlg" << endmsg;
+    return GaudiAlgorithm::initialize();
+}
+
+StatusCode Edm4hepReadAlg::execute()
+{
+    debug() << "begin execute Edm4hepReadAlg" << endmsg;
+
+    auto headers = m_headerCol.get();
+    auto header = headers->at(0);
+    auto mcCol = m_mcParCol.get();
+
+    info() << "Run " << header.getRunNumber() << " Event " << header.getEventNumber() << " { ";
+    for ( auto p : *mcCol ) {
+        info() << p.getObjectID().index << " : [";
+        for ( auto it = p.daughters_begin(), end = p.daughters_end(); it != end; ++it ) {
+            info() << " " << it->getObjectID().index;
+        }
+        info() << " ]; ";
+    }
+    info() << "}" << endmsg;
+
+    return StatusCode::SUCCESS;
+}
+
+StatusCode Edm4hepReadAlg::finalize()
+{
+    debug() << "begin finalize Edm4hepReadAlg" << endmsg;
+    return GaudiAlgorithm::finalize();
+}

Examples/src/Edm4hepTest/Edm4hepReadAlg.h

+#ifndef TEST_EDM4HEP_WRITE_ALG_H
+#define TEST_EDM4HEP_WRITE_ALG_H
+
+#include "FWCore/DataHandle.h"
+#include "GaudiAlg/GaudiAlgorithm.h"
+
+namespace edm4hep {
+    class EventHeaderCollection;
+    class MCParticleCollection;
+}
+
+class Edm4hepReadAlg : public GaudiAlgorithm
+{
+
+    public :
+
+        Edm4hepReadAlg(const std::string& name, ISvcLocator* svcLoc);
+
+        virtual StatusCode initialize();
+        virtual StatusCode execute();
+        virtual StatusCode finalize();
+
+    private :
+
+        DataHandle<edm4hep::EventHeaderCollection> m_headerCol{"EventHeader", Gaudi::DataHandle::Reader, this};
+        DataHandle<edm4hep::MCParticleCollection> m_mcParCol{"MCParticle", Gaudi::DataHandle::Reader, this};
+
+};
+
+#endif  // TEST_EDM4HEP_WRITE_ALG_H

Examples/src/Edm4hepTest/Edm4hepWriteAlg.cpp

+#include "Edm4hepWriteAlg.h"
+#include "edm4hep/EventHeaderCollection.h"
+#include "edm4hep/MCParticleCollection.h"
+
+DECLARE_COMPONENT(Edm4hepWriteAlg)
+
+Edm4hepWriteAlg::Edm4hepWriteAlg(const std::string& name, ISvcLocator* svcLoc)
+    : GaudiAlgorithm(name, svcLoc)
+{
+    declareProperty("HeaderCol", m_headerCol);
+    declareProperty("OutputCol", m_mcParCol, "MCParticle collection (output)");
+}
+
+StatusCode Edm4hepWriteAlg::initialize()
+{
+    debug() << "begin initialize Edm4hepWriteAlg" << endmsg;
+    return GaudiAlgorithm::initialize();
+}
+
+StatusCode Edm4hepWriteAlg::execute()
+{
+    debug() << "begin execute Edm4hepWriteAlg" << endmsg;
+
+    static int evtNo = 0;
+
+    auto headers = m_headerCol.createAndPut();
+    auto header = headers->create();
+    header.setRunNumber(-999);
+    header.setEventNumber(evtNo++);
+    // header.setDetectorName("TEST");
+
+    //auto mcCol = new edm4hep::MCParticleCollection;
+    //m_mcParCol.put(mcCol);
+    auto mcCol = m_mcParCol.createAndPut();
+
+    auto p1 = mcCol->create();
+    auto p2 = mcCol->create();
+
+    for ( int i = 0; i < 4; ++i ) {
+        auto d = mcCol->create();
+        d.addParent(p1);
+        d.addParent(p2);
+        p1.addDaughter(d);
+        p2.addDaughter(d);
+    }
+
+    return StatusCode::SUCCESS;
+}
+
+StatusCode Edm4hepWriteAlg::finalize()
+{
+    debug() << "begin finalize Edm4hepWriteAlg" << endmsg;
+    return GaudiAlgorithm::finalize();
+}

Examples/src/Edm4hepTest/Edm4hepWriteAlg.h

+#ifndef TEST_EDM4HEP_WRITE_ALG_H
+#define TEST_EDM4HEP_WRITE_ALG_H
+
+#include "FWCore/DataHandle.h"
+#include "GaudiAlg/GaudiAlgorithm.h"
+
+namespace edm4hep {
+    class EventHeaderCollection;
+    class MCParticleCollection;
+}
+
+class Edm4hepWriteAlg : public GaudiAlgorithm
+{
+
+    public :
+
+        Edm4hepWriteAlg(const std::string& name, ISvcLocator* svcLoc);
+
+        virtual StatusCode initialize();
+        virtual StatusCode execute();
+        virtual StatusCode finalize();
+
+    private :
+
+        DataHandle<edm4hep::EventHeaderCollection> m_headerCol{"EventHeader", Gaudi::DataHandle::Writer, this};
+        DataHandle<edm4hep::MCParticleCollection> m_mcParCol{"MCParticle", Gaudi::DataHandle::Writer, this};
+
+};
+
+#endif  // TEST_EDM4HEP_WRITE_ALG_H

Simulation/DetSimAna/src/Edm4hepWriterAnaElemTool.cpp

+#include "Edm4hepWriterAnaElemTool.h"
+
+#include "G4Event.hh"
+#include "G4THitsCollection.hh"
+
+#include "DD4hep/Detector.h"
+#include "DD4hep/Plugins.h"
+#include "DDG4/Geant4Converter.h"
+#include "DDG4/Geant4Mapping.h"
+#include "DDG4/Geant4HitCollection.h"
+#include "DDG4/Geant4Data.h"
+#include "DDG4/Geant4Hits.h"
+
+DECLARE_COMPONENT(Edm4hepWriterAnaElemTool)
+
+void
+Edm4hepWriterAnaElemTool::BeginOfRunAction(const G4Run*) {
+    G4cout << "Begin Run of detector simultion..." << G4endl;
+}
+
+void
+Edm4hepWriterAnaElemTool::EndOfRunAction(const G4Run*) {
+    G4cout << "End Run of detector simultion..." << G4endl;
+}
+
+void
+Edm4hepWriterAnaElemTool::BeginOfEventAction(const G4Event* anEvent) {
+    msg() << "Event " << anEvent->GetEventID() << endmsg;
+}
+
+void
+Edm4hepWriterAnaElemTool::EndOfEventAction(const G4Event* anEvent) {
+    auto mcCol = m_mcParCol.get();
+    msg() << "mcCol size: " << mcCol->size() << endmsg;
+    // save all data
+
+    // create collections.
+    auto trackercols = m_trackerCol.createAndPut();
+    auto calorimetercols = m_calorimeterCol.createAndPut();
+    auto calocontribcols = m_caloContribCol.createAndPut();
+
+    auto vxdcols = m_VXDCol.createAndPut();
+    auto ftdcols = m_FTDCol.createAndPut();
+    auto sitcols = m_SITCol.createAndPut();
+    auto tpccols = m_TPCCol.createAndPut();
+    auto setcols = m_SETCol.createAndPut();
+
+    // readout defined in DD4hep
+    auto lcdd = &(dd4hep::Detector::getInstance());
+    auto allReadouts = lcdd->readouts();
+
+    for (auto& readout : allReadouts) {
+        info() << "Readout " << readout.first << endmsg;
+    }
+
+    // retrieve the hit collections
+    G4HCofThisEvent* collections = anEvent->GetHCofThisEvent();
+    if (!collections) {
+        warning() << "No collections found. " << endmsg;
+        return;
+    }
+    int Ncol = collections->GetNumberOfCollections();
+    for (int icol = 0; icol < Ncol; ++icol) {
+        G4VHitsCollection* collect = collections->GetHC(icol);
+        if (!collect) {
+            warning() << "Collection iCol " << icol << " is missing" << endmsg;
+            continue;
+        }
+        size_t nhits = collect->GetSize();
+        info() << "Collection " << collect->GetName()
+               << " #" << icol
+               << " has " << nhits << " hits."
+               << endmsg;
+        if (nhits==0) {
+            // just skip this collection.
+            continue;
+        }
+
+        edm4hep::SimTrackerHitCollection* tracker_col_ptr = nullptr;
+        edm4hep::SimCalorimeterHitCollection* calo_col_ptr = nullptr;
+
+        // the mapping between hit collection and the data handler
+        if (collect->GetName() == "VXDCollection") {
+            tracker_col_ptr = vxdcols;
+        } else if (collect->GetName() == "FTDCollection") {
+            tracker_col_ptr = ftdcols;
+        } else if (collect->GetName() == "SITCollection") {
+            tracker_col_ptr = sitcols;
+        } else if (collect->GetName() == "TPCCollection") {
+            tracker_col_ptr = tpccols;
+        } else if (collect->GetName() == "SETCollection") {
+            tracker_col_ptr = setcols;
+        } else if (collect->GetName() == "SETCollection") {
+            tracker_col_ptr = setcols;
+        } else if (collect->GetName() == "CaloHitsCollection") {
+            calo_col_ptr = calorimetercols;
+        } else {
+            warning() << "Unknown collection name: " << collect->GetName()
+                      << ". The SimTrackerCol will be used. " << endmsg;
+            tracker_col_ptr = trackercols;
+        }
+
+
+
+        // There are different types (new and old)
+
+        dd4hep::sim::Geant4HitCollection* coll = dynamic_cast<dd4hep::sim::Geant4HitCollection*>(collect);
+        if (coll) {
+            info() << " cast to dd4hep::sim::Geant4HitCollection. " << endmsg;
+            for(size_t i=0; i<nhits; ++i)   {
+
+                dd4hep::sim::Geant4HitData* h = coll->hit(i);
+
+                dd4hep::sim::Geant4Tracker::Hit* trk_hit = dynamic_cast<dd4hep::sim::Geant4Tracker::Hit*>(h);
+                if ( 0 != trk_hit )   {
+                    dd4hep::sim::Geant4HitData::Contribution& t = trk_hit->truth;
+                    int trackID = t.trackID;
+                    // t.trackID = m_truth->particleID(trackID);
+                }
+                // Geant4Calorimeter::Hit* cal_hit = dynamic_cast<Geant4Calorimeter::Hit*>(h);
+                // if ( 0 != cal_hit )   {
+                //     Geant4HitData::Contributions& c = cal_hit->truth;
+                //     for(Geant4HitData::Contributions::iterator j=c.begin(); j!=c.end(); ++j)  {
+                //         Geant4HitData::Contribution& t = *j;
+                //         int trackID = t.trackID;
+                //         // t.trackID = m_truth->particleID(trackID);
+                //     }
+                // }
+            }
+            continue;
+        }
+
+        typedef G4THitsCollection<dd4hep::sim::Geant4Hit> HitCollection;
+        HitCollection* coll2 = dynamic_cast<HitCollection*>(collect);
+
+        if (coll2) {
+            info() << " cast to G4THitsCollection<dd4hep::sim::Geant4Hit>. " << endmsg;
+
+            int n_trk_hit = 0;
+            int n_cal_hit = 0;
+
+            for(size_t i=0; i<nhits; ++i)   {
+                dd4hep::sim::Geant4Hit* h = dynamic_cast<dd4hep::sim::Geant4Hit*>(coll2->GetHit(i));
+                if (!h) {
+                    warning() << "Failed to cast to dd4hep::sim::Geant4Hit. " << endmsg;
+                    continue;
+                }
+
+                dd4hep::sim::Geant4TrackerHit* trk_hit = dynamic_cast<dd4hep::sim::Geant4TrackerHit*>(h);
+                if (trk_hit) {
+                    ++n_trk_hit;
+                    // auto edm_trk_hit = trackercols->create();
+                    auto edm_trk_hit = tracker_col_ptr->create();
+
+                    edm_trk_hit.setCellID(trk_hit->cellID);
+                    edm_trk_hit.setEDep(trk_hit->energyDeposit/CLHEP::GeV);
+                    edm_trk_hit.setTime(trk_hit->truth.time/CLHEP::ns);
+                    edm_trk_hit.setPathLength(trk_hit->length/CLHEP::mm);
+                    // lc_hit->setMCParticle(lc_mcp);
+                    double pos[3] = {trk_hit->position.x()/CLHEP::mm,
+                                     trk_hit->position.y()/CLHEP::mm,
+                                     trk_hit->position.z()/CLHEP::mm};
+                    edm_trk_hit.setPosition(edm4hep::Vector3d(pos));
+
+                    float mom[3] = {trk_hit->momentum.x()/CLHEP::GeV,
+                                    trk_hit->momentum.y()/CLHEP::GeV,
+                                    trk_hit->momentum.z()/CLHEP::GeV};
+                    edm_trk_hit.setMomentum(edm4hep::Vector3f(mom));
+                }
+
+                dd4hep::sim::Geant4CalorimeterHit* cal_hit = dynamic_cast<dd4hep::sim::Geant4CalorimeterHit*>(h);
+                if (cal_hit) {
+                    ++n_cal_hit;
+                    auto edm_calo_hit = calo_col_ptr->create();
+                    edm_calo_hit.setCellID(cal_hit->cellID);
+                    edm_calo_hit.setEnergy(cal_hit->energyDeposit);
+                    float pos[3] = {cal_hit->position.x()/CLHEP::mm,
+                                    cal_hit->position.y()/CLHEP::mm,
+                                    cal_hit->position.z()/CLHEP::mm};
+                    edm_calo_hit.setPosition(edm4hep::Vector3f(pos));
+
+                    // contribution
+                    typedef dd4hep::sim::Geant4CalorimeterHit::Contributions Contributions;
+                    typedef dd4hep::sim::Geant4CalorimeterHit::Contribution Contribution;
+                    for (Contributions::const_iterator j = cal_hit->truth.begin();
+                         j != cal_hit->truth.end(); ++j) {
+                        const Contribution& c = *j;
+                        // The legacy Hit object does not contains positions of contributions.
+                        // float contrib_pos[] = {float(c.x/mm), float(c.y/mm), float(c.z/mm)};
+                        auto edm_calo_contrib = calocontribcols->create();
+                        edm_calo_contrib.setPDG(c.pdgID);
+                        edm_calo_contrib.setEnergy(c.deposit/CLHEP::GeV);
+                        edm_calo_contrib.setTime(c.time/CLHEP::ns);
+                        edm_calo_contrib.setStepPosition(edm4hep::Vector3f(pos));
+                        edm_calo_contrib.setParticle(mcCol->at(0)); // todo
+                        edm_calo_hit.addContribution(edm_calo_contrib);
+                    }
+                }
+
+            }
+
+            info() << n_trk_hit << " hits cast to dd4hep::sim::Geant4TrackerHit. " << endmsg;
+            info() << n_cal_hit << " hits cast to dd4hep::sim::Geant4CalorimeterHit. " << endmsg;
+
+
+            continue;
+        }
+
+        warning() << "Failed to convert to collection "
+                  << collect->GetName()
+                  << endmsg;
+        
+    }
+}
+
+void
+Edm4hepWriterAnaElemTool::PreUserTrackingAction(const G4Track*) {
+
+}
+
+void
+Edm4hepWriterAnaElemTool::PostUserTrackingAction(const G4Track*) {
+
+}
+
+void
+Edm4hepWriterAnaElemTool::UserSteppingAction(const G4Step*) {
+
+}
+
+StatusCode
+Edm4hepWriterAnaElemTool::initialize() {
+    StatusCode sc;
+
+    return sc;
+}
+
+StatusCode
+Edm4hepWriterAnaElemTool::finalize() {
+    StatusCode sc;
+
+    return sc;
+}
+
+

Simulation/DetSimAna/src/Edm4hepWriterAnaElemTool.h

+#ifndef Edm4hepWriterAnaElemTool_h
+#define Edm4hepWriterAnaElemTool_h
+
+#include "GaudiKernel/AlgTool.h"
+#include "FWCore/DataHandle.h"
+#include "DetSimInterface/IAnaElemTool.h"
+
+#include "edm4hep/MCParticleCollection.h"
+#include "edm4hep/SimTrackerHitCollection.h"
+#include "edm4hep/SimCalorimeterHitCollection.h"
+#include "edm4hep/CaloHitContributionCollection.h"
+
+class Edm4hepWriterAnaElemTool: public extends<AlgTool, IAnaElemTool> {
+
+public:
+
+    using extends::extends;
+
+    /// IAnaElemTool interface
+    // Run
+    virtual void BeginOfRunAction(const G4Run*) override;
+    virtual void EndOfRunAction(const G4Run*) override;
+
+    // Event
+    virtual void BeginOfEventAction(const G4Event*) override;
+    virtual void EndOfEventAction(const G4Event*) override;
+
+    // Tracking
+    virtual void PreUserTrackingAction(const G4Track*) override;
+    virtual void PostUserTrackingAction(const G4Track*) override;
+
+    // Stepping
+    virtual void UserSteppingAction(const G4Step*) override;
+
+
+    /// Overriding initialize and finalize
+    StatusCode initialize() override;
+    StatusCode finalize() override;
+
+private:
+    // In order to associate MCParticle with contribution, we need to access MC Particle.
+    DataHandle<edm4hep::MCParticleCollection> m_mcParCol{"MCParticle", 
+            Gaudi::DataHandle::Writer, this};
+
+    // Generic collections for Tracker and Calorimeter
+    DataHandle<edm4hep::SimTrackerHitCollection> m_trackerCol{"SimTrackerCol", 
+            Gaudi::DataHandle::Writer, this};
+    DataHandle<edm4hep::SimCalorimeterHitCollection> m_calorimeterCol{"SimCalorimeterCol", 
+            Gaudi::DataHandle::Writer, this};
+    DataHandle<edm4hep::CaloHitContributionCollection> m_caloContribCol{"SimCaloContributionCol", 
+            Gaudi::DataHandle::Writer, this};
+
+    // Dedicated collections for CEPC
+    DataHandle<edm4hep::SimTrackerHitCollection> m_VXDCol{"VXDCollection", 
+            Gaudi::DataHandle::Writer, this};
+    DataHandle<edm4hep::SimTrackerHitCollection> m_FTDCol{"FTDCollection", 
+            Gaudi::DataHandle::Writer, this};
+    DataHandle<edm4hep::SimTrackerHitCollection> m_SITCol{"SITCollection", 
+            Gaudi::DataHandle::Writer, this};
+    DataHandle<edm4hep::SimTrackerHitCollection> m_TPCCol{"TPCCollection", 
+            Gaudi::DataHandle::Writer, this};
+    DataHandle<edm4hep::SimTrackerHitCollection> m_SETCol{"SETCollection", 
+            Gaudi::DataHandle::Writer, this};
+
+};
+
+#endif