Lesson 2 Project: Using the Position-Salary Dataset on Kaggle

# Run this in your terminal or command prompt
pip install numpy==1.19.4 pandas matplotlib scikit-learn

# Import libraries
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.linear_model import LinearRegression
from sklearn.preprocessing import PolynomialFeatures
from sklearn.metrics import r2_score

# Load the dataset
dataset = pd.read_csv('Position_Salaries.csv')
print(dataset.head())

# Basic information about the dataset
print(dataset.info())
print(dataset.describe())

# Check column names to ensure correct access
print("Column names:", dataset.columns.tolist())

# Visualize the data
plt.figure(figsize=(10, 6))
plt.scatter(dataset.iloc[:, 1], dataset.iloc[:, 2], color='red')  # Using column positions instead of names
plt.title('Salary vs Position Level')
plt.xlabel('Position Level')
plt.ylabel('Salary')
plt.show()

# Separate features and target
X = dataset.iloc[:, 1:2].values  # Position level (column index 1)
y = dataset.iloc[:, 2].values   # Salary (column index 2)

# Fit linear regression model
linear_reg = LinearRegression()
linear_reg.fit(X, y)

# Predict with linear model
linear_pred = linear_reg.predict(X)

# Calculate R score for linear model
linear_r2 = r2_score(y, linear_pred)
print(f"Linear Regression R score: {linear_r2}")

# Create polynomial features (try different degrees)
degree = 4  # You can experiment with different values
poly_features = PolynomialFeatures(degree=degree)
X_poly = poly_features.fit_transform(X)

# Fit polynomial regression model
poly_reg = LinearRegression()
poly_reg.fit(X_poly, y)

# Predict with polynomial model
poly_pred = poly_reg.predict(X_poly)

# Calculate R score for polynomial model
poly_r2 = r2_score(y, poly_pred)
print(f"Polynomial Regression (degree {degree}) R score: {poly_r2}")

# Create a more continuous set of points for smoother curve visualization
X_grid = np.arange(min(X), max(X), 0.1).reshape(-1, 1)

# Plot the original data points
plt.figure(figsize=(12, 8))
plt.scatter(X, y, color='red', label='Data Points')

# Plot the linear regression line
plt.plot(X, linear_pred, color='blue', label=f'Linear Regression (R = {linear_r2:.3f})')

# Plot the polynomial regression curve
plt.plot(X_grid, poly_reg.predict(poly_features.transform(X_grid)), 
         color='green', label=f'Polynomial Regression degree {degree} (R = {poly_r2:.3f})')

plt.title('Salary vs Position Level: Linear vs Polynomial Regression')
plt.xlabel('Position Level')
plt.ylabel('Salary')
plt.legend()
plt.grid(True)
plt.show()

# Predict salary for position level 6.5
new_level = np.array([[6.5]])

# Linear model prediction
linear_prediction = linear_reg.predict(new_level)
print(f"Linear model predicts salary for level 6.5: ${linear_prediction[0]:.2f}")

# Polynomial model prediction
poly_prediction = poly_reg.predict(poly_features.transform(new_level))
print(f"Polynomial model predicts salary for level 6.5: ${poly_prediction[0]:.2f}")