linearregression_pytorch.py

# -*- coding: utf-8 -*-
"""LinearRegression_pytorch.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1MmABwkx8IKTZOJAFWHjI6KmnqHYgwySP

**Keerthana Golla**(NetID-kg58)

# Homework 1: Problem 4: Linear regression in PyTorch
## Predicting home prices with the California housing dataset

This file contains code that helps you get started on linear regression as a neural network with one input layer with N units corresponding to the number of predictors, and an output layer with one unit. You will implement linear regression to predict the median value of a home in a census block in California from  eight predictor variables. This notebook has already been set up to load this data for you and convert it into a Pandas dataframe.
"""

# Commented out IPython magic to ensure Python compatibility.
from sklearn.datasets import fetch_california_housing
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import torch
import torch.nn as nn
from sklearn.utils import shuffle
from sklearn.model_selection import train_test_split, KFold
from sklearn.preprocessing import StandardScaler
from sklearn import metrics
import numpy as np

# %matplotlib inline
plt.rcParams['figure.figsize'] = (8.0,5.0) # set default size of plots
plt.rcParams['image.interpolation'] = 'nearest'
plt.rcParams['image.cmap'] = 'viridis'

# Some more magic so that the notebook will reload external python modules;
# see http://stackoverflow.com/questions/1907993/autoreload-of-modules-in-ipython

# %load_ext autoreload
# %autoreload 2

"""## Reading data and plotting

We download the California housing dataset, and load it into a pandas dataframe. We have data on median home price in census block groups in California. We want to predict the median home price on the basis of eight numeric predictors at the block group level.

- MedInc (median income in block group)

- HouseAge (median house age in block group)

- AveRooms (average number of rooms per household)

- AveBedrms (average number of bedrooms per household)

- Population (block group population)

- AveOccup (average number of household members)

- Latitude (block group latitude)

- Longitude (block group longitude)

Before starting on any task, it is often useful to understand the data by visualizing it. For this
dataset, we will show you how to use histograms and scatter plots to visualize the data.
"""

# read the California Housing data and place it in a pandas's dataframe

cdata = fetch_california_housing()
predictors = pd.DataFrame(data = cdata.data, columns = cdata.feature_names)
y = pd.DataFrame(cdata.target,columns = ['price'])
cal_df = pd.concat([predictors,y],axis=1)
cal_df.head()

"""# Data visualization
- plot the distribution of the target variable
    - note that all home prices over 5 are assigned a value of 5.
    - distribution is not Gaussian, right skew with outliers at the high end.
    
- plot the predictor variables
    - AveRooms, AveBedrooms, Population, AveOccupancy have extreme values; the range of values is high, quite a few outliers. Unlikely to be useful predictors, because almost all homes cluster at the very low end of values!
    - MedInc, HouseAge, Latitude, Longitude seem more promising as predictors, with nice distributions.
"""

# plot distribution of the target variable
sns.histplot(cal_df['price']);

# look at the distributions of the individual predictor variables
plt.style.use('seaborn')

cal_df.hist(figsize=(12, 10), bins=30, edgecolor="black")
plt.subplots_adjust(hspace=0.7, wspace=0.4)

"""# Spatial visualization
- prices have been discretized by the scatterplot function to simplify the legend.
- most expensive homes are along the coast and in Sacramento, as expected.
- Latitude and Longitude are very likely to be good predictors of home value. (location, location, location!)

"""

# visualizing the spatial distribution of homes
sns.scatterplot(
    data=cal_df,
    x="Longitude",
    y="Latitude",
    size="price",
    hue="price",
    palette="viridis",
    alpha=0.5,
)
plt.legend(title="Price", bbox_to_anchor=(1.05, 0.95), loc="upper left")
plt.title("Median home price as a function of spatial location");

