Merge pull request #95 from rest-for-physics/lobis-merge

Macro to merge simulation files
rest-for-physics · Aug 20, 2023 · c79d9ed · c79d9ed
2 parents 9ba07e6 + 014e28e
commit c79d9ed
Show file tree

Hide file tree

Showing 6 changed files with 183 additions and 10 deletions.
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -59,6 +59,7 @@ install(FILES ${MAC} DESTINATION ./macros/geant4)
 
 set(rest_macros ${rest_macros} "restGeant4_ViewEvent")
 set(rest_macros ${rest_macros} "restGeant4_ViewGeometry")
+set(rest_macros ${rest_macros} "restGeant4_MergeRestG4Files")
 set(rest_macros
     ${rest_macros}
     PARENT_SCOPE)

diff --git a/inc/TRestGeant4Metadata.h b/inc/TRestGeant4Metadata.h
@@ -56,6 +56,9 @@ class TRestGeant4Metadata : public TRestMetadata {
     void ReadDetector();
     void ReadBiasing();
 
+    // Metadata is the result of a merge of other metadata
+    bool fIsMerge = false;
+
     /// Class used to store and retrieve geometry info
     TRestGeant4GeometryInfo fGeant4GeometryInfo;
 
@@ -306,7 +309,7 @@ class TRestGeant4Metadata : public TRestMetadata {
 
     /// Returns the probability per event to register (write to disk) hits in a GDML volume given its geometry
     /// name.
-    Double_t GetStorageChance(TString volume);
+    Double_t GetStorageChance(const TString& volume);
 
     Double_t GetMaxStepSize(const TString& volume);
 
@@ -364,20 +367,25 @@ class TRestGeant4Metadata : public TRestMetadata {
     /// Returns the world magnetic field in Tesla
     inline TVector3 GetMagneticField() const { return fMagneticField; }
 
-    Int_t GetActiveVolumeID(TString name);
+    Int_t GetActiveVolumeID(const TString& name);
 
     Bool_t isVolumeStored(const TString& volume) const;
 
     void SetActiveVolume(const TString& name, Double_t chance, Double_t maxStep = 0);
 
     void PrintMetadata() override;
 
+    void Merge(const TRestGeant4Metadata&);
+
     TRestGeant4Metadata();
     TRestGeant4Metadata(const char* configFilename, const std::string& name = "");
 
     ~TRestGeant4Metadata();
 
-    ClassDefOverride(TRestGeant4Metadata, 13);
+    TRestGeant4Metadata(const TRestGeant4Metadata& metadata);
+    TRestGeant4Metadata& operator=(const TRestGeant4Metadata& metadata);
+
+    ClassDefOverride(TRestGeant4Metadata, 14);
 
     // Allow modification of otherwise inaccessible / immutable members that shouldn't be modified by the user
     friend class SteppingAction;

diff --git a/inc/TRestGeant4PhysicsInfo.h b/inc/TRestGeant4PhysicsInfo.h
@@ -11,7 +11,7 @@
 class G4VProcess;
 
 class TRestGeant4PhysicsInfo {
-    ClassDef(TRestGeant4PhysicsInfo, 2);
+    ClassDef(TRestGeant4PhysicsInfo, 3);
 
    private:
     std::map<Int_t, TString> fProcessNamesMap = {};
@@ -22,8 +22,6 @@ class TRestGeant4PhysicsInfo {
 
     std::map<TString, TString> fProcessTypesMap = {};  // process name -> process type
 
-    std::mutex fMutex;  //!
-
    public:
     TString GetProcessName(Int_t id) const;
     Int_t GetProcessID(const TString& processName) const;

diff --git a/macros/REST_Geant4_MergeRestG4Files.C b/macros/REST_Geant4_MergeRestG4Files.C
@@ -0,0 +1,121 @@
+#include <string>
+#include <vector>
+
+#include "TRestGeant4Event.h"
+#include "TRestGeant4Metadata.h"
+#include "TRestTask.h"
+
+#ifndef RestTask_Geant4_MergeRestG4Files
+#define RestTask_Geant4_MergeRestG4Files
+
+//*******************************************************************************************************
+//***
+//*** Your HELP is needed to verify, validate and document this macro.
+//*** This macro might need update/revision.
+//***
+//*******************************************************************************************************
+
+/*
+ * Author: Luis Obis ([email protected])
+ * Description: Macro used to merge simulation files, before any analysis is performed
+ * All files are understood to have the same configuration (same generator, same detector, etc.)
+ * They can only differ in the random seed, run number or number of events.
+ */
+
+// Usage:
+// restGeant4_MergeRestG4Files merge_result.root /path/to/directory/with/files/*.root
+
+using namespace std;
+
+void REST_Geant4_MergeRestG4Files(const char* outputFilename, const char* inputFilesDirectory) {
+    // TODO: use glob pattern instead of directory. Already tried this but conflicts with TRestTask...
+
+    cout << "Output file: " << outputFilename << endl;
+    // print input
+    cout << "Input files directory: " << inputFilesDirectory << endl;
+
+    // find all .root files in the directory
+    vector<string> inputFiles = TRestTools::GetFilesMatchingPattern(string(inputFilesDirectory) + "/*.root");
+
+    cout << "Number of input files: " << inputFiles.size() << endl;
+    for (auto file : inputFiles) {
+        cout << "  - " << file << endl;
+    }
+
+    if (inputFiles.size() == 0) {
+        cerr << "No input files found" << endl;
+        exit(1);
+    }
+
+    // open the first file
+    TRestGeant4Metadata mergeMetadata;
+
+    auto mergeRun = new TRestRun();
+    mergeRun->SetName("run");
+    mergeRun->SetOutputFileName(outputFilename);
+    mergeRun->FormOutputFile();
+    mergeRun->GetOutputFile()->cd();
+    mergeRun->SetRunType("restG4");
+
+    TRestGeant4Event* mergeEvent = nullptr;
+    auto mergeEventTree = mergeRun->GetEventTree();
+    mergeEventTree->Branch("TRestGeant4EventBranch", "TRestGeant4Event", &mergeEvent);
+
+    set<Int_t> eventIds;
+
+    // iterate over all other files
+    for (int i = 0; i < inputFiles.size(); i++) {
+        cout << "Processing file " << i + 1 << "/" << inputFiles.size() << endl;
+
+        map<Int_t, Int_t>
+            eventIdUpdates;  // repeatedId -> newId. Make sure if there are repeated event ids in a file
+                             // (because of sub-events) they keep the same event id after modification
+        auto run = TRestRun(inputFiles[i].c_str());
+        auto metadata = (TRestGeant4Metadata*)run.GetMetadataClass("TRestGeant4Metadata");
+        if (i == 0) {
+            mergeMetadata = *metadata;
+        } else {
+            mergeMetadata.Merge(*metadata);
+        }
+        TRestGeant4Event* event = nullptr;
+        auto eventTree = run.GetEventTree();
+        eventTree->SetBranchAddress("TRestGeant4EventBranch", &event);
+        for (int j = 0; j < eventTree->GetEntries(); j++) {
+            eventTree->GetEntry(j);
+            *mergeEvent = *event;
+
+            Int_t eventId = mergeEvent->GetID();
+            if (eventIdUpdates.find(eventId) != eventIdUpdates.end()) {
+                eventId = eventIdUpdates[eventId];
+                cout << "WARNING: event ID " << mergeEvent->GetID() << " with sub-event ID "
+                     << mergeEvent->GetSubID() << " already exists. It was already changed to " << eventId
+                     << endl;
+            } else if (eventIds.find(eventId) != eventIds.end()) {
+                const Int_t maxEventId = *max_element(eventIds.begin(), eventIds.end());
+                eventId = maxEventId + 1;
+                eventIdUpdates[mergeEvent->GetID()] = eventId;
+                cout << "WARNING: event ID " << mergeEvent->GetID() << " with sub-event ID "
+                     << mergeEvent->GetSubID() << " already exists. Changing to " << eventId << endl;
+            }
+            mergeEvent->SetID(eventId);
+            eventIds.insert(mergeEvent->GetID());
+
+            mergeEventTree->Fill();
+            mergeRun->GetAnalysisTree()->Fill();
+        }
+    }
+
+    cout << "Output filename: " << mergeRun->GetOutputFileName() << endl;
+    cout << "Output file: " << mergeRun->GetOutputFile() << endl;
+
+    mergeRun->GetOutputFile()->cd();
+
+    gGeoManager->Write("Geometry", TObject::kOverwrite);
+
+    mergeMetadata.SetName("geant4Metadata");
+    mergeMetadata.Write();
+    mergeRun->UpdateOutputFile();
+    mergeRun->CloseFile();
+}
+
+#endif
diff --git a/src/TRestGeant4Metadata.cxx b/src/TRestGeant4Metadata.cxx
@@ -1485,7 +1485,7 @@ void TRestGeant4Metadata::PrintMetadata() {
 
 ///////////////////////////////////////////////
 /// \brief Returns the id of an active volume giving as parameter its name.
-Int_t TRestGeant4Metadata::GetActiveVolumeID(TString name) {
+Int_t TRestGeant4Metadata::GetActiveVolumeID(const TString& name) {
     Int_t id;
     for (id = 0; id < (Int_t)fActiveVolumes.size(); id++) {
         if (fActiveVolumes[id] == name) return id;
@@ -1536,7 +1536,7 @@ Bool_t TRestGeant4Metadata::isVolumeStored(const TString& volume) const {
 ///////////////////////////////////////////////
 /// \brief Returns the probability of an active volume being stored
 ///
-Double_t TRestGeant4Metadata::GetStorageChance(TString volume) {
+Double_t TRestGeant4Metadata::GetStorageChance(const TString& volume) {
     Int_t id;
     for (id = 0; id < (Int_t)fActiveVolumes.size(); id++) {
         if (fActiveVolumes[id] == volume) return fChance[id];
@@ -1571,3 +1571,46 @@ size_t TRestGeant4Metadata::GetGeant4VersionMajor() const {
     }
     return std::stoi(majorVersion.Data());
 }
+
+void TRestGeant4Metadata::Merge(const TRestGeant4Metadata& metadata) {
+    fIsMerge = true;
+
+    fNEvents += metadata.fNEvents;
+    fNRequestedEntries += metadata.fNRequestedEntries;
+    fSeed = 0;  // seed makes no sense in a merged file
+}
+
+TRestGeant4Metadata::TRestGeant4Metadata(const TRestGeant4Metadata& metadata) { *this = metadata; }
+
+TRestGeant4Metadata& TRestGeant4Metadata::operator=(const TRestGeant4Metadata& metadata) {
+    fIsMerge = metadata.fIsMerge;
+    fGeant4GeometryInfo = metadata.fGeant4GeometryInfo;
+    fGeant4PhysicsInfo = metadata.fGeant4PhysicsInfo;
+    fGeant4PrimaryGeneratorInfo = metadata.fGeant4PrimaryGeneratorInfo;
+    fGeant4Version = metadata.fGeant4Version;
+    fGdmlReference = metadata.fGdmlReference;
+    fMaterialsReference = metadata.fMaterialsReference;
+    fEnergyRangeStored = metadata.fEnergyRangeStored;
+    fActiveVolumes = metadata.fActiveVolumes;
+    fChance = metadata.fChance;
+    fMaxStepSize = metadata.fMaxStepSize;
+    // fParticleSource = metadata.fParticleSource; // segfaults (pointers!)
+    fNBiasingVolumes = metadata.fNBiasingVolumes;
+    fBiasingVolumes = metadata.fBiasingVolumes;
+    fMaxTargetStepSize = metadata.fMaxTargetStepSize;
+    fSubEventTimeDelay = metadata.fSubEventTimeDelay;
+    fFullChain = metadata.fFullChain;
+    fSensitiveVolumes = metadata.fSensitiveVolumes;
+    fNEvents = metadata.fNEvents;
+    fNRequestedEntries = metadata.fNRequestedEntries;
+    fSimulationMaxTimeSeconds = metadata.fSimulationMaxTimeSeconds;
+    fSeed = metadata.fSeed;
+    fSaveAllEvents = metadata.fSaveAllEvents;
+    fRemoveUnwantedTracks = metadata.fRemoveUnwantedTracks;
+    fRemoveUnwantedTracksKeepZeroEnergyTracks = metadata.fRemoveUnwantedTracksKeepZeroEnergyTracks;
+    fRemoveUnwantedTracksVolumesToKeep = metadata.fRemoveUnwantedTracksVolumesToKeep;
+    fKillVolumes = metadata.fKillVolumes;
+    fRegisterEmptyTracks = metadata.fRegisterEmptyTracks;
+    fMagneticField = metadata.fMagneticField;
+    return *this;
+}
diff --git a/src/TRestGeant4PhysicsInfo.cxx b/src/TRestGeant4PhysicsInfo.cxx
@@ -7,6 +7,8 @@ using namespace std;
 
 ClassImp(TRestGeant4PhysicsInfo);
 
+std::mutex insertMutex;
+
 set<TString> TRestGeant4PhysicsInfo::GetAllParticles() const {
     set<TString> particles = {};
     for (const auto& [_, name] : fParticleNamesMap) {
@@ -59,7 +61,7 @@ void TRestGeant4PhysicsInfo::InsertProcessName(Int_t id, const TString& processN
     if (fProcessNamesMap.count(id) > 0) {
         return;
     }
-    std::lock_guard<std::mutex> lock(fMutex);
+    std::lock_guard<std::mutex> lock(insertMutex);
     fProcessNamesMap[id] = processName;
     fProcessNamesReverseMap[processName] = id;
 
@@ -70,7 +72,7 @@ void TRestGeant4PhysicsInfo::InsertParticleName(Int_t id, const TString& particl
     if (fParticleNamesMap.count(id) > 0) {
         return;
     }
-    std::lock_guard<std::mutex> lock(fMutex);
+    std::lock_guard<std::mutex> lock(insertMutex);
     fParticleNamesMap[id] = particleName;
     fParticleNamesReverseMap[particleName] = id;
 }