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

# # Supervised learning (without calculus!)
# This tutorial wants to show the (possibly) most naive way to train the most naive neural network on the [MNIST database](https://en.wikipedia.org/wiki/MNIST_database). For our model we only need to known matrix multiplication, for loops. Importantly, no knowledge of calculus  (gradient descent) is required and coding syntax is explained in detail. The reader expect to:
#
# - present key ideas in machine learning in simple setting,
# - be impressed by the power of simple matrix multiplication,
# - build intuition for issues such as overfitting and
# - have plenty of extensions of the model to experiment.
#
# (This tutorial was originally made as a concrete application of the linear algebra course at Otago University.)
#
# ## Dataset and packages
# We first import some basic packages along with the training/testing and test dataset and we reduce the size of the training dataset. (Note that tensorflow is only used to import the data.)

# +
# We import the packages that we will use
import matplotlib.pyplot as plt
import tensorflow as tf
import numpy as np
from numpy.linalg import norm as Euclid_norm
from time import time

#We import the dataset
mnist = tf.keras.datasets.mnist
(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0
size_train_small = 200
x_train, y_train = x_train[:size_train_small], y_train[:size_train_small]
# -

# ## Data visualisation
#
# Our aim is to "train" a "machine"  capable of recognising the ten digits $\{0,1,2,3,4,5,6,7,8,9\}$ when hand-written. It is natural to show the machine a bunch of these digits and then to train it to choose correctly by rewarding the machine  when correct, and penalising  it otherwise (much like humans do). Once trained we send the machine into the unkown territory of the testing data set, and see how well it performs. Lukily, images are matrices (tabular arrays) where each position represent a pixel and the value of this entry is the color (and machines love matrices). Let's see this.

fig, ax = plt.subplots(1,5,figsize=(10,4))
fig.suptitle('.jpg files as 28x28 matrices with the respective labels', fontsize=15)
for j in range(5):
    ax[j].imshow(x_train[j:j+1][0])
    ax[j].set_title(f'Label: {y_train[j]}');
plt.show();
print("To see the first image in matrix notation type x_train[0] in a cell and run it. ");

# ## The neural network
#
# So an image is a $28\times 28$ matrix, or equivalenty a $28^2\times 1$ vector/matrix, obtained by flattening it (stack all rows in one line). For each image/vector the machine needs to return a number between 0 and 9 (according to what the image is understood to be). What is the easiest (reasonable) way to obtain ten evaluations out of a $28^2\times 1$ vector $x$. Of course it is matrix multiplication by a $10\times 28^2$ matrix `W`. The matrix `W` is called the weights matrix (it weights  `x` to to measure its value out of 10 different options) and our model follows the simple decision ruled
#
#
# **PREDICTION**: `x` is the digit equal to the index of the biggest entry in `Wx`.
#
#
# For example if `Wx`  $ = [0,0,0,1,0,3,-\pi, 7,100]^T$, then the machine tells us that `x` is the picture of a 9. That simple.
#
#
# **TRAINING** = finding a good `W`. What's "good"? It's up to us to decide.  For example, an option is to tell the machine that $W x$ needs to have a very high third entry if the label of $x$ is 2, this is the same as asking that 
#
# - $\|W x-y \|_2$ is small,
# - $y = [0,0,C,0,0,0,0,0,0,0]$,
# - $C>0$ is very big and 
#
# where $\|$    `x`  $ \|_2=(x_0^2+...+x^2_9)^{\frac12}$ is the 10-$d$ Euclidean norm. So we require $\|W x-y \|_2$ to be samll for all images in the traning set by requiring that the loss function
#
# $$
# L(W)= \frac1n\sum_{i=1}^n\|Wx_i - y_i\|_2
# $$
#
# is small, where $\{(x_i,y_i)\}_{i=1}^n$ are the pairs of images with their vector label and $n$ is `size_train_small`. Mathematically, we want a `W` that solves the optimisation problem $\arg\min_W L(W)$. SoHow do we train this model? With a 
#
# **Trivialized gradient descent alghoritym**: 
#
# Initialize at the entry `(i, j) = (0, 0)`,  then 
#
# 1. Take the current `W` and create a copy `W_temp` that equals `W` apart from `W_temp[i,j] = W[i,j] + 1`,
# 2. if `L(W_temp) < L(W)` make `W_temp` the new `W`.
# 3. Fix a new entry `(i,j)`  to move  horizontally in the $10\times 28^2$ matrix `W` (skipping lines once at the last column and staring back from `(0,0)` once all entries have been visited)
# 4. start back at 1.

# ## Model construction
#
# We first flatten all images and construct labeld vectors in `y_label`.

#Flatten images in training set
x_train = x_train.reshape(x_train.shape[0],28*28)
x_test = x_test.reshape(x_test.shape[0],28**2)
#Create 10x1 vectors with a 1-entry corresponding the respective
#image label
y_label = np.zeros((y_train.shape[0], 10))
for i in range(y_train.shape[0]):
    y_label[i,y_train[i]] = 1


# The following is the class cotaining our model. It performs three main tasks:
#
# I. `nl = NL()`: the model is created and named `nl`, and the size of each image input vector is set to 28*28, the output vector size to 10, the weight matrix `W` is initialised to a $10\times28*28$  matrix where each entry is a 0, the number of times we run through all the entries of `W` in our gradiant descent (in `self.steps`).
#
# II. `nl.fit(x_train,y_label)`: this is the trivialized gradient descent, check the description of the steps!
#
# III. `nl.predict(x_test):` the decision making process. The `nl` computes the matrix product `W@x` with the current `W` for each image `x` in `x_test` and returns its decisions, i.e. the index of the  biggest entry in `W@x`.
#
# The remaining two functions simply compute the average correct predictions of `nl` and allow to return the current weights matrix `W`. 

class NL:
    def __init__(self):
        self.size_in = 28*28 #Lenght of a flattened image
        self.size_out = 10 #Ten possible choices for the possible digit in the image
        self.weights = np.zeros((self.size_out, self.size_in)) #Initialize W as a zero matrix
        self.steps = 10 #How many times the descent algorythm is performed on all entries of W
        
    def fit(self,x_train,y_label):
        """Train the neural network by performing the trivialized 
        gradient descent."""
        t_start = time() #Time at start of training
        steps = self.steps*(self.size_in*self.size_out) #Number of iteration of the descent algo
        W = self.weights
        h, C = 2., 10**5. #Set the learning step h and the big constant C to scale y_label
        y_label = C*y_label
        
        #Algorythm starts here:
        for i in range(steps): #Steps 3 and 4 happen here
            
            #Step 1: make a copy of W
            W_temp = W.copy()
            #Change the current entry of W_temp by h
            j = i % (self.size_in*self.size_out -1)
            W_temp.reshape(np.prod(W_temp.shape))[j] =  W_temp.reshape(np.prod(W_temp.shape))[j] + h
            
            #Step 2: Create a matrix contaning all vectors Wx - y
            Wx_minus_y = np.ones(x_train.shape[0]*10).reshape(x_train.shape[0],10)
            for  m in range(x_train.shape[0]):
                Wx_minus_y[m] = W@x_train[m]- y_label[m]
            #and create a matrix contaning all vectors W_temp x - y    
            W_temp_x_minus_y = np.ones(x_train.shape[0]*10).reshape(x_train.shape[0],10)
            for  m in range(x_train.shape[0]):
                W_temp_x_minus_y[m] = W_temp@x_train[m] - y_label[m]
            #and then compute the loss under the current and temporary weight matrices
            Loss_W = Euclid_norm(Wx_minus_y,axis=0).mean()
            Loss_W_temp = Euclid_norm(W_temp_x_minus_y,axis=0).mean()
            #and decide whether we descend with W_temp
            if Loss_W_temp < Loss_W:
                W = W_temp
        
        t_end = time() #Time at end of training
        self.weights = W  # The training has finished and the last W is set to be the model weight matrix
        print("Training time:", np.round(t_end-t_start,2),"s")
    
    def predict(self,x_test):
        """Predict the digit in each image x in x_test by returning
        the position where W@x is the highest.
        """
        y_pred = np.ones(x_test.shape[0])
        for i in range(x_test.shape[0]):
            y_pred[i] = np.where(self.weights@x_test[i]
                        == np.max(self.weights@x_test[i]))[0][0]
        return y_pred

    def return_weights(self):
        return self.weights
    
    def accuracy(self,y_test,y_pred):
        return len(np.where((y_pred - y_test) == 0.)[0])/len(y_test)


# ## Training and predictions
#
# So initialise/instantiate the model and check its predictions before training, which are of course goiong to be pretty bad and expected to be around 10% accuracy (why?).

nl = NL()
#Compute prediction accuracy without training
y_pred = nl.predict(x_test)
print("Probability of correct prediction before training:\n", nl.accuracy(y_pred,y_test))

# Finally, we are ready to train (or fit) the model. So lets feed to `nl` the training images with the respective labels and let do its homework.

nl.fit(x_train,y_label)

# Now that the model is trained, let's check its predictions.

y_pred_in_sample = nl.predict(x_train)
print("Probability of correct prediction after training\n\n- on traning set:",
      nl.accuracy(y_pred_in_sample,y_train))
y_pred_out_of_sample = nl.predict(x_test)
print("\n- on test set:",
      nl.accuracy(y_pred_out_of_sample,y_test))

# ## Discussion and exercises 
#
# So we trained the machine to recognise digits with more than 60% accuracy with rather trivial mathematics!!! This is very exciting because 60% is so much higher than 10%, meaning that the machine did learn a good deal about digits during its first five minutes of training.
#
# However, note that on the training set the machine, or better, the weight matrix performs much better, close to 90%. This is a classic case of overfitting, mainly due to the small sized training set that we provided. So we leave the reader with some questions to think about:
#
# - How can we interpret the rows of `W`?
# - How can you fix the overfitting? 
# - Can you speed up the calculations to allow for bigger size traning sets (would dropouts make sense here)?
# - Can you improve the model by allowing $h$ and $-h$ learning steps?

# ## Tensorflow equivalent
#
# We now train our model using TensorFlow on the same dataset and then on the full dataset. We first reimport the data and create a small and a large training set. Then instantiate the first neural network, display its description and its performance before traning. (Notice that the tf model performs the flattening for us.)

# + tags=[]
#reimport data and create a big and small dataset
(x_train_big, y_train_big), (x_test, y_test) = mnist.load_data()
size_train_big = 20000
x_train_big, y_train_big = x_train_big[:size_train_big], y_train_big[:size_train_big]
x_train_small, y_train_small = x_train[:size_train_small], y_train[:size_train_small]

#Model instantiation and summary
model = tf.keras.models.Sequential([
  tf.keras.layers.Flatten(input_shape=(28, 28)),
  tf.keras.layers.Dense(10,use_bias=False)
])
print("This is the model we instantiated:\n")
print(model.summary())
#Choose a loss function
loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
model.compile(optimizer='adam',
              loss=loss_fn,
              metrics=['accuracy'])
#Check accuracy before training
print('\n\nUntrained accuracy:',
np.round(model.evaluate(x_test,  y_test, verbose=0)[1],4))
#Train the model
print(f"\n\nNow we fit/train the model on the small data set",
      f"with {size_train_small} images.\n")
model.fit(x_train, y_train, epochs=5)
#Check accuracy in/out-of-sample
print("\n\nProbability of correct prediction after training\n\n- on traning set:", 
       np.round(model.evaluate(x_train,  y_train, verbose=0)[1],4))
print("\n- on test set:",
       np.round(model.evaluate(x_test,  y_test, verbose=0)[1],4),"\n");
# -

# So the TensorFlow training, though blindingly fast and using gradient descent, yields a model performs similarly to our naively trained model! In particular it seems to be also suffering from overfitting. Lets try now to train the same model on a bigger set.

# +
model = tf.keras.models.Sequential([
  tf.keras.layers.Flatten(input_shape=(28, 28)),
  tf.keras.layers.Dense(10,use_bias=False)
])
model.compile(optimizer='adam',
              loss=loss_fn,
              metrics=['accuracy'])

print(f"\n\nWe fit/train the model on the bit data set",
      f"with {size_train_big} images.\n\n")
model.fit(x_train_big, y_train_big, epochs=5)
print("\n\nProbability of correct prediction after training\n\n- on traning set:",
       np.round(model.evaluate(x_train_big,  y_train_big, verbose=0)[1],4))
print("\n- on test set:",
       np.round(model.evaluate(x_test,  y_test, verbose=0)[1],4));
# -

# Great news! The in-sample and out-of-sample accuracy are high and very close, so we can conclude that the model is not overfitting the traning set.
#
#
# Have fun!
#
#
# ---
# Author: <NAME>
#
# Github: [ltoniazzi](https://github.com/ltoniazzi)

# ### Future additions
#
# - Speed up this code with JAX.

StarcoderdataJupyterScriptsDedupFiltered