Skip to content

AOD: Prepare Thinner for TPC-only V0s #12510

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
f3sch merged 1 commit into from
Jan 22, 2024
Merged

AOD: Prepare Thinner for TPC-only V0s #12510

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
187 changes: 137 additions & 50 deletions Framework/AODMerger/src/aodThinner.cxx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,27 +9,29 @@
// granted to it by virtue of its status as an Intergovernmental Organization
// or submit itself to any jurisdiction.

#include <map>
#include <list>
#include <fstream>
#include <unordered_map>
#include <getopt.h>

#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 <TGrid.h>
#include <TMap.h>
#include <TLeaf.h>
#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
Expand All @@ -38,38 +40,60 @@ 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 <inputfile.root> Contains input file path to the file to be thinned. Default: %s\n", inputFileName.c_str());
printf(" --output <outputfile.root> 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 <inputfile.root> Contains input file path to the file to be thinned. Default: %s\n", inputFileName.c_str());
printf(" --output/-o <outputfile.root> 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;
}
}

printf("AOD reduction started with:\n");
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) {
Expand All @@ -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;
}
Expand Down Expand Up @@ -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;
}
Expand All @@ -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<int, bool> 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<int> acceptedTracks(trackExtraTree->GetEntries(), -1);
std::vector<bool> hasCollision(trackExtraTree->GetEntries(), false);
std::vector<int> 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);
Expand All @@ -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);

Expand All @@ -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;
Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -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;
Expand Down