Merge pull request #487 from rest-for-physics/dSMergeIssue

Dataset merge issue
rest-for-physics · Nov 6, 2023 · 8dc9f22 · 8dc9f22
2 parents 9c2fae7 + 8c48a98
commit 8dc9f22
Show file tree

Hide file tree

Showing 2 changed files with 65 additions and 18 deletions.
diff --git a/source/framework/core/inc/TRestDataSet.h b/source/framework/core/inc/TRestDataSet.h
@@ -169,6 +169,7 @@ class TRestDataSet : public TRestMetadata {
     inline void SetQuantity(const std::map<std::string, RelevantQuantity>& quantity) { fQuantity = quantity; }
 
     TRestDataSet& operator=(TRestDataSet& dS);
+    Bool_t Merge(const TRestDataSet& dS);
     void Import(const std::string& fileName);
     void Import(std::vector<std::string> fileNames);
     void Export(const std::string& filename);

diff --git a/source/framework/core/src/TRestDataSet.cxx b/source/framework/core/src/TRestDataSet.cxx
@@ -894,12 +894,35 @@ TRestDataSet& TRestDataSet::operator=(TRestDataSet& dS) {
     fFilterEqualsTo = dS.GetFilterEqualsTo();
     fQuantity = dS.GetQuantity();
     fTotalDuration = dS.GetTotalTimeInSeconds();
-    fFileSelection = dS.GetFileSelection();
     fCut = dS.GetCut();
 
     return *this;
 }
 
+///////////////////////////////////////////////
+/// \brief This function merge different TRestDataSet
+/// metadata in current dataSet
+///
+Bool_t TRestDataSet::Merge(const TRestDataSet& dS) {
+    auto obsNames = GetObservablesList();
+    for (const auto& obs : fObservablesList) {
+        if (std::find(obsNames.begin(), obsNames.end(), obs) != obsNames.end()) {
+            RESTError << "Cannot merge dataSets with different observable list " << RESTendl;
+            return false;
+        }
+    }
+
+    if (fStartTime > dS.GetStartTime()) fStartTime = dS.GetStartTime();
+    if (fEndTime < dS.GetEndTime()) fEndTime = dS.GetEndTime();
+
+    auto fileSelection = dS.GetFileSelection();
+    fFileSelection.insert(fFileSelection.end(), fileSelection.begin(), fileSelection.end());
+
+    fTotalDuration += dS.GetTotalTimeInSeconds();
+
+    return true;
+}
+
 ///////////////////////////////////////////////
 /// \brief This function imports metadata from a root file
 /// it import metadata info from the previous dataSet
@@ -956,25 +979,48 @@ void TRestDataSet::Import(std::vector<std::string> fileNames) {
             return;
         }
 
-    if (fileNames.size() == 0) return;
-
-    TFile* file = TFile::Open(fileNames[0].c_str(), "READ");
-    if (file != nullptr) {
-        TIter nextkey(file->GetListOfKeys());
-        TKey* key;
-        while ((key = (TKey*)nextkey())) {
-            std::string kName = key->GetClassName();
-            if (REST_Reflection::GetClassQuick(kName.c_str()) != nullptr &&
-                REST_Reflection::GetClassQuick(kName.c_str())->InheritsFrom("TRestDataSet")) {
-                TRestDataSet* dS = file->Get<TRestDataSet>(key->GetName());
-                if (GetVerboseLevel() >= TRestStringOutput::REST_Verbose_Level::REST_Info)
-                    dS->PrintMetadata();
-                *this = *dS;
+    int count = 0;
+    auto it = fileNames.begin();
+    while (it != fileNames.end()) {
+        std::string fileName = *it;
+        TFile* file = TFile::Open(fileName.c_str(), "READ");
+        bool isValid = false;
+        if (file != nullptr) {
+            TIter nextkey(file->GetListOfKeys());
+            TKey* key;
+            while ((key = (TKey*)nextkey())) {
+                std::string kName = key->GetClassName();
+                if (REST_Reflection::GetClassQuick(kName.c_str()) != nullptr &&
+                    REST_Reflection::GetClassQuick(kName.c_str())->InheritsFrom("TRestDataSet")) {
+                    TRestDataSet* dS = file->Get<TRestDataSet>(key->GetName());
+                    if (GetVerboseLevel() >= TRestStringOutput::REST_Verbose_Level::REST_Info)
+                        dS->PrintMetadata();
+
+                    if (count == 0) {
+                        *this = *dS;
+                        isValid = true;
+                    } else {
+                        isValid = Merge(*dS);
+                    }
+
+                    if (isValid) count++;
+                }
             }
+        } else {
+            RESTError << "Cannot open " << fileName << RESTendl;
         }
-    } else {
-        RESTError << "Cannot open " << fileNames[0] << RESTendl;
-        exit(1);
+
+        if (!isValid) {
+            RESTError << fileName << " is not a valid dataSet skipping..." << RESTendl;
+            it = fileNames.erase(it);
+        } else {
+            ++it;
+        }
+    }
+
+    if (fileNames.empty()) {
+        RESTError << "File selection is empty, dataSet will not be imported " << RESTendl;
+        return;
     }
 
     RESTInfo << "Opening list of files. First file: " << fileNames[0] << RESTendl;