Two changes in the plot_result() method.

ahmedfgad · web-flow · commit 517612c4a5b6 · 2021-06-14T23:14:19.000-04:00
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
@@ -2530,7 +2530,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, save_as=None):
+    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.
@@ -2540,7 +2540,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_as: save image instead of showing it.
+            save_dir: Directory to save the figure.
 
         Returns the figure.
         """
@@ -2557,10 +2557,10 @@ def plot_result(self, title="PyGAD - Iteration vs. Fitness", xlabel="Generation"
         matplotlib.pyplot.xlabel(xlabel)
         matplotlib.pyplot.ylabel(ylabel)
 
-        if save_as is None:
-            matplotlib.pyplot.show()        # show plot
-        else:
-            fig.savefig(f'{save_as}.png')   # save as image
+        if not save_dir is None:
+            matplotlib.pyplot.savefig(fname=save_dir, 
+                                      bbox_inches='tight')
+        matplotlib.pyplot.show()
 
         return fig
 
@@ -2589,4 +2589,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
