# ---
# 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
# ---

# # Exercise 4.2 - Solution
# ## Linear regression
# In this task we will design and train a linear model using [Keras](https://keras.io/).
#
# ### Tasks
# 1. Complete the implemetation of the `LinearLayer`
# 2. Define a meaningful objective
# 3. Implement gradient descent and train the linear model for 80 epochs.

# +
import tensorflow as tf
from tensorflow import keras
import numpy as np
import matplotlib.pyplot as plt

layers = keras.layers
# -

# ### Simulation of data
# Let's first simulate some noisy data

np.random.seed(1904)
x = np.float32(np.linspace(-1, 1, 100)[:,np.newaxis])
y = np.float32(2 * x[:,0] + 0.3 * np.random.randn(100))
print("x.shape:", x.shape)
print("y.shape:", y.shape)


# ### Implement linear model

# Now, we have to design a <b>linear layer</b> that maps from the input $x$ to the output $y$ using a single adaptive weight $w$:
#  
# $$y = w \cdot x$$
#
# ### Task 1
# Complete the implementation of the `LinearLayer` by adding the linear transformation in the `call` function.

class LinearLayer(layers.Layer):

    def __init__(self, units=1, input_dim=1):  # when intializing the layer the weights have to be initialized
        super(LinearLayer, self).__init__()
        w_init = tf.random_normal_initializer()
        self.w = tf.Variable(initial_value=w_init(shape=(input_dim, units), dtype="float32"),
                             trainable=True)

    def call(self, inputs):  # when calling the layer the linear transformation has to be performed
        return tf.matmul(inputs, self.w)


# Build a model using the implemented layer.

model = keras.models.Sequential()
model.add(LinearLayer(units=1, input_dim=1))

model.build((None, 1))
print(model.summary())

# ### Performance before the training
# Plot data and model before the training

# +
y_pred = model(x)

fig, ax = plt.subplots(1)
ax.plot(x, y, 'bo', label='data')
ax.plot(x, y_pred, 'r-', label='model')
ax.set(xlabel='$x$', ylabel='$y$')
ax.grid()
ax.legend(loc='lower right')
plt.tight_layout()


# -

# ### Task 2: Define the objective function
# Define a meaningful objective here (regression task).  
# Note that you can use `tf.reduce_mean()` to average your loss estimate over the full data set (100 points).

def loss(x, y):
    return tf.reduce_mean((tf.squeeze(x)-tf.squeeze(y))**2)


# ### Task 3 - Train the model using gradient descent
# 'Train' the linear model for 80 epochs (or iterations) with a meaningful learning rate and implement gradient descent.  
# Hint: you can access the adaptive parameters using `model.trainable_weights` and perform $w' \rightarrow w-z$ using `w.assign_sub(z)`

# +
epochs = 80  # number of epochs
lr = 0.1  # learning rate

for epoch in range(epochs):

    with tf.GradientTape() as tape:
        output = model(x, training=True)
        # Compute loss value
        loss_value = loss(tf.convert_to_tensor(y), output)
        grads = tape.gradient(loss_value, model.trainable_weights)
        
        for weight, grad in zip(model.trainable_weights, grads):
            weight.assign_sub(lr * grad)

        print("Current loss at epoch %d: %.4f" % (epoch, float(loss_value)))
# -

# ### Performance of the fitted model
# Plot data and model after the training

# +
fig, ax = plt.subplots(1)

y_pred = model(x)

ax.plot(x, y, 'bo', label='data')
ax.plot(x, y_pred, 'r-', label='model')
ax.set(xlabel='$x$', ylabel='$y$')
ax.grid()
ax.legend(loc='lower right')
plt.tight_layout()

StarcoderdataJupyterScriptsDedupFiltered