# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---
# # Linear Regression - Sklearn, Statsmodel, TensorFlow, Keras, PyTorch
# Author: <NAME>
#
# Date: September 22, 2019
#
# Updated: September 24, 2019
# The goal of this notebook is to compute a Linear Regression using different frameworks. The frameworks that will be investigated are:
# - Sklearn
# - Statsmodel
# - TensorFlow 1.0
# - Keras
# - PyTorch
# - TensorFlow 2.0
# The Ordinary Least Squares equation is the following:
#
# \begin{equation}
# Y_i = \beta_0 + \beta_1 X_i + \epsilon_i
# \end{equation}
#
# The data will be created according to this equation.
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

def plot_output(X, y, predicted=None):
    '''
    Plots Regression Line Charts
    
    Arguments:
    X - Accepts a DataFrame or numpy array for the X data
    y - Accepts a DataFrame or numpy array for the y data
    predicted - Accepts a DataFrame or numpy array for the predicted y data. If None, this function only plots the scatter plot.
    '''
    
    plt.clf()
    plt.plot(X, y, 'go', label='True Data', alpha=0.5)
    if predicted is not None:
        plt.plot(X, predicted, '--', label='Predictions', alpha=0.5)
        plt.legend(loc='best')
    plt.show();

# +
np.random.seed(1) #Setting a seed to recreate experiment.
num_points = 250 #The number of data points
X = 2.5 * np.random.randn(num_points) + 1.5   # Array of 100 values with mean = 1.5, stddev = 2.5
res = 0.5 * np.random.randn(num_points)       # Generate 100 residual terms
y = 2 + 0.3 * X + res                  # Actual values of Y
# -
plot_output(X, y)
# Reshaping X and Y necessary when the data has one column so that sklearn and pytorch will work properly. In addition, the numpy array is forced to have dtype np.float32 so that pytorch will accept these values.
X = np.array(X, dtype=np.float32).reshape(-1,1)
y = np.array(y, dtype=np.float32).reshape(-1,1)
# Scaling the data so that it helps with convergence when running gradient descent in TensorFlow, Keras and Pytorch.
# +
from sklearn.preprocessing import StandardScaler
X_scaler = StandardScaler()
X_scaled = X_scaler.fit_transform(X)
y_scaler = StandardScaler()
y_scaled = y_scaler.fit_transform(y)

# -
def plot_output(X, y, predicted=None):
    '''
    Plots Regression Line Charts
    
    Arguments:
    X - Accepts a DataFrame or numpy array for the X data
    y - Accepts a DataFrame or numpy array for the y data
    predicted - Accepts a DataFrame or numpy array for the predicted y data
    '''
    
    plt.clf()
    plt.plot(X, y, 'go', label='True Data', alpha=0.5)
    if predicted is not None:
        plt.plot(X, predicted, '--', label='Predictions', alpha=0.5)
        plt.legend(loc='best')
    plt.show();

# ### Sklearn
# +
from sklearn.linear_model import LinearRegression
lin_model = LinearRegression() #Set the model
lin_model.fit(X,y) #Fit the model
predicted = lin_model.predict(X)
plot_output(X, y, predicted)
# -
# ### Statsmodel
# +
import statsmodels.api as sm # For constants
from statsmodels.regression import linear_model
#Compared to SKLearn, statsmodel doesn't include the bias term.
X_c = sm.add_constant(X) 
linear_model_sm = linear_model.OLS(y, X_c).fit()  #Fit the model
predicted = linear_model_sm.predict(X_c)
plot_output(X, y, predicted)
# -
# ### Standardized Settings for TensorFlow, Keras and PyTorch
input_dim = 1        # takes variable 'x' 1 input variable.
output_dim = 1       # takes variable 'y' 1 output variable.
learning_rate = 0.01 
epochs = 1000
# ### TensorFlow 1.0
# Steps to calculate a linear regression using TensorFlow. The goal is to create the placeholders and variables to satisfy the following linear equation:
# $$ y = X^TW + b$$
import tensorflow as tf
print('Version:', tf.__version__)
# Create placeholders for X and y, because these values will be supplied in the opimization step using feed_dict.
X_tf = tf.placeholder(dtype=tf.float32, shape=[None, input_dim])
y_tf = tf.placeholder(dtype=tf.float32, shape=[None, output_dim])
# Create variables for W and b. Initializing W with Xavier/Glorot Initialization and b with zeros. W and B are updated during gradient descent so that its why they are variables instead of placeholders
W_tf = tf.Variable(tf.glorot_uniform_initializer()((input_dim, output_dim)), dtype=tf.float32, name='W')
b_tf = tf.Variable(tf.zeros_initializer()((output_dim)), dtype=tf.float32, name='b')
# Define the y predicted equation
#WX + b
y_pred_tf = tf.add(tf.matmul(X_tf, W_tf), b_tf)
# Define the Cost Function. 
# In the case of Linear Regression, the Mean Squared Error is being minimized:
# $$MSE = \frac{1}{n} \sum_{i=1}^n (\hat{y_i} - y_i)^2$$
#The cost function
cost_tf = tf.losses.mean_squared_error(y_tf, y_pred_tf)
# Set the optimizer.
#The optimizer
optimizer_tf = tf.train.GradientDescentOptimizer(learning_rate).minimize(cost_tf)
# Set the Global Variables Initializer. This step is mandatory.
#Global Variables Initilializer -- Mandatory
init = tf.global_variables_initializer()
# Running the computation graph.
#The Tensorflow Session
with tf.Session() as sess:
    sess.run(init)
    for epoch in range(epochs):
        
        #The dictionary to feed in the variables for hte optimizer
        feed_dict = {X_tf: X_scaled, y_tf: y_scaled}
        
        #Run the optimizer at every epoch
        #Also saves the weight and bias.
        _, weight, bias = sess.run([optimizer_tf, W_tf, b_tf], feed_dict)
            
        if (epoch + 1) % 50 == 0:
            #Calculate the cost at every epoch
            print("Epoch", (epoch + 1), ": cost =", cost_tf.eval(feed_dict))
