Done apart from running it on the cluster, making the plot and make the code nice