Gate_Multi-detector_Post-processing/runPreselection.cpp at main · mraedler/Gate_Multi-detector_Post-processing · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
#include <iomanip>
#include <iostream>
#include <TFile.h>
#include "include/utils.h"
#include "include/CASToRTools.h"
#include "include/eventSelection.h"

// Globals
bool g_verbose;
int g_minSectorDifference;
TString g_treeName;
std::vector<TString> g_fullPaths;
TString g_lutName;
std::vector<LutEntry> g_lut;


void processSingleFile(const size_t idx) {
    std::cout << "Processing: " << g_fullPaths[idx] << std::endl;

    TFile* file = getTFile(g_fullPaths[idx], "UPDATE", g_verbose);
    TTree* tree = getTTree(file, g_treeName, g_verbose);

    if (!checkIfLeafExists(tree, "scatterTest") || !checkIfLeafExists(tree, "trueness")) {
        std::cerr << "Error: coincidence attributes leafs missing. Run addCoincidenceAttributes first." << std::endl;
        std::exit(1);
    }

    // Remove branches before cloning the tree that do not need to be saved in the new file
    tree->SetBranchStatus("runID", false);
    tree->SetBranchStatus("axialPos", false);
    tree->SetBranchStatus("rotationAngle", false);
    tree->SetBranchStatus("sinogramTheta", false);
    tree->SetBranchStatus("sinogramS", false);
    tree->SetBranchStatus("comptVolName1", false);
    tree->SetBranchStatus("comptVolName2", false);
    tree->SetBranchStatus("RayleighVolName1", false);
    tree->SetBranchStatus("RayleighVolName2", false);

    // Set up new file containing only the preselected data
    TString newFullPath = g_fullPaths[idx];
    newFullPath.Remove(newFullPath.Last('.'));
    newFullPath += "_new.root";
    TFile* newFile = TFile::Open(newFullPath, "RECREATE");
    //newFile->cd();  // Ensure the new file is linked
    TTree* newTree = tree->CloneTree(0);

    // Get necessary leafs
    // todo: Can also be moved to the respective event selection functions
    TLeaf* gantryID1 = tree->GetLeaf("gantryID1");
    TLeaf* gantryID2 = tree->GetLeaf("gantryID2");
    TLeaf* rsectorID1 = tree->GetLeaf("rsectorID1");
    TLeaf* rsectorID2 = tree->GetLeaf("rsectorID2");
    TLeaf* trueness = tree->GetLeaf("trueness");
    TLeaf* scatterTest = tree->GetLeaf("scatterTest");

    // Add boolean branch for the preselection
    Bool_t passMinSectorDifferenceTest;
    TBranch* b1 = tree->Branch("passMinSectorDifferenceTest", &passMinSectorDifferenceTest, "passMinSectorDifferenceTest/O");
    runMinSectorDifferenceTest(tree, gantryID1, gantryID2, rsectorID1, rsectorID2, passMinSectorDifferenceTest, b1, g_lutName, g_minSectorDifference, g_verbose);

    Bool_t passScatterTest;
    TBranch* b2 = tree->Branch("passScatterTest", &passScatterTest, "passScatterTest/O");
    runScatterTest(tree, scatterTest, gantryID1, gantryID2, trueness, passMinSectorDifferenceTest, passScatterTest, b2, g_verbose);

    Long64_t nEntries = tree->GetEntries();
    Long64_t progressStep = nEntries / 10;

    for (Long64_t ii = 0; ii < nEntries; ++ii) {
        tree->GetEntry(ii);

        // Print progress (adding nEntries / 2 corresponds to rounding in integer division)
        if (ii % progressStep == 0) {std::cout << (ii * 100 + nEntries / 2) / nEntries << " %" << std::endl;}

        if (!passMinSectorDifferenceTest) continue;
        if (!passScatterTest) continue;

        newTree->Fill();
    }

    std::cout << "Writing: " << newFullPath << std::endl;

    newTree->Write("", TObject::kOverwrite);
    newFile->Close();
    file->Close();

    std::remove(g_fullPaths[idx]);
    std::rename(newFullPath, g_fullPaths[idx]);  // rename new file
}


int main(int argc, char* argv[]) {
    // Check for the required input
    if (argc < 2) {
        printf("Error: no path provided.\n");
        return 1;
    }
    const char *path = argv[1];
    //printf("Path: %s\n", path);

    // Set globals
    g_verbose = false;
    g_minSectorDifference = 2;
    g_treeName = "MergedCoincidences";
    g_fullPaths = getListOfRootFilePaths(path, g_verbose);

    g_lutName = "TB_J-PET_7th_gen_brain_insert_WHR_4_18_1_mm";

    std::map<std::string, ArgumentOptions> argOpts = {
        {"lut", {{"TB_J-PET_7th_gen_brain_insert_WHR_4_18_1_mm", "TB_J-PET_7th_gen_brain_insert_WHR_6_30_1_mm"}, g_lutName}}
    };

    parseArguments(argc, argv, argOpts);
    checkArguments(argOpts);

    g_lut = readLutBinary("/data/local1/raedler/J-PET/CASToR/castor/config/scanner/" + g_lutName + ".lut");

    // runSequentially(g_fullPaths.size(), processSingleFile);
    runInSeparateProcesses(g_fullPaths.size(), processSingleFile, 128);

    return 0;
}