# 8. Calculate the y predicted values. Note: Remember to unscale your data.
predicted = weight * X_scaled + bias
predicted = y_scaler.inverse_transform(predicted) 
plot_output(X, y, predicted)
# ### Keras
# Keras is a high-level API that uses TensorFlow as its backend. Creating a model is simpler with the Keras API. Essentially this is similar to putting together lego bricks.
import tensorflow as tf
from tensorflow.keras import Sequential
from tensorflow.keras.layers import Dense
print('Version:', tf.keras.__version__)
# Creating a Sequential model in Keras. A Sequential model is when layers are linearly stacked.
# - For a Linear Regression, the activation layer will be 'linear' and the output dimension being 1 due to a value being a regression.
# - The input shape needs to be defined.
# - This layer is essentially y = Wx + b, with an activation of 'linear'. 'linear' means that no additional calculations will be done and will return the identity y = Wx + b.
model = Sequential([
    Dense(output_dim, activation='linear',  input_shape=(input_dim,)),       
])
# Setting the loss function, optimizer and then compiling the model. This step configures the model for training
# +
optimizer = tf.keras.optimizers.SGD(learning_rate)
model.compile(loss='mse', optimizer=optimizer, metrics=['mae', 'mse'])
# -
model.summary()
# Training the model using fit()
history = model.fit(x = X_scaled, y = y_scaled, epochs = epochs, verbose=0)
# Calculate the y predicted values using the predict() function on the model. Note: Remember to unscale your data.
predicted = model.predict(X_scaled)
predicted = y_scaler.inverse_transform(predicted)
plot_output(X, y, predicted)
# ### PyTorch
# PyTorch is a framework created by Facebook. Implementing the model is slightly different. The most apparent differences is the way forward propagation is coded and GPU implementation is not seamless.
import torch
print('Version:', torch.__version__)

# Create the model.
# - If manually creating the tensors, must wrap with <a href="https://pytorch.org/docs/stable/nn.html">torch.nn.Parameter</a>
# - <a href="https://pytorch.org/docs/stable/nn.html#torch.nn.Linear">torch.nn.Linear</a> is a layer that applies a linear transformation to y = Wx + b
# - The forward propagation passes x to the linear transformation
#
# Generally, the layer method is used in tutorial.
# +
class LinearRegressionModel_PyT(torch.nn.Module):
    def __init__(self, input_dim, output_dim):
        super(LinearRegressionModel_PyT, self).__init__()
        self.W = torch.nn.Parameter(torch.nn.init.xavier_uniform_(torch.empty(input_dim, output_dim)))
        self.b = torch.nn.Parameter(torch.zeros(output_dim))
    def forward(self, x):
        out = torch.matmul(x, self.W) + self.b
        return out
    
class LinearRegressionModel_PyT_Layers(torch.nn.Module):
    def __init__(self, input_dim, output_dim):
        super(LinearRegressionModel_PyT_Layers, self).__init__()
        self.linear = torch.nn.Linear(input_dim, output_dim)
    
    def forward(self, x):
        out = self.linear(x)
        return out

# -
# Initiate the model.
# +
model = LinearRegressionModel_PyT_Layers(input_dim, output_dim)
####### For GPU #######
if torch.cuda.is_available():
    model.cuda()
# -
# Initiate the cost function and the optimizer.
criterion = torch.nn.MSELoss()
optimiser = torch.optim.SGD(model.parameters(), lr = learning_rate)
# Converting the numpy data to torch tensors.
# Converting X inputs and y labels to torch tensors
if torch.cuda.is_available():
    ####### For GPU #######
    X_data = torch.from_numpy(X_scaled).cuda()
    y_data = torch.from_numpy(y_scaled).cuda()
