# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     name: python3
# ---

# + [markdown] id="7Q6NHKUCo8Aq"
# In this notebook, we will run a ready-made network starting from some ATLAS data, which is already normalized. There is also an alternative to train the network from scratch.  

# + [markdown] id="ONFAOUa5o8Au"
# ## Look into the dataset

# + [markdown] id="EzZcwOnUo8Av"
# First we need to make sure that Python 3.8 is used in the notebook. It is required in order to open this certain .pkl-file. 

# + colab={"base_uri": "https://localhost:8080/", "height": 35} id="hNqX-DTuo8Av" outputId="8c09a2d9-ee57-41fc-9602-9019deeb7122"
import sys
sys.version

# + [markdown] id="5rTdMGQYRrbq"
# Then, we will import all necessary modules.

# + id="U_VvHO_DRaP1"
import os
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import urllib.request
from sklearn.model_selection import train_test_split
np.warnings.filterwarnings('ignore', category=np.VisibleDeprecationWarning)

# + [markdown] id="Yw9PyQjszdMV"
# Then, we will read the dataset and open into Pandas. Note that we have to change the paths to the directory where the dataset files are.
#
#

# + colab={"base_uri": "https://localhost:8080/", "height": 643} id="qlS5HxLLR2tw" outputId="8d446dff-dec8-4fc1-b373-81ad53a6c347"
df = pd.DataFrame()

filename = "monojet_Zp2000.0_DM_50.0_chan3.csv"

if not os.path.exists(filename):
    print(f"Downloading {filename}")
    url = "https://www.dropbox.com/s/elclfkfeg2s7vgs/monojet_Zp2000.0_DM_50.0_chan3.csv?dl=1"
    urllib.request.urlretrieve(url, filename)
    print(f"Downloaded {filename}")
else:
    print(f"{filename} already exists")

with open(filename, 'r') as datasetfile:
    for line in datasetfile:
        df = pd.concat([df, pd.DataFrame([tuple(line.strip().split(';'))])], ignore_index = True)

df

# + [markdown] id="Fd7Kg61ti5Py"
# Extracting Jet particles only by checking every objects in all events.

# + colab={"base_uri": "https://localhost:8080/", "height": 419} id="Pe51__rjSKif" outputId="f596dc31-1ff5-406e-edce-582763e6a731"
dataset = df.values
X = []

# Remove all NaN and Null values
for i in range(len(dataset)):
    X.append([j for j in dataset[i] if j == j and j != ""])

X = np.array(X)
data = pd.DataFrame(columns=['E', 'pt', 'eta', 'phi'])

for i in range(len(X)):
    for j in range(len(X[i])):
        # Checking if j particle
        if X[i][j][0] == 'j':
            four = list(map(float,X[i][j].split(',')[1:]))
            data = data.append({'E': four[0], 'pt': four[1],'eta': four[2],'phi': four[3]}, ignore_index=True)
        
data

# + [markdown] id="EkK7DUB8SuAI"
# Then, we'll normalise the data.

# + colab={"base_uri": "https://localhost:8080/"} id="6poFYIqnSoen" outputId="643cb510-45ad-481c-8a05-ba1ce8dfe3f3"
print('All samples:')
print(data)

normalized_data = pd.DataFrame(columns=['E', 'pt', 'eta', 'phi'])

normalized_data['E'] = np.log10(data['E'])
normalized_data['pt'] = np.log10(data['pt'])
normalized_data['eta'] = data['eta'] / 5
normalized_data['phi'] = data['phi'] / 3

print('\nSamples after normalization:')
print(normalized_data)

# + [markdown] id="YdvirQz9yyLE"
# ### Split the data into Train and Test set

# + [markdown] id="Gv5W27YgTGsX"
# Then, we'll split the data into train and test set.
#

# + colab={"base_uri": "https://localhost:8080/"} id="9uuTlDIbS3QC" outputId="a608a1de-2f02-4ae9-bafe-cb40deccf0f8"
train, test = train_test_split(normalized_data, test_size=0.2, random_state = 42)

print('\nTraining samples:')
print(train)

print('\nTesting samples:')
print(test)

# + colab={"base_uri": "https://localhost:8080/", "height": 1000} id="ClQQcIc9TOef" outputId="862120c9-5c06-4162-ba1f-1bf8e2993a76"
unit_list = ['[log(GeV)]', '[log(GeV)]', '[rad/5]', '[rad/3]']
variable_list = [r'$E$', r'$p_T$', r'$\eta$', r'$\phi$']

