2.6 - Noah's Notes - First predictive model with scikit-learn

In this notebook, we will build our first predictive model using Python's powerful scikit-learn library. We will focus on Linear Regression, a fundamental algorithm in machine learning, and explore how scikit-learn simplifies the process of training, evaluating, and using models.

It is the companion of the notebook Linear Regression from Scratch: there we wrote the prediction rule, the cost and the gradient descent loop ourselves; here the library does that work for us. To be able to check the result, we will use data generated from a line we choose ourselves, $y = 4 + 3x$, plus some random noise. The model will only see the noisy points, so if it works it should recover a slope close to 3 and an intercept close to 4.

Throughout this notebook, we will cover the following steps:

1. Import the required libraries: We will load the scikit-learn tools we need.
2. Load and prepare the data: We will generate synthetic data for our example.
3. Split the data into training and test sets: We will separate the data to train and effectively evaluate the model.
4. Train the model: We will fit the linear regression model to our training data.
5. Evaluate the model: We will measure the model's performance using standard metrics.
6. Visualize the results: We will plot the data and the regression line for visual understanding.
7. Use the model to make predictions: We will make predictions with the trained model.

Step 1: Import the required libraries

---

In this section, we will import the essential scikit-learn functions and classes we need to build, train, evaluate, and use our linear regression model.

# Import LinearRegression to build the model
from sklearn.linear_model import LinearRegression

# Import train_test_split to split the data into training and test sets
from sklearn.model_selection import train_test_split

# Import mean_squared_error and r2_score to evaluate the model
from sklearn.metrics import mean_squared_error, r2_score

Step 2: Data loading and preparation

---

For this example, we will generate a synthetic dataset that simulates a linear relationship between an independent variable (X) and a dependent variable (y), adding a bit of noise to make the example more realistic. This allows us to control the characteristics of the data and better understand how the model works.

The line we plant is $y = 4 + 3x$ (slope 3, intercept 4), and the noise is drawn from a standard normal distribution (mean 0, standard deviation 1). Note the shape (m, 1): scikit-learn expects X as a table with one row per example and one column per feature.

import numpy as np

# Generate synthetic data
m = 50  # Number of samples
np.random.seed(0)  # For reproducibility: the same "random" numbers on every run
X = 10 * np.random.rand(m, 1)  # m data points between 0 and 10, shape (50, 1)
y = 4 + 3 * X + np.random.randn(m, 1)  # The planted line y = 4 + 3x, plus noise

Step 3: Split the data into training and test sets

---

It is standard practice in machine learning to split the data into two sets:

• Training set: to train the model.
• Test set: to evaluate its performance on unseen data.

Common splits are 70/30, 75/25, and 80/20. We will use an 80% training and 20% test split.

This helps avoid overfitting and obtain a more realistic estimate of the model's generalization ability. We will use scikit-learn's train_test_split function to easily perform this split.

# Split the data into training and test sets
# test_size=0.2 -> 20% of the rows (10 of 50) go to the test set
# random_state=0 -> the same shuffle, and therefore the same split, on every run
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=0)

print(f"Training set size: {X_train.shape[0]}")
print(f"Test set size: {X_test.shape[0]}")

Training set size: 40
Test set size: 10

Step 4: Model training

---

Here is where we use scikit-learn's LinearRegression class. We will create an instance of this model and "train" it using the training set (X_train and y_train). The fit() method will find the best coefficients for the regression line that best fits these data.

Unlike the from-scratch version, there is no learning rate and no number of iterations to choose: for linear regression, scikit-learn computes the least-squares solution directly. What the model has learned is stored in attributes whose names end with an underscore: coef_ (the weight $w$, i.e. the slope) and intercept_ (the bias $b$).

# Create a Linear Regression model instance
model = LinearRegression()

# Train the model with the training data
model.fit(X_train, y_train)

print("Model trained successfully.")
print(f"Coefficient (weight): {model.coef_[0][0]:.4f}")
print(f"Intercept (bias): {model.intercept_[0]:.4f}")

Model trained successfully.
Coefficient (weight): 2.9820
Intercept (bias): 4.0371

From 40 noisy points, the model has recovered the planted line almost exactly: a slope of about 2.98 (the true value is 3) and an intercept of about 4.04 (the true value is 4). The small differences are caused by the noise.

Step 5: Model evaluation

---

Once the model has been trained, we need to evaluate how well it performs on data it has not seen during training: the test set. For regression, two common metrics are:

