Allow test ratio to zero in WildFireSplitter (#63)

* Allow test ratio to zero in WildFireSplitter * fix flake8 * add a test to improve code coverage * fix test * put back the warning on FireIds * fix style
pyronear · May 19, 2020 · 065286d · 065286d
1 parent 61e6410
commit 065286d
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 6 deletions.
diff --git a/pyronear/datasets/wildfire/split_strategy.py b/pyronear/datasets/wildfire/split_strategy.py
@@ -69,12 +69,22 @@ def split(self, dataset, ratios, seed=42):
         self._fire_id_to_size_to_exhaust = copy.deepcopy(self._fire_id_to_size)
 
         # Let's get
-        fire_ids = {'train': self._get_fire_ids_for_one_split(n_samples_train),
-                    'val': self._get_fire_ids_for_one_split(n_samples_val)}
-        fire_ids['test'] = [id_ for id_ in self._fire_id_to_size if id_ not in (fire_ids['train'] + fire_ids['val'])]
-        # Finish exhaustion
-        for fire_id_test in fire_ids['test']:
-            del self._fire_id_to_size_to_exhaust[fire_id_test]
+        if ratios['test'] > 0:
+            fire_ids = {'train': self._get_fire_ids_for_one_split(n_samples_train),
+                        'val': self._get_fire_ids_for_one_split(n_samples_val)}
+            fire_ids['test'] = [id_ for id_ in self._fire_id_to_size
+                                if id_ not in (fire_ids['train'] + fire_ids['val'])]
+            # Finish exhaustion
+            for fire_id_test in fire_ids['test']:
+                del self._fire_id_to_size_to_exhaust[fire_id_test]
+
+        else:
+            fire_ids = {'train': self._get_fire_ids_for_one_split(n_samples_train)}
+            fire_ids['val'] = [id_ for id_ in self._fire_id_to_size if id_ not in fire_ids['train']]
+            # Finish exhaustion
+            for fire_id_test in fire_ids['val']:
+                del self._fire_id_to_size_to_exhaust[fire_id_test]
+            fire_ids['test'] = []
 
         n_samples_remaining = len(self._fire_id_to_size_to_exhaust)
         if n_samples_remaining != 0:

diff --git a/test/test_datasets_wildfire_split.py b/test/test_datasets_wildfire_split.py
@@ -93,6 +93,15 @@ def test_inconsistent_ratios_raise_exception(self):
         with self.assertRaises(ValueError):
             WildFireSplitter(ratios)
 
+    def test_splitting_with_test_to_zero(self):
+        ratios = {'train': 0.81, 'val': 0.19, 'test': 0}
+
+        splitter = WildFireSplitter(ratios, seed=42)
+        splitter.fit(self.wildfire)
+
+        for (set_, ratio_) in splitter.ratios_.items():
+            self.assertAlmostEqual(ratio_, ratios[set_], places=2)
+
     def test_splitting_gives_good_splits_size(self):
         n_samples_expected = {'train': 684, 'val': 147, 'test': 143}
         ratios = {'train': 0.7, 'val': 0.15, 'test': 0.15}