diff --git a/DG_Approximation.py b/DG_Approximation.py
index fc43cdb04b34e9d804b39afb524af6185feeea4a..274b2df6dd5b8edf2e08222a50e575a8137d4995 100644
--- a/DG_Approximation.py
+++ b/DG_Approximation.py
@@ -10,8 +10,12 @@ TODO: Contemplate how to make shock tubes comparable
 TODO: Check whether 'projection' is always a np.array()
 TODO: Check whether all instance variables sensible
 TODO: Use cfl_number for updating, not just time
+TODO: Remove Training_Data.py -> Done
+TODO: Remove unnecessary comments -> Done
+TODO: Combine saving of plots -> Done
 
 """
+import os
 import numpy as np
 from sympy import Symbol
 import math
@@ -44,6 +48,7 @@ class DGScheme(object):
         self._left_bound = kwargs.pop('left_bound', -1)
         self._right_bound = kwargs.pop('right_bound', 1)
         self._verbose = kwargs.pop('verbose', False)
+        self._plot_dir = kwargs.pop('plot_dir', 'testing')
         self._history_threshold = kwargs.pop('history_threshold', math.ceil(0.2/self._cfl_number))
         self._detector = detector
         self._detector_config = kwargs.pop('detector_config', {})
@@ -130,7 +135,18 @@ class DGScheme(object):
             + str(self._final_time) + '__wave_speed_' + str(self._wave_speed) + '__number_of_cells_' \
             + str(self._num_grid_cells) + '__polynomial_degree_' + str(self._polynomial_degree)
 
-        self._detector.save_plots(name)
+        # Set paths for plot files if not existing already
+        if not os.path.exists(self._plot_dir):
+            os.makedirs(self._plot_dir)
+
+        # Save plots
+        for identifier in plt.get_figlabels():
+            # Set path for figure directory if not existing already
+            if not os.path.exists(self._plot_dir + '/' + identifier):
+                os.makedirs(self._plot_dir + '/' + identifier)
+
+            plt.figure(identifier)
+            plt.savefig(self._plot_dir + '/' + identifier + '/' + name + '.pdf')
 
     def _reset(self):
         # Set additional necessary instance variables
diff --git a/Troubled_Cell_Detector.py b/Troubled_Cell_Detector.py
index 6d487a699174b0c39910ee5aafe11701e07a0a08..35c4d241427bc7304b2f87a6f3b40a6a193cef4b 100644
--- a/Troubled_Cell_Detector.py
+++ b/Troubled_Cell_Detector.py
@@ -3,9 +3,9 @@
 @author: Laura C. Kühle, Soraya Terrab (sorayaterrab)
 
 TODO: Move plotting to separate file (try to adjust for different equations)
+TODO: Improve figure identifiers -> Done
 
 """
-import os
 import numpy as np
 import matplotlib.pyplot as plt
 import seaborn as sns
@@ -35,7 +35,6 @@ class TroubledCellDetector(object):
         self._quadrature = quadrature
 
         # Set parameters from config if existing
-        self._plot_dir = config.pop('plot_dir', 'fig')
         self._colors = config.pop('colors', {})
 
         self._check_colors()
@@ -165,36 +164,6 @@ class TroubledCellDetector(object):
 
         return np.reshape(np.array(approx), (1, len(approx) * num_points))
 
-    def save_plots(self, name):
-        # Set paths for plot files if not existing already
-        if not os.path.exists(self._plot_dir):
-            os.makedirs(self._plot_dir)
-
-        if not os.path.exists(self._plot_dir + '/exact_and_approx'):
-            os.makedirs(self._plot_dir + '/exact_and_approx')
-
-        if not os.path.exists(self._plot_dir + '/semilog_error'):
-            os.makedirs(self._plot_dir + '/semilog_error')
-
-        if not os.path.exists(self._plot_dir + '/error'):
-            os.makedirs(self._plot_dir + '/error')
-
-        if not os.path.exists(self._plot_dir + '/shock_tube'):
-            os.makedirs(self._plot_dir + '/shock_tube')
-
-        # Save plots
-        plt.figure('exact_and_approx')
-        plt.savefig(self._plot_dir + '/exact_and_approx/' + name + '.pdf')
-
-        plt.figure('semilog_error')
-        plt.savefig(self._plot_dir + '/semilog_error/' + name + '.pdf')
-
-        plt.figure('error')
-        plt.savefig(self._plot_dir + '/error/' + name + '.pdf')
-
-        plt.figure('shock_tube')
-        plt.savefig(self._plot_dir + '/shock_tube/' + name + '.pdf')
-
 
 class NoDetection(TroubledCellDetector):
     def get_cells(self, projection):
@@ -335,16 +304,6 @@ class WaveletDetector(TroubledCellDetector):
 
         return max_error
 
-    def save_plots(self, name):
-        super().save_plots(name)
-
-        # Set path for details plot files if not existing already
-        if not os.path.exists(self._plot_dir + '/coeff_details'):
-            os.makedirs(self._plot_dir + '/coeff_details')
-
-        plt.figure('coeff_details')
-        plt.savefig(self._plot_dir + '/coeff_details/' + name + '.pdf')
-
     def _plot_coarse_mesh(self, projection):
         coarse_cell_len = 2*self._cell_len
         coarse_mesh = np.arange(self._left_bound - (0.5*coarse_cell_len), self._right_bound + (1.5*coarse_cell_len),