CaTS/src/RootIO.cc

0001 //
0002 // ********************************************************************
0003 // * License and Disclaimer                                           *
0004 // *                                                                  *
0005 // * The  Geant4 software  is  copyright of the Copyright Holders  of *
0006 // * the Geant4 Collaboration.  It is provided  under  the terms  and *
0007 // * conditions of the Geant4 Software License,  included in the file *
0008 // * LICENSE and available at  http://cern.ch/geant4/license .  These *
0009 // * include a list of copyright holders.                             *
0010 // *                                                                  *
0011 // * Neither the authors of this software system, nor their employing *
0012 // * institutes,nor the agencies providing financial support for this *
0013 // * work  make  any representation or  warranty, express or implied, *
0014 // * regarding  this  software system or assume any liability for its *
0015 // * use.  Please see the license in the file  LICENSE  and URL above *
0016 // * for the full disclaimer and the limitation of liability.         *
0017 // *                                                                  *
0018 // * This  code  implementation is the result of  the  scientific and *
0019 // * technical work of the GEANT4 collaboration.                      *
0020 // * By using,  copying,  modifying or  distributing the software (or *
0021 // * any work based  on the software)  you  agree  to acknowledge its *
0022 // * use  in  resulting  scientific  publications,  and indicate your *
0023 // * acceptance of all terms of the Geant4 Software license.          *
0024 // ********************************************************************
0025 //
0026 // CaTS (Calorimetry and Tracking Simulation)
0027 //
0028 // Authors: Hans Wenzel and Soon Yung Jun
0029 //          (Fermi National Accelerator Laboratory)
0030 //
0031 // History:
0032 // October  18th, 2021 : first implementation
0033 // November 10th, 2021 : implement writing one file per worker thread
0034 //                       which are then merged
0035 // ********************************************************************
0036 //
0037 /// \file RootIO.cc
0038 /// \brief Implementation of the CaTS::RootIO class
0039
0040 #ifdef WITH_ROOT
0041 // Project headers
0042 #include "RootIO.hh"
0043 #include "ConfigurationManager.hh"
0044 #include "Event.hh"
0045 #include "PhotonSD.hh"
0046 #include "PhotonHit.hh"
0047 #include "InteractionHit.hh"
0048 #include "lArTPCHit.hh"
0049 #include "TrackerHit.hh"
0050 #include "MscHit.hh"
0051 #include "CalorimeterHit.hh"
0052 #include "DRCalorimeterHit.hh"
0053 // Geant4 headers
0054 #include "G4Event.hh"
0055 #include "G4ios.hh"
0056 #include "G4Threading.hh"
0057 #include "G4RunManager.hh"
0058 #include "G4HCofThisEvent.hh"
0059 #include "G4VHit.hh"
0060 #include "G4VHitsCollection.hh"
0061 #include "G4SDManager.hh"
0062 // Root headers
0063 #include "TBranch.h"
0064 #include "TFile.h"
0065 #include "TObject.h"
0066 #include "TROOT.h"
0067 #include "TSystem.h"
0068 #include "TTree.h"
0069
0070 // G4Mutex in a file scope
0071 namespace
0072 {
0073   // Mutex to lock RootIO constructor
0074   G4Mutex RootIOMutex = G4MUTEX_INITIALIZER;
0075 } // namespace
0076
0077 //....oooOO0OOooo........oooOO0OOooo........oooOO0OOooo........oooOO0OOooo......
0078 RootIO::RootIO()
0079 {
0080   G4AutoLock lock(&RootIOMutex);
0081
0082   ROOT::EnableThreadSafety();
0083   gSystem->Load("libCaTSClassesDict");
0084
0085   fFileName = ConfigurationManager::getInstance()->getfname() + ".root";
0086
0087   if (G4Threading::IsWorkerThread())
0088   {
0089     fFileName += std::to_string(G4Threading::G4GetThreadId());
0090   }
0091
0092   if (G4Threading::IsWorkerThread() ||
0093       !G4Threading::IsMultithreadedApplication())
0094   {
0095     G4cout << "RootIO:: Opening File: " << fFileName << G4endl;
0096     fFile = TFile::Open(fFileName.c_str(), "RECREATE");
0097     fTree = new TTree("Events", "ROOT tree containing Hit collections");
0098   }
0099 }
0100
0101 //....oooOO0OOooo........oooOO0OOooo........oooOO0OOooo........oooOO0OOooo......
0102 RootIO* RootIO::GetInstance()
0103 {
0104   static G4ThreadLocal RootIO instance;
0105   return &instance;
0106 }
0107
0108 //....oooOO0OOooo........oooOO0OOooo........oooOO0OOooo........oooOO0OOooo......
0109 void RootIO::Write(const G4Event* event)
0110 {
0111   G4bool verbose = ConfigurationManager::getInstance()->isEnable_verbose();
0112
0113   G4HCofThisEvent* HCE = event->GetHCofThisEvent();
0114   if (HCE == nullptr)
0115   {
0116     return;
0117   }
0118
0119   // Record hits from sensitive detectors into the CaTS event data
0120   Event* CaTSEvt = new Event();
0121   CaTSEvt->SetEventNr(event->GetEventID());
0122
0123   if (verbose)
0124   {
0125     G4cout << "Number of collections:  " << HCE->GetNumberOfCollections()
0126            << G4endl;
0127   }
0128
0129   for (int i = 0; i < HCE->GetNumberOfCollections(); i++)
0130   {
0131     G4VHitsCollection* hc = HCE->GetHC(i);
0132     std::vector<G4String> y = Split(hc->GetName(), '_');
0133     G4String Classname = y[1];
0134
0135     if (verbose)
0136     {
0137       G4cout << "Classname: " << Classname << " Collection size: "
0138          <<  hc->GetSize() << G4endl;
0139     }
0140
0141     if (Classname == "lArTPC")
0142     {
0143       AddHits<lArTPCHit>(hc, CaTSEvt);
0144     }
0145     else if (Classname == "PhotonDetector")
0146     {
0147       AddHits<PhotonHit>(hc, CaTSEvt);
0148     }
0149     else if (Classname == "Target")
0150     {
0151       AddHits<InteractionHit>(hc, CaTSEvt);
0152     }
0153     else if (Classname == "Tracker")
0154     {
0155       AddHits<TrackerHit>(hc, CaTSEvt);
0156     }
0157     else if (Classname == "Msc")
0158     {
0159       AddHits<MscHit>(hc, CaTSEvt);
0160     }
0161     else if (Classname == "Calorimeter")
0162     {
0163       AddHits<CalorimeterHit>(hc, CaTSEvt);
0164     }
0165     else if (Classname == "DRCalorimeter")
0166     {
0167       AddHits<DRCalorimeterHit>(hc, CaTSEvt);
0168     }
0169     else
0170     {
0171       G4cout << "SD type: " << Classname << " unknown" << G4endl;
0172     }
0173   }
0174
0175   if (verbose)
0176   {
0177     G4cout << "writing Event: " << CaTSEvt->GetEventNumber() << G4endl;
0178   }
0179
0180   G4AutoLock lock(&RootIOMutex);
0181   if (!fEvtBranch)
0182   {
0183     Int_t bufsize = 64000;
0184     fEvtBranch = fTree->Branch("event.", &CaTSEvt, bufsize, 0);
0185   }
0186   else
0187   {
0188     fEvtBranch->SetAddress(&CaTSEvt);
0189   }
0190
0191   fTree->Fill();
0192   fFile->Write("", TObject::kOverwrite);
0193
0194   CaTSEvt->Reset();
0195   delete CaTSEvt;
0196 }
0197
0198 //....oooOO0OOooo........oooOO0OOooo........oooOO0OOooo........oooOO0OOooo......
0199 void RootIO::Close()
0200 {
0201   G4cout << " Closing File: " << fFileName << G4endl;
0202   fFile->Close();
0203   fEvtBranch = nullptr;
0204 }
0205
0206 //....oooOO0OOooo........oooOO0OOooo........oooOO0OOooo........oooOO0OOooo......
0207 void RootIO::Merge()
0208 {
0209   if (G4Threading::IsMasterThread())
0210   {
0211     auto nthreads = G4RunManager::GetRunManager()->GetNumberOfThreads();
0212     G4cout << "RootIO::Merging: " << nthreads << " threads" << G4endl;
0213
0214     std::vector<TFile*> files;
0215     std::vector<TTree*> trees;
0216     TList* list = new TList;
0217
0218     for (int i = 0; i < nthreads; i++)
0219     {
0220       G4String fileName = fFileName + std::to_string(i);
0221       files.push_back(TFile::Open(fileName.c_str()));
0222       trees.push_back((TTree*) (files[i]->Get("Events")));
0223       list->Add(trees[i]);
0224
0225       if (i == nthreads - 1)
0226       {
0227         auto* file = TFile::Open(fFileName.c_str(), "RECREATE");
0228
0229         auto tree = TTree::MergeTrees(list);
0230         tree->SetName("Events");
0231
0232         Event* event = new Event();
0233         tree->SetBranchAddress("event.", &event);
0234         TBranch* branch = tree->GetBranch("event.");
0235         int nevents = branch->GetEntries();
0236         G4cout << "Number of Events:  " << nevents << G4endl;
0237
0238         tree->Write();
0239         file->Close();
0240       }
0241       // Delete the merged file
0242       if(std::remove(fileName.c_str()) !=0)
0243       {
0244     G4cout << "Error deleting file: " << fileName << G4endl;
0245       }
0246     }
0247   }
0248 }
0249
0250 //....oooOO0OOooo........oooOO0OOooo........oooOO0OOooo........oooOO0OOooo......
0251 std::vector<G4String> RootIO::Split(const G4String& s, char delim)
0252 {
0253     std::stringstream ss(s);
0254     G4String item;
0255     std::vector<G4String> elems;
0256
0257     while (std::getline(ss, item, delim))
0258     {
0259         elems.push_back(item);
0260     }
0261     return elems;
0262 }
0263
0264 #endif