TRestComponentDataSet.cxx

/*************************************************************************
 * This file is part of the REST software framework.                     *
 *                                                                       *
 * Copyright (C) 2016 GIFNA/TREX (University of Zaragoza)                *
 * For more information see https://gifna.unizar.es/trex                 *
 *                                                                       *
 * REST is free software: you can redistribute it and/or modify          *
 * it under the terms of the GNU General Public License as published by  *
 * the Free Software Foundation, either version 3 of the License, or     *
 * (at your option) any later version.                                   *
 *                                                                       *
 * REST is distributed in the hope that it will be useful,               *
 * but WITHOUT ANY WARRANTY; without even the implied warranty of        *
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the          *
 * GNU General Public License for more details.                          *
 *                                                                       *
 * You should have a copy of the GNU General Public License along with   *
 * REST in $REST_PATH/LICENSE.                                           *
 * If not, see https://www.gnu.org/licenses/.                            *
 * For the list of contributors see $REST_PATH/CREDITS.                  *
 *************************************************************************/

#include "TRestComponentDataSet.h"

#include <numeric>

#include <TKey.h>
#include <TLatex.h>

ClassImp(TRestComponentDataSet);

TRestComponentDataSet::TRestComponentDataSet() { Initialize(); }

TRestComponentDataSet::~TRestComponentDataSet() {}

TRestComponentDataSet::TRestComponentDataSet(const char* cfgFileName, const std::string& name)
    : TRestComponent(cfgFileName) {
    Initialize();

    LoadConfigFromFile(fConfigFileName, name);

    if (GetVerboseLevel() >= TRestStringOutput::REST_Verbose_Level::REST_Info) PrintMetadata();
}

void TRestComponentDataSet::Initialize() {
    TRestComponent::Initialize();

    SetSectionName(this->ClassName());
}

Double_t TRestComponentDataSet::GetRate(std::vector<Double_t> point) {
    if (!HasNodes()) {
        RESTError << "TRestComponentDataSet::GetRate. The component has no nodes!" << RESTendl;
        RESTError << "Try calling TRestComponentDataSet::LoadDataSets" << RESTendl;

        RESTInfo << "Trying to load datasets" << RESTendl;
        LoadDataSets();
        if (IsDataSetLoaded())
            RESTInfo << "Sucess!" << RESTendl;
        else
            return 0;
    }

    if (HasNodes() && fActiveNode == -1) {
        RESTError << "TRestComponentDataSet::GetRate. Active node has not been defined" << RESTendl;
        return 0;
    }

    Double_t density =
        GetDensityForActiveNode()->GetBinContent(GetDensityForActiveNode()->GetBin(point.data()));

    Double_t norm = 1;

    // Perhaps this value could be stored internally
    // for (size_t n = 0; n < fNbins.size(); n++) norm = norm * (fRanges[n].Y() - fRanges[n].X()) / fNbins[n];

    return norm * density;
}

Double_t TRestComponentDataSet::GetTotalRate() {
    THnD* dHist = GetDensityForActiveNode();

    Double_t integral = 0;
    if (dHist != nullptr) integral = dHist->ComputeIntegral();

    // Perhaps this value could be stored internally
    for (size_t n = 0; n < fNbins.size(); n++)
        integral = integral * (fRanges[n].Y() - fRanges[n].X()) / fNbins[n];

    return integral;
}

