diff --git a/ANN_Data_Generator.py b/ANN_Data_Generator.py
index 37fbe93ddd08eee3165c4d7f17964b4d55a6d6f7..0a9bebc857aeffa3e9ff305e903291bbfaa757be 100644
--- a/ANN_Data_Generator.py
+++ b/ANN_Data_Generator.py
@@ -8,6 +8,7 @@ TODO: Adapt code to generate both normalized and non-normalized data -> Done
 TODO: Improve verbose output -> Done
 TODO: Change order of methods -> Done
 TODO: Fix bug in initialization of input matrix -> Done
+TODO: Improve function selection (more even distribution) -> Done
 
 """
 
@@ -88,13 +89,13 @@ class TrainingDataGenerator(object):
         print('Samples to complete:', num_samples)
         tic = timeit.default_timer()
 
-        num_function_samples = num_samples//len(initial_conditions)
-        function_id = 0
         input_data = np.zeros((num_samples, self._stencil_length+2))
+        num_init_cond = len(initial_conditions)
 
         count = 0
         for i in range(num_samples):
             # Pick a Function here
+            function_id = i % num_init_cond
             initial_condition = initial_conditions[function_id]['function']
             initial_condition.randomize(initial_conditions[function_id]['config'])
 
@@ -122,11 +123,6 @@ class TrainingDataGenerator(object):
                 input_data[i] = dg_scheme.build_training_data(
                     centers[self._stencil_length//2], self._stencil_length, initial_condition)
 
-            # Update Function ID
-            if (i % num_function_samples == num_function_samples - 1) \
-                    and (function_id != len(initial_conditions)-1):
-                function_id = function_id + 1
-
             count += 1
             if count % 100 == 0:
                 print(str(count) + ' samples completed.')