ENH Sets xlabel only if it is not defined in PDP (#15610)

thomasjpfan · adrinjalali · commit c170be1c73f9 · 2019-11-19T11:54:01.000+01:00
diff --git a/sklearn/inspection/_partial_dependence.py b/sklearn/inspection/_partial_dependence.py
@@ -933,9 +933,12 @@ def plot(self, ax=None, n_cols=3, line_kw=None, contour_kw=None):
             ylim = axi.get_ylim()
             axi.vlines(self.deciles[fx[0]], 0, 0.05, transform=trans,
                        color='k')
-            axi.set_xlabel(self.feature_names[fx[0]])
             axi.set_ylim(ylim)
 
+            # Set xlabel if it is not already set
+            if not axi.get_xlabel():
+                axi.set_xlabel(self.feature_names[fx[0]])
+
             if len(values) == 1:
                 if n_cols is None or i % n_cols == 0:
                     axi.set_ylabel('Partial dependence')
diff --git a/sklearn/inspection/tests/test_plot_partial_dependence.py b/sklearn/inspection/tests/test_plot_partial_dependence.py
@@ -256,6 +256,24 @@ def test_plot_partial_dependence_with_same_axes(pyplot, clf_boston, boston):
                                 feature_names=boston.feature_names, ax=ax)
 
 
+def test_plot_partial_dependence_feature_name_reuse(pyplot, clf_boston,
+                                                    boston):
+    # second call to plot does not change the feature names from the first
+    # call
+
+    feature_names = boston.feature_names
+    disp = plot_partial_dependence(clf_boston, boston.data,
+                                   [0, 1],
+                                   grid_resolution=10,
+                                   feature_names=feature_names)
+
+    plot_partial_dependence(clf_boston, boston.data, [0, 1],
+                            grid_resolution=10, ax=disp.axes_)
+
+    for i, ax in enumerate(disp.axes_.ravel()):
+        assert ax.get_xlabel() == feature_names[i]
+
+
 def test_plot_partial_dependence_multiclass(pyplot):
     grid_resolution = 25
     clf_int = GradientBoostingClassifier(n_estimators=10, random_state=1)
