Does this also allow to pass a ax that was drawn on before? Or will it be cleared like currently in master?

For this specific case, if the object was drawn on by a PartialDependenceDisplay object, it will not error as shown above.

If it was drawn on by a normal plot function, it will error. For example:

fig, ax = pyplot.subplots()
ax.plot([0, 1, 2], [1, 2, 3])
plot_partial_dependence(..., ax=ax)

Is it reasonable to store the display object as an attribute of the ax?

The display objects stores all the data which can get pretty big. You would expect that the data can be GCed if the Display object isn't there anymore.

But with this PR, the data exists as long as ax exists.

If we can, it would be ideal to also support that case

fig, ax = pyplot.subplots()
ax.plot([0, 1, 2], [1, 2, 3])
plot_partial_dependence(..., ax=ax)

This is typically useful if you want to plot the data points and then the PDPs on top, as done in #15582 (comment)

Yes I have to make sure the axes is a weak reference.

As in the axes has a weak reference to the display object.

Sorry if I'm missing something but I don't see where the weak ref is being created? (this should also be tested)

I was just thinking about the weakref on the train. I am going to move this back to WIP for a little bit.

thank you for the early reviews.

If we can, it would be ideal to also support that case

fig, ax = pyplot.subplots()
ax.plot([0, 1, 2], [1, 2, 3])
plot_partial_dependence(..., ax=ax)

This is typically useful if you want to plot the data points and then the PDPs on top, as done in #15582 (comment)

Supporting this can run into an unclear API. For example:

fig, ax = pyplot.subplots()
ax.plot([0, 1, 2], [1, 2, 3])

plot_partial_dependence(est, features=[0, 1, 2], ax=ax)

For this case, I do not think a user would want the first call to plot to be used in all 3 partial dependence plots.

The only case where this can make sense is when there is only one feature:

fig, ax = pyplot.subplots()
ax.plot([0, 1, 2], [1, 2, 3])

plot_partial_dependence(est, features=[0], ax=ax)

I am -0.2 on this. It will lead to another code path to handle len(features)==1 and isinstance(ax, plt.Axes).

Please ping when ready @thomasjpfan .

Also, I was looking at pandas internal: when you pass ax to df.plot() but they need to clear the figure, they raise a UserWarning. Maybe we should do that too to avoid surprises. To reproduce:

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

d = {'a': np.arange(10),
     'b': np.arange(10) + 23 * 4}

fig, ax = plt.subplots()

df = pd.DataFrame(d)
df.plot(subplots=True, ax=ax)  # ask for subplots but also for a single ax => warning
plt.show()

@NicolasHug This is ready for another review.

I slightly prefer the current behavior where we do not clear the axes and raise an error if the axes has been used. If we relax this restriction and automatically clear the axes with a warning, there will be code out there that depends on this behavior.

please refresh my memory: if I pass an array-like as ax, with stuff already drawn-on, do we raise an error?

In retrospect I'm not even sure we want to raise errors/warnings.

If there as an array like of ax with stuff already drawn on, it will raise an error if the number of axes in the array like is not the expected number of axes. If it has the expected number of axes, then they will be drawn on directly.

#15760 provides a better solution to how plot_partial_dependence checks the axes.