WpZ_analysis/WZPlots.cpp at master · kdlong/WpZ_analysis · 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
124
125
126
127
128
#include "WZPlots.h"
#include <iostream>

WZPlots::WZPlots(const std::string& rootFileName, const std::string& baseFolder,
                 std::vector<std::string> weightNames)

{
    rootFile_ = new TFile(rootFileName.c_str(), "RECREATE");
    rootFileName_ = rootFileName;
    if(baseFolder != "")
        rootFile_->mkdir(baseFolder.c_str());
    weightNames_ = weightNames;

    for(const auto& weightName : weightNames_)
    {
        if(rootFile_->mkdir(weightName.c_str()) == NULL)
            std::cerr << "Error making weight directory!";
    }
    plotSets_.resize(weightNames_.size());


}
void WZPlots::addHist(const std::string& directory, const std::string& subdirectory,
                      const std::string& title, const std::string& xAxisTitle,
                      const std::string& yAxisTitle, const int numBins,
                      const float xMin, const float xMax)
{
    int i = 0;
    for (auto& weightName : weightNames_)
    {
        std::string path = weightName + "/" + directory;
        if (!rootFile_->GetDirectory(path.c_str()))
            rootFile_->mkdir(path.c_str());
        rootFile_->cd(path.c_str());

        std::string weightTitle = title + " (" + weightName + ")";
        TH1F* hist = new TH1F(subdirectory.c_str(), weightTitle.c_str(),
                              numBins, xMin, xMax);
        hist->GetXaxis()->SetTitle(xAxisTitle.c_str());
        hist->GetYaxis()->SetTitle(yAxisTitle.c_str());
        plotSets_[i][directory][subdirectory] = hist;
        i++;
    }
}
void WZPlots::add2DHist(const std::string& directory, const std::string& subdirectory,
                        const std::string& title, const std::string& xAxisTitle,
                        const std::string& yAxisTitle, const int numBinsX,
                        const float xMin, const float xMax, const int numBinsY,
                        const float yMin, const float yMax)
{
    int i = 0;
    for (auto& weightName : weightNames_)
    {
        std::string path = weightName + "/" + directory;
        if (!rootFile_->GetDirectory(path.c_str()))
            rootFile_->mkdir(path.c_str());
        rootFile_->cd(path.c_str());

        std::string weightTitle = title + " (" + weightName + ")";
        TH2F* hist = new TH2F(subdirectory.c_str(), weightTitle.c_str(),
                              numBinsX, xMin, xMax, numBinsY, yMin, yMax);
        hist->GetXaxis()->SetTitle(xAxisTitle.c_str());
        hist->GetYaxis()->SetTitle(yAxisTitle.c_str());
        plotSets_[i][directory][subdirectory] = hist;
        i++;
    }
}
void WZPlots::fillHist(const std::string& plotGroup, const std::string plot,
                       std::vector<float>& values, const std::vector<float>& weights,
                       const float luminosity)
{
    if (weights.size() != plotSets_.size())
    {
        std::cout << "Error, too many weights in histogram fill function";
        std::cout << "\nweights.size() = " << weights.size();
        std::cout << "\nplotSets_.size() = " << plotSets_.size();
        std::cout << "\nExiting do to error.\n";
        exit(0);
    }
    //std::cout << "Filling plot " << plotGroup << " " << plot << std::endl;
    for (unsigned int i = 0; i < weights.size(); i++)
    {
        if (values.size() == 1)
        {
            TH1F* hist = static_cast<TH1F*>(plotSets_[i][plotGroup][plot]);
            hist->Fill(values[0], weights[i]*luminosity);
        }
        else if (values.size() == 2)
        {
            TH2F* hist = static_cast<TH2F*>(plotSets_[i][plotGroup][plot]);
            hist->Fill(values[0], values[1], weights[i]*luminosity);
        }
        else
        {
            std::cout <<  "Not set up to plot > 2D histograms yet.";
            exit(0);
        }
    }
}
void WZPlots::writeToFile()
{
    int i = 0;
    for (auto& plotSet : plotSets_)
    {
        std::string weightPath = rootFileName_ + ":" + weightNames_[i++];
        for (auto& plotGroup : plotSet)
        {
            std::string path = weightPath + "/" + plotGroup.first;
            rootFile_->cd(path.c_str());
            for (auto& plot : plotGroup.second)
            {
                plot.second->Write();
            }
        }
    }
}
WZPlots::~WZPlots()
{
    for (auto& plotSet : plotSets_)
    {
        for (auto& particlePlots : plotSet)
        {
            for (auto& plot : particlePlots.second)
                delete plot.second;
        }
    }
    delete rootFile_;
}