diff --git a/Framework/AODMerger/src/aodThinner.cxx b/Framework/AODMerger/src/aodThinner.cxx index 6ed8463cae3e3..811e907aa79ab 100644 --- a/Framework/AODMerger/src/aodThinner.cxx +++ b/Framework/AODMerger/src/aodThinner.cxx @@ -9,27 +9,29 @@ // granted to it by virtue of its status as an Intergovernmental Organization // or submit itself to any jurisdiction. -#include -#include -#include +#include #include #include "TSystem.h" +#include "TStopwatch.h" +#include "TString.h" +#include "TRegexp.h" #include "TFile.h" #include "TTree.h" +#include "TBranch.h" #include "TList.h" #include "TKey.h" #include "TDirectory.h" #include "TObjString.h" -#include -#include -#include +#include "TGrid.h" +#include "TMap.h" +#include "TLeaf.h" #include "aodMerger.h" // AOD reduction tool // Designed for the 2022 pp data with specific selections: -// - Remove all TPC only tracks +// - Remove all TPC only tracks, optionally keep TPC-only V0 tracks // - Remove all V0s which refer to any removed track // - Remove all cascade which refer to any removed V0 // - Remove all ambiguous track entries which point to a track with collision @@ -38,30 +40,45 @@ int main(int argc, char* argv[]) { std::string inputFileName("AO2D.root"); std::string outputFileName("AO2D_thinned.root"); - int exitCode = 0; // 0: success, >0: failure + int exitCode = 0; // 0: success, !=0: failure + bool bOverwrite = false; int option_index = 1; + + const char* const short_opts = "i:o:KOh"; static struct option long_options[] = { - {"input", required_argument, nullptr, 0}, - {"output", required_argument, nullptr, 1}, - {"help", no_argument, nullptr, 2}, + {"input", required_argument, nullptr, 'i'}, + {"output", required_argument, nullptr, 'o'}, + {"overwrite", no_argument, nullptr, 'O'}, + {"help", no_argument, nullptr, 'h'}, {nullptr, 0, nullptr, 0}}; while (true) { - int c = getopt_long(argc, argv, "", long_options, &option_index); - if (c == -1) { - break; - } else if (c == 0) { - inputFileName = optarg; - } else if (c == 1) { - outputFileName = optarg; - } else if (c == 2) { - printf("AO2D thinning tool. Options: \n"); - printf(" --input Contains input file path to the file to be thinned. Default: %s\n", inputFileName.c_str()); - printf(" --output Target output ROOT file. Default: %s\n", outputFileName.c_str()); - return -1; - } else { - return -2; + const auto opt = getopt_long(argc, argv, short_opts, long_options, &option_index); + if (opt == -1) { + break; // use defaults + } + switch (opt) { + case 'i': + inputFileName = optarg; + break; + case 'o': + outputFileName = optarg; + break; + case 'O': + bOverwrite = true; + printf("Overwriting existing output file if existing\n"); + break; + case 'h': + case '?': + default: + printf("AO2D thinning tool. Options: \n"); + printf(" --input/-i Contains input file path to the file to be thinned. Default: %s\n", inputFileName.c_str()); + printf(" --output/-o Target output ROOT file. Default: %s\n", outputFileName.c_str()); + printf("\n"); + printf(" Optional Arguments:\n"); + printf(" --overwrite/-O Overwrite existing output file\n"); + return -1; } } @@ -69,7 +86,14 @@ int main(int argc, char* argv[]) printf(" Input file: %s\n", inputFileName.c_str()); printf(" Ouput file name: %s\n", outputFileName.c_str()); - auto outputFile = TFile::Open(outputFileName.c_str(), "RECREATE", "", 501); + TStopwatch clock; + clock.Start(kTRUE); + + auto outputFile = TFile::Open(outputFileName.c_str(), (bOverwrite) ? "RECREATE" : "CREATE", "", 501); + if (outputFile == nullptr) { + printf("Error: File %s exists or cannot be created!\n", outputFileName.c_str()); + return 1; + } TDirectory* outputDir = nullptr; if (inputFileName.find("alien:") == 0) { @@ -87,18 +111,21 @@ int main(int argc, char* argv[]) keyList->Sort(); for (auto key1 : *keyList) { + // Keep metaData if (((TObjString*)key1)->GetString().EqualTo("metaData")) { auto metaData = (TMap*)inputFile->Get("metaData"); outputFile->cd(); metaData->Write("metaData", TObject::kSingleKey); } + // Keep parentFiles if (((TObjString*)key1)->GetString().EqualTo("parentFiles")) { auto parentFiles = (TMap*)inputFile->Get("parentFiles"); outputFile->cd(); parentFiles->Write("parentFiles", TObject::kSingleKey); } + // Skip everything else, except 'DF_*' if (!((TObjString*)key1)->GetString().BeginsWith("DF_")) { continue; } @@ -131,15 +158,27 @@ int main(int argc, char* argv[]) } } + // Scan versions e.g. 001 or 002 ... + TString v0Name{"O2v0_???"}, trkExtraName{"O2trackextra*"}; + TRegexp v0Re(v0Name, kTRUE), trkExtraRe(trkExtraName, kTRUE); + for (TObject* obj : *treeList) { + TString st = obj->GetName(); + if (st.Index(v0Re) != kNPOS) { + v0Name = st; + } else if (st.Index(trkExtraRe) != kNPOS) { + trkExtraName = st; + } + } + // Certain order needed in order to populate vectors of skipped entries - auto v0Entry = (TObject*)treeList->FindObject("O2v0_001"); + auto v0Entry = (TObject*)treeList->FindObject(v0Name); treeList->Remove(v0Entry); treeList->AddFirst(v0Entry); // Prepare maps for track skimming - auto trackExtraTree = (TTree*)inputFile->Get(Form("%s/%s", dfName, "O2trackextra")); // for example we can use this line to access the track table + auto trackExtraTree = (TTree*)inputFile->Get(Form("%s/%s", dfName, trkExtraName.Data())); if (trackExtraTree == nullptr) { - printf("O2trackextra table not found\n"); + printf("%s table not found\n", trkExtraName.Data()); exitCode = 6; break; } @@ -149,27 +188,57 @@ int main(int argc, char* argv[]) exitCode = 7; break; } - auto v0s = (TTree*)inputFile->Get(Form("%s/%s", dfName, "O2v0_001")); + auto v0s = (TTree*)inputFile->Get(Form("%s/%s", dfName, v0Name.Data())); if (v0s == nullptr) { - printf("O2v0_001 table not found\n"); + printf("%s table not found\n", v0Name.Data()); exitCode = 8; break; } + // We need to loop over the V0s once and flag the prong indices + int trackIdxPos = 0, trackIdxNeg = 0; + std::unordered_map keepV0TPCs; + v0s->SetBranchAddress("fIndexTracks_Pos", &trackIdxPos); + v0s->SetBranchAddress("fIndexTracks_Neg", &trackIdxNeg); + auto nV0s = v0s->GetEntriesFast(); + for (int i{0}; i < nV0s; ++i) { + v0s->GetEntry(i); + keepV0TPCs[trackIdxPos] = true; + keepV0TPCs[trackIdxNeg] = true; + } + std::vector acceptedTracks(trackExtraTree->GetEntries(), -1); std::vector hasCollision(trackExtraTree->GetEntries(), false); std::vector keepV0s(v0s->GetEntries(), -1); uint8_t tpcNClsFindable = 0; + bool bTPClsFindable = false; uint8_t ITSClusterMap = 0; + bool bITSClusterMap = false; uint8_t TRDPattern = 0; + bool bTRDPattern = false; float_t TOFChi2 = 0; - // char16_t fTPCNClsFindableMinusFound = 0; - trackExtraTree->SetBranchAddress("fTPCNClsFindable", &tpcNClsFindable); - trackExtraTree->SetBranchAddress("fITSClusterMap", &ITSClusterMap); - trackExtraTree->SetBranchAddress("fTRDPattern", &TRDPattern); - trackExtraTree->SetBranchAddress("fTOFChi2", &TOFChi2); - // trackExtraTree->SetBranchAddress("fTPCNClsFindableMinusFound", &fTPCNClsFindableMinusFound); + bool bTOFChi2 = false; + + // Test if track properties exist + TBranch* br = nullptr; + TIter next(trackExtraTree->GetListOfBranches()); + while ((br = (TBranch*)next())) { + TString brName = br->GetName(); + if (brName == "fTPCNClsFindable") { + trackExtraTree->SetBranchAddress("fTPCNClsFindable", &tpcNClsFindable); + bTPClsFindable = true; + } else if (brName == "fITSClusterMap") { + trackExtraTree->SetBranchAddress("fITSClusterMap", &ITSClusterMap); + bITSClusterMap = true; + } else if (brName == "fTRDPattern") { + trackExtraTree->SetBranchAddress("fTRDPattern", &TRDPattern); + bTRDPattern = true; + } else if (brName == "fTOFChi2") { + trackExtraTree->SetBranchAddress("fTOFChi2", &TOFChi2); + bTOFChi2 = true; + } + } int fIndexCollisions = 0; track_iu->SetBranchAddress("fIndexCollisions", &fIndexCollisions); @@ -181,28 +250,34 @@ int main(int argc, char* argv[]) trackExtraTree->GetEntry(i); track_iu->GetEntry(i); - // Remove TPC only tracks - if (tpcNClsFindable > 0. && ITSClusterMap == 0 && TRDPattern == 0 && TOFChi2 < -1.) { + // Flag collisions + hasCollision[i] = (fIndexCollisions >= 0); + + // Remove TPC only tracks, if (opt.) they are not assoc. to a V0 + if ((!bTPClsFindable || tpcNClsFindable > 0.) && + (!bITSClusterMap || ITSClusterMap == 0) && + (!bTRDPattern || TRDPattern == 0) && + (!bTOFChi2 || TOFChi2 < -1.) && + (keepV0TPCs.find(i) == keepV0TPCs.end())) { counter++; } else { acceptedTracks[i] = i - counter; } - hasCollision[i] = (fIndexCollisions >= 0); } for (auto key2 : *treeList) { - auto treeName = ((TObjString*)key2)->GetString().Data(); - - auto inputTree = (TTree*)inputFile->Get(Form("%s/%s", dfName, treeName)); - printf(" Processing tree %s with %lld entries with total size %lld\n", treeName, inputTree->GetEntries(), inputTree->GetTotBytes()); + TString treeName = ((TObjString*)key2)->GetString().Data(); // Connect trees but do not copy entries (using the clone function) // NOTE Basket size etc. are copied in CloneTree() - if (!outputDir) { + if (outputDir == nullptr) { outputDir = outputFile->mkdir(dfName); printf("Writing to output folder %s\n", dfName); } outputDir->cd(); + + auto inputTree = (TTree*)inputFile->Get(Form("%s/%s", dfName, treeName.Data())); + printf(" Processing tree %s with %lld entries with total size %lld\n", treeName.Data(), inputTree->GetEntries(), inputTree->GetTotBytes()); auto outputTree = inputTree->CloneTree(0); outputTree->SetAutoFlush(0); @@ -213,7 +288,7 @@ int main(int argc, char* argv[]) for (int i = 0; i < branches->GetEntriesFast(); ++i) { TBranch* br = (TBranch*)branches->UncheckedAt(i); TString branchName(br->GetName()); - TString tableName(getTableName(branchName, treeName)); + TString tableName(getTableName(branchName, treeName.Data())); // register index of track index ONLY if (!tableName.EqualTo("O2track")) { continue; @@ -243,10 +318,10 @@ int main(int argc, char* argv[]) } } - bool processingTracks = memcmp(treeName, "O2track", 7) == 0; // matches any of the track tables - bool processingCascades = memcmp(treeName, "O2cascade", 9) == 0; - bool processingV0s = memcmp(treeName, "O2v0", 4) == 0; - bool processingAmbiguousTracks = memcmp(treeName, "O2ambiguoustrack", 16) == 0; + bool processingTracks = treeName.BeginsWith("O2track"); // matches any of the track tables + bool processingCascades = treeName.BeginsWith("O2cascade"); + bool processingV0s = treeName.BeginsWith("O2v0"); + bool processingAmbiguousTracks = treeName.BeginsWith("O2ambiguoustrack"); auto indexV0s = -1; if (processingCascades) { @@ -337,6 +412,18 @@ int main(int argc, char* argv[]) gSystem->Unlink(outputFile->GetName()); } + clock.Stop(); + + // Report savings + auto sBefore = inputFile->GetSize(); + auto sAfter = outputFile->GetSize(); + if (sBefore <= 0 || sAfter <= 0) { + printf("Warning: Empty input or output file after thinning!\n"); + exitCode = 9; + } + auto spaceSaving = (1 - ((double)sAfter) / ((double)sBefore)) * 100; + printf("Stats: After=%lld / Before=%lld Bytes ---> Saving %.1f%% diskspace!\n", sAfter, sBefore, spaceSaving); + printf("Timing: CPU=%.2f (s); Real=%.2f (s)\n", clock.CpuTime(), clock.RealTime()); printf("End of AOD thinning.\n"); return exitCode;