Plot iris petal scatter with regression with python

iris.py

import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np
from matplotlib.patches import Ellipse
from sklearn.linear_model import LinearRegression
from sklearn.metrics import r2_score

def plot_iris_petal_scatter_with_regression():
    """
    Downloads Iris dataset, extracts petal data, creates a scatter plot with regression,
    overlays a filled ellipse for each class based on petal distribution.
    """
    try:
        # Download the Iris dataset from the UCI Machine Learning Repository
        url = "https://archive.ics.uci.edu/ml/machine-learning-databases/iris/iris.data"
        column_names = ["sepal_length", "sepal_width", "petal_length", "petal_width", "class"]
        df = pd.read_csv(url, names=column_names)

        # Extract petal data (petal length and petal width)
        petal_data = df[["petal_length", "petal_width"]]

        # Prepare data for linear regression
        X = petal_data[["petal_length"]]
        y = petal_data["petal_width"]

        # Create and train the linear regression model
        model = LinearRegression()
        model.fit(X, y)

        # Calculate the R-squared value
        y_pred = model.predict(X)
        r2 = r2_score(y, y_pred)
        print(f"R-squared (R^2): {r2}")

        # Generate predictions for plotting the regression line
        x_range = pd.DataFrame({
            "petal_length": [petal_data["petal_length"].min(), petal_data["petal_length"].max()]
        })
        y_range = model.predict(x_range)

        # Create scatter plot with class differentiation using seaborn
        plt.figure(figsize=(10, 6))
        sns.scatterplot(data=df, x="petal_length", y="petal_width", hue="class", s=100)

        # Plot the linear regression line
        plt.plot(x_range, y_range, color="red", linewidth=2, label="Regression Line")

        # Overlay a half-transparent filled ellipse for each class
        ax = plt.gca()
        unique_classes = df["class"].unique()
        palette = sns.color_palette("deep", n_colors=len(unique_classes))
        color_map = dict(zip(unique_classes, palette))

        for cls in unique_classes:
            sub_df = df[df["class"] == cls]
            x_vals = sub_df["petal_length"].values
            y_vals = sub_df["petal_width"].values
            # Compute the mean of x and y
            mean_x, mean_y = x_vals.mean(), y_vals.mean()
            # Compute the covariance matrix
            cov = np.cov(x_vals, y_vals)
            # Compute the eigenvalues and eigenvectors
            eigvals, eigvecs = np.linalg.eigh(cov)
            # Compute the angle of the ellipse
            angle = np.degrees(np.arctan2(*eigvecs[:, 0][::-1]))
            # Scale the eigenvalues by 4 to enclose most of the data distribution
            width, height = 4 * np.sqrt(eigvals)
            ellipse = Ellipse((mean_x, mean_y),
                              width=width,
                              height=height,
                              angle=angle,
                              facecolor=color_map[cls],
                              alpha=0.3,
                              edgecolor=None,
                              label=f"{cls} region")
            ax.add_patch(ellipse)

        plt.title("Iris Petal Scatter Plot by Class with Regression and Regions")
        plt.xlabel("Petal Length")
        plt.ylabel("Petal Width")
        plt.legend(title="Class")
        plt.grid(True)
        plt.show()

    except Exception as e:
        print(f"An error occurred: {e}")
        print("Please check your internet connection or the data download.")

# Run the function to create and show the plot.
plot_iris_petal_scatter_with_regression()