Files
learning_science/iris_flowers/linear_regression.py
T
2026-05-11 07:37:50 +01:00

116 lines
3.2 KiB
Python

import matplotlib.pyplot as plt
import pandas as pd
from sklearn import datasets, linear_model
from sklearn.metrics import mean_squared_error, r2_score
from sklearn.model_selection import train_test_split
import os
def load_petal_data():
"""
Loads the Iris dataset and extracts petal dimensions.
Returns:
X (DataFrame): Feature matrix containing 'petal length (cm)'.
y (Series): Target vector containing 'petal width (cm)'.
"""
iris = datasets.load_iris(as_frame=True)
df = iris.frame
# We want to predict Petal Width based on Petal Length
# Reshaping X to be 2D array as required by sklearn (n_samples, n_features)
X = df[['petal length (cm)']]
y = df['petal width (cm)']
return X, y
def train_linear_model(X_train, y_train):
"""
Trains a Linear Regression model.
Args:
X_train: Training features.
y_train: Training target.
Returns:
model: Trained LinearRegression model.
"""
model = linear_model.LinearRegression()
model.fit(X_train, y_train)
return model
def evaluate_model_performance(model, X_test, y_test):
"""
Evaluates the model using Mean Squared Error and R-squared score.
Args:
model: Trained model.
X_test: Test features.
y_test: Test target.
"""
y_pred = model.predict(X_test)
mse = mean_squared_error(y_test, y_pred)
r2 = r2_score(y_test, y_pred)
print("Model Performance:")
print(f" Coefficients: {model.coef_[0]:.4f}")
print(f" Intercept: {model.intercept_:.4f}")
print(f" Mean Squared Error (MSE): {mse:.4f}")
print(f" Coefficient of Determination (R^2): {r2:.4f}")
def visualize_regression_line(model, X, y):
"""
Plots the data points and the regression line.
Args:
model: Trained model.
X: All features (for plotting).
y: All targets (for plotting).
"""
if not os.path.exists('plots'):
os.makedirs('plots')
plt.figure(figsize=(10, 6))
# Plot actual data points
plt.scatter(X, y, color='black', label='Actual Data')
# Plot regression line
# We predict on the full X range to draw the line
y_pred_full = model.predict(X)
plt.plot(X, y_pred_full, color='blue', linewidth=3, label='Regression Line')
plt.xlabel('Petal Length (cm)')
plt.ylabel('Petal Width (cm)')
plt.title('Linear Regression: Petal Length vs Petal Width')
plt.legend()
plt.grid(True, alpha=0.3)
output_path = 'plots/linear_regression.png'
plt.savefig(output_path)
plt.close()
print(f"Regression plot saved to {output_path}")
def main():
print("Starting Linear Regression Analysis...")
# 1. Load Data
X, y = load_petal_data()
# 2. Split Data, 80/20 rule
X_train, X_test, y_train, y_test = train_test_split(X, y)#, test_size=0.2)
# 3. Train Model
model = train_linear_model(X_train, y_train)
# 4. Evaluate
evaluate_model_performance(model, X_test, y_test)
# 5. Visualize
visualize_regression_line(model, X, y)
print("Analysis Complete.")
if __name__ == "__main__":
main()