branches=["E","pt","eta","phi"]

n_bins = 100

plt.figure(figsize=(14, 14), dpi=300)

for i in range(0,4):
    plt.subplot(2,2,i+1)
    n_hist_data, bin_edges, _ = plt.hist(train[branches[i]], bins=n_bins)
    plt.xlabel(xlabel=variable_list[i] + ' ' + unit_list[i])
    plt.ylabel('# of events')

plt.suptitle('Four momentum values for training instances.', size = 20, y=1.02)
plt.tight_layout()
plt.savefig("fourmomentum_training", bbox_inches='tight', dpi=300)
plt.show()

# + [markdown] id="nIC0k7pno8A1"
# ## Setting up the network

# + [markdown] id="yl9IRE8-Wl_x"
# ### Installing prerequisites

# + [markdown] id="vmb6Yg1vVix_"
# Install FastAI. Access to google drive will be requested.

# + id="ynfxHLrsVQUo"
# !pip install torchtext==0.8.1
# !pip install -U fastbook
import fastbook
fastbook.setup_book()

# + [markdown] id="8jhxz2sko8A1"
# ### Preparing the data

# + [markdown] id="XgSqVfPyo8A1"
# Adding the two datasets as TensorDatasets to PyTorch (also loading all other classes we'll need later)

# + id="rHR5D4OoVVKv"
import torch
import torch.nn as nn
import torch.optim as optim
import torch.utils.data
from torch.autograd import Variable

from torch.utils.data import TensorDataset
from torch.utils.data import DataLoader

from fastai import learner
from fastai.data import core

train_x = train
test_x = test
train_y = train_x  # y = x since we are building an autoencoder
test_y = test_x

# Constructs a tensor object of the data and wraps them in a TensorDataset object.
train_ds = TensorDataset(torch.tensor(train_x.values, dtype=torch.float), torch.tensor(train_y.values, dtype=torch.float))
valid_ds = TensorDataset(torch.tensor(test_x.values, dtype=torch.float), torch.tensor(test_y.values, dtype=torch.float))

# + [markdown] id="YO7XPNyeo8A2"
# We now set things up to load the data, and we use a batch size that was optimized by previous students...note also that this is fastai v2, migration thanks to <NAME>.

# + id="q-9XUIk0o8A3"
batchsize = 256

# Converts the TensorDataset into a DataLoader object and combines into one DataLoaders object (a basic wrapper
# around several DataLoader objects). 
train_dl = DataLoader(train_ds, batch_size=batchsize, shuffle=True)
valid_dl = DataLoader(valid_ds, batch_size=batchsize * 2)
dls = core.DataLoaders(train_dl, valid_dl)


# + [markdown] id="N31Ydu4xo8A3"
# ### Preparing the network

# + [markdown] id="OhyfYI2to8A3"
# Here we have an example network. Details aren't too important, as long as they match what was already trained for us...in this case we have a LeakyReLU, tanh activation function, and a number of layers that goes from 4 to 200 to 20 to 3 (number of features in the hidden layer that we pick for testing compression) and then back all the way to 4. 

# + id="VoEmFVeGo8A4" colab={"base_uri": "https://localhost:8080/"} outputId="685d1ffa-fed9-4634-da13-70105af9e730"
class AE_3D_200_LeakyReLU(nn.Module):
    def __init__(self, n_features=4):
        super(AE_3D_200_LeakyReLU, self).__init__()
        self.en1 = nn.Linear(n_features, 200)
        self.en2 = nn.Linear(200, 200)
        self.en3 = nn.Linear(200, 20)
        self.en4 = nn.Linear(20, 3)
        self.de1 = nn.Linear(3, 20)
        self.de2 = nn.Linear(20, 200)
        self.de3 = nn.Linear(200, 200)
        self.de4 = nn.Linear(200, n_features)
        self.tanh = nn.Tanh()

    def encode(self, x):
        return self.en4(self.tanh(self.en3(self.tanh(self.en2(self.tanh(self.en1(x)))))))

    def decode(self, x):
        return self.de4(self.tanh(self.de3(self.tanh(self.de2(self.tanh(self.de1(self.tanh(x))))))))

    def forward(self, x):
        z = self.encode(x)
        return self.decode(z)

    def describe(self):
        return 'in-200-200-20-3-20-200-200-out'

