diff --git a/pygad.py b/pygad.py index 7ee11ca..5b2595b 100644 --- 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_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,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_dir: Directory to save the figure. 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) + + if not save_dir is None: + matplotlib.pyplot.savefig(fname=save_dir, + bbox_inches='tight') matplotlib.pyplot.show() + return fig def save(self, filename): @@ -2583,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 \ No newline at end of file + return ga_in