else:
    X_data = torch.from_numpy(X_scaled)
    y_data = torch.from_numpy(y_scaled)
# Training the model.
#  - Clear gradient buffer
#  - Forward propagation
#  - Get loss 
#  - Get gradients from backward propagation
#  - Update parameters for gradient descent
#  - Rinse and repeat
for epoch in range(epochs):
    # Clear gradient buffers in the optimiser so that it doesn't carry forward
    optimiser.zero_grad()
    #The forward propagation step.
    y_pred = model(X_data)
    
    # Get loss for the predicted output
    loss = criterion(y_pred, y_data)
    
    # Get gradients w.r.t to parameters. This is the backward propagation step.
    loss.backward()
    # Update parameters for gradient descent
    optimiser.step()
    if (epoch + 1) % 50 == 0:
        print('epoch {}, loss {}'.format(epoch + 1, loss.item()))    
# Calculate the y predicted by passing in the data into the model.Note: Remember to unscale your data.
# +
with torch.no_grad():
    if torch.cuda.is_available():
        predicted = model(torch.from_numpy(X_scaled).cuda()).cpu().data.numpy()
    else:
        predicted = model(torch.from_numpy(X_scaled)).data.numpy()
predicted = y_scaler.inverse_transform(predicted)
# -
plot_output(X, y, predicted)
# ## Note: The environment has been switched at this point of the notebook
# ### Tensorflow 2.0
# The API has changed drastically in TensorFlow 2.0. The structure is very similar to PyTorch.
import tensorflow as tf
print('Version:', tf.__version__)
print('Version:', tf.keras.__version__)

# Creating the model.
#
# Differences from TensorFlow 1.0:
# - `tf.placeholders()`, `tf.Session()` has been removed
# - `tf.glorot_uniform_initializer()` is gone, and has become `tf.keras.initializers.GlorotUniform()`
#
# Alternatively, the model class can utilize the layers in Keras. This is analogous to PyTorch.
# +
class LinearRegressionModel_TF2(tf.keras.Model):
    def __init__(self, input_dim, output_dim):
        super(LinearRegressionModel_TF2, self).__init__()
        self.W = tf.Variable(tf.keras.initializers.GlorotUniform()((input_dim, output_dim)), dtype=tf.float32, name='W')
        self.b = tf.Variable(tf.zeros_initializer()((output_dim)), dtype=tf.float32, name='b')
    
    #In PyTorch the function will be 'forward' instead.
    def call(self, x):
        out = tf.matmul(x, self.W) + self.b
        return out
    
class LinearRegressionModel_TF2_Layers(tf.keras.Model):
    def __init__(self, input_dim, output_dim):
        super(LinearRegressionModel_TF2_Keras, self).__init__()
        self.d1 = tf.keras.layers.Dense(output_dim, activation='linear', input_dim=(input_dim,))
        
    def call(self, x):
        x = self.d1(x)
        return x

# -
# Initiate the model.
model = LinearRegressionModel_TF2(input_dim, output_dim)

# Initiate the loss function, gradient function, and optimzer.
# +
@tf.function
def loss(model, inputs, targets):
    loss_object = tf.keras.losses.MeanSquaredError()
    
    #Forward propagation
    predictions = model(inputs)
    
    #Get loss for the predicted output
    output = loss_object(targets, predictions)
    return output
@tf.function
def grad(model, inputs, targets):
    #For automatic differentiation
    with tf.GradientTape() as tape:
        loss_value = loss(model, inputs, targets)
        
    #This calculates the gradient.
    return tape.gradient(loss_value, model.trainable_variables)

optimizer = tf.keras.optimizers.SGD(learning_rate)
# -
for epoch in range(epochs):
    
    #This is the backward propagation step to get the parameters
    grads = grad(model, X_scaled, y_scaled)
    
    #Update parameters for gradient descent
    optimizer.apply_gradients(zip(grads, model.trainable_variables))
    
    if (epoch+1) %  50 == 0:
        template = 'Epoch {}, Loss: {}'
        print(template.format(epoch+1, loss(model, X_scaled, y_scaled)))
        
# +
predicted = y_scaler.inverse_transform(model(X_scaled))
plot_output(X, y, predicted)
# -
# Alternatively, the model can be ran through a .compile(), .fit(), and .predict() from the Keras API.
model = LinearRegressionModel_TF2(input_dim, output_dim)
model.compile(loss=tf.keras.losses.MeanSquaredError(), optimizer=optimizer)
model.fit(x=X_scaled, y=y_scaled, epochs=epochs, verbose=0)
plot_output(X, y, y_scaler.inverse_transform(model.predict(X_scaled)))
# Conclusion: The code for TensorFlow 2.0 is a lot cleaner than TensorFlow 1.0