Change behaviour when all files are annotated

IPI_bombyx.py

@@ -10,12 +10,14 @@ import sys
 
 
 class BetterIter(object):
-    def __init__(self, to_wrap):
+    def __init__(self, to_wrap, save, file_list):
         self.wrapped = iter(to_wrap)
         self.max_num = len(to_wrap)
         self.to_wrap = to_wrap
         self.overlap = False
         self.curr_len = None
+        self.save = save
+        self.file_list = file_list
 
     def __iter__(self):
         return self
@@ -33,9 +35,10 @@ class BetterIter(object):
                 val = next(self.wrapped)
 
             self.done_file.seek(0)
-            self.list_pass = [int(v) for v in self.done_file.read().split('\n') if len(v)]
+            self.list_pass = set(int(v) for v in self.done_file.read().split('\n') if len(v))
+            self.usr_pass = set(v for f in self.save for v in f.passage.unique() if v in self.file_list)
             self._test_done()
-            if (val not in self.list_pass) or self.overlap:
+            if (val not in self.list_pass) or (self.overlap and (val not in self.usr_pass)):
                 break
         self.current = val
         return val
@@ -87,7 +90,7 @@ def main(args):
         overlap = True
     if args.nb_ind != -1:
         df = df[df.nbindiv <= args.nb_ind]
-    samples_order = BetterIter(np.random.choice(len(df), len(df), replace=False))
+    samples_order = BetterIter(np.random.choice(len(df), len(df), replace=False), save, file_list)
     samples_order.done_file = done_file
     samples_order.overlap = overlap
     ind = samples_order.__next__()