# %%
from sklearn.datasets import load_iris

iris = load_iris(as_frame=True)
X = iris.data[["sepal width (cm)", "petal width (cm)"]]
y = iris.target

# %%
from sklearn.model_selection import train_test_split

X_train, X_test, y_train, y_test = train_test_split(
    X, y, random_state=42
)

# %%
from sklearn.linear_model import LogisticRegression

model = LogisticRegression().fit(X_train, y_train)

# %%
from sklearn.inspection import DecisionBoundaryDisplay

display = DecisionBoundaryDisplay.from_estimator(model, X_train, alpha=0.5)
scatter = display.ax_.scatter(
    X_train["sepal width (cm)"], X_train["petal width (cm)"], c=y_train, edgecolor="k"
)
display.ax_.legend(scatter.legend_elements()[0], iris.target_names)