Extended color options.

DG_Approximation.py

@@ -14,9 +14,9 @@ TODO: Contemplate moving basis/wavelet matrices to Vectors_of_Polynomials
 TODO: Contemplate how to make shock tubes comparable
 TODO: Improve saving of plots -> Done (moved to Troubled_Cell_Detector)
 TODO: Added option to set plot directory -> Done
-TODO: Extend color options
+TODO: Extend color options -> Done
 TODO: Implement type check for all kwargs and configs
-TODO: Add option to not save plots
+TODO: Add option to not save plots -> Done (not call function in Main)
 
 """
 import numpy as np
@@ -123,7 +123,7 @@ class DGScheme(object):
             current_time += time_step
 
         # Plot exact/approximate results, errors, shock tubes and any detector-dependant plots
-        self._detector.plot_results(projection, troubled_cell_history, time_history, 'k-', 'y')
+        self._detector.plot_results(projection, troubled_cell_history, time_history)
 
         if self._verbose:
             plt.show()

Troubled_Cell_Detector.py

@@ -33,11 +33,17 @@ class TroubledCellDetector(object):
         self._init_cond = init_cond
         self._quadrature = quadrature
 
-        # Set plot directory from config if existing
+        # Set parameters from config if existing
         self._plot_dir = config.pop('plot_dir', 'fig')
+        self._colors = config.pop('colors', {})
 
+        self._check_colors()
         self._reset(config)
 
+    def _check_colors(self):
+        self._colors['exact'] = self._colors.get('exact', 'k-')
+        self._colors['approx'] = self._colors.get('approx', 'y')
+
     def _reset(self, config):
         pass
 
@@ -47,9 +53,9 @@ class TroubledCellDetector(object):
     def get_cells(self, projection):
         pass
 
-    def plot_results(self, projection, troubled_cell_history, time_history, color_exact, color_approx):
+    def plot_results(self, projection, troubled_cell_history, time_history):
         self._plot_shock_tube(troubled_cell_history, time_history)
-        max_error = self._plot_mesh(projection, color_exact, color_approx)
+        max_error = self._plot_mesh(projection)
 
         print("p =", self._polynom_degree)
         print("N =", self._num_grid_cells)
@@ -66,14 +72,14 @@ class TroubledCellDetector(object):
         plt.ylabel('Time')
         plt.title('Shock Tubes')
 
-    def _plot_mesh(self, projection, color_exact, color_approx):
+    def _plot_mesh(self, projection):
         grid, exact = self._calculate_exact_solution(self._mesh[2:-2], self._cell_len)
         approx = self._calculate_approximate_solution(projection[:, 1:-1])
 
         pointwise_error = np.abs(exact-approx)
         max_error = np.max(pointwise_error)
 
-        self._plot_solution_and_approx(grid, exact, approx, color_exact, color_approx)
+        self._plot_solution_and_approx(grid, exact, approx, self._colors['exact'], self._colors['approx'])
         plt.legend(['Exact', 'Approx'])
         self._plot_semilog_error(grid, pointwise_error)
         self._plot_error(grid, exact, approx)
@@ -177,6 +183,12 @@ class NoDetection(TroubledCellDetector):
 
 
 class WaveletDetector(TroubledCellDetector):
+    def _check_colors(self):
+        self._colors['fine_exact'] = self._colors.get('fine_exact', 'k-.')
+        self._colors['fine_approx'] = self._colors.get('fine_approx', 'b-.')
+        self._colors['coarse_exact'] = self._colors.get('coarse_exact', 'k-')
+        self._colors['coarse_approx'] = self._colors.get('coarse_approx', 'y')
+
     def _reset(self, config):
         # Set fixed basis and wavelet vectors
         self._basis = OrthonormalLegendre(self._polynom_degree).get_vector(x)
@@ -214,9 +226,9 @@ class WaveletDetector(TroubledCellDetector):
     def _get_cells(self, multiwavelet_coeffs, projection):
         return []
 
-    def plot_results(self, projection, troubled_cell_history, time_history, color_exact, color_approx):
+    def plot_results(self, projection, troubled_cell_history, time_history):
         self._plot_details(projection)
-        super().plot_results(projection, troubled_cell_history, time_history, color_exact, color_approx)
+        super().plot_results(projection, troubled_cell_history, time_history)
 
     def _plot_details(self, projection):
         fine_mesh = self._mesh[2:-2]
@@ -308,15 +320,15 @@ class WaveletDetector(TroubledCellDetector):
             matrix.append(row)
         return matrix
 
-    def _plot_mesh(self, projection, color_exact, color_approx):
+    def _plot_mesh(self, projection):
         grid, exact = self._calculate_exact_solution(self._mesh[2:-2], self._cell_len)
         approx = self._calculate_approximate_solution(projection[:, 1:-1])
 
         pointwise_error = np.abs(exact-approx)
         max_error = np.max(pointwise_error)
 
-        self._plot_coarse_mesh(projection, color_exact, color_approx)
-        self._plot_solution_and_approx(grid, exact, approx, 'k-.', 'b-.')
+        self._plot_coarse_mesh(projection)
+        self._plot_solution_and_approx(grid, exact, approx, self._colors['fine_exact'], self._colors['fine_approx'])
         plt.legend(['Exact (Coarse)', 'Approx (Coarse)', 'Exact (Fine)', 'Approx (Fine)'])
         self._plot_semilog_error(grid, pointwise_error)
         self._plot_error(grid, exact, approx)
@@ -333,7 +345,7 @@ class WaveletDetector(TroubledCellDetector):
         plt.figure(4)
         plt.savefig(self._plot_dir + '/coeff_details/' + name + '.pdf')
 
-    def _plot_coarse_mesh(self, projection, color_exact, color_approx):
+    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),
                                 coarse_cell_len)
@@ -343,7 +355,7 @@ class WaveletDetector(TroubledCellDetector):
         # Plot exact and approximate solutions for coarse mesh
         grid, exact = self._calculate_exact_solution(coarse_mesh[1:-1], coarse_cell_len)
         approx = self._calculate_approximate_solution(coarse_projection)
-        self._plot_solution_and_approx(grid, exact, approx, color_exact, color_approx)
+        self._plot_solution_and_approx(grid, exact, approx, self._colors['coarse_exact'], self._colors['coarse_approx'])
 
 
 class Boxplot(WaveletDetector):