Linear Regression in Python

A complete, from-data-to-diagnostics linear regression workflow in Python — using scikit-learn for the standard path, and a from-scratch Normal Equation implementation to show exactly what's happening underneath.

The Full scikit-learn Workflow

import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_squared_error, r2_score
import matplotlib.pyplot as plt

# 1. Dataset
df = pd.DataFrame({
    "hours_studied": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
    "marks":         [52, 58, 62, 68, 75, 78, 84, 88, 92, 97],
})

# 2. Features and target
X = df[["hours_studied"]]
y = df["marks"]

# 3. Train/test split
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 4. Model creation and training
model = LinearRegression()
model.fit(X_train, y_train)

# 5. Prediction
y_pred = model.predict(X_test)

# 6. Evaluation
print("Intercept:", model.intercept_)
print("Coefficient:", model.coef_)
print("MSE:", mean_squared_error(y_test, y_pred))
print("R²:", r2_score(y_test, y_pred))

# 7. Visualization
plt.scatter(X_train, y_train, label="Training data", alpha=0.6)
plt.plot(X_train, model.predict(X_train), color="red", label="Fitted line")
plt.scatter(X_test, y_test, color="green", label="Test data")
plt.xlabel("Hours studied"); plt.ylabel("Marks"); plt.legend()
plt.show()

Expected output: a near-perfect fit given this clean, close-to-linear toy data — R² close to 1.0. On real data, expect R² well below 1.0; treat a suspiciously perfect fit as a signal to check for data leakage, not a reason to celebrate.

The Normal Equation, From Scratch

import numpy as np

X_raw = np.array([[1], [2], [3], [4], [5]])
y = np.array([52, 58, 62, 68, 75])

X_b = np.c_[np.ones((len(X_raw), 1)), X_raw]   # add a column of 1s for the intercept term

b = np.linalg.inv(X_b.T @ X_b) @ X_b.T @ y
print(b)   # [46.2  5.6] -- matches scikit-learn exactly

The column of 1s prepended to \(X\) is what lets a single matrix formula solve for both the intercept and the slope(s) simultaneously — the intercept is treated as the coefficient on a feature that's always 1.

Diagnosing the Fit

residuals = y_test - y_pred

plt.scatter(y_pred, residuals)
plt.axhline(0, color="red", linestyle="--")
plt.xlabel("Predicted"); plt.ylabel("Residuals")
plt.show()
# See Linear Regression Assumptions for how to read this plot

Saving and Reusing the Model

import joblib

joblib.dump(model, "linear_regression_model.pkl")
loaded_model = joblib.load("linear_regression_model.pkl")
print(loaded_model.predict([[11]]))

See ML Model Deployment for the security considerations around loading serialized models from untrusted sources.

Common Mistakes

• Reporting R² or MSE computed on the training set as if it were the model's real-world performance — always evaluate on the held-out test set.
• Forgetting .reshape(-1, 1) when passing a single feature as a plain 1D array, triggering scikit-learn's "expected 2D array" error.
• Skipping the residual plot step — a model can have a reasonable R² while still clearly violating assumptions that a quick plot would reveal.

Interview Relevance

Q: "Implement linear regression without using scikit-learn." The Normal Equation code above is exactly this answer — being able to write \((X^TX)^{-1}X^Ty\) from memory, and explain the column-of-1s trick for the intercept, is a common interview signal that you understand the algorithm, not just the library call.

Practice Question

Modify the from-scratch Normal Equation code to work with two features instead of one (hint: add a second column to X_raw before building X_b).

Related ML Notes

• Linear Regression with Gradient Descent
• Linear Regression Example
• R² Score