"""# PyTorch model building (10 points)
- Create the model (Read [3.5.1](https://d2l.ai/chapter_linear-regression/linear-regression-concise.html#defining-the-model))
    - define the network structure: 1 layer network (4 points)
       - set the learning rate lr, and number of inputs num_inputs
       - use [nn.Linear](https://pytorch.org/docs/stable/generated/torch.nn.Linear.html#torch.nn.Linear) with num_inputs inputs and one output to set up self.net
       - initialize the weights of network to numbers drawn from a zero-mean normal distribution with std = noise (use self.net.weight.data.normal_(mean,std))
       - initialize bias weight to 0 (use self.net.bias.data.fill_(val))
    - define the forward function (Read 3.5.1) (2 points)
       - propagate the input X through the net
    - define the loss function (mean squared error) (Read 3.5.2) (2 points)
       - [nn.functional.mse_loss](https://pytorch.org/docs/stable/generated/torch.nn.functional.mse_loss.html)
    - configure the optimizer (Read 3.5.3) (2 points)
       - use [torch.optim.SGD()](https://pytorch.org/docs/stable/optim.html) with self.parameters() and set learning to be the learning rate of the network
"""

import torch.optim as optim
import torch.nn.functional as F
class RegressionModel(nn.Module):
    def __init__(self, num_inputs, lr, noise=0.01):
        super(RegressionModel, self).__init__()

        # Define the network structure
        self.net = nn.Linear(num_inputs, 1)  # 1 layer with num_inputs inputs and one output

        # Initialize the weights of the network with a zero-mean normal distribution
        self.net.weight.data.normal_(0, noise)

        # Initialize bias weight to 0
        self.net.bias.data.fill_(0)

        # Define the optimizer (Stochastic Gradient Descent)
        self.optimizer = optim.SGD(self.parameters(), lr=lr)

    def forward(self, X):
        return self.net(X)

    def loss(self, yhat, y):
        # Calculate the mean squared error loss
        return F.mse_loss(yhat, y)

    def configure_optimizers(self):
        return self.optimizer

"""## Test your model set up code
- run this code after completing the functions above.
- you should be able to view your model's structure and its weights and bias
"""

# Create an instance of the RegressionModel
model = RegressionModel(cal_df.shape[1] - 1, 1)

# Print the model's structure
print(model)

# Access model's weights and bias
print("Model Weights:")
print(model.net.weight)

print("\nModel Bias:")
print(model.net.bias)

"""# Splitting and standardizing training data (6 points)
- split training data (X,y) into a training set (XXtrain,yytrain) and (Xtest,ytest) using sklearn's [train_test_split](https://scikit-learn.org/stable/modules/generated/sklearn.model_selection.train_test_split.html) (with test_size = 0.2)
- If standardize is True, use sklearn.preprocessing StandardScaler() and fit to XXtrain. Then apply the scaler to transform both XXtrain and Xtest.
- Split (XXtrain,yytrain) further into (Xtrain,ytrain) and (Xval,yval) using using sklearn's train_test_split (with test_size = 0.2)
- Convert Xtrain,Xval,Xtest,ytrain,yval,ytest into Torch tensors, using [torch.Tensor()](https://pytorch.org/docs/stable/tensors.html)
- Return tensors Xtrain,Xval,Xtest,ytrain,yval,ytest
   - the shapes of the y tensors should be of the form torch.Size([N,1]) [**reshape** if needed]
"""