TCanvas* TRestComponentDataSet::DrawComponent(std::vector<std::string> drawVariables,
                                              std::vector<std::string> scanVariables, Int_t binScanSize,
                                              TString drawOption) {
    if (drawVariables.size() > 2 || drawVariables.size() == 0) {
        RESTError << "TRestComponentDataSet::DrawComponent. The number of variables to be drawn must "
                     "be 1 or 2!" << RESTendl;
        return fCanvas;
    }

    if (scanVariables.size() > 2 || scanVariables.size() == 0) {
        RESTError << "TRestComponentDataSet::DrawComponent. The number of variables to be scanned must "
                     "be 1 or 2!" << RESTendl;
        return fCanvas;
    }

    std::vector<int> scanIndexes;
    for (const auto& x : scanVariables) scanIndexes.push_back(GetVariableIndex(x));

    Int_t nPlots = 1;
    size_t n = 0;
    for (const auto& x : scanIndexes) {
        if (fNbins[x] % binScanSize != 0) {
            RESTWarning << "The variable " << scanVariables[n] << " contains " << fNbins[x]
                        << " bins and it doesnt match with a bin size " << binScanSize << RESTendl;
            RESTWarning << "The bin size must be a multiple of the number of bins." << RESTendl;
            RESTWarning << "Redefining bin size to 1." << RESTendl;
            binScanSize = 1;
        }
        nPlots *= fNbins[x] / binScanSize;
        n++;
    }

    Int_t nPlotsX = 0;
    Int_t nPlotsY = 0;

    if (scanIndexes.size() == 2) {
        nPlotsX = fNbins[scanIndexes[0]] / binScanSize;
        nPlotsY = fNbins[scanIndexes[1]] / binScanSize;
    } else {
        nPlotsX = TRestTools::CanvasDivisions(nPlots)[1];
        nPlotsY = TRestTools::CanvasDivisions(nPlots)[0];
    }

    RESTInfo << "Number of plots to be generated: " << nPlots << RESTendl;
    RESTInfo << "Canvas size : " << nPlotsX << " x " << nPlotsY << RESTendl;

    if (fCanvas != nullptr) {
        delete fCanvas;
        fCanvas = nullptr;
    }

    fCanvas = new TCanvas(this->GetName(), this->GetName(), 0, 0, nPlotsX * 640, nPlotsY * 480);
    fCanvas->Divide(nPlotsX, nPlotsY, 0.01, 0.01);

    std::vector<int> variableIndexes;
    for (const auto& x : drawVariables) variableIndexes.push_back(GetVariableIndex(x));

    for (int n = 0; n < nPlotsX; n++)
        for (int m = 0; m < nPlotsY; m++) {
            fCanvas->cd(n * nPlotsY + m + 1);

            THnD* hnd = GetDensity();

            int binXo = binScanSize * n + 1;
            int binXf = binScanSize * n + binScanSize;
            int binYo = binScanSize * m + 1;
            int binYf = binScanSize * m + binScanSize;

            if (scanVariables.size() == 2) {
                hnd->GetAxis(scanIndexes[0])->SetRange(binXo, binXf);
                hnd->GetAxis(scanIndexes[1])->SetRange(binYo, binYf);
            } else if (scanVariables.size() == 1) {
                binXo = binScanSize * nPlotsY * n + binScanSize * m + 1;
                binXf = binScanSize * nPlotsY * n + binScanSize * m + binScanSize;
                hnd->GetAxis(scanIndexes[0])->SetRange(binXo, binXf);
            }

            if (variableIndexes.size() == 1) {
                TH1D* h1 = hnd->Projection(variableIndexes[0]);
                std::string hName;

                if (scanIndexes.size() == 2)
                    hName = scanVariables[0] + "(" + IntegerToString(binXo) + ", " + IntegerToString(binXf) +
                            ") " + scanVariables[1] + "(" + IntegerToString(binYo) + ", " +
                            IntegerToString(binYf) + ") ";

                if (scanIndexes.size() == 1)
                    hName = scanVariables[0] + "(" + IntegerToString(binXo) + ", " + IntegerToString(binXf) +
                            ") ";

                TH1D* newh = (TH1D*)h1->Clone(hName.c_str());
                newh->SetTitle(hName.c_str());
                newh->SetStats(false);
                newh->GetXaxis()->SetTitle((TString)drawVariables[0]);
                newh->SetMarkerStyle(kFullCircle);
                newh->Draw("PLC PMC");

                TString entriesStr = "Entries: " + IntegerToString(newh->GetEntries());
                TLatex* textLatex = new TLatex(0.62, 0.825, entriesStr);
                textLatex->SetNDC();
                textLatex->SetTextColor(1);
                textLatex->SetTextSize(0.05);
                textLatex->Draw("same");
                delete h1;
            }

            if (variableIndexes.size() == 2) {
                TH2D* h2 = hnd->Projection(variableIndexes[0], variableIndexes[1]);

                std::string hName;
                if (scanIndexes.size() == 2)
                    hName = scanVariables[0] + "(" + IntegerToString(binXo) + ", " + IntegerToString(binXf) +
                            ") " + scanVariables[1] + "(" + IntegerToString(binYo) + ", " +
                            IntegerToString(binYf) + ") ";

                if (scanIndexes.size() == 1)
                    hName = scanVariables[0] + "(" + IntegerToString(binXo) + ", " + IntegerToString(binXf) +
                            ") ";

                TH2D* newh = (TH2D*)h2->Clone(hName.c_str());
                newh->SetStats(false);
                newh->GetXaxis()->SetTitle((TString)drawVariables[0]);
                newh->GetYaxis()->SetTitle((TString)drawVariables[1]);
                newh->SetTitle(hName.c_str());
                newh->Draw(drawOption);

                TString entriesStr = "Entries: " + IntegerToString(newh->GetEntries());
                TLatex* textLatex = new TLatex(0.62, 0.825, entriesStr);
                textLatex->SetNDC();
                textLatex->SetTextColor(1);
                textLatex->SetTextSize(0.05);
                textLatex->Draw("same");
                delete h2;
            }
        }

    return fCanvas;
}

