Newer
Older
#include "Edm4hepWriterAnaElemTool.h"
#include "G4Step.hh"
#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*) {
auto msglevel = msgLevel();
if (msglevel == MSG::VERBOSE || msglevel == MSG::DEBUG) {
verboseOutput = true;
}
G4cout << "Begin Run of detector simultion..." << G4endl;
// access geometry service
m_geosvc = service<IGeomSvc>("GeomSvc");
if (m_geosvc) {
dd4hep::Detector* dd4hep_geo = m_geosvc->lcdd();
// try to get the constants
R = dd4hep_geo->constant<double>("tracker_region_rmax")/dd4hep::mm*CLHEP::mm;
Z = dd4hep_geo->constant<double>("tracker_region_zmax")/dd4hep::mm*CLHEP::mm;
info() << "Tracker Region "
<< " R: " << R
<< " Z: " << Z
<< endmsg;
} else {
error() << "Failed to find GeomSvc." << endmsg;
}
}
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());
if (verboseOutput) {
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());
// Add parent-daughter relation
for(auto imc = 0; imc<mcGenParticle.parents_size(); imc++){
newparticle.addToParents(mcGenParticle.getParents(imc));
}
for(auto imc = 0; imc<mcGenParticle.daughters_size(); imc++){
newparticle.addToDaughters(mcGenParticle.getDaughters(imc));
}
msg() << "mcCol size (original) : " << mcCol->size() << endmsg;
// reset
m_track2primary.clear();
}
void
Edm4hepWriterAnaElemTool::EndOfEventAction(const G4Event* anEvent) {
msg() << "mcCol size (after simulation) : " << 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();
auto otkendcapcols = m_OTKEndCapCol.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 lumicalcol = m_LumicalCol.createAndPut();
auto lumicalconribcols = m_LumicalContributionCol.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();
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
// 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() == "OTKBarrelCollection") {
tracker_col_ptr = otkbarrelcols;
} else if (collect->GetName() == "OTKEndCapCollection") {
tracker_col_ptr = otkendcapcols;
} 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() == "LumicalCollection"){
calo_col_ptr = lumicalcol;
calo_contrib_col_ptr = lumicalconribcols;
} 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;
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
}
// 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);
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
// ===================================================================
// 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) {
debug() << "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);
debug() << " 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
}
Edm4hepWriterAnaElemTool::UserSteppingAction(const G4Step* aStep) {
auto aTrack = aStep->GetTrack();
// try to get user track info
auto trackinfo = dynamic_cast<CommonUserTrackInfo*>(aTrack->GetUserInformation());
// ========================================================================
// Note:
// if there is no track info, then do nothing.
// ========================================================================
if (trackinfo) {
// back scattering is defined as following:
// - pre point is not in tracker
// - post point is in tracker
// For details, look at Mokka's source code:
// https://llrforge.in2p3.fr/trac/Mokka/browser/trunk/source/Kernel/src/SteppingAction.cc
const auto& pre_pos = aStep->GetPreStepPoint()->GetPosition();
const auto& post_pos = aStep->GetPostStepPoint()->GetPosition();
bool is_pre_in_tracker = pre_pos.perp() < R && std::fabs(pre_pos.z()) < Z;
bool is_post_in_tracker = post_pos.perp() < R && std::fabs(post_pos.z()) < Z;
if ((!is_pre_in_tracker) and is_post_in_tracker) {
// set back scattering
auto idxedm4hep = trackinfo->idxEdm4hep();
mcCol->at(idxedm4hep).setBackscatter(true);
debug() << " set back scatter for MCParticle "
<< " (ID: " << idxedm4hep << ")"
<< endmsg;
}
}