def make_train_val_test_split(X, y, standardize=True):
    # Split X, y into a training and test set in an 80/20 ratio.
    Xtrain, Xtest, ytrain, ytest = train_test_split(X, y, test_size=0.2, random_state=42)

    # Standardize training and test set if the flag is True
    if standardize:
        scaler = StandardScaler()
        Xtrain = scaler.fit_transform(Xtrain)
        Xtest = scaler.transform(Xtest)

    # Split the training set further into train and val sets in an 80/20 ratio
    Xtrain, Xval, ytrain, yval = train_test_split(Xtrain, ytrain, test_size=0.2, random_state=42)

    # Convert train data to tensors
    Xtrain_tensor = torch.Tensor(Xtrain)
    ytrain_tensor = torch.Tensor(ytrain).view(-1, 1)  # Reshape ytrain if needed

    # Convert val data to tensors
    Xval_tensor = torch.Tensor(Xval)
    yval_tensor = torch.Tensor(yval).view(-1, 1)  # Reshape yval if needed

    # Convert test data to tensors
    Xtest_tensor = torch.Tensor(Xtest)
    ytest_tensor = torch.Tensor(ytest).view(-1, 1)  # Reshape ytest if needed

    return Xtrain_tensor, Xval_tensor, Xtest_tensor, ytrain_tensor, yval_tensor, ytest_tensor

"""## Test the train/val/test split function
- run the cell below after implementing the function in the previous cell.
"""

X = cal_df.iloc[:,:-1].values # predictor matrix
y = cal_df['price'].values    # target vector

Xtrain_tensor,Xval_tensor,Xtest_tensor,ytrain_tensor,yval_tensor,ytest_tensor = \
        make_train_val_test_split(X,y,standardize=True)
print('Train data: ', Xtrain_tensor.shape, ytrain_tensor.shape)
print('Val data: ', Xval_tensor.shape, yval_tensor.shape)
print('Test data: ', Xtest_tensor.shape, ytest_tensor.shape)

"""# Train the model
The train function takes the predictor matrix X and the target vector y, the learning rate lr, and the number of epochs, num_epochs,  to train as inputs. Instead of using D2L's trainer class (explained in Section 3.2), we actually write out the training loop explicitly to be able to understand the steps in the training process. Here is the algorithm.

1. Construct tensors Xtrain_tensor,Xval_tensor,Xtest_tensor,ytrain_tensor,yval_tensor,ytest_tensor by calling make_train_val_test_split() with predictor vector X and target vector y (and standardize=True).
2. Set up the model by instantiating RegressionModel with the number of predictors in Xtrain_tensor, and learning rate lr
3. Configure the optimizer for the model by calling the configure_optimizers() method defined in the RegressionModel class.
4. Initialize torch tensors of length num_epochs to store train set loss and val set loss (for plotting)
5. For epoch in range(num_epochs)
   - zero the parameter gradients in the optimizer (use zero_grad())
   - perform a forward pass, and calculate loss on train set (model.loss()). Store train set loss in train set loss tensor.
   - do a backward step to calculate gradients (loss.backward())
   - get the optimizer to take a step in direction of the gradient (optimizer.step())
   - calculate loss on val set and store in val set loss tensor
6. Plot train set loss and val set loss on y-axis and epoch number on x-axis. These are the training curves.
7. return the regression model object, and the Xtest tensor and ytest tensor. The last two are for model evaluation.
"""

def create_train_test_model(X, y, model, num_epochs):
    # Split into train, val, test sets
    Xtrain_tensor, Xval_tensor, Xtest_tensor, ytrain_tensor, yval_tensor, ytest_tensor = \
        make_train_val_test_split(X, y, standardize=True)

    # Configure optimizer for the model
    optimizer = model.configure_optimizers()

    # Initialize train loss and val loss tensors for plotting
    train_loss_history = []
    val_loss_history = []

    for epoch in range(num_epochs):
        # Zero the parameter gradients
        optimizer.zero_grad()

        # Forward pass
        ytrain_pred = model(Xtrain_tensor)

        # Calculate and store train loss
        train_loss = model.loss(ytrain_pred, ytrain_tensor)
        train_loss.backward()

        # Optimize step
        optimizer.step()

        # Calculate loss on val set
        with torch.no_grad():
            yval_pred = model(Xval_tensor)
            val_loss = model.loss(yval_pred, yval_tensor)

        # Store train and val loss
        train_loss_history.append(train_loss.item())
        val_loss_history.append(val_loss.item())

    # Plot train set loss and val set loss as a function of epoch
    plt.plot(range(1, num_epochs + 1), train_loss_history, label='Train Loss')
    plt.plot(range(1, num_epochs + 1), val_loss_history, label='Validation Loss')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.legend()
    plt.show()

    # Print values of train set loss and val set loss on the last epoch
    print(f'Final Train Loss: {train_loss_history[-1]}')
    print(f'Final Validation Loss: {val_loss_history[-1]}')

    return model, Xtest_tensor, ytest_tensor