model = AE_3D_200_LeakyReLU()
model.to('cpu')

# + [markdown] id="fL55Au3Mo8A5"
# We now have to pick a loss function - MSE loss is appropriate for a compression autoencoder since it reflects the [(input-output)/input] physical quantity that we want to minimize. 

# + id="BdnUzozOo8A5"
from fastai.metrics import mse

loss_func = nn.MSELoss()

#bn_wd = False  # Don't use weight decay for batchnorm layers
#true_wd = True  # weight decay will be used for all optimizers
wd = 1e-6

recorder = learner.Recorder()
learn = learner.Learner(dls, model=model, wd=wd, loss_func=loss_func, cbs=recorder)
#was: learn = basic_train.Learner(data=db, model=model, loss_func=loss_func, wd=wd, callback_fns=ActivationStats, bn_wd=bn_wd, true_wd=true_wd)

# + [markdown] id="FlPY150to8A6"
# ## Alternative 1: Running a pre-trained network

# + [markdown] id="CXcKfd2Zo8A6"
# Now we load the pre-trained network. 

# + id="jtUw7pkDo8A6" colab={"base_uri": "https://localhost:8080/"} outputId="15685bae-8421-44d0-ad4d-25370a337551"
# To load pre-trained network, uncomment next line
#learn.load("AE_3D_200")

# + [markdown] id="-ovcjQ4do8A7"
# Then we evaluate the MSE on this network - it should be of the order of 0.001 or less if all has gone well...if it has not trained as well (note the pesky 0-mass peak above...) then it's going to be a bit higher.

# + id="rvFEVY90o8A7" colab={"base_uri": "https://localhost:8080/", "height": 34} outputId="b4cfaf76-5e4c-4661-f0e5-3a42fb5eb818"
#learn.validate()

# + [markdown] id="nUx44PKXo8A9"
# ## Alternative 2: Training a new network

# + [markdown] id="9M9CAWR8o8A-"
# Instead of using a pre-trained network, an alternative is to train a new network and use that instead. 

# + [markdown] id="j3ol98V1o8A-"
# First, we want to find the best learning rate. The learning rate is a hyper-paramater that sets how much the weights of the network will change each step with respect to the loss gradient.
#
# Then we plot the loss versus the learning rates. We're interested in finding a good order of magnitude of learning rate, so we plot with a log scale.
#
# A good value for the learning rates is then either:
# - one tenth of the minimum before the divergence
# - when the slope is the steepest

# + id="R95TqfhDo8A_" colab={"base_uri": "https://localhost:8080/", "height": 323} outputId="54c74051-7360-4634-c66b-48ee49ffc73c"
from fastai.callback import schedule

lr_min, lr_steep = learn.lr_find()

print('Learning rate with the minimum loss:', lr_min)
print('Learning rate with the steepest gradient:', lr_steep)

# + [markdown] id="7rs07-DGo8A_"
# Now we want to run the training!
#
# User-chosen variables:
# - n_epoch: The number of epochs, i.e how many times the to run through all of the training data once (i.e the 1266046 entries, see cell 2)
# - lr: The learning rate. Either choose lr_min, lr_steep from above or set your own.
#

# + id="qJqmUnZno8A_" colab={"base_uri": "https://localhost:8080/", "height": 1000} outputId="24705c24-7c68-4860-84f6-e89d6d2f9193"
import time

start = time.perf_counter() # Starts timer
learn.fit_one_cycle(n_epoch=100)
end = time.perf_counter() # Ends timer
delta_t = end - start
print('Training took', delta_t, 'seconds')

# + [markdown] id="GbmCaFGSo8BA"
# Then we plot the loss as a function of batches and epochs to check if we reach a plateau.

# + id="OI3vHyzEo8BA" colab={"base_uri": "https://localhost:8080/", "height": 270} outputId="6e163b49-0dd5-47f9-c134-1656eccc29a0"
recorder.plot_loss()

# + [markdown] id="V7qUJsFGo8BC"
# Then we evaluate the MSE on this network - it should be of the order of 0.001 or less if all has gone well...if it has not trained as well (note the pesky 0-mass peak above...) then it's going to be a bit higher.

# + id="btex604vo8BD" colab={"base_uri": "https://localhost:8080/", "height": 51} outputId="1ee88431-f1d8-48a0-ea8f-43cda5f16757"
learn.validate()
trainingMSE = learn.validate(dl=train_dl)
validationMSE = learn.validate(dl=valid_dl)
print(f"Training MSE: {trainingMSE} \nValidation MSE: {validationMSE}")

