Update pygad.py

rengel8 · rengel8 · commit ccfe679c7881 · 2021-06-14T23:01:09.000+02:00
Allows to save the plot into a file as an option, since linux might throw an error showing it.
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):
+    def plot_result(self, title="PyGAD - Iteration vs. Fitness", xlabel="Generation", ylabel="Fitness", linewidth=3, save_as=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,6 +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.
 
         Returns the figure.
         """
@@ -2555,7 +2556,12 @@ def plot_result(self, title="PyGAD - Iteration vs. Fitness", xlabel="Generation"
         matplotlib.pyplot.title(title)
         matplotlib.pyplot.xlabel(xlabel)
         matplotlib.pyplot.ylabel(ylabel)
-        matplotlib.pyplot.show()
+
+        if save_as is None:
+            matplotlib.pyplot.show()        # show plot
+        else:
+            fig.savefig(f'{save_as}.png')   # save as image
+
         return fig
 
     def save(self, filename):