"""# Analyze model performance (10 points)
- get histogram of errors (actual price - predicted price) on set aside test data (Xtest_tensor,ytest_tensor) (3 points)
- get scatterplot of actual price vs predicted price on set aside test data (3 points)
- compute R2 of concordance between actual and predicted values using sklearn.metrics [r2_score()](https://scikit-learn.org/stable/modules/generated/sklearn.metrics.r2_score.html) function. (4 points)
"""

from sklearn.metrics import r2_score
# Calculate errors (actual price - predicted price) on test data
num_epochs = 100  # Set the number of training epochs
model = RegressionModel(Xtrain_tensor.shape[1], lr=0.01)  # Adjust the model and learning rate
trained_model, Xtest_tensor, ytest_tensor = create_train_test_model(X, y, model, num_epochs)

errors = (ytest_tensor - model(Xtest_tensor)).detach().numpy()

# Plot histogram of errors
plt.hist(errors, bins=50, alpha=0.5, color='blue')
plt.xlabel('Error (Actual Price - Predicted Price)')
plt.ylabel('Frequency')
plt.title('Histogram of Errors')
plt.show()

# Create scatter plot of predicted vs actual price
plt.scatter(ytest_tensor, model(Xtest_tensor).detach().numpy(), alpha=0.5)
plt.xlabel('Actual Price')
plt.ylabel('Predicted Price')
plt.title('Scatterplot of Actual Price vs Predicted Price')
plt.show()

# Assess model quality using R-squared (R2)
r2 = r2_score(ytest_tensor, model(Xtest_tensor).detach().numpy())
print(f'R-squared (R2) Value: {r2}')

"""# Visualize model parameters
- run code in the following two cells
"""

# Model parameters
print('Model weights:', model.net.weight)
print('Model bias:', model.net.bias)

# plot model coefficients
coefs = pd.DataFrame(model.net.weight.data.detach().numpy(),columns=cal_df.columns[:-1])
sns.barplot(y=coefs.columns, x=coefs.iloc[0].values,orient='h')

"""# Cross-validation to assess model quality (10 points + 10 points)

Up to now, we have assessed model quality with a single train/test split. Now we use a more systematic approach based on 5-fold crossvalidation.  The $R^2$ value is used as a measure of model quality. The closer $R^2$ is to 1., the better the model is.
"""

from sklearn.preprocessing import StandardScaler

def make_train_val_test_cv_split(X, y, train_index, test_index, standardize=True):
    # Construct the training and test subsets of X, y
    Xtrain, Xtest = X[train_index], X[test_index]
    ytrain, ytest = y[train_index], y[test_index]

    # If the standardize flag is set, fit a Standard Scaler on training data and transform the training and test X
    if standardize:
        scaler = StandardScaler()
        Xtrain = scaler.fit_transform(Xtrain)
        Xtest = scaler.transform(Xtest)

    # Split the training X, y set into train and val sets
    Xtrain, Xval, ytrain, yval = train_test_split(Xtrain, ytrain, test_size=0.2, random_state=42)

    # Convert train data to tensors
    Xtrain_tensor = torch.Tensor(Xtrain)
    ytrain_tensor = torch.Tensor(ytrain).view(-1, 1)  # Reshape ytrain if needed

    # Convert val data to tensors
    Xval_tensor = torch.Tensor(Xval)
    yval_tensor = torch.Tensor(yval).view(-1, 1)  # Reshape yval if needed

    # Convert test data to tensors
    Xtest_tensor = torch.Tensor(Xtest)
    ytest_tensor = torch.Tensor(ytest).view(-1, 1)  # Reshape ytest if needed

    return Xtrain_tensor, Xval_tensor, Xtest_tensor, ytrain_tensor, yval_tensor, ytest_tensor

