|
| 1 | +#include <string> |
| 2 | +#include <vector> |
| 3 | + |
| 4 | +#include "TRestGeant4Event.h" |
| 5 | +#include "TRestGeant4Metadata.h" |
| 6 | +#include "TRestTask.h" |
| 7 | + |
| 8 | +#ifndef RestTask_Geant4_MergeRestG4Files |
| 9 | +#define RestTask_Geant4_MergeRestG4Files |
| 10 | + |
| 11 | +//******************************************************************************************************* |
| 12 | +//*** |
| 13 | +//*** Your HELP is needed to verify, validate and document this macro. |
| 14 | +//*** This macro might need update/revision. |
| 15 | +//*** |
| 16 | +//******************************************************************************************************* |
| 17 | + |
| 18 | +/* |
| 19 | + * Author: Luis Obis (lobis@unizar.es) |
| 20 | + * Description: Macro used to merge simulation files, before any analysis is performed |
| 21 | + * All files are understood to have the same configuration (same generator, same detector, etc.) |
| 22 | + * They can only differ in the random seed, run number or number of events. |
| 23 | + */ |
| 24 | + |
| 25 | +// Usage: |
| 26 | +// restGeant4_MergeRestG4Files merge_result.root /path/to/directory/with/files/*.root |
| 27 | + |
| 28 | +using namespace std; |
| 29 | + |
| 30 | +void REST_Geant4_MergeRestG4Files(const char* outputFilename, const char* inputFilesDirectory) { |
| 31 | + // TODO: use glob pattern instead of directory. Already tried this but conflicts with TRestTask... |
| 32 | + |
| 33 | + cout << "Output file: " << outputFilename << endl; |
| 34 | + // print input |
| 35 | + cout << "Input files directory: " << inputFilesDirectory << endl; |
| 36 | + |
| 37 | + // find all .root files in the directory |
| 38 | + vector<string> inputFiles = TRestTools::GetFilesMatchingPattern(string(inputFilesDirectory) + "/*.root"); |
| 39 | + |
| 40 | + cout << "Number of input files: " << inputFiles.size() << endl; |
| 41 | + for (auto file : inputFiles) { |
| 42 | + cout << " - " << file << endl; |
| 43 | + } |
| 44 | + |
| 45 | + if (inputFiles.size() == 0) { |
| 46 | + cerr << "No input files found" << endl; |
| 47 | + exit(1); |
| 48 | + } |
| 49 | + |
| 50 | + // open the first file |
| 51 | + TRestGeant4Metadata mergeMetadata; |
| 52 | + |
| 53 | + auto mergeRun = new TRestRun(); |
| 54 | + mergeRun->SetName("run"); |
| 55 | + mergeRun->SetOutputFileName(outputFilename); |
| 56 | + mergeRun->FormOutputFile(); |
| 57 | + mergeRun->GetOutputFile()->cd(); |
| 58 | + mergeRun->SetRunType("restG4"); |
| 59 | + |
| 60 | + TRestGeant4Event* mergeEvent = nullptr; |
| 61 | + auto mergeEventTree = mergeRun->GetEventTree(); |
| 62 | + mergeEventTree->Branch("TRestGeant4EventBranch", "TRestGeant4Event", &mergeEvent); |
| 63 | + |
| 64 | + set<Int_t> eventIds; |
| 65 | + |
| 66 | + // iterate over all other files |
| 67 | + for (int i = 0; i < inputFiles.size(); i++) { |
| 68 | + cout << "Processing file " << i + 1 << "/" << inputFiles.size() << endl; |
| 69 | + |
| 70 | + map<Int_t, Int_t> |
| 71 | + eventIdUpdates; // repeatedId -> newId. Make sure if there are repeated event ids in a file |
| 72 | + // (because of sub-events) they keep the same event id after modification |
| 73 | + auto run = TRestRun(inputFiles[i].c_str()); |
| 74 | + auto metadata = (TRestGeant4Metadata*)run.GetMetadataClass("TRestGeant4Metadata"); |
| 75 | + if (i == 0) { |
| 76 | + mergeMetadata = *metadata; |
| 77 | + } else { |
| 78 | + mergeMetadata.Merge(*metadata); |
| 79 | + } |
| 80 | + TRestGeant4Event* event = nullptr; |
| 81 | + auto eventTree = run.GetEventTree(); |
| 82 | + eventTree->SetBranchAddress("TRestGeant4EventBranch", &event); |
| 83 | + for (int j = 0; j < eventTree->GetEntries(); j++) { |
| 84 | + eventTree->GetEntry(j); |
| 85 | + *mergeEvent = *event; |
| 86 | + |
| 87 | + Int_t eventId = mergeEvent->GetID(); |
| 88 | + if (eventIdUpdates.find(eventId) != eventIdUpdates.end()) { |
| 89 | + eventId = eventIdUpdates[eventId]; |
| 90 | + cout << "WARNING: event ID " << mergeEvent->GetID() << " with sub-event ID " |
| 91 | + << mergeEvent->GetSubID() << " already exists. It was already changed to " << eventId |
| 92 | + << endl; |
| 93 | + } else if (eventIds.find(eventId) != eventIds.end()) { |
| 94 | + const Int_t maxEventId = *max_element(eventIds.begin(), eventIds.end()); |
| 95 | + eventId = maxEventId + 1; |
| 96 | + eventIdUpdates[mergeEvent->GetID()] = eventId; |
| 97 | + cout << "WARNING: event ID " << mergeEvent->GetID() << " with sub-event ID " |
| 98 | + << mergeEvent->GetSubID() << " already exists. Changing to " << eventId << endl; |
| 99 | + } |
| 100 | + mergeEvent->SetID(eventId); |
| 101 | + eventIds.insert(mergeEvent->GetID()); |
| 102 | + |
| 103 | + mergeEventTree->Fill(); |
| 104 | + mergeRun->GetAnalysisTree()->Fill(); |
| 105 | + } |
| 106 | + } |
| 107 | + |
| 108 | + cout << "Output filename: " << mergeRun->GetOutputFileName() << endl; |
| 109 | + cout << "Output file: " << mergeRun->GetOutputFile() << endl; |
| 110 | + |
| 111 | + mergeRun->GetOutputFile()->cd(); |
| 112 | + |
| 113 | + gGeoManager->Write("Geometry", TObject::kOverwrite); |
| 114 | + |
| 115 | + mergeMetadata.SetName("geant4Metadata"); |
| 116 | + mergeMetadata.Write(); |
| 117 | + mergeRun->UpdateOutputFile(); |
| 118 | + mergeRun->CloseFile(); |
| 119 | +} |
| 120 | + |
| 121 | +#endif |
0 commit comments