DeepTop/extract.cpp at master · mjagkow/DeepTop · 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
#include "TTree.h"
#include "TFile.h"
#include <iostream>
#include <fstream>
#include <sys/stat.h>


using namespace std;


void print_usage() {
    cout << "Usage:\n"
            "  extract <input.root> [output.csv]" << endl;
}

int main(int argc, char **argv) {

    if (argc < 2) {
        cout << "Too few parameters, please specify ROOT file path.\n";
        print_usage();
        return 1;
    }


    {
        struct stat buffer;
        if ( stat(argv[1], &buffer) != 0 ) {
            cout << "Error: cannot open file '" << argv[1] << "', check path and permissions.\n";
            return 1;
        }
    }


    string outputFileName;
    if (argc >= 3)
        outputFileName = argv[2];
    else {
        outputFileName = argv[1];
        auto pos = outputFileName.find_last_of('.');
        if (pos == string::npos)
            outputFileName += ".csv";
        else {
            auto slashPos = outputFileName.find_last_of('/');
            if (slashPos == string::npos)
                slashPos = 0;
            else
                slashPos ++;
            outputFileName = outputFileName.substr(slashPos, pos - slashPos) + ".csv";
        }
    }


    TFile *f = TFile::Open(argv[1], "r");
    TTree *t = dynamic_cast<TTree*>(f->Get("Vars"));

    Float_t Lepton_PX, Lepton_PY, Lepton_PZ, Lepton_E;
    Float_t MET_PX, MET_PY, MET_PZ, MET_E;
    Float_t MtW, MET, Pt_Lep, DPhi_LepNu;

    t->SetBranchAddress("Lepton_PX", &Lepton_PX);
    t->SetBranchAddress("Lepton_PY", &Lepton_PY);
    t->SetBranchAddress("Lepton_PZ", &Lepton_PZ);
    t->SetBranchAddress("Lepton_E" , &Lepton_E );
    t->SetBranchAddress("MET_PX", &MET_PX);
    t->SetBranchAddress("MET_PY", &MET_PY);
    t->SetBranchAddress("MET_PZ", &MET_PZ);
    t->SetBranchAddress("MET_E", &MET_E );
    t->SetBranchAddress("MtW", &MtW);
    t->SetBranchAddress("MET", &MET);
    t->SetBranchAddress("Pt_Lep", &Pt_Lep);
    t->SetBranchAddress("DPhi_LepNu", &DPhi_LepNu);

    ofstream of(outputFileName, ios_base::out | ios_base::trunc);
    of << "MtW, MET, Pt_Lep, DPhi_LepNu, Lepton_PX, Lepton_PY, Lepton_PZ, Lepton_E, MET_PX, MET_PY, MET_E\n";

    for (size_t i = 0; i < t->GetEntries(); i++) {
        t->GetEntry(i);
        of << MtW << ", " << MET << ", " << Pt_Lep << ", " << DPhi_LepNu << ", " << Lepton_PX << ", " << Lepton_PY << ", " << Lepton_PZ << ", " << Lepton_E << ", " << MET_PX << ", " << MET_PY << ", " << MET_E << '\n';
    }

    of.close();
    cout << " input: " << argv[1] << endl;
    cout << "output: " << outputFileName << endl;

    cout << t->GetEntries() << endl;

    return 0;
}