mala_tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "# Machine learning for computational materials science and chemistry with MALA"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "## What is MALA?\n",
    "\n",
    "A short summary about MALA and how it works is given in `mala_background.pdf`."
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Setting up MALA\n",
    "\n",
    "For this tutorial, a Google Collab enviroment will be provided that includes all necessary packages. Generally, MALA is an open source framework that can be obtained [here](https://github.com/mala-project/mala). Detailled (installation) instructions can be found [here](https://mala-project.github.io/mala/).\n",
    "\n",
    "A few examples at the end of the notebook tackle advanced applications. The necessary backends are, for the ease of installation, not bundled with the Google Collab environment. Interested readers may install them themselves on their machines, for in presence workshops, they will be demonstrated by the host, as are sosme aspects of data generation."
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Loading the modules\n",
    "\n",
    "These modules will be necessary for the tutorials discussed here."
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# MALA itself.\n",
    "!git clone https://github.com/mala-project/mala.git mala_repo; cd mala_repo; pip install .\n",
    "import mala\n",
    "\n",
    "# We would like to visualize simple plots.\n",
    "# The font size can sometimes be a bit small for Jupyter Notebooks.\n",
    "\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "font = {'size'   : 22}\n",
    "matplotlib.rc('font', **font)\n",
    "\n",
    "# For the data paths.\n",
    "from os.path import join as pj\n",
    "\n",
    "# Only for the prediction down below.\n",
    "from ase.io import read\n",
    "\n",
    "# To do some timings down below.\n",
    "from time import time\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Data I: Performing simulations (presentation only)\n",
    "\n",
    "(See presenter screen)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "# The MALA interface (hands-on)\n",
    "\n",
    "Within extended ML frameworks, a big problem is reproducibility. Models often depend on a lot of so called hyperparameters, that characterize model behavior. This may include, but in no way be limited to, neural network layer sizes and number, training procedures, description of data specifics (e.g.: how is the local density of states sampled?), etc.\n",
    "There is a number of ways to efficiently handle this problem. A lot of frameworks rely on command line arguments to deal with this, i.e., the user provides a, potentially extensive, list of command line arguments upon runtime. Another good way to handle all this is the usage of input files. This is consistent with the way standard computational science simulation codes work.\n",
    "An obvious downside to this is that one has to have a framework at hand to prepare these input files.\n",
    "\n",
    "MALA follows a route sort of in between the two approaches. The central quantity is the `Parameters()` object. It holds ALL (hyper)parameters one could use in the course of a MALA run. It is structured by the subtasks of MALA."
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "parameters = mala.Parameters()\n",
    "\n",
    "# All parameters related to how data is handled in general.\n",
    "parameters.data\n",
    "\n",
    "# \"Targets\" always refers to the quantity being learned.\n",
    "parameters.targets\n",
    "\n",
    "# \"Descriptors\" always refers to the quantity from which we learn.\n",
    "parameters.descriptors\n",
    "\n",
    "# \"Data generation\" refers to useful routines for creating training data.\n",
    "# These routines are mostly experimental at the moment and will not be discussed here in detail\n",
    "parameters.datageneration\n",
    "\n",
    "# \"Network\" refers to everything related to neural network creation and training.\n",
    "# In the future, support for more models is planned, and this collection of parameters\n",
    "# will be updated to reflect this.\n",
    "parameters.network\n",
    "\n",
    "# \"Hyperparameters\" means hyperparameter optimization. This is the process of finding the optimal\n",
    "# hyperparameters for MALA model training, and we will come back to this process later.\n",
    "parameters.hyperparameters"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Individual parameter objects can be printed to see what's hidden inside. This also works on the main object as well.\n",
    "\n",
    "Try printing a different parameter subset."
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "parameters.data.show()"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Finally, there are some high-level parameters one needs when performing ML at scale."
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Whether or not to use a GPU for model training and inference.\n",
    "parameters.use_gpu\n",
    "\n",
    "# Whether or not to use MPI parallel CPU inference (no training supported).\n",
    "# This option is either for pre-processing or production runs of trained models.\n",
    "# More on that later.\n",
    "parameters.use_mpi\n",
    "\n",
    "# Manual seeds can be used to fix the Pseudo RNG to re-create models with the exact\n",
    "# same model weights.\n",
    "parameters.manual_seed\n",
    "\n",
    "# A comment may be useful to distinguish between sets of parameters.\n",
    "parameters.comment\n",
    "\n",
    "# This is useful for adjusting the output level of MALA.\n",
    "parameters.verbosity"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "All of this does not explain how MALA handles reproducibility. Write hundreds of lines of parameter statement in each python script is not exactly maintainable.\n",
    "Therefore MALA provides a .json interface."
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "parameters.save(...)\n",
    "# Loading can be done via Parameters.load_from_file(...)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Have a look at the .json file that was just created. You will see that it is structured in the same ways as the python object, allowing fast access. You will see some parameters that are not dicussed here since they exceed the scope of this tutorial. For a first excercise, try to modify parameters both in python and in json and see whether loading will recover those changes. The comment and manual seed are good first examples for this."
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Set a comment\n",
    "...\n",
    "\n",
    "# Save.\n",
    "..."
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Edit something in the file (e.g. the manual_seed), reload and print.\n",
    "..."
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "For extended experiments, it is very useful to operate with such input files and only use the in-python parameter editing when absolutely necessary. Further, concluded experiments can be saved in this way for future reference.\n",
    "In the following, we will use a combination of both approaches for the sake of transparency."
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Data II: Data preprocessing (presentation-only)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "(See presenter screen)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "!git clone https://github.com/mala-project/test-data.git mala_test_data;\n",
    "data_path = \"/content/mala_test_data/Be2\""
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Visualizing and reproducing output data (hands-on)\n",
    "\n",
    "Before we train a model, it is a good idea to think about which metric is important, i.e., how do we test if a model is good?\n",
    "\n",
    "In essence, the advantage of MALA is the access to multiple observables. Two easily accesible metrics are the density of states (DOS) and the band energy. We will now learn how to calculate them from the LDOS (the actual DFT LDOS in this case) so we can do the same after model training to test our models.\n",
    "\n",
    "For this, we first have to make sure the correct LDOS parameters are used."
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "parameters.targets.target_type = \"LDOS\"\n",
    "parameters.targets.ldos_gridsize = 11\n",
    "parameters.targets.ldos_gridspacing_ev = 2.5\n",
    "parameters.targets.ldos_gridoffset_ev = -5"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Now we can create an LDOS calculator and directly populate it with the LDOS data from the data set."
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "ldos_calculator = mala.LDOS.from_numpy_file(parameters, pj(data_path, \"Be_snapshot0.out.npy\"))"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Afterwards, we have to read in some additional information from the simulation data (size of the real space grid, temperature, etc.)."
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "ldos_calculator.read_additional_calculation_data(pj(data_path, \"Be_snapshot0.out\"))"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Now we can access the DOS and the band energy as properties of the calculator object."
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Electronic structure data is provided as properties of calculator object,\n",
    "# i.e. .energy_grid (x-axis) and .density_of_states (y-axis).\n",
    "# Band energy can be accessed the same way.\n",
    "# Plot with e.g. matplotlib."
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Training a model (hands-on)\n",
    "\n",
    "Finally, we can train a neural network based model for the electronic structure. First, let us review which parameters we need for this. In the following we will slowly adapt the parameters until we get a good model out of it.\n",
    "\n",
    "Since we want to learn about the inner workings of MALA, we want full output. We will also fix the manual seed, so that all the models are comparable."
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "parameters = mala.Parameters()\n",
    "parameters.verbosity = 2\n",
    "\n",
    "# Choose whatever you like.\n",
    "parameters.manual_seed = ..."
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Since MALA provides quite a few reasonable default values, in the simplest case the only thing we have to decide upon is the data we want to learn on and the architecture of the neural network.\n",
    "\n",
    "For each training we have to specify training (`\"tr\"`) and validation (`\"va\"`) snapshots. The former are used to actually tune the network weights, the latter monitor model performance during training. They will become very relevant as we optimize the training process.\n",
    "\n",
    "Deciding on the layer sizes is usually done AFTER the data is loaded, since the first and last layer need to match up with the data provided."
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "data_handler = mala.DataHandler(parameters)\n",
    "data_handler.add_snapshot(input_file_name, path_to_input_file,\n",
    "                          output_file_name, path_to_output_file, \"tr\")\n",
    "# What about validation data?\n",
    "\n",
    "# This already loads data into RAM!\n",
    "data_handler.prepare_data()\n",
    "\n",
    "# This statement is best kept AFTER the data handling\n",
    "parameters.network.layer_sizes = [data_handler.input_dimension,\n",
    "                                  100,\n",
    "                                  data_handler.output_dimension]"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Now we can actually train a network."
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "network = mala.Network(parameters)\n",
    "trainer = mala.Trainer(parameters, network, data_handler)\n",
    "trainer.train_network()"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Well, how was that? Do we have a good model now, can we predict the LDOS with this? That is not an easy question to answer from this output alone. First of all, we see loss values being printed, and those look all nice, but they are not trivially related to physical/chemical accuracy, which we are actually looking for.\n",
    "\n",
    "We can test this by using the `Tester` class. The class works similar to the Trainer class. We add data, push them through the model, and then use the results to perform calculations.\n",
    "\n",
    "We just have to make sure that the LDOS is correctly integrated by setting the appropriate parameters. Then we can add data to test. We should always test on data different from the one we trained on. Also, we now have to specify the corresponding calculation output, since we may need this for integration."
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "parameters.targets.ldos_gridsize = ...\n",
    "parameters.targets.ldos_gridspacing_ev = ...\n",
    "parameters.targets.ldos_gridoffset_ev = ...\n",
    "\n",
    "# The data handler still holds the data from training.\n",
    "# What happens when we instead create a new one?\n",
    "# Try adding more then one snapshot.\n",
    "# We have four in total, trained on one, validated on one,\n",
    "# that leaves one for testing.\n",
    "data_handler.clear_data()\n",
    "data_handler.add_snapshot(input_file_name, path_to_input_file,\n",
    "                          output_file_name, path_to_output_file, what_goes_here?,\n",
    "                          calculation_output_file=what_is_this?)\n",
    "\n",
    "# What happens when you select \"True\" here?\n",
    "data_handler.prepare_data(reparametrize_scaler=False)\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Now comes the actual object with which to test. We simply tell it which observables to test for and off we go. The results are given as a dictionary and in the units of meV/atom."
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Other observables could be \"number_of_electrons\" or \"total_energy\". Feel free to try them out!\n",
    "tester = mala.Tester(parameters, network, data_handler, observables_to_test=[\"band_energy\"])\n",
    "results = tester.test_all_snapshots()\n",
    "print(results)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "That already looks pretty solid. On the machine this notebook was tested on, errors below 0.1 meV/atom were reported, i.e., the network reproduces the band energy excellently. We may also visualize the DOS, and will do so later, but for now let's focus on the band energy. Can we get even better results? Of course, for simple data like this, we already seem to be doing a good job. But larger and more diverse data sets can be tricky.\n",
    "\n",
    "So let us see if we can improve the training routine even further. To do so, first let's rewrite the code above into functions so that we can toy around with them."
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "def training(parameters):\n",
    "    ...\n",
    "\n",
    "    # Beware that the network layers still have to be treated after data loading!\n",
    "\n",
    "    return parameters, data_handler, network\n",
    "\n",
    "def testing(parameters, data_handler, network):\n",
    "    ...\n",
    "\n",
    "    return results\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Now let's see what happens if we start changing up the parameters. Let's start with those that influence the network architecture first. If we define the functions in an appropriate way, we can easily manipulate layer sizes of the neural networks for instance. What happens if we choose a ridiculously small number of neurons for the hidden layer? What happens if we choose a large number? Try e.g. 2 or 1000 as number for the hidden layer (the one in between)."
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "parameters = mala.Parameters()\n",
    "parameters.verbosity = 2\n",
    "parameters.manual_seed = ...\n",
    "parameters.network.layer_sizes = [...]\n",
    "\n",
    "parameters, data_handler, network = training(parameters)\n",
    "print(testing(parameters, data_handler, network))\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Ok, so this clearly has an effect. We can build better and worse models by simply adjusting this number. Models that are too small don't have enough information capacity - they simply cannot capture the information present in the data. This is called **underfitting**.\n",
    "\n",
    "But also big models (if you haven't tried - try e.g. 1000 neurons!) do not yield immediate improvement. In my case, 1000 neurons in the middle gave noticeably worse performance.\n",
    "\n",
    "Why is that? Large models possess a large capacity to store information. If they are not carefully fitted, they may very soon begin to **overfit**, i.e. learn to predict the training data - and only the training data. This is obviously not what is supposed to happen.\n",
    "\n",
    "This is what the aforementioned validation data is for. Validation data is data unseen by the model during the weight optimization, but checked after each epoch. The idea is simple: Since the model is not directly fitted on the validation data, we can use it to monitor fitting progress. If validation accuracy is drastically different from training accuracy, something is off.\n",
    "\n",
    "If you check your outputs, you will realize that both losses behave differently for different models. In the 1000 neuron case at the end you will encounter a factor of 10 between both of them! For a smaller model, it was only 2.\n",
    "\n",
    "There is one way to alleviate this: Early stopping. The idea is to end training prematurely if we see validation accuracy stagnating for a certain amount of time, i.e., the model starting to overfit. MALA supports this with a simple command. Let's try it out with 1000 neurons."
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "parameters = mala.Parameters()\n",
    "parameters.verbosity = 2\n",
    "parameters.manual_seed = ...\n",
    "parameters.network.layer_sizes = [...]\n",
    "parameters.running.early_stopping_epochs = ...\n",
    "\n",
    "parameters, data_handler, network = training(parameters)\n",
    "print(testing(parameters, data_handler, network))\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "In my case, that helped a tiny bit with the accuracy of one snapshot. So we're getting there. Before we overthinking this specific aspect, let's introduce some more hyperparameters and test them out. We can toy with the length of the training, the learning rate (i.e. the aggressivness of the gradient updates of the neural networks, large learning rates mean larger updates per epoch, which can lead to oscillations, small learning rates can lead to stagnation) and the optimizer being used."
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "parameters = mala.Parameters()\n",
    "parameters.verbosity = 2\n",
    "parameters.manual_seed = ...\n",
    "parameters.network.layer_sizes = [...]\n",
    "parameters.running.early_stopping_epochs = ...\n",
    "# Default value is 100 - try out longer trainings as well. Or shorter!\n",
    "parameters.running.max_number_epochs = 100\n",
    "# Default value is 0.5 - try out a few!\n",
    "parameters.running.learning_rate = 0.5\n",
    "# Default value is SGD - try out Adam instead!\n",
    "parameters.running.trainingtype = \"SGD\"\n",
    "\n",
    "parameters, data_handler, network = training(parameters)\n",
    "print(testing(parameters, data_handler, network))\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "You may beginning to feel a bit overwhelmed by all the choices one can make here. We'll address this very real struggle in a second, but we'll just add a bit more gasoline to the fire while were at it.\n",
    "\n",
    "There are two aspects of ML that can be very important (and that MALA supports) that we have not yet talked about.\n",
    "\n",
    "The first is data scaling. The idea behind data scaling is to make data more uniform - vectorial fields as in our case may contain larger and smaller components per data point and it can be hard for a network to correctly learn this. By normalizing or standardizing data, performance is often improved. MALA supports normalization feature-wise or per entire dataset. Empirically, we have made good experiences with the following. (Be aware: if you change the scaling, the absolute loss values will inevitably change - this is why we monitor the band energy.)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "parameters = mala.Parameters()\n",
    "parameters.verbosity = 2\n",
    "parameters.manual_seed = 2023\n",
    "parameters.network.layer_sizes = [100]\n",
    "\n",
    "# Options are \"None\", \"normal\" and \"standard\". The two latter ones can alsoe\n",
    "# be used in combination with the phrase \"feature-wise-\".\n",
    "parameters.data.input_rescaling_type = \"feature-wise-standard\"\n",
    "parameters.data.output_rescaling_type = \"normal\"\n",
    "\n",
    "\n",
    "parameters, data_handler, network = training(parameters)\n",
    "print(testing(parameters, data_handler, network))\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "In my case, adding the data scaling improved performance drastically. Feel free to try out other forms of scaling (for no scaling, just specify `\"None\"`).\n",
    "\n",
    "Secondly, we can adapt the activation functions of the neural net, i.e., the linear transformation at the end of each layer. In MALA, this is realized by a list as well. We can specify individual layers, or simply give only one type to be used at each layer."
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "parameters = mala.Parameters()\n",
    "parameters.verbosity = 2\n",
    "parameters.manual_seed = 2023\n",
    "parameters.network.layer_sizes = [100]\n",
    "parameters.data.input_rescaling_type = \"feature-wise-standard\"\n",
    "parameters.data.output_rescaling_type = \"normal\"\n",
    "\n",
    "# Sigmoid is default, try out ReLU or LeakyReLU!\n",
    "parameters.network.layer_activations = [\"Sigmoid\"]\n",
    "\n",
    "\n",
    "parameters, data_handler, network = training(parameters)\n",
    "print(testing(parameters, data_handler, network))\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "As you will have noticed, the result seem to vary quite a bit.\n",
    "\n",
    "So what do we make of all of this?\n",
    "We can build different models, and they give different accuracy. But that's also dependent on the data scaling to be used. And that is also dependent on the way we train them. What now?\n",
    "\n",
    "Well, welcome to hell. This complexity is what makes (deep) ML hard, just as the obvious speed-up makes it worthwhile.\n",
    "\n",
    "There are two important concepts at play here:\n",
    "\n",
    "1. Intuition and experience - understanding your data. Your data has inherent structure that can be exploited in several ways. As a rule of thumb activation functions may be chosen in a way that reflects the data itself, i.e. if we have a lot of smooth curves, Sigmoid functions may work well (they do here). LeakyReLU on the other hand is a standard choice for more complicated data. Model size should reflect the complexity in your data set, and a simple data set does not warrant a huge model. We have seen this just now. Training routines should be optimized for the model.\n",
    "2. Good, old fashioned non-linear optimization. Essentially we have a N-dimensional hypersphere with some distinct or float valued choices and can simply test out all the combinations.\n",
    "\n",
    "The best possible route is a combination of the two. By adding as much prior experience from 1. into 2. we can perform a limited optimization and still end up with the best possible model."
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Tuning the Hyperparameters (hands-on)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "MALA provides a convenient way to tune the hyperparameters of models. There are multiple algorithms implemented, but we will focus on the `optuna` library, to which MALA provides an interface. The idea is easy: we give MALA a set of hyperparameters to optimize, fix the others, give it some data and let optuna do its work. Optuna internally uses elaborate algorithms to determine optimal hyperparameters from observed data.\n",
    "\n",
    "First, let us think about which hyperparameters to fix. We want to fix as many as possible and reasonable to get the best trade-off between optimal results and minimal computational effort.\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "parameters = mala.Parameters()\n",
    "parameters.verbosity = 2\n",
    "parameters.manual_seed = 2023\n",
    "parameters.data.input_rescaling_type = \"feature-wise-standard\"\n",
    "parameters.data.output_rescaling_type = \"normal\"\n",
    "parameters.running.max_number_epochs = 100\n",
    "parameters.running.mini_batch_size = 40\n",
    "parameters.running.trainingtype = \"SGD\""
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Now we can specify data to be used.\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "data_handler = mala.DataHandler(parameters)\n",
    "..."
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Before we instantiate the hyperparameter optimizer, two important parameters to set are\n",
    "\n",
    "1. The number of trials (test networks) to train\n",
    "2. How often to train a test network - if we train each proposed network a number of times and evaluate the average performant, we can discard unrobust outliers. In the interest of time, let's keep it at 1 here."
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "parameters.hyperparameters.n_trials = 20\n",
    "parameters.hyperparameters.number_training_per_trial = 1\n",
    "\n",
    "hyperoptimizer = mala.HyperOptOptuna(parameters, data_handler)\n",
    "\n",
    "hyperoptimizer.add_hyperparameter(\"categorical\", \"learning_rate\",\n",
    "                                     choices=[0.1, 0.2, 0.5])\n",
    "# Which other hyperparameters should we maybe add?\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "The syntax is a bit clunky at times, since we are trying to cramp a complicated system into a few function calls. Now let's run. This may take a while. Grab a coffee and sit back :)\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "hyperoptimizer.perform_study()\n",
    "\n",
    "# Will save the results directly to the parameters.\n",
    "hyperoptimizer.set_optimal_parameters()\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Let's train the best network one final time and see where we stand. In my case, the network that Optuna suggest is bigger and deeper then the ones we have used before, but that may differ. Keep in mind we are performing a very\n",
    " limited search here in the interest of time.\n",
    "\n",
    "Also, please keep in mind that the identified network already includes a first and last layer."
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "...\n",
    "trainer.train_network()\n",
    "print(testing(parameters, data_handler, network))"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "The results may not necessarily be better with this limited search then what we had above. In my case they are very good, but again, this is a limited search compared to reasonable defaults. For actual production runs, one would set up a thorough, longer search before training production level models."
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Excurse: How to tune bispectrum hyperparameters (presentation only)\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "(See presenter screen)\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Inference and prediction (hands-on / presentation)\n",
    "\n",
    "We now understand how we can build optimal models with MALA. Once this is done, we would like to use these models to predict properties of interest. For this, let us first see how MALA can be used to save and load models.\n",
    "\n",
    "If you haven't trained a model after the hyperparameter optimization, train one now, elsewise you can skip this part.\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "...\n",
    "trainer.train_network()\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Saving such a model now becomes a one-liner. Models will be saved as `.zip` archives containing the network weights, parameter json file and coefficients for the data scaling."
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "trainer.save_run(\"Be_model\")"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Loading is equally simple, but dependent on the task. Two possible task can be done: Testing (as we have done above, to verify whether a model is performing well) or predicting (using the model to predict the electronic structure of arbitrary atomic configurations).\n",
    "\n",
    "We will have a quick look at the testing interface again to plot the density of states, and see how MALA can directly give us access to the electronic structure of a material.\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "parameters, network, data_handler, tester = mala.Tester.load_run(\"Be_model\")\n",
    "...\n",
    "# Can you predict for more then one snapshot at a time?\n",
    "actual_ldos, predicted_ldos = tester.predict_targets(0)\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Now we can visualize the DOS via the internal LDOS calculator of the DataHandler object."
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "ldos_calculator = data_handler.target_calculator\n",
    "ldos_calculator.read_additional_calculation_data(data_handler.get_snapshot_calculation_output(0))\n",
    "# In the same way we had used \"from_numpy_file\" above, we can read the LDOS from memory with\n",
    "# \"read_from_array\"\n",
    "...\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "This finally leads us to the prediction part - the part, to which such efforts eventually should culminate too. Since this part again requires LAMMPS, it will not be hands-on.\n",
    "\n",
    "Please see the presenter screen for the prediction part."
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Making MALA performant (presentation only)\n",
    "\n",
    "(See presenter screen)\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}