Merge pull request #549 from rest-for-physics/aezq_dsgmImprovements

TRestDataSet and TRestDataSetGainMap improvements

Author: AlvaroEzq

source/framework/analysis/src/TRestDataSetGainMap.cxx

@@ -213,6 +213,8 @@ void TRestDataSetGainMap::GenerateGainMap() {
 /// \brief Function to calibrate a dataset with this gain map.
 ///
 /// \param dataSetFileName the name of the root file where the TRestDataSet to be calibrated is stored.
+/// If the file is not a TRestDataSet, it will be treated as a file pattern for several TRestRun files
+/// to generate a temporary TRestDataSet with the needed observables.
 /// \param outputFileName the name of the output (root) file where the calibrated TRestDataSet will be
 /// exported. If empty, the output file will be named as the input file plus the name of the
 /// TRestDataSetGainMap. E.g. "data/myDataSet.root" -> "data/myDataSet_<gmName>.root".
@@ -230,7 +232,17 @@ void TRestDataSetGainMap::CalibrateDataSet(const std::string& dataSetFileName, s
 
     TRestDataSet dataSet;
     dataSet.EnableMultiThreading(true);
-    dataSet.Import(dataSetFileName);
+
+    if (TRestTools::isDataSet(dataSetFileName)) {
+        dataSet.Import(dataSetFileName);
+    } else {
+        RESTWarning << dataSetFileName << " is not a dataset. Generating a temporal one..." << RESTendl;
+        // generate the dataset with the needed observables
+        dataSet.SetFilePattern(dataSetFileName);
+        dataSet.SetObservablesList({"*"});  // get all observables
+        dataSet.GenerateDataSet();
+    }
+
     auto dataFrame = dataSet.GetDataFrame();
 
     // Define a new column with the identifier (pmID) of the module for each row (event)
@@ -294,17 +306,9 @@ void TRestDataSetGainMap::CalibrateDataSet(const std::string& dataSetFileName, s
     }
 
     // Export dataset. Exclude columns if requested.
-    std::set<std::string> excludeCol;  // vector with the explicit column names to be excluded
     auto columns = dataSet.GetDataFrame().GetColumnNames();
-    // Get the columns to be excluded from the list of columns. It accepts wildcards "*" and "?"
-    for (auto& eC : excludeColumns) {
-        if (eC.find("*") != std::string::npos || eC.find("?") != std::string::npos) {
-            for (auto& c : columns)
-                if (MatchString(c, eC)) excludeCol.insert(c);
-        } else if (std::find(columns.begin(), columns.end(), eC) != columns.end())
-            excludeCol.insert(eC);
-    }
-    // Remove the calibObsName, calibObsNameFullSpc and pmIDname from the list of columns to be excluded
+    std::set<std::string> excludeCol = TRestTools::GetMatchingStrings(columns, excludeColumns);
+    // Never exclude the calibObsName, calibObsNameFullSpc and pmIDname
     excludeCol.erase(calibObsName);
     excludeCol.erase(calibObsNameFullSpc);
     excludeCol.erase(pmIDname);

source/framework/core/inc/TRestDataSet.h

@@ -60,9 +60,6 @@ class TRestDataSet : public TRestMetadata {
     /// It contains a list of the observables that will be added to the final tree or exported file
     std::vector<std::string> fObservablesList;  //<
 
-    /// It contains a list of the process where all observables should be added
-    std::vector<std::string> fProcessObservablesList;  //<
-
     /// A list of metadata members where filters will be applied
     std::vector<std::string> fFilterMetadata;  //<
 
@@ -172,7 +169,6 @@ class TRestDataSet : public TRestMetadata {
     inline auto GetFilePattern() const { return fFilePattern; }
     inline auto GetObservablesList() const { return fObservablesList; }
     inline auto GetFileSelection() const { return fFileSelection; }
-    inline auto GetProcessObservablesList() const { return fProcessObservablesList; }
     inline auto GetFilterMetadata() const { return fFilterMetadata; }
     inline auto GetFilterContains() const { return fFilterContains; }
     inline auto GetFilterGreaterThan() const { return fFilterGreaterThan; }
@@ -215,6 +211,6 @@ class TRestDataSet : public TRestMetadata {
     TRestDataSet(const char* cfgFileName, const std::string& name = "");
     ~TRestDataSet();
 
-    ClassDefOverride(TRestDataSet, 8);
+    ClassDefOverride(TRestDataSet, 9);
 };
 #endif

source/framework/core/src/TRestDataSet.cxx

@@ -359,29 +359,14 @@ void TRestDataSet::GenerateDataSet() {
 
     ///// Disentangling process observables --> producing finalList
     TRestRun run(fFileSelection.front());
-    std::vector<std::string> finalList;
-    finalList.push_back("runOrigin");
-    finalList.push_back("eventID");
-    finalList.push_back("timeStamp");
+    std::set<std::string> finalList;
+    finalList.insert("runOrigin");
+    finalList.insert("eventID");
+    finalList.insert("timeStamp");
 
     auto obsNames = run.GetAnalysisTree()->GetObservableNames();
-    for (const auto& obs : fObservablesList) {
-        if (std::find(obsNames.begin(), obsNames.end(), obs) != obsNames.end()) {
-            finalList.push_back(obs);
-        } else {
-            RESTWarning << " Observable " << obs << " not found in observable list, skipping..." << RESTendl;
-        }
-    }
-
-    for (const auto& name : obsNames) {
-        for (const auto& pcs : fProcessObservablesList) {
-            if (name.find(pcs) == 0) finalList.push_back(name);
-        }
-    }
-
-    // Remove duplicated observables if any
-    std::sort(finalList.begin(), finalList.end());
-    finalList.erase(std::unique(finalList.begin(), finalList.end()), finalList.end());
+    auto obsFromList = TRestTools::GetMatchingStrings(obsNames, fObservablesList);
+    finalList.insert(obsFromList.begin(), obsFromList.end());
 
     if (fMT)
         ROOT::EnableImplicitMT();
@@ -397,11 +382,11 @@ void TRestDataSet::GenerateDataSet() {
     // Adding new user columns added to the dataset
     for (const auto& [cName, cExpression] : fColumnNameExpressions) {
         RESTInfo << "Adding column to dataset: " << cName << RESTendl;
-        finalList.emplace_back(cName);
+        finalList.emplace(cName);
         fDataFrame = DefineColumn(cName, cExpression);
     }
 
-    RegenerateTree(finalList);
+    RegenerateTree(std::vector<std::string>(finalList.begin(), finalList.end()));
 
     RESTInfo << " - Dataset generated!" << RESTendl;
 }
@@ -672,21 +657,13 @@ void TRestDataSet::PrintMetadata() {
     RESTMetadata << "  " << RESTendl;
 
     if (!fObservablesList.empty()) {
-        RESTMetadata << " Single observables added:" << RESTendl;
+        RESTMetadata << " Observables added:" << RESTendl;
         RESTMetadata << " -------------------------" << RESTendl;
         for (const auto& l : fObservablesList) RESTMetadata << " - " << l << RESTendl;
 
         RESTMetadata << "  " << RESTendl;
     }
 
-    if (!fProcessObservablesList.empty()) {
-        RESTMetadata << " Process observables added: " << RESTendl;
-        RESTMetadata << " -------------------------- " << RESTendl;
-        for (const auto& l : fProcessObservablesList) RESTMetadata << " - " << l << RESTendl;
-
-        RESTMetadata << "  " << RESTendl;
-    }
-
     if (!fFilterMetadata.empty()) {
         RESTMetadata << " Metadata filters: " << RESTendl;
         RESTMetadata << " ----------------- " << RESTendl;
@@ -811,7 +788,10 @@ void TRestDataSet::InitFromConfigFile() {
 
         std::vector<std::string> obsList = REST_StringHelper::Split(observables, ",");
 
-        for (const auto& l : obsList) fProcessObservablesList.push_back(l);
+        for (const auto& l : obsList) {
+            std::string processObsPattern = l + "_*";
+            fObservablesList.push_back(processObsPattern);
+        }
 
         obsProcessDefinition = GetNextElement(obsProcessDefinition);
     }
@@ -1033,7 +1013,6 @@ TRestDataSet& TRestDataSet::operator=(TRestDataSet& dS) {
     fFilePattern = dS.GetFilePattern();
     fObservablesList = dS.GetObservablesList();
     fFileSelection = dS.GetFileSelection();
-    fProcessObservablesList = dS.GetProcessObservablesList();
     fFilterMetadata = dS.GetFilterMetadata();
     fFilterContains = dS.GetFilterContains();
     fFilterGreaterThan = dS.GetFilterGreaterThan();

source/framework/tools/inc/TRestTools.h

@@ -29,6 +29,7 @@
 
 #include <map>
 #include <memory>
+#include <set>
 #include <string>
 #include <vector>
 
@@ -85,6 +86,8 @@ class TRestTools {
     static std::string GetFileNameRoot(const std::string& fullname);
     static std::vector<std::string> GetObservablesInString(const std::string& observablesStr,
                                                            bool removeDuplicates = true);
+    static std::set<std::string> GetMatchingStrings(const std::vector<std::string>& stack,
+                                                    const std::vector<std::string>& wantedStrings);
 
     static int GetBinaryFileColumns(std::string fname);
 

source/framework/tools/src/TRestTools.cxx

@@ -861,6 +861,31 @@ std::vector<std::string> TRestTools::GetObservablesInString(const std::string& o
     return obsList;
 }
 
+////////////////////////////////////////////////////////////
+/// \brief Returns a set of strings that match the wanted strings from the stack.
+/// The wanted strings can contain wildcards "*" and "?".
+/// \param stack: vector of strings to be searched
+/// \param wantedStrings: vector of strings with the wanted strings to be matched.
+/// \return a set of strings that match the wanted strings
+/// e.g.
+/// Input: stack = {"x1", "x2", "x11", "y1", "y2", "y11", "z1", "z2"},
+/// wantedStrings = {"x*", "y?", "z1"},
+/// Output: {"x1", "x11", "x2", "y1", "y2", "z1"}
+///
+std::set<std::string> TRestTools::GetMatchingStrings(const std::vector<std::string>& stack,
+                                                     const std::vector<std::string>& wantedStrings) {
+    std::set<std::string> result;
+    for (auto& ws : wantedStrings) {
+        if (ws.find("*") != std::string::npos || ws.find("?") != std::string::npos) {
+            for (auto& c : stack)
+                if (MatchString(c, ws)) result.insert(c);
+        } else if (std::find(stack.begin(), stack.end(), ws) != stack.end())
+            result.insert(ws);
+    }
+    // return std::vector<std::string>(result.begin(), result.end()); // convert to vector
+    return result;
+}
+
 ///////////////////////////////////////////////
 /// \brief Returns the input string but without multiple slashes ("/")
 ///