"""# Training in crossvalidation mode

X,y are predictor matrix and target vector. K is the number of crossvalidation folds. Models is a list of length K of regression model objects. num_epochs is the number of epochs to train each model

1. Create a KFold() object called kfolds with n_splits = K
2. For (train_index,test_index) in kfolds.split(X)

   - Construct tensors Xtrain_tensor,Xval_tensor,Xtest_tensor,ytrain_tensor,yval_tensor,ytest_tensor by calling make_train_val_test_split() with predictor vector X and target vector y (and standardize=True).
   - Extract model from models for that fold
   - Configure the optimizer for the model for that fold by calling the configure_optimizers() method defined in the RegressionModel class.
   - Initialize torch tensors of length num_epochs to store train set loss and val set loss (for plotting)
   - For epoch in range(num_epochs)
       - zero the parameter gradients in the optimizer (use zero_grad())
       - perform a forward pass, and calculate loss on train set (model.loss()). Store train set loss in train set loss tensor.
       - do a backward step to calculate gradients (loss.backward())
       - get the optimizer to take a step in direction of the gradient (optimizer.step())
       - calculate loss on val set and store in val set loss tensor
   - Plot train set loss and val set loss on y-axis and epoch number on x-axis. These are the training curves for that fold.
3. return the regression models list, and the Xtest tensor and ytest tensor. The last two are for model evaluation.
"""

from sklearn.model_selection import KFold

def create_train_test_cv_model(X, y, K, models, num_epochs):
    kfolds = KFold(n_splits=K)
    all_model_predictions = []
    r2_scores = []

    for train_index, test_index in kfolds.split(X):
        Xtrain_tensor, Xval_tensor, Xtest_tensor, ytrain_tensor, yval_tensor, ytest_tensor = make_train_val_test_cv_split(X, y, train_index, test_index, standardize=True)
        #ytrain_tensor = ytrain_tensor.reshape(-1,1)
        #yval_tensor = yval_tensor.reshape(-1,1)
        #ytest_tensor = ytest_tensor.reshape(-1,1)
        model = models.pop(0)  # Pop the first model from the list

        optimizer = model.configure_optimizers()

        train_losses = torch.zeros(num_epochs)
        val_losses = torch.zeros(num_epochs)

        for epoch in range(num_epochs):
            optimizer.zero_grad()

            yhat_train = model(Xtrain_tensor)
            yhat_val = model(Xval_tensor)

            train_loss = model.loss(yhat_train, ytrain_tensor)
            val_loss = model.loss(yhat_val, yval_tensor)

            train_losses[epoch] = train_loss
            val_losses[epoch] = val_loss

            train_loss.backward()
            optimizer.step()

        with torch.no_grad():
            yhat_test = model(Xtest_tensor).detach()

        r2 = r2_score(ytest_tensor.detach().numpy(), yhat_test.numpy())
        r2_scores.append(r2)

        all_model_predictions.append(yhat_test)

        # Plot train set loss and val set loss as a function of epoch
        plt.figure(figsize=(10, 5))
        # plt.plot(range(num_epochs), train_losses.detach.numpy(), label='Train Loss')
        plt.plot(range(num_epochs), train_losses.detach().numpy(), label='Train Loss')
        plt.plot(range(num_epochs), val_losses.detach().numpy(), label='Validation Loss')
        plt.xlabel('Epoch')
        plt.ylabel('Loss')
        plt.legend()
        plt.title(f'Training and Validation Curves - Fold {len(r2_scores)}')

        print(f'Final Train Loss (Fold {len(r2_scores)}): {train_losses[-1]:.4f}')
        print(f'Final Validation Loss (Fold {len(r2_scores)}): {val_losses[-1]:.4f}')
        print(f'R-squared (R2) Score (Fold {len(r2_scores)}): {r2:.4f}')

    return models,Xtest_tensor,ytest_tensor, all_model_predictions, r2_scores

