Cannot Set Ylim For Sklearn Partial Dependence Plot

Contents hide

1 Cannot set ylim for sklearn partial dependence plot

Problem Description:

I am following this example on sklearn documentation

I want to change the limits of y axis so I can visually compare results from different models.

import pandas as pd
import matplotlib.pyplot as plt
from sklearn.datasets import load_diabetes
from sklearn.tree import DecisionTreeRegressor
from sklearn.inspection import PartialDependenceDisplay

diabetes = load_diabetes()
X = pd.DataFrame(diabetes.data, columns=diabetes.feature_names)
y = diabetes.target

tree = DecisionTreeRegressor()
tree.fit(X, y)

fig, ax = plt.subplots(figsize=(12, 6))
ax.set_ylim(50,300)
tree_disp = PartialDependenceDisplay.from_estimator(tree, X, ["age", "bmi"], ax=ax)

However, it seems that ax.set_ylim get ignored no matter what I specify. On the other hand, ax.set_title given in example works fine.

Solution – 1

PartialDependenceDisplay have an axes_ attribute that represents both matplotlib’s axes of the figure.

You can modify them as follow:

tree_disp = PartialDependenceDisplay.from_estimator(tree, X, ["age", "bmi"], ax=ax)
tree_disp.axes_[0][0].set_ylim(50,300)
tree_disp.axes_[0][1].set_ylim(50,300)

This will output the following plot:

Rate this post