fixed test and stratified splitter function (finally)

Author: Maren Pielka

src/data_stack/dataset/splitter.py

@@ -110,16 +110,19 @@ def _determine_split_indices(self, dataset_length: int, split_config: Dict, data
         # split the training set until the desired number of splits is reached
         if "train" in split_config.keys():
             train_indices, test_indices, train_targets, test_targets = train_test_split(indices, targets,
-                                                    train_size = int(len(indices)*split_config["train"]))
+                                                    train_size = int(len(indices)*split_config["train"]),
+                                                    stratify=targets)
             split_indices.append(train_indices)
         else:
             logging.error("Training split was not provided")
             sys.exit(1)
-        for split in split_config.keys():
+        for split in list(split_config.keys())[:-1]:
             if split != "train":
                 train_indices, test_indices, train_targets, test_targets = train_test_split(test_indices, test_targets,
-                                                    train_size = int(len(indices)*split_config[split]))
+                                                    train_size = int(len(indices)*split_config[split]),
+                                                    stratify=test_targets)
                 split_indices.append(train_indices)
+        # any remaining indices are added to the last split
         split_indices.append(test_indices)
         return split_indices
 

unittests/dataset/test_splitter.py

@@ -29,7 +29,8 @@ def dataset_iterator(self) -> DatasetIteratorIF:
 
     @pytest.fixture
     def dataset_iterator_stratifiable(self) -> DatasetIteratorIF:
-        return SequenceDatasetIterator(dataset_sequences=[list(range(10)), [0,0,0,1,1,0,0,1,0,1]])
+        return SequenceDatasetIterator(dataset_sequences=[list(range(20)), list(np.ones(8, dtype=int))+
+                                                                            list(np.zeros(12, dtype=int))])
 
     def test_random_splitter(self, ratios: List[int], dataset_iterator: DatasetIteratorIF):
         splitter_impl = RandomSplitterImpl(ratios=ratios, seed=100)
@@ -51,9 +52,9 @@ def test_stratification(self, split_config: Dict[str, int], dataset_iterator_str
         iterator_splits = splitter.split(dataset_iterator_stratifiable)
 
         # target distribution should be equal among all splits
-        assert(sum([sample[1] for sample in iterator_splits[0]]) == 2)
-        assert(sum([sample[1] for sample in iterator_splits[1]]) == 1)
-        assert(sum([sample[1] for sample in iterator_splits[1]]) == 1)
+        assert(sum([sample[1] for sample in iterator_splits[0]]) == 4)
+        assert(sum([sample[1] for sample in iterator_splits[1]]) == 2)
+        assert(sum([sample[1] for sample in iterator_splits[2]]) == 2)
 
     def test_seeding(self):
         ratios = [0.4, 0.6]