"""## Test your crossvalidation trainer
- run the cell below after completing the two functions above
"""

X = cal_df.iloc[:,:-1].values # predictor matrix
y = cal_df['price'].values    # target vector
K = 5
lr = 0.1
num_epochs = 100

models = [RegressionModel(Xtrain_tensor.shape[1],lr=lr) for i in range(K)]

models,Xtest_tensor,ytest_tensor, all_model_predictions, r2_scores = create_train_test_cv_model(X,y,K,models,num_epochs)

"""# Assessing crossvalidation performance"""

#code here
# assess model quality using r2 on test data on all K models
# report mean and std on r2 (4 lines)
print(np.asarray(r2_scores).mean())
print(np.asarray(r2_scores).std())

"""# Visualizing coefficients of CV models
- run the cell below to visualize your model coefficients
"""

cal_df.columns[:-1]
models = [RegressionModel(Xtrain_tensor.shape[1],lr=lr) for i in range(K)]
coefs = pd.DataFrame()
for i in range(K):
    coef = pd.DataFrame(models[i].net.weight.data.detach().numpy(),columns=cal_df.columns[:-1])
    coefs = pd.concat([coefs,coef],axis=0)
coefs = coefs.reset_index(drop=True)

color = {"whiskers": "black", "medians": "black", "caps": "black"}
coefs.plot.box(vert=False, color=color)
plt.axvline(x=0, ymin=-1, ymax=1, color="black", linestyle="--")
_ = plt.title("Coefficients of linear regression model via 5-fold cross-validation")

"""# Adding weight decay and polynomial basis function expansion
- define a new class RegressionModelL2 by copying the RegressionModel class and adding a new argument to the init function: wd (for weight decay or lambda)
- fill in the methods below
- you need to modify the torch.optim.SGD call to include weight_decay (see [3.7](https://d2l.ai/chapter_linear-regression/weight-decay.html#concise-implementation))
- 5 points for the class definition below
"""

import torch.nn.functional as F

class RegressionModelL2(nn.Module):
    def __init__(self, num_inputs, lr, noise=0.01, wd=0.1):
        super(RegressionModelL2, self).__init__()

        # Define the network structure with one linear layer
        self.net = nn.Linear(num_inputs, 1)

        # Initialize the weights with a zero-mean normal distribution and std = noise
        self.net.weight.data.normal_(mean=0, std=noise)

        # Initialize the bias weight to 0
        self.net.bias.data.fill_(0)

        # Store the weight decay parameter
        self.wd = wd

    def forward(self, X):
        # Forward pass through the network
        return self.net(X)

    def loss(self, yhat, y):
        # Calculate mean squared error loss with L2 regularization
        mse_loss = F.mse_loss(yhat, y)

        # Calculate L2 regularization term
        l2_reg = 0.5 * self.wd * sum(p.pow(2).sum() for p in self.parameters())

        # Total loss includes both MSE loss and L2 regularization term
        total_loss = mse_loss + l2_reg
        return total_loss

    def configure_optimizers(self):
        # Configure the optimizer with stochastic gradient descent (SGD)
        optimizer = torch.optim.SGD(self.parameters(), lr=lr, weight_decay=self.wd)
        return optimizer

"""# Set up L2 regression model and train using a simple train/val/test split
- 5 points for running the code below and assessing model quality as a function of the value of wd
- try different values of wd and see its impact on r2 as well as on the coefficient magnitudes and signs
"""

