Merge pull request ahmedfgad#47 from rengel8/save-matplot-option

ahmedfgad · web-flow · commit c2e78b2a3ea6 · 2021-06-15T16:24:13.000-04:00
Two changes in the `plot_result()` method.
1) Rename the parameter to `save_dir`. If it is not None, then save.
2) Keep the `.show()` function called regardless of saving the figure or not.
diff --git a/pygad.py b/pygad.py
@@ -2792,7 +2792,7 @@ def best_solution(self, pop_fitness=None):
 
         return best_solution, best_solution_fitness, best_match_idx
 
-    def plot_result(self, title="PyGAD - Iteration vs. Fitness", xlabel="Generation", ylabel="Fitness", linewidth=3):
+    def plot_result(self, title="PyGAD - Iteration vs. Fitness", xlabel="Generation", ylabel="Fitness", linewidth=3, save_dir=None):
 
         """
         Creates and shows a plot that summarizes how the fitness value evolved by generation. Can only be called after completing at least 1 generation. If no generation is completed, an exception is raised.
@@ -2802,6 +2802,7 @@ def plot_result(self, title="PyGAD - Iteration vs. Fitness", xlabel="Generation"
             xlabel: Label on the X-axis.
             ylabel: Label on the Y-axis.
             linewidth: Line width of the plot.
+            save_dir: Directory to save the figure.
 
         Returns the figure.
         """
@@ -2817,7 +2818,12 @@ def plot_result(self, title="PyGAD - Iteration vs. Fitness", xlabel="Generation"
         matplotlib.pyplot.title(title)
         matplotlib.pyplot.xlabel(xlabel)
         matplotlib.pyplot.ylabel(ylabel)
+
+        if not save_dir is None:
+            matplotlib.pyplot.savefig(fname=save_dir, 
+                                      bbox_inches='tight')
         matplotlib.pyplot.show()
+
         return fig
 
     def save(self, filename):
@@ -2845,4 +2851,4 @@ def load(filename):
         raise FileNotFoundError("Error reading the file {filename}. Please check your inputs.".format(filename=filename))
     except:
         raise BaseException("Error loading the file. Please check if the file exists.")
-    return ga_in
+    return ga_in