void TRestComponentDataSet::PrintMetadata() {
    TRestComponent::PrintMetadata();

    if (!fDataSetFileNames.empty()) {
        RESTMetadata << " " << RESTendl;
        RESTMetadata << " == Dataset filenames ==" << RESTendl;

        for (const auto& x : fDataSetFileNames) RESTMetadata << "- " << x << RESTendl;

        RESTMetadata << " " << RESTendl;
    }

    if (!fParameter.empty() && fParameterizationNodes.empty()) {
        RESTMetadata << "This component has no nodes!" << RESTendl;
        RESTMetadata << " Use: LoadDataSets() to initialize the nodes" << RESTendl;
    }

    if (!fWeights.empty()) {
        RESTMetadata << " " << RESTendl;
        RESTMetadata << " == Weights ==" << RESTendl;

        for (const auto& x : fWeights) RESTMetadata << "- " << x << RESTendl;

        RESTMetadata << " " << RESTendl;
    }

    RESTMetadata << " Use : PrintStatistics() to check node statistics" << RESTendl;
    RESTMetadata << "----" << RESTendl;
}

void TRestComponentDataSet::PrintStatistics() {
    if (fNSimPerNode.empty() && IsDataSetLoaded()) fNSimPerNode = ExtractNodeStatistics();

    if (!HasNodes() && !IsDataSetLoaded()) {
        RESTWarning << "TRestComponentDataSet::PrintStatistics. Empty nodes and no dataset loaded!"
                    << RESTendl;
        RESTWarning << "Invoking TRestComponentDataSet::LoadDataSets might solve the problem" << RESTendl;
        return;
    }

    auto result = std::accumulate(fNSimPerNode.begin(), fNSimPerNode.end(), 0);
    std::cout << "Total counts : " << result << std::endl;
    std::cout << std::endl;

    std::cout << " Parameter node statistics (" << fParameter << ")" << std::endl;
    int n = 0;
    for (const auto& p : fParameterizationNodes) {
        std::cout << " - Value : " << p << " Counts: " << fNSimPerNode[n] << std::endl;
        n++;
    }
}

void TRestComponentDataSet::InitFromConfigFile() {
    TRestComponent::InitFromConfigFile();

    auto ele = GetElement("dataset");
    while (ele != nullptr) {
        fDataSetFileNames.push_back(GetParameter("filename", ele, ""));
        ele = GetNextElement(ele);
    }
}

void TRestComponentDataSet::FillHistograms() {
    fNSimPerNode = ExtractNodeStatistics();

    if (!IsDataSetLoaded()) {
        RESTError << "TRestComponentDataSet::FillHistograms. Dataset has not been initialized!" << RESTendl;
        return;
    }

    if (fParameterizationNodes.empty()) {
        RESTWarning << "Nodes have not been defined" << RESTendl;
        RESTWarning << "The full dataset will be used to generate the density distribution" << RESTendl;
        fParameterizationNodes.push_back(-137);
    }

    RESTInfo << "Generating N-dim histograms" << RESTendl;
    int nIndex = 0;
    for (const auto& node : fParameterizationNodes) {
        ROOT::RDF::RNode df = ROOT::RDataFrame(0);
        if (fParameterizationNodes.size() == 1 && node == -137) {
            RESTInfo << "Creating component with no parameters (full dataset used)" << RESTendl;
            df = fDataSet.GetDataFrame();
            fParameterizationNodes.clear();
        } else {
            RESTInfo << "Creating THnD for parameter " << fParameter << ": " << DoubleToString(node)
                     << RESTendl;
            std::string filter = fParameter + " == " + DoubleToString(node);
            df = fDataSet.GetDataFrame().Filter(filter);
        }

        Int_t* bins = new Int_t[fNbins.size()];
        Double_t* xmin = new Double_t[fNbins.size()];
        Double_t* xmax = new Double_t[fNbins.size()];

        for (size_t n = 0; n < fNbins.size(); n++) {
            bins[n] = fNbins[n];
            xmin[n] = fRanges[n].X();
            xmax[n] = fRanges[n].Y();
        }

        TString hName = fParameter + "_" + DoubleToString(node);
        if (fParameterizationNodes.empty()) hName = "full";

        std::vector<std::string> varsAndWeight = fVariables;

        if (!fWeights.empty()) {
            std::string weightsStr = "";
            for (size_t n = 0; n < fWeights.size(); n++) {
                if (n > 0) weightsStr += "*";

                weightsStr += fWeights[n];
            }
            df = df.Define("componentWeight", weightsStr);
            varsAndWeight.push_back("componentWeight");
        }

        auto hn = df.HistoND({hName, hName, (int)fNbins.size(), bins, xmin, xmax}, varsAndWeight);
        THnD* hNd = new THnD(*hn);
        hNd->Scale(1. / fNSimPerNode[nIndex]);

        fNodeDensity.push_back(hNd);
        fActiveNode = nIndex;
        nIndex++;
    }
}

