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 <TKey.h>
 
#include <numeric>
 
ClassImp(TRestComponentDataSet);
 
TRestComponentDataSet::TRestComponentDataSet() {}
 
TRestComponentDataSet::~TRestComponentDataSet() {}
 
TRestComponentDataSet::TRestComponentDataSet(const char* cfgFileName, const std::string& name)
    : TRestComponent(cfgFileName) {
    LoadConfigFromFile(fConfigFileName, name);
 
    if (GetVerboseLevel() >= TRestStringOutput::REST_Verbose_Level::REST_Info) PrintMetadata();
}
 
void TRestComponentDataSet::Initialize() {
    TRestComponent::Initialize();
 
    SetSectionName(this->ClassName());
 
    LoadDataSets();
}
 
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::Initialize() might solve the problem" << RESTendl;
        return;
    }
 
    auto result = std::accumulate(fNSimPerNode.begin(), fNSimPerNode.end(), 0);
    RESTInfo << "Total counts : " << result << RESTendl;
    std::cout << std::endl;
 
    RESTInfo << " Parameter node statistics (" << fParameter << ")" << RESTendl;
    int n = 0;
    for (const auto& p : fParameterizationNodes) {
        RESTInfo << " - Value : " << p << " Counts: " << fNSimPerNode[n] << RESTendl;
        n++;
    }
}
 
void TRestComponentDataSet::InitFromConfigFile() {
    TRestComponent::InitFromConfigFile();
 
    auto ele = GetElement("dataset");
    while (ele != nullptr) {
        fDataSetFileNames.push_back(GetParameter("filename", ele, ""));
        ele = GetNextElement(ele);
    }
 
    if (!fDataSetFileNames.empty()) Initialize();
}
 
void TRestComponentDataSet::FillHistograms() {
    if (!fNodeDensity.empty()) return;
 
    if (fNbins.size() == 0) {
        RESTError
            << "TRestComponentDataSet::FillHistograms. Trying to fill histograms but no variables found!"
            << RESTendl;
        return;
    }
 
    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) {
        Int_t from = 0;
        Int_t to = 0;
        if (fSamples > 0 && fTotalSamples[nIndex] - fSamples > 0) {
            from = fRandom->Integer(fTotalSamples[nIndex] - fSamples);
            to = from + fSamples;
            fNSimPerNode[nIndex] = fSamples;
        }
 
        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().Range(from, to);
            fParameterizationNodes.clear();
        } else {
            RESTInfo << "Creating THnD for parameter " << fParameter << ": " << DoubleToString(node)
                     << RESTendl;
            Double_t pUp = node * (1 + fPrecision / 2);
            Double_t pDown = node * (1 - fPrecision / 2);
            std::string filter = fParameter + " < " + DoubleToString(pUp) + " && " + fParameter + " > " +
                                 DoubleToString(pDown);
            df = fDataSet.GetDataFrame().Filter(filter).Range(from, to);
        }
 
        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++;
    }
}
 
void TRestComponentDataSet::RegenerateActiveNodeDensity() {
    if (fActiveNode >= 0 && fNodeDensity[fActiveNode]) {
        delete fNodeDensity[fActiveNode];
    } else {
        RESTError << "TRestComponentDataSet::RegenerateActiveNode. Active node undefined!" << RESTendl;
        return;
    }
 
    Int_t from = 0;
    Int_t to = 0;
    if (fSamples > 0 && fTotalSamples[fActiveNode] - fSamples > 0) {
        from = fRandom->Integer(fTotalSamples[fActiveNode] - fSamples);
        to = from + fSamples;
        fNSimPerNode[fActiveNode] = fSamples;
    }
 
    Double_t node = GetActiveNodeValue();
    RESTInfo << "Creating THnD for parameter " << fParameter << ": " << DoubleToString(node) << RESTendl;
 
    ROOT::RDF::RNode df = ROOT::RDataFrame(0);
    Double_t pUp = node * (1 + fPrecision / 2);
    Double_t pDown = node * (1 - fPrecision / 2);
    std::string filter =
        fParameter + " < " + DoubleToString(pUp) + " && " + fParameter + " > " + DoubleToString(pDown);
    df = fDataSet.GetDataFrame().Filter(filter).Range(from, to);
 
    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[fActiveNode]);
 
    fNodeDensity[fActiveNode] = hNd;
}
 
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;
 
    fTotalSamples.clear();
 
    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) {
        Double_t pUp = p * (1 + fPrecision / 2);
        Double_t pDown = p * (1 - fPrecision / 2);
        std::string filter =
            fParameter + " < " + DoubleToString(pUp) + " && " + fParameter + " > " + DoubleToString(pDown);
        RESTInfo << "Counting stats for : " << fParameter << " = " << p << RESTendl;
        auto nEv = fDataSet.GetDataFrame().Filter(filter).Count();
        fTotalSamples.push_back(*nEv);
        RESTInfo << "Total entries for " << fParameter << ":" << p << " = " << *nEv << RESTendl;
        if (fSamples != 0) {
            nEv = fDataSet.GetDataFrame().Filter(filter).Range(fSamples).Count();
        }
 
        if ((Int_t)*nEv < fSamples) {
            RESTWarning << "The number of requested samples (" << fSamples
                        << ") is higher than the number of dataset entries (" << *nEv << ")" << RESTendl;
        }
        RESTInfo << "Samples to be used for " << fParameter << ":" << p << " = " << *nEv << RESTendl;
        stats.push_back(*nEv);
    }
    return stats;
}
 
Bool_t TRestComponentDataSet::LoadDataSets() {
    if (fDataSetFileNames.empty()) {
        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::Initialize()" << 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;
}