• Mean Squared Error (MSE): the average of the squared differences between the predictions and the true values. The lower, the better. Squaring makes large errors count more.
• Coefficient of Determination (R²): $R^2 = 1 - \frac{\sum (y - \hat{y})^2}{\sum (y - \bar{y})^2}$. It compares the errors of the model with the spread of the values around their mean. $R^2 = 1$ is a perfect fit and $R^2 = 0$ means the model is no better than always predicting the mean.

# Make predictions on the test set
y_pred = model.predict(X_test)

# Compute Mean Squared Error (MSE)
mse = mean_squared_error(y_test, y_pred)

# Compute Coefficient of Determination (R²)
r2 = r2_score(y_test, y_pred)

print(f"Mean Squared Error (MSE): {mse:.4f}")
print(f"Coefficient of Determination (R²): {r2:.4f}")

Mean Squared Error (MSE): 0.8376
Coefficient of Determination (R²): 0.9757

# Compare with the error on the training data: similar values mean the model generalizes
train_mse = mean_squared_error(y_train, model.predict(X_train))

print(f"Training MSE: {train_mse:.4f}")
print(f"Test MSE:     {mse:.4f}")

Training MSE: 0.8775
Test MSE:     0.8376

The test MSE (about 0.84) is close to the variance of the noise we added (1, because its standard deviation is 1): the errors of the model are mostly just that noise, which no straight line can remove. The R² of about 0.98 on unseen data, and a training error of the same size as the test error, tell us that the model has captured the relationship without overfitting.

Step 6: Visualization of the results

---

Visualization is a powerful tool to understand the model fit. We will plot the original data points (training and test) and the regression line learned by the model. This will allow us to visually see how well the line represents the relationship between X and y, and how it compares to the test data.

import matplotlib.pyplot as plt

# Visualize the training/test data and the regression line
plt.scatter(X_train, y_train, color='blue', label='Training Data')
plt.scatter(X_test, y_test, color='red', label='Test Data')
plt.plot(X, model.predict(X), color='green', label='Regression Line')
plt.xlabel('X')
plt.ylabel('y')
plt.title('Linear Regression with scikit-learn')
plt.legend()
plt.show()

Step 7: Use the model to make predictions

---

Finally, we use the trained model to make predictions on new values of the independent variable (X) that the model has not seen before. This simulates how the model would be used in a real-world scenario to predict future outcomes.

predict() expects the same shape as X: one row per example. Pay attention to the range: the training data only covered $x$ between 0 and 10.

• For $x = 2.5$, inside that range, the model is interpolating.
• For $x = 15$, outside it, the model is extrapolating: it simply extends the line. Here the answer is close to the true value ($4 + 3 \cdot 15 = 49$) only because we planted a perfectly straight line; with real data nothing guarantees that the pattern continues beyond the range we trained on.

# Use the model to predict on new data
X_new = np.array([[2.5], [15]])  # Two new X values, one row each
y_new_pred = model.predict(X_new)

for x, y_hat in zip(X_new, y_new_pred):
    print(f"Prediction for X = {x[0]} is {y_hat[0]:.2f}")
print("(15 is outside the 0-10 range we trained on: extrapolation)")

Prediction for X = 2.5 is 11.49
Prediction for X = 15.0 is 48.77
(15 is outside the 0-10 range we trained on: extrapolation)

Conclusion

In this notebook, we have implemented a linear regression model using the scikit-learn library. Unlike a manual implementation, scikit-learn greatly simplifies the process, providing efficient and optimized tools for:

• Data splitting: train_test_split makes it easy to separate the data into training and test sets, which is essential for good training and realistic evaluation.
• Model training: LinearRegression abstracts the mathematical and computational details of training, allowing us to fit the model with a simple call to fit() — no learning rate, no loop, no gradients to write.
• Model evaluation: Metrics such as mean_squared_error and r2_score give us a quantitative measure of the model's performance on unseen data.

The model recovered the line we planted (slope ≈ 3, intercept ≈ 4), its errors on the test set are of the size of the noise, and its training and test errors are similar. Visualizing the results confirms that the regression line runs through both the training and the test points.

In summary, scikit-learn is a powerful and easy-to-use tool for building and evaluating machine learning models, which speeds up development and allows us to focus on interpreting results and improving the model. A good next step is to try the same few calls on your own data, with more than one feature.