import numpy as np


def _consecutive_index_generator(length, offset=0):
    """Generate pair of ids of consecutive images.

    Offset is the distance between the images.
    """
    offset += 1
    for i in range(length - offset):
        yield (i, i + offset)


def generate_learning_set(array, random_permutation=True, offset=0):
    """Generate learning set of consecutive scans

    Parameters
    ----------
    array: numpy array of shape n_scans x n_voxels
        Array of masked scans

    random_permutation: boolean
        If True, consecutive scans are switched with a probability of .5

    offset: int
        Distance between two consecutive scans

    Returns
    -------
    learning_set: list of (img_a, img_b, label)
        Two consecutive images, label is 1 if images are ordered, 0 otherwise
    """
    np.random.seed()

    learning_set = []
    for (ia, ib) in _consecutive_index_generator(array.shape[0], offset=offset):
        label = np.random.randint(0, 2)
        if label == 0:
            ia, ib = ib, ia
        learning_set.append((ia, ib, label))

    return learning_set


if __name__ == '__main__':
    array = np.arange(5)
    res = generate_learning_set(array)
    assert(len(res) == 4)
    for ia, ib, label in res:
        if label == 0:
            ia, ib = ib, ia

    res = generate_learning_set(array, random_permutation=False, offset=1)
    assert(len(res) == 3)
    for ia, ib, label in res:
        assert(ia < ib)
        assert(label == 1)
    print('Basic testing is OK')

@@ -706,16 +706,17 @@
 et: 
+(
 list of 
 (img
@@ -715,22 +715,33 @@
  of 
-(
 img
-_a,
+, list of
  img
-_b
 ,
+ list of
  lab
@@ -756,11 +756,16 @@
     
-Two
+Lists of
  con
@@ -778,14 +778,37 @@
 ve i
-mages,
+ndices and list of labels. If
  lab
@@ -817,17 +817,16 @@
  is 
-1 if imag
+0, indic
 es a
@@ -839,21 +839,21 @@
 ered
-, 0 otherwise
+ decreasingly
 %0A   
@@ -883,27 +883,59 @@
 ()%0A%0A    
-learning_se
+ia_list = %5B%5D%0A    ib_list = %5B%5D%0A    label_lis
 t = %5B%5D%0A 
@@ -1029,16 +1029,69 @@
  label =
+ 1%0A        if random_permutation:%0A            label =
  np.rand
@@ -1107,32 +1107,36 @@
 t(0, 2)%0A        
+   
+ 
 if label == 0:%0A 
@@ -1126,32 +1126,36 @@
  if label == 0:%0A
+    
             ia, 
@@ -1174,27 +1174,22 @@
         
-learning_se
+ia_lis
 t.append
@@ -1193,48 +1193,111 @@
 end(
-(
 ia
-, ib, label))%0A%0A    return learning_se
+)%0A        ib_list.append(ib)%0A        label_list.append(label)%0A%0A    return ia_list, ib_list, label_lis
 t%0A%0A%0A
@@ -1401,24 +1401,27 @@
 sert(len(res
+%5B0%5D
 ) == 4)%0A    
@@ -1433,35 +1433,41 @@
 a, ib, label in 
+zip(*
 res
+)
 :%0A        if lab
@@ -1503,16 +1503,39 @@
  ib, ia%0A
+        assert(ia %3C ib)
 %0A    res
@@ -1620,16 +1620,19 @@
 (len(res
+%5B0%5D
 ) == 3)%0A
@@ -1656,19 +1656,25 @@
 abel in 
+zip(*
 res
+)
 :%0A      
@@ -1678,31 +1678,34 @@
      assert(
-ia %3C ib
+label == 1
 )%0A        as
@@ -1705,34 +1705,31 @@
      assert(
-label == 1
+ia %3C ib
 )%0A    print(

Fix return type