Cleaned up Snakefile.

Snakefile

-configfile: 'config.yaml'
 import ANN_Data_Generator, Initial_Condition, ANN_Training
 from ANN_Training import evaluate_models
 import numpy as np
 
-def replace_none(list):
-    return {} if list is None else list
+configfile: 'config.yaml'
 
 DIR = config['data_directory']
 MODELS = config['models']
-
 if config['random_seed'] is not None:
     np.random.seed(config['random_seed'])
 
@@ -34,11 +31,13 @@ rule test_model:
             trainer= ANN_Training.ModelTrainer({'model_name': model, 'dir': DIR,
                                             'model_dir': DIR, **MODELS[model]})
             models[model] = trainer
-        evaluate_models(models, DIR, 2, params.colors, params.compare_normalization)
+        evaluate_models(models=models, directory=DIR, num_iterations=2, colors=params.colors,
+            compare_normalization=params.compare_normalization)
 
 rule generate_data:
     output:
         DIR+'/input_data.npy',
+        DIR+'/normalized_input_data.npy',
         DIR+'/output_data.npy'
     params:
         left_bound = config['left_boundary'],
@@ -55,28 +54,26 @@ rule generate_data:
             initial_conditions.append({
                 'function': getattr(Initial_Condition, function)(
                     params.left_bound, params.right_bound, {}),
-                'config': replace_none(config['functions'][function])})
+                'config': {} if function in config['functions'] else config['functions'][function]
+            })
 
-        generator = ANN_Data_Generator.TrainingDataGenerator(initial_conditions,
+        generator = ANN_Data_Generator.TrainingDataGenerator(initial_conditions=initial_conditions,
             left_bound=params.left_bound, right_bound=params.right_bound, balance=params.balance,
             stencil_length=params.stencil_length, directory=DIR)
-        data = generator.build_training_data(params.sample_number)
-        # print(data[0])
+        data = generator.build_training_data(num_samples=params.sample_number)
 
 rule train_model:
     input:
         DIR+'/input_data.npy',
+        DIR+'/normalized_input_data.npy',
         DIR+'/output_data.npy'
-    # params:
-    #     # parameter = MODELS[{model}],
-    #     lambda wildcards: [MODELS[wildcards.model]]
     log:
         DIR+'/log/train_model_{model}.log'
     output:
         DIR+'/trained models/model__{model}.pt',
         DIR+'/trained models/loss__{model}.pt'
     run:
-        trainer= ANN_Training.ModelTrainer({'model_name': wildcards.model, 'dir': DIR,
+        trainer= ANN_Training.ModelTrainer(config={'model_name': wildcards.model, 'dir': DIR,
                                         'model_dir': DIR, **MODELS[wildcards.model]})
         trainer.epoch_training()
         trainer.save_model()
\ No newline at end of file