Fix stratified folding mistake in remainders

src/Platform/CMakeLists.txt

@@ -5,4 +5,4 @@ include_directories(${BayesNet_SOURCE_DIR}/lib/mdlp)
 include_directories(${BayesNet_SOURCE_DIR}/lib/argparse/include)
 include_directories(${BayesNet_SOURCE_DIR}/lib/json/include)
 add_executable(main main.cc Folding.cc platformUtils.cc Experiment.cc Datasets.cc)
-target_link_libraries(main BayesNet ArffFiles mdlp "${TORCH_LIBRARIES} ")
+target_link_libraries(main BayesNet ArffFiles mdlp "${TORCH_LIBRARIES}")

src/Platform/Datasets.cc

@@ -89,6 +89,10 @@ namespace platform {
         }
         return datasets[name]->getTensors();
     }
+    bool Datasets::isDataset(string name)
+    {
+        return datasets.find(name) != datasets.end();
+    }
     Dataset::Dataset(Dataset& dataset)
     {
         path = dataset.path;

src/Platform/Datasets.h

@@ -58,6 +58,7 @@ namespace platform {
         pair<vector<vector<float>>&, vector<int>&> getVectors(string name);
         pair<vector<vector<int>>&, vector<int>&> getVectorsDiscretized(string name);
         pair<torch::Tensor&, torch::Tensor&> getTensors(string name);
+        bool isDataset(string name);
     };
 };
 

src/Platform/Folding.cc

@@ -72,7 +72,7 @@ void StratifiedKFold::build()
         }
         while (remainder_samples_to_take > 0) {
             int fold = (rand() % static_cast<int>(k));
-            if (stratified_indices[fold].size() == fold_size) {
+            if (stratified_indices[fold].size() == fold_size + 1) {
                 continue;
             }
             auto it = next(class_indices[label].begin(), 1);

src/Platform/main.cc

@@ -81,6 +81,10 @@ int main(int argc, char** argv)
     vector<string> filesToProcess;
     auto datasets = platform::Datasets(path, true, platform::ARFF);
     if (file_name != "") {
+        if (!datasets.isDataset(file_name)) {
+            cerr << "Dataset " << file_name << " not found" << endl;
+            exit(1);
+        }
         filesToProcess.push_back(file_name);
     } else {
         filesToProcess = platform::Datasets(path, true, platform::ARFF).getNames();