X = cal_df.iloc[:,:-1].values # predictor matrix
y = cal_df['price'].values    # target vector
lr = 0.1
num_epochs = 100
wd = 0.01

l2_model = RegressionModelL2(X.shape[1],lr=lr,wd=wd)
l2_model,Xtest_tensor,ytest_tensor = create_train_test_model(X,y,l2_model,num_epochs)

# assess model quality using r2 (1 line)
r2_score(ytest_tensor.detach().numpy(), l2_model(Xtest_tensor).detach().numpy())

# plot model coefficients
coefs = pd.DataFrame(l2_model.net.weight.data.detach().numpy(),columns=cal_df.columns[:-1])
sns.barplot(y=coefs.columns, x=coefs.iloc[0].values,orient='h')

"""# Polynomial basis function expansion
- use sklearn [PolynomialFeatures](https://scikit-learn.org/stable/modules/generated/sklearn.preprocessing.PolynomialFeatures.html)
- 5 points for the code in this cell
"""

from sklearn.preprocessing import PolynomialFeatures

X = cal_df.iloc[:,:-1]        # predictor matrix
y = cal_df['price'].values    # target vector

# make a PolynomialFeatures constructor (make include_bias = False) for polynomial of degree 2
# fit and transform X using that constructor, call that Xpoly
# (2 lines)

# get the feature names associated with Xpoly using the method [get_feature_names_out](https://scikit-learn.org/stable/modules/generated/sklearn.preprocessing.PolynomialFeatures.html) of the constructor
# print the shape Xpoly
from sklearn.preprocessing import PolynomialFeatures

# Create PolynomialFeatures constructor with degree 2
poly = PolynomialFeatures(degree=2, include_bias=False)

# Fit and transform X using the constructor
Xpoly = poly.fit_transform(X)

# Getting the feature names
Xpoly_cols = poly.get_feature_names_out(input_features=X.columns)

# Get the feature names associated with Xpoly
feature_names = poly.get_feature_names_out(input_features=X.columns)

# Print the shape of Xpoly
print("Shape of Xpoly:", Xpoly.shape)

# understand the correlation structure of Xpoly
Xdf = pd.DataFrame(Xpoly,columns=Xpoly_cols)
corr_matrix = Xdf.corr().abs()
plt.imshow(corr_matrix)
plt.rcParams['axes.grid'] = False
plt.colorbar()

# drop the features that have correlations greater than 0.95 with other predictive features
# get upper triangle of correlation matrix
upper = corr_matrix.where(np.triu(np.ones(corr_matrix.shape), k=1).astype(np.bool_))

# find features with correlation greater than 0.95
to_drop = [column for column in upper.columns if any(upper[column] > 0.95)]

# drop highly correlated features
Xdf.drop(to_drop, axis=1, inplace=True)
print(Xdf.shape)

# now Xdf has just the selected features for model building

"""# Train the L2 model with the selected features (15 points)
- run the cell below
- try different weights for wd and assess the effect on model coefficients and model performance
- what is the best setting for wd?
- also test variation with changes in lr and num_epochs
- does the model with expanded basis functions do better than the one without? Why or why not?
"""

lr = 0.1
num_epochs = 50
wd = 0.1

l2_model = RegressionModelL2(Xdf.shape[1],lr=lr,wd=wd)
l2_model,Xtest_tensor,ytest_tensor = create_train_test_model(Xdf,y,l2_model,num_epochs)

# compute r2 score for l2_model after training (1 line)
r2_score_after_training = r2_score(ytest_tensor, l2_model(Xtest_tensor).detach().numpy())
print(r2_score_after_training)

# plot model coefficients
coefs = pd.DataFrame(l2_model.net.weight.data.detach().numpy(),columns=Xdf.columns)
sns.barplot(y=coefs.columns, x=coefs.iloc[0].values,orient='h')