DDG4/src/Geant4Output2ROOT.cpp

0001 //==========================================================================
0002 //  AIDA Detector description implementation
0003 //--------------------------------------------------------------------------
0004 // Copyright (C) Organisation europeenne pour la Recherche nucleaire (CERN)
0005 // All rights reserved.
0006 //
0007 // For the licensing terms see $DD4hepINSTALL/LICENSE.
0008 // For the list of contributors see $DD4hepINSTALL/doc/CREDITS.
0009 //
0010 // Author     : M.Frank
0011 //
0012 //==========================================================================
0013
0014 // Framework include files
0015 #include <DD4hep/Printout.h>
0016 #include <DD4hep/Primitives.h>
0017 #include <DD4hep/InstanceCount.h>
0018 #include <DDG4/Geant4HitCollection.h>
0019 #include <DDG4/Geant4Output2ROOT.h>
0020 #include <DDG4/Geant4Particle.h>
0021 #include <DDG4/Geant4Data.h>
0022
0023 // Geant4 include files
0024 #include <G4HCofThisEvent.hh>
0025 #include <G4ParticleTable.hh>
0026 #include <G4Run.hh>
0027
0028 // ROOT include files
0029 #include <TFile.h>
0030 #include <TTree.h>
0031 #include <TBranch.h>
0032 #include <TSystem.h>
0033
0034
0035 using namespace dd4hep::sim;
0036
0037 /// Standard constructor
0038 Geant4Output2ROOT::Geant4Output2ROOT(Geant4Context* ctxt, const std::string& nam)
0039   : Geant4OutputAction(ctxt, nam), m_file(nullptr), m_tree(nullptr) {
0040   declareProperty("Section",              m_section = "EVENT");
0041   declareProperty("HandleMCTruth",        m_handleMCTruth = true);
0042   declareProperty("DisabledCollections",  m_disabledCollections);
0043   declareProperty("DisableParticles",     m_disableParticles);
0044   declareProperty("FilesByRun",           m_filesByRun = false);
0045   InstanceCount::increment(this);
0046 }
0047
0048 /// Default destructor
0049 Geant4Output2ROOT::~Geant4Output2ROOT() {
0050   closeOutput();
0051   InstanceCount::decrement(this);
0052 }
0053
0054 /// Close current output file
0055 void Geant4Output2ROOT::closeOutput()   {
0056   if (m_file) {
0057     TDirectory::TContext ctxt(m_file);
0058     Sections::iterator i = m_sections.find(m_section);
0059     info("+++ Closing ROOT output file %s", m_file->GetName());
0060     if ( i != m_sections.end() )
0061       m_sections.erase(i);
0062     m_branches.clear();
0063     m_tree->Write();
0064     m_file->Close();
0065     m_tree = nullptr;
0066     detail::deletePtr (m_file);
0067   }
0068 }
0069
0070 /// Create/access tree by name
0071 TTree* Geant4Output2ROOT::section(const std::string& nam) {
0072   Sections::const_iterator i = m_sections.find(nam);
0073   if (i == m_sections.end()) {
0074     TDirectory::TContext ctxt(m_file);
0075     TTree* t = new TTree(nam.c_str(), ("Geant4 " + nam + " information").c_str());
0076     m_sections.emplace(nam, t);
0077     return t;
0078   }
0079   return (*i).second;
0080 }
0081
0082 /// Callback to store the Geant4 run information
0083 void Geant4Output2ROOT::beginRun(const G4Run* run) {
0084   std::string fname = m_output;
0085   if ( m_filesByRun )    {
0086     size_t idx = m_output.rfind(".");
0087     if ( m_file )  {
0088       closeOutput();
0089     }
0090     fname  = m_output.substr(0, idx);
0091     fname += _toString(run->GetRunID(), ".run%08d");
0092     if ( idx != std::string::npos )
0093       fname += m_output.substr(idx);
0094   }
0095   if ( !m_file && !fname.empty() ) {
0096     TDirectory::TContext ctxt(TDirectory::CurrentDirectory());
0097     if ( !gSystem->AccessPathName(fname.c_str()) )  {
0098       gSystem->Unlink(fname.c_str());
0099     }
0100     std::unique_ptr<TFile> file(TFile::Open(fname.c_str(), "RECREATE", "dd4hep Simulation data"));
0101     if ( !file )  {
0102       file.reset(TFile::Open((fname+".1").c_str(), "RECREATE", "dd4hep Simulation data"));
0103     }
0104     if ( !file )  {
0105       except("Failed to create ROOT output file:'%s'", fname.c_str());
0106     }
0107     if (file->IsZombie()) {
0108       detail::deletePtr (m_file);
0109       except("Failed to open ROOT output file:'%s'", fname.c_str());
0110     }
0111     m_file = file.release();
0112     m_tree = section(m_section);
0113   }
0114   Geant4OutputAction::beginRun(run);
0115 }
0116
0117 /// Fill single EVENT branch entry (Geant4 collection data)
0118 int Geant4Output2ROOT::fill(const std::string& nam, const ComponentCast& type, void* ptr) {
0119   if (m_file) {
0120     TBranch* b = 0;
0121     Branches::const_iterator i = m_branches.find(nam);
0122     if (i == m_branches.end()) {
0123       const std::type_info& typ = type.type();
0124       TClass* cl = TBuffer::GetClass(typ);
0125       if (cl) {
0126         b = m_tree->Branch(nam.c_str(), cl->GetName(), (void*) 0);
0127         b->SetAutoDelete(false);
0128         m_branches.emplace(nam, b);
0129       }
0130       else {
0131         throw std::runtime_error("No ROOT TClass object availible for object type:" + typeName(typ));
0132       }
0133     }
0134     else {
0135       b = (*i).second;
0136     }
0137     Long64_t evt = b->GetEntries(), nevt = b->GetTree()->GetEntries(), num = nevt - evt;
0138     if (nevt > evt) {
0139       b->SetAddress(0);
0140       while (num > 0) {
0141         b->Fill();
0142         --num;
0143       }
0144     }
0145     b->SetAddress(&ptr);
0146     int nbytes = b->Fill();
0147     if (nbytes < 0) {
0148       throw std::runtime_error("Failed to write ROOT collection:" + nam + "!");
0149     }
0150     return nbytes;
0151   }
0152   return 0;
0153 }
0154
0155 /// Commit data at end of filling procedure
0156 void Geant4Output2ROOT::commit(OutputContext<G4Event>& ctxt) {
0157   if (m_file) {
0158     TObjArray* a = m_tree->GetListOfBranches();
0159     Long64_t evt = m_tree->GetEntries() + 1;
0160     Int_t nb = a->GetEntriesFast();
0161     /// Fill NULL pointers to all branches, which have less entries than the Event branch
0162     for (Int_t i = 0; i < nb; ++i) {
0163       TBranch* br_ptr = (TBranch*) a->UncheckedAt(i);
0164       Long64_t br_evt = br_ptr->GetEntries();
0165       if (br_evt < evt) {
0166         Long64_t num = evt - br_evt;
0167         br_ptr->SetAddress(0);
0168         while (num > 0) {
0169           br_ptr->Fill();
0170           --num;
0171         }
0172       }
0173     }
0174     m_tree->SetEntries(evt);
0175   }
0176   Geant4OutputAction::commit(ctxt);
0177 }
0178
0179 /// Callback to store the Geant4 event
0180 void Geant4Output2ROOT::saveEvent(OutputContext<G4Event>& /* ctxt */) {
0181   if ( !m_disableParticles )  {
0182     Geant4ParticleMap* parts = context()->event().extension<Geant4ParticleMap>();
0183     if ( parts )   {
0184       typedef Geant4HitWrapper::HitManipulator Manip;
0185       typedef Geant4ParticleMap::ParticleMap ParticleMap;
0186       Manip* manipulator = Geant4HitWrapper::manipulator<Geant4Particle>();
0187       G4ParticleTable* table = G4ParticleTable::GetParticleTable();
0188       const ParticleMap& pm = parts->particles();
0189       std::vector<void*> particles;
0190       particles.reserve(pm.size());
0191       for ( const auto& i : pm )   {
0192         auto* p = i.second;
0193         G4ParticleDefinition* def = table->FindParticle(p->pdgID);
0194         p->charge = int(3.0 * (def ? def->GetPDGCharge() : -1.0)); // Assume e-/pi-
0195         particles.emplace_back((ParticleMap::mapped_type*)p);
0196       }
0197       fill("MCParticles",manipulator->vec_type,&particles);
0198     }
0199   }
0200 }
0201
0202 /// Callback to store each Geant4 hit collection
0203 void Geant4Output2ROOT::saveCollection(OutputContext<G4Event>& /* ctxt */, G4VHitsCollection* collection) {
0204   Geant4HitCollection* coll = dynamic_cast<Geant4HitCollection*>(collection);
0205   std::string hc_nam = collection->GetName();
0206   for(const auto& n : m_disabledCollections)  {
0207     if ( n == hc_nam )   {
0208       return;
0209     }
0210   }
0211   if (coll) {
0212     std::vector<void*> hits;
0213     coll->getHitsUnchecked(hits);
0214     size_t nhits = coll->GetSize();
0215     if ( m_handleMCTruth && m_truth && nhits > 0 )   {
0216       hits.reserve(nhits);
0217       try  {
0218         for(size_t i=0; i<nhits; ++i)   {
0219           Geant4HitData* h = coll->hit(i);
0220           Geant4Tracker::Hit* trk_hit = dynamic_cast<Geant4Tracker::Hit*>(h);
0221           if ( 0 != trk_hit )   {
0222             Geant4HitData::Contribution& t = trk_hit->truth;
0223             int trackID = t.trackID;
0224             t.trackID = m_truth->particleID(trackID);
0225           }
0226           Geant4Calorimeter::Hit* cal_hit = dynamic_cast<Geant4Calorimeter::Hit*>(h);
0227           if ( 0 != cal_hit )   {
0228             Geant4HitData::Contributions& c = cal_hit->truth;
0229             for(Geant4HitData::Contributions::iterator j=c.begin(); j!=c.end(); ++j)  {
0230               Geant4HitData::Contribution& t = *j;
0231               int trackID = t.trackID;
0232               t.trackID = m_truth->particleID(trackID);
0233             }
0234           }
0235         }
0236       }
0237       catch(...)   {
0238         error("+++ Exception while saving collection %s.",hc_nam.c_str());
0239       }
0240     }
0241     fill(hc_nam, coll->vector_type(), &hits);
0242   }
0243 }