# + colab={"base_uri": "https://localhost:8080/"} id="CLqHDZFWvgSG" outputId="fdedc6b1-d78f-4ac0-9832-b83ad23b0131"
modelpath = 'AE_3D_200'
learn.save(modelpath)


# + [markdown] id="kyjq7962o8BD"
# Let's plot all of this, with ratios (thanks to code by <NAME>)

# + [markdown] id="65ZUWBzeo8BD"
# ## Plotting the outputs of the network

# + [markdown] id="TCbLMnVMo8BE"
# A function in case we want to un-normalize and get back to physical quantities...

# + id="KyV5Llt6o8BE"
def custom_unnormalize(df):
    df['eta'] = df['eta'] * 5
    df['phi'] = df['phi'] * 3
    df['E'] = 10**df['E']
    df['pt'] = 10**(df['pt'])
    return df


# + [markdown] id="D5O7RnWfo8BF"
# Make the histograms from the dataset...

# + id="_xvEBcB-o8BF" colab={"base_uri": "https://localhost:8080/", "height": 1000} outputId="e270624e-a632-4659-fd0e-5a9068bf6a77"
import numpy as np

plt.close('all')
unit_list = ['[GeV]', '[GeV]', '[rad]', '[rad]']
variable_list = [r'$E$', r'$p_T$', r'$\eta$', r'$\phi$']
line_style = ['--', '-']
colors = ['orange', 'c']
markers = ['*', 's']

model.to('cpu')

# Histograms
idxs = (0, 100000)  # Choose events to compare
data = torch.tensor(test[idxs[0]:idxs[1]].values, dtype=torch.float)
pred = model(data)
pred = pred.detach().numpy()
data = data.detach().numpy()

data_df = pd.DataFrame(data, columns=test.columns)
pred_df = pd.DataFrame(pred, columns=test.columns)

unnormalized_data_df = custom_unnormalize(data_df)
unnormalized_pred_df = custom_unnormalize(pred_df)    
    
alph = 0.8
n_bins = 200

plt.figure(figsize=(14,14), dpi=300)

for i in np.arange(4):
    plt.subplot(2,2,i+1)
    n_hist_data, bin_edges, _ = plt.hist(data[:, i], color=colors[1], label='Input', alpha=1, bins=n_bins)
    n_hist_pred, _, _ = plt.hist(pred[:, i], color=colors[0], label='Output', alpha=alph, bins=bin_edges)
    plt.suptitle(test.columns[i])
    plt.xlabel(test.columns[i])
    plt.ylabel('Number of events')
    plt.yscale('log')
    plt.legend()

plt.suptitle('Results of Compression.', size = 20, y = 1.02)
plt.tight_layout()
plt.savefig("compressionresults", bbox_inches='tight', dpi = 300)
plt.show()



# + [markdown] id="7PQcal5DyRy_"
# As we can see, autoencoder is able to retain most (if not all) data even after reducing 4 variables to 3.
#
# We will now compute average error.

# + id="1rlIfxOPo8BG" colab={"base_uri": "https://localhost:8080/"} outputId="96d3d118-cce3-4547-d875-17f5e9fa49e1"
def getRatio(bin1,bin2):
    bins = []
    for b1,b2 in zip(bin1,bin2):
        if b1==0 and b2==0:
            bins.append(0.)
        elif b2==0:
            bins.append(None)
        else:
            bins.append((float(b2)-float(b1))/b1)
    return bins   

rat = getRatio(n_hist_data,n_hist_pred)
print(f'Average Error: {np.mean(rat)}')

# + [markdown] id="qKjpBFVKymOl"
# Then, plot average error by bins.

# + colab={"base_uri": "https://localhost:8080/", "height": 515} id="SqNJrwz3Mi5H" outputId="9326c946-83d0-4e43-d02b-c1d9f2e7a59e"
plt.figure(figsize=[16,8])
plt.plot(rat)
plt.xlabel('Bins', size=16)
plt.ylabel('Error', size=16)
plt.title("Relative Error by Bins", size=20)
plt.grid(axis='y')
plt.tight_layout()
plt.savefig("relativeerror", bbox_inches='tight', dpi = 300)
plt.show()

StarcoderdataJupyterScriptsDedupFiltered