THnD* TRestComponentDataSet::GetDensityForNode(Double_t node) {
    int n = 0;
    for (const auto& x : fParameterizationNodes) {
        if (x == node) {
            return fNodeDensity[n];
        }
        n++;
    }

    RESTError << "Parametric node : " << node << " was not found in component" << RESTendl;
    PrintNodes();
    return nullptr;
}

THnD* TRestComponentDataSet::GetDensityForActiveNode() {
    if (fActiveNode >= 0) return fNodeDensity[fActiveNode];

    RESTError << "The active node is invalid" << RESTendl;
    PrintNodes();
    return nullptr;
}

TH1D* TRestComponentDataSet::GetHistogram(Double_t node, std::string varName) {
    SetActiveNode(node);
    return GetHistogram(varName);
}

TH1D* TRestComponentDataSet::GetHistogram(std::string varName) {
    if (fActiveNode < 0) return nullptr;

    Int_t v1 = GetVariableIndex(varName);

    if (v1 >= 0 && GetDensityForActiveNode()) return GetDensityForActiveNode()->Projection(v1);

    return nullptr;
}

TH2D* TRestComponentDataSet::GetHistogram(Double_t node, std::string varName1, std::string varName2) {
    SetActiveNode(node);
    return GetHistogram(varName1, varName2);
}

TH2D* TRestComponentDataSet::GetHistogram(std::string varName1, std::string varName2) {
    if (fActiveNode < 0) return nullptr;

    Int_t v1 = GetVariableIndex(varName1);
    Int_t v2 = GetVariableIndex(varName2);

    if (v1 >= 0 && v2 >= 0)
        if (GetDensityForActiveNode()) return GetDensityForActiveNode()->Projection(v1, v2);

    return nullptr;
}

TH3D* TRestComponentDataSet::GetHistogram(Double_t node, std::string varName1, std::string varName2,
                                          std::string varName3) {
    SetActiveNode(node);
    return GetHistogram(varName1, varName2, varName3);
}

TH3D* TRestComponentDataSet::GetHistogram(std::string varName1, std::string varName2, std::string varName3) {
    if (fActiveNode < 0) return nullptr;

    Int_t v1 = GetVariableIndex(varName1);
    Int_t v2 = GetVariableIndex(varName2);
    Int_t v3 = GetVariableIndex(varName3);

    if (v1 >= 0 && v2 >= 0 && v3 >= 0)
        if (GetDensityForActiveNode()) return GetDensityForActiveNode()->Projection(v1, v2, v3);

    return nullptr;
}

std::vector<Double_t> TRestComponentDataSet::ExtractParameterizationNodes() {
    if (!fParameterizationNodes.empty()) return fParameterizationNodes;

    RESTInfo << "Extracting parameterization nodes" << RESTendl;

    std::vector<double> vs;
    if (!IsDataSetLoaded()) {
        RESTError << "TRestComponentDataSet::ExtractParameterizationNodes. Dataset has not been initialized!"
                  << RESTendl;
        return vs;
    }

    auto parValues = fDataSet.GetDataFrame().Take<double>(fParameter);
    for (const auto v : parValues) vs.push_back(v);

    std::vector<double>::iterator ip;
    ip = std::unique(vs.begin(), vs.begin() + vs.size());
    vs.resize(std::distance(vs.begin(), ip));
    std::sort(vs.begin(), vs.end());
    ip = std::unique(vs.begin(), vs.end());
    vs.resize(std::distance(vs.begin(), ip));

    return vs;
}

