Newer
Older
#include "Edm4hepWriterAnaElemTool.h"
#include "G4Event.hh"
#include "G4THitsCollection.hh"
#include "G4EventManager.hh"
#include "G4TrackingManager.hh"
#include "G4SteppingManager.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 "DetSimSD/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;
// event info
m_userinfo = nullptr;
if (anEvent->GetUserInformation()) {
m_userinfo = dynamic_cast<CommonUserEventInfo*>(anEvent->GetUserInformation());
m_userinfo->dumpIdxG4Track2Edm4hep();
}
auto mcGenCol = m_mcParGenCol.get();
mcCol = m_mcParCol.createAndPut();
// copy the MC particle first
for (auto mcGenParticle: *mcGenCol) {
auto newparticle = mcCol->create();
newparticle.setPDG (mcGenParticle.getPDG());
newparticle.setGeneratorStatus(mcGenParticle.getGeneratorStatus());
newparticle.setSimulatorStatus(mcGenParticle.getSimulatorStatus());
newparticle.setCharge (mcGenParticle.getCharge());
newparticle.setTime (mcGenParticle.getTime());
newparticle.setMass (mcGenParticle.getMass());
newparticle.setVertex (mcGenParticle.getVertex());
newparticle.setEndpoint (mcGenParticle.getEndpoint());
newparticle.setMomentum (mcGenParticle.getMomentum());
newparticle.setMomentumAtEndpoint(mcGenParticle.getMomentumAtEndpoint());
newparticle.setSpin (mcGenParticle.getSpin());
newparticle.setColorFlow (mcGenParticle.getColorFlow());
}
msg() << "mcCol size (original) : " << mcCol->size() << endmsg;
// reset
m_track2primary.clear();
}
void
Edm4hepWriterAnaElemTool::EndOfEventAction(const G4Event* anEvent) {
// 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();
auto ecalbarrelcol = m_EcalBarrelCol.createAndPut();
auto ecalbarrelcontribcols = m_EcalBarrelContributionCol.createAndPut();
auto ecalendcapscol = m_EcalEndcapsCol.createAndPut();
auto ecalendcapscontribcols = m_EcalEndcapsContributionCol.createAndPut();
auto ecalendcapringcol = m_EcalEndcapRingCol.createAndPut();
auto ecalendcapringcontribcol = m_EcalEndcapRingContributionCol.createAndPut();
auto hcalbarrelcol = m_HcalBarrelCol.createAndPut();
auto hcalbarrelcontribcols = m_HcalBarrelContributionCol.createAndPut();
auto hcalendcapscol = m_HcalEndcapsCol.createAndPut();
auto hcalendcapscontribcols = m_HcalEndcapsContributionCol.createAndPut();
auto hcalendcapringcol = m_HcalEndcapRingCol.createAndPut();
auto hcalendcapringcontribcol = m_HcalEndcapRingContributionCol.createAndPut();
auto coilcols = m_COILCol.createAndPut();
auto muonbarrelcol = m_MuonBarrelCol.createAndPut();
auto muonbarrelcontribcols = m_MuonBarrelContributionCol.createAndPut();
auto muonendcapscol = m_MuonEndcapsCol.createAndPut();
auto muonendcapscontribcols = m_MuonEndcapsContributionCol.createAndPut();
auto driftchamberhitscol = m_DriftChamberHitsCol.createAndPut();
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
// 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;
edm4hep::CaloHitContributionCollection* calo_contrib_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;
calo_contrib_col_ptr = calocontribcols;
} else if (collect->GetName() == "EcalBarrelCollection") {
calo_col_ptr = ecalbarrelcol;
calo_contrib_col_ptr = ecalbarrelcontribcols;
} else if (collect->GetName() == "EcalEndcapsCollection") {
calo_col_ptr = ecalendcapscol;
calo_contrib_col_ptr = ecalendcapscontribcols;
} else if (collect->GetName() == "EcalEndcapRingCollection") {
calo_col_ptr = ecalendcapringcol;
calo_contrib_col_ptr = ecalendcapringcontribcol;
} else if (collect->GetName() == "HcalBarrelCollection") {
calo_col_ptr = hcalbarrelcol;
calo_contrib_col_ptr = hcalbarrelcontribcols;
} else if (collect->GetName() == "HcalEndcapsCollection") {
calo_col_ptr = hcalendcapscol;
calo_contrib_col_ptr = hcalendcapscontribcols;
} else if (collect->GetName() == "HcalEndcapRingCollection") {
calo_col_ptr = hcalendcapringcol;
calo_contrib_col_ptr = hcalendcapringcontribcol;
} else if (collect->GetName() == "COILCollection") {
tracker_col_ptr = coilcols;
} else if (collect->GetName() == "MuonBarrelCollection") {
calo_col_ptr = muonbarrelcol;
calo_contrib_col_ptr = muonbarrelcontribcols;
} else if (collect->GetName() == "MuonEndcapsCollection") {
calo_col_ptr = muonendcapscol;
calo_contrib_col_ptr = muonendcapscontribcols;
} else if (collect->GetName() == "DriftChamberHitsCollection") {
tracker_col_ptr = driftchamberhitscol;
} else {
warning() << "Unknown collection name: " << collect->GetName()
<< ". Please register in Edm4hepWriterAnaElemTool. " << endmsg;
continue;
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
}
// 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);
++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));
// get the truth or contribution
auto& truth = trk_hit->truth;
int trackID = truth.trackID;
int pritrkid = m_track2primary[trackID];
if (pritrkid <= 0) {
error() << "Failed to find the primary track for trackID #" << trackID << endmsg;
pritrkid = 1;
}
if (m_userinfo) {
auto idxedm4hep = m_userinfo->idxG4Track2Edm4hep(pritrkid);
if (idxedm4hep != -1) {
edm_trk_hit.setMCParticle(mcCol->at(idxedm4hep));
}
}
if (pritrkid != trackID) {
// If the track is a secondary, then the primary track id and current track id is different
edm_trk_hit.setProducedBySecondary(true);
}
}
dd4hep::sim::Geant4CalorimeterHit* cal_hit = dynamic_cast<dd4hep::sim::Geant4CalorimeterHit*>(h);
++n_cal_hit;
auto edm_calo_hit = calo_col_ptr->create();
edm_calo_hit.setCellID(cal_hit->cellID);
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 = calo_contrib_col_ptr->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));
// from the track id, get the primary track
int pritrkid = m_track2primary[c.trackID];
if (pritrkid<=0) {
error() << "Failed to find the primary track for trackID #" << c.trackID << endmsg;
pritrkid = 1;
}
if (m_userinfo) {
auto idxedm4hep = m_userinfo->idxG4Track2Edm4hep(pritrkid);
if (idxedm4hep != -1) {
edm_calo_contrib.setParticle(mcCol->at(idxedm4hep)); // todo
}
}
edm_calo_hit.addToContributions(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* track) {
int curtrkid = track->GetTrackID();
int curparid = track->GetParentID();
int pritrkid = curparid;
// try to find the primary track id from the parent track id.
if (curparid) {
auto it = m_track2primary.find(curparid);
if (it == m_track2primary.end()) {
error() << "Failed to find primary track for track id " << curparid << endmsg;
} else {
pritrkid = it->second;
}
} else {
// curparid is 0, it is primary
pritrkid = curtrkid;
}
m_track2primary[curtrkid] = pritrkid;
lintao@ihep.ac.cn
committed
Edm4hepWriterAnaElemTool::PostUserTrackingAction(const G4Track* track) {
int curtrkid = track->GetTrackID(); // starts from 1
int curparid = track->GetParentID();
lintao@ihep.ac.cn
committed
if (m_userinfo) {
idxedm4hep = m_userinfo->idxG4Track2Edm4hep(curtrkid);
}
if (curparid == 0) { // Primary Track
lintao@ihep.ac.cn
committed
// select the primary tracks (parentID == 0)
lintao@ihep.ac.cn
committed
if (idxedm4hep == -1) {
error () << "Failed to get idxedm4hep according to the g4track id " << curtrkid << endmsg;
return;
}
if (idxedm4hep>=mcCol->size()) {
lintao@ihep.ac.cn
committed
error() << "out of range: curtrkid is " << curtrkid
<< " while the MCParticle size is " << mcCol->size() << endmsg;
return;
}
auto primary_particle = mcCol->at(idxedm4hep);
lintao@ihep.ac.cn
committed
const G4ThreeVector& stop_pos = track->GetPosition();
edm4hep::Vector3d endpoint(stop_pos.x()/CLHEP::mm,
stop_pos.y()/CLHEP::mm,
stop_pos.z()/CLHEP::mm);
primary_particle.setEndpoint(endpoint);
const G4ThreeVector& stop_mom = track->GetMomentum();
edm4hep::Vector3f mom_endpoint(stop_mom.x()/CLHEP::GeV,
stop_mom.y()/CLHEP::GeV,
stop_mom.z()/CLHEP::GeV);
primary_particle.setMomentumAtEndpoint(mom_endpoint);
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
// ===================================================================
// Update the flags of the primary particles
// ===================================================================
// processes
static G4VProcess* proc_decay = nullptr;
if (!proc_decay) {
G4ProcessManager* pm
= track->GetDefinition()->GetProcessManager();
G4int nprocesses = pm->GetProcessListLength();
G4ProcessVector* pv = pm->GetProcessList();
for(G4int i=0; i<nprocesses; ++i){
if((*pv)[i]->GetProcessName()=="Decay"){
proc_decay = (*pv)[i];
}
}
}
// flags
bool is_decay = false;
G4TrackingManager* tm = G4EventManager::GetEventManager()
-> GetTrackingManager();
G4TrackVector* secondaries = tm->GimmeSecondaries();
if(secondaries) {
size_t nSeco = secondaries->size();
for (size_t i = 0; i < nSeco; ++i) {
G4Track* sectrk = (*secondaries)[i];
G4ParticleDefinition* secparticle = sectrk->GetDefinition();
const G4VProcess* creatorProcess = sectrk->GetCreatorProcess();
// select the necessary processes
if (creatorProcess==proc_decay) {
info() << "Creator Process is Decay for secondary particle: "
<< " idx: " << i
<< " trkid: " << sectrk->GetTrackID() // not valid until track
<< " particle: " << secparticle->GetParticleName()
<< " pdg: " << secparticle->GetPDGEncoding()
<< " at position: " << sectrk->GetPosition() //
<< " time: " << sectrk->GetGlobalTime()
<< " momentum: " << sectrk->GetMomentum() //
<< endmsg;
is_decay = true;
// create secondaries in MC particles
// todo: convert the const collection to non-const
// auto mcCol = const_cast<edm4hep::MCParticleCollection*>(m_mcParCol.get());
auto mcp = mcCol->create();
mcp.setPDG(secparticle->GetPDGEncoding());
mcp.setGeneratorStatus(0); // not created by Generator
mcp.setCreatedInSimulation(1);
mcp.setCharge(secparticle->GetPDGCharge());
mcp.setTime(sectrk->GetGlobalTime()/CLHEP::ns); // todo
mcp.setMass(secparticle->GetPDGMass());
const G4ThreeVector& sec_init_pos = sectrk->GetPosition();
double x=sec_init_pos.x()/CLHEP::mm;
double y=sec_init_pos.y()/CLHEP::mm;
double z=sec_init_pos.z()/CLHEP::mm;
const G4ThreeVector& sec_init_mom = sectrk->GetMomentum();
double px=sec_init_mom.x()/CLHEP::GeV;
double py=sec_init_mom.y()/CLHEP::GeV;
double pz=sec_init_mom.z()/CLHEP::GeV;
mcp.setVertex(edm4hep::Vector3d(x,y,z)); // todo
mcp.setEndpoint(edm4hep::Vector3d(x,y,z)); // todo
mcp.setMomentum(edm4hep::Vector3f(px,py,pz)); // todo
mcp.setMomentumAtEndpoint(edm4hep::Vector3f(px,py,pz)); //todo
mcp.addToParents(primary_particle);
primary_particle.addToDaughters(mcp);
// store the edm4hep obj idx in track info.
// using this idx, the MCParticle object could be modified later.
auto trackinfo = new CommonUserTrackInfo();
trackinfo->setIdxEdm4hep(mcp.getObjectID().index);
sectrk->SetUserInformation(trackinfo);
info() << " Appending MCParticle: (id: "
<< mcp.getObjectID().index << ")"
<< endmsg;
}
}
}
// now update
if (is_decay) {
primary_particle.setDecayedInTracker(is_decay);
primary_particle.setDecayedInCalorimeter(is_decay);
}
lintao@ihep.ac.cn
committed
} else {
// TODO: select other interested tracks
}
}
void
Edm4hepWriterAnaElemTool::UserSteppingAction(const G4Step*) {
}
StatusCode
Edm4hepWriterAnaElemTool::initialize() {
StatusCode sc;
return sc;
}
StatusCode
Edm4hepWriterAnaElemTool::finalize() {
StatusCode sc;
return sc;
}