Merge pull request #18 from rodrigo-arenas/0.5.xdev

plot fitness evolution metric selection

demo/Boson_Houses_decision_tree.py

@@ -57,7 +57,7 @@
 print("r-squared: ", "{:.2f}".format(r_squared))
 
 print("Best k solutions: ", evolved_estimator.hof)
-plot = plot_fitness_evolution(evolved_estimator)
+plot = plot_fitness_evolution(evolved_estimator, metric="fitness_sd")
 plt.show()
 
 plot_search_space(evolved_estimator)

sklearn_genetic/_version.py

@@ -1,3 +1 @@
-
-__version__ = "0.4.1"
-
+__version__ = "0.5.0dev0"

sklearn_genetic/plots.py

@@ -1,36 +1,43 @@
 import seaborn as sns
 
 from .utils import logbook_to_pandas
+from .parameters import Metrics
 
 
 """
 This module contains some useful function to explore the results of the optimization routines
 """
 
 
-def plot_fitness_evolution(estimator):
+def plot_fitness_evolution(estimator, metric="fitness"):
     """
     Parameters
     ----------
     estimator: estimator object
         A fitted estimator from :class:`~sklearn_genetic.GASearchCV`
+    metric: {"fitness", "fitness_std", "fitness_max", "fitness_min"}, default="fitness"
+        Logged metric into the estimator history to plot
 
     Returns
     -------
     Lines plot with the fitness value in each generation
 
     """
+
+    if metric not in Metrics.list():
+        raise ValueError(f"metric must be one of {Metrics.list()}, but got {metric} instead")
+
     sns.set_style("white")
 
-    fitness_history = estimator.history["fitness"]
+    fitness_history = estimator.history[metric]
 
     palette = sns.color_palette("rocket")
     sns.set(rc={"figure.figsize": (10, 10)})
 
     ax = sns.lineplot(
         x=range(len(estimator)), y=fitness_history, markers=True, palette=palette
     )
-    ax.set_title("Fitness average evolution over generations")
+    ax.set_title(f"{metric.capitalize()} average evolution over generations")
 
     ax.set(xlabel="generations", ylabel=f"fitness ({estimator.scoring})")
     return ax

sklearn_genetic/tests/test_plots.py

@@ -1,3 +1,4 @@
+import pytest
 from sklearn.datasets import load_boston
 from sklearn.model_selection import train_test_split
 from sklearn.tree import DecisionTreeRegressor
@@ -44,8 +45,15 @@
 def test_plot_evolution():
     plot = plot_fitness_evolution(evolved_estimator)
 
+    with pytest.raises(Exception) as excinfo:
+        plot = plot_fitness_evolution(evolved_estimator, metric="accuracy")
+
+    assert str(excinfo.value) == "metric must be one of ['fitness', 'fitness_std', 'fitness_max', 'fitness_min'], " \
+                                 "but got accuracy instead"
+
 
 def test_plot_space():
+    plot = plot_search_space(evolved_estimator)
     plot = plot_search_space(evolved_estimator)
     plot = plot_search_space(
         evolved_estimator, features=["ccp_alpha", "max_depth", "min_samples_split"]