std::vector<Int_t> TRestComponentDataSet::ExtractNodeStatistics() {
    if (!fNSimPerNode.empty()) return fNSimPerNode;

    std::vector<Int_t> stats;
    if (!IsDataSetLoaded()) {
        RESTError << "TRestComponentDataSet::ExtractNodeStatistics. Dataset has not been initialized!"
                  << RESTendl;
        return stats;
    }

    RESTInfo << "Counting statistics for each node ..." << RESTendl;
    RESTInfo << "Number of nodes : " << fParameterizationNodes.size() << RESTendl;
    for (const auto& p : fParameterizationNodes) {
        std::string filter = fParameter + " == " + DoubleToString(p);
        auto nEv = fDataSet.GetDataFrame().Filter(filter).Count();
        stats.push_back(*nEv);
    }
    return stats;
}

Bool_t TRestComponentDataSet::LoadDataSets() {
    if (fDataSetFileNames.empty()) {
        RESTWarning << "Dataset filename was not defined. You may still use "
                       "TRestComponentDataSet::LoadDataSet( filename );" << RESTendl;
        fDataSetLoaded = false;
        return fDataSetLoaded;
    }

    RESTInfo << "Loading datasets" << RESTendl;

    std::vector<std::string> fullFileNames;
    for (const auto& name : fDataSetFileNames) {
        // TODO we get here a list of files. However, we will need to weight each dataset with a factor
        // to consider the contribution of each background component.
        // Of course, we could previously take a factor into account already in the dataset, through the
        // definition of a new column. But being this way would allow us to play around with the
        // background model without having to regenerate the dataset.
        std::string fileName = SearchFile(name);
        if (fileName.empty()) {
            RESTError << "TRestComponentDataSet::LoadDataSet. Error loading file : " << name << RESTendl;
            RESTError << "Does the file exist?" << RESTendl;
            RESTError << "You may use `<globals> <searchPath ...` to indicate the path location" << RESTendl;
            return false;
        }
        fullFileNames.push_back(fileName);
    }

    fDataSet.Import(fullFileNames);
    fDataSetLoaded = true;

    if (fDataSet.GetTree() == nullptr) {
        RESTError << "Problem loading dataset from file list :" << RESTendl;
        for (const auto& f : fDataSetFileNames) RESTError << " - " << f << RESTendl;
        return false;
    }

    if (GetVerboseLevel() >= TRestStringOutput::REST_Verbose_Level::REST_Info) fDataSet.PrintMetadata();

    if (VariablesOk() && WeightsOk()) {
        fParameterizationNodes = ExtractParameterizationNodes();
        FillHistograms();
        return fDataSetLoaded;
    }

    return fDataSetLoaded;
}

Bool_t TRestComponentDataSet::VariablesOk() {
    Bool_t ok = true;
    std::vector cNames = fDataSet.GetDataFrame().GetColumnNames();

    for (const auto var : fVariables)
        if (std::count(cNames.begin(), cNames.end(), var) == 0) {
            RESTError << "Variable ---> " << var << " <--- NOT found on dataset" << RESTendl;
            ok = false;
        }
    return ok;
}

Bool_t TRestComponentDataSet::WeightsOk() {
    Bool_t ok = true;
    std::vector cNames = fDataSet.GetDataFrame().GetColumnNames();

    for (const auto var : fWeights)
        if (std::count(cNames.begin(), cNames.end(), var) == 0) {
            RESTError << "Weight ---> " << var << " <--- NOT found on dataset" << RESTendl;
            ok = false;
        }
    return ok;
}

Bool_t TRestComponentDataSet::ValidDataSet() {
    if (!IsDataSetLoaded()) {
        RESTWarning << "TRestComponentDataSet::ValidDataSet. Dataset has not been loaded" << RESTendl;
        RESTWarning << "Try calling TRestComponentDataSet::LoadDataSets" << RESTendl;

        RESTInfo << "Trying to load datasets" << RESTendl;
        LoadDataSets();
        if (IsDataSetLoaded()) {
            RESTInfo << "Sucess!" << RESTendl;
        } else {
            RESTError << "Failed loading datasets" << RESTendl;
            return false;
        }
    }

    if (HasNodes() && fActiveNode == -1) {
        RESTError << "TRestComponentDataSet::ValidDataSet. Active node has not been defined" << RESTendl;
        return false;
    }
    return true;
}