Improved directory selection.

ANN_Data_Generator.py

@@ -16,6 +16,7 @@ TODO: Fix bug in 'input_data' calculation -> Done
 TODO: Change order of imports -> Done
 TODO: Improve calculation of normalization -> Done
 TODO: Improve calculation of random point for stencil -> Done
+TODO: Improve directory selection -> Done
 
 """
 
@@ -47,7 +48,7 @@ class TrainingDataGenerator(object):
 
     """
     def __init__(self, initial_conditions, left_bound=-1, right_bound=1, balance=0.5,
-                 stencil_length=3, directory=None):
+                 stencil_length=3, directory='test_data'):
         """Initializes TrainingDataGenerator.
 
         Parameters
@@ -85,8 +86,6 @@ class TrainingDataGenerator(object):
                 self._troubled_functions.append(function)
 
         # Set directory
-        self._data_dir = 'test_data'
-        if directory is not None:
         self._data_dir = directory
         if not os.path.exists(self._data_dir):
             os.makedirs(self._data_dir)