Address Guillaume's comments

Author: ArturoAmorQ

python_scripts/trees_sol_01.py

@@ -149,40 +149,40 @@
 # plots in the figure.
 
 # %% tags=["solution"]
-import numpy as np
+# import numpy as np
 from matplotlib import cm
 
-classes = np.unique(tree.classes_)
 _, axs = plt.subplots(ncols=3, nrows=1, sharey=True, figsize=(12, 5))
 plt.suptitle("Predicted probabilities for decision tree model", y=1.05)
 plt.subplots_adjust(bottom=0.45)
 
-for idx, class_of_interest in enumerate(classes):
-    axs[idx].set_title(f"Class {class_of_interest}")
-    disp = DecisionBoundaryDisplay.from_estimator(
+for idx, (class_of_interest, ax) in enumerate(zip(tree.classes_, axs)):
+    ax.set_title(f"Class {class_of_interest}")
+    DecisionBoundaryDisplay.from_estimator(
         tree,
         data_test,
         response_method="predict_proba",
         class_of_interest=class_of_interest,
-        ax=axs[idx],
+        ax=ax,
         vmin=0,
         vmax=1,
     )
-    axs[idx].scatter(
+    ax.scatter(
         data_test["Culmen Length (mm)"].loc[target_test == class_of_interest],
         data_test["Culmen Depth (mm)"].loc[target_test == class_of_interest],
         marker="o",
         c="w",
         edgecolor="k",
     )
-    axs[idx].set_xlabel("Culmen Length (mm)")
-    axs[idx].set_ylabel("Culmen Depth (mm)" if idx == 0 else None)
+    ax.set_xlabel("Culmen Length (mm)")
+    if idx == 0:
+        ax.set_ylabel("Culmen Depth (mm)")
 
 ax = plt.axes([0.15, 0.14, 0.7, 0.05])
 plt.colorbar(
     cm.ScalarMappable(cmap="viridis"), cax=ax, orientation="horizontal"
 )
-_ = plt.title("Probability")
+_ = ax.set_title("Class probability")
 
 # %% [markdown] tags=["solution"]
 # ```{note}