machine_learning_example_pytorch.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "ml_example_qml_pytorch.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cDY7k0i-0e4n",
        "colab_type": "text"
      },
      "source": [
        "Before we start, let's make sure we got access to a GPU:\n",
        "\n",
        "Edit->Notebook Setting->Hardware Accelerator"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "g243ZlQjdBK2",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from tensorflow.python.client import device_lib\n",
        "device_lib.list_local_devices()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QFcc452jt9Cb",
        "colab_type": "text"
      },
      "source": [
        "# ML Coding example #2: Atomization Energies with PyTorch (Feed-Forward Neural Network)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qJXOJ3Ix1qtJ",
        "colab_type": "text"
      },
      "source": [
        "The QM7 dataset contains XYZ structures for 7101 small molecules, with up to 7 atoms of type CNO, saturated with H.\n",
        "\n",
        "Some of them look like this:\n",
        "![alt text](https://i.imgur.com/vQymf39.png)\n",
        "\n",
        "Attempt to map all organic molecules with up to 7 CNO-atoms.\n",
        "\n",
        "For each molecule you will be given the atomization energy calculated using a QM method (PBE0/def2-TZVP).\n",
        "\n",
        "\n",
        "*   Instead of the raw 7101 XYZ files, we will use the Bag-of-Bonds features for each molecule, along with the atomization energies.\n",
        "\n",
        "References: \n",
        "\n",
        "1.   Rupp et al. (2012) Phys Rev Lett https://doi.org/10.1103/PhysRevLett.108.058301\n",
        "2.   Hansen et al. (2015) J Phys Chem Lett https://doi.org/10.1021/acs.jpclett.5b00831\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "aQvcmL2fuMH6",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!wget -O bob.npy https://www.dropbox.com/s/vyiwza2uy4jkczg/bob.npy\n",
        "!wget -O hof.npy https://www.dropbox.com/s/zy717f8mwxaegff/hof.npy"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5e_aHFh22uNq",
        "colab_type": "text"
      },
      "source": [
        "Load the data into numpy arrays and plot the distribution of energies:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TAI3PeVqvKHW",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import numpy as np\n",
        "\n",
        "Xall = np.load(\"bob.npy\")\n",
        "Yall = np.load(\"hof.npy\")\n",
        "\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "plt.hist(Yall)\n",
        "plt.xlabel(\"Atomization Energy [kcal/mol]\")\n",
        "plt.ylabel(\"Counts\")\n",
        "plt.show()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HGToKpH-3T8l",
        "colab_type": "text"
      },
      "source": [
        "The values are very large and negative, which can be a problem for neural networks. To avoid this, we can normalize the data or scale all the target values by a factor.\n",
        "\n",
        "Scale the values by the mean of the dataset and divide into Training and Test sets:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MQ30rjC62xng",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "n_train = 4000\n",
        "n_valid = 1000\n",
        "n_test  = 2000\n",
        "\n",
        "Yall = np.load(\"hof.npy\")\n",
        "y_mean = np.mean(Yall)\n",
        "print(y_mean)\n",
        "Yall /= y_mean\n",
        "\n",
        "Xtrain = Xall[:n_train]\n",
        "Xvalid = Xall[n_train:n_valid+n_train]\n",
        "Xtest  = Xall[-n_test:]\n",
        "\n",
        "Ytrain = Yall[:n_train]\n",
        "Yvalid = Yall[n_train:n_valid+n_train]\n",
        "Ytest  = Yall[-n_test:]"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LoEpzct14Iwm",
        "colab_type": "text"
      },
      "source": [
        "Plotting again!"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xl_36Z8G3yx9",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "plt.hist(Yall, label=\"All\")\n",
        "plt.hist(Ytrain, label=\"Training\")\n",
        "plt.hist(Ytest, label=\"Test\")\n",
        "plt.hist(Yvalid, label=\"Validation\")\n",
        "\n",
        "plt.xlabel(\"Target value\")\n",
        "plt.ylabel(\"Counts\")\n",
        "plt.legend()\n",
        "plt.show()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Jp1rU5_oAe_g",
        "colab_type": "text"
      },
      "source": [
        "Wondering what our representations look like? Print a row in `Xtrain`!"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hLeEftEJAj8m",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "print(Xtrain[1])\n",
        "print(len(Xtrain[1]))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "87Yn9bus6s6V",
        "colab_type": "text"
      },
      "source": [
        "## Defining out Architecture:\n",
        "Four layers (including input and output), so this qualifies for the buzzword \"Deep Learning\":\n",
        "<img src=\"https://i.imgur.com/dsEGsVg.png\" width=\"400\">\n",
        "\n",
        "Fully-connected neural network with two hidden layer in PyTorch:\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "V47ghWBwuSA2",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "\n",
        "class NeuralNet(nn.Module):\n",
        "\n",
        "  def __init__(self, input_size, hidden1_size, hidden2_size):\n",
        "\n",
        "    super(NeuralNet, self).__init__()\n",
        "\n",
        "    # Define the operations in the layers\n",
        "    # Here \"Linear\", i.e: y = a*X.T + b\n",
        "    self.fc1 = nn.Linear(input_size,hidden1_size)\n",
        "    self.fc2 = nn.Linear(hidden1_size,hidden2_size)\n",
        "    self.fc3 = nn.Linear(hidden2_size, 1)\n",
        "\n",
        "  def forward(self, x):\n",
        "\n",
        "    # Define the activation functions,\n",
        "    # i.e. the \"forward\" passes\n",
        "    x = torch.sigmoid(self.fc1(x))\n",
        "    x = torch.sigmoid(self.fc2(x))\n",
        "    x = self.fc3(x)\n",
        "\n",
        "    return x\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rbSqeUOU8IKX",
        "colab_type": "text"
      },
      "source": [
        "Sigmoid functions again:\\begin{equation}\n",
        "\\sigma\\left(x\\right) = \\frac{1}{1 + \\exp\\left(-x\\right)}\n",
        "\\end{equation}\n",
        "\n",
        "Linear connections between the layers:\n",
        "\\begin{equation}\n",
        "y = b + \\sum_i x_i  w_i\n",
        "\\end{equation}\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dYWQLYl6SUgq",
        "colab_type": "text"
      },
      "source": [
        "Since we are using a GPU, all the training and test data must be moved to the GPU:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mm1enx-0u62k",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "dtype = torch.float\n",
        "device = torch.device(\"cuda:0\")\n",
        "\n",
        "x = torch.from_numpy(Xtrain).to(device, dtype)\n",
        "y = torch.from_numpy(Ytrain.reshape((n_train,1))).to(device, dtype)\n",
        "\n",
        "xv = torch.from_numpy(Xvalid).to(device, dtype)\n",
        "yv = torch.from_numpy(Yvalid.reshape((n_valid,1))).to(device, dtype)\n",
        "\n",
        "xs = torch.from_numpy(Xtest).to(device, dtype)\n",
        "ys = torch.from_numpy(Ytest.reshape((n_test,1))).to(device, dtype)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ey3JsNSOAJOQ",
        "colab_type": "text"
      },
      "source": [
        "## Optimizing the Neural Network weights:"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jwHdGJE5HKIn",
        "colab_type": "text"
      },
      "source": [
        "First, lets define the neural network, and the sizes of the hidden layers. Secondly, we have to send it to the GPU."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-3wt6qZFwR-B",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "input_size = 465\n",
        "\n",
        "# 30 & 5 are good values\n",
        "# 64 & 50 also worth a try\n",
        "hidden1_size = 30\n",
        "hidden2_size = 5\n",
        "\n",
        "# Define our model\n",
        "model = NeuralNet(input_size, hidden1_size, hidden2_size)\n",
        "\n",
        "# Send it to the GPU\n",
        "model = model.to(device)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DuNV3vQkKXgM",
        "colab_type": "text"
      },
      "source": [
        "Check parameters:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qzbCwuAhKZh7",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "for param in model.parameters():\n",
        "  print(param.shape)\n",
        "  print(param)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IH0X_gFMHVmL",
        "colab_type": "text"
      },
      "source": [
        "Next, we have to select a loss function. For $y = sin(x)$ we minimized the squared error:\n",
        "\n",
        "\\begin{equation}\n",
        "L_2\\left(\\mathbf{w}\\right) = \\sum_{i=1}^{N} \\left(y_i^\\mathrm{true} -  y_i^\\mathrm{predicted} \\right)^{2}\n",
        "\\end{equation}\n",
        "From experience, the L1-loss function (mean absolute error) works better for this experiment:\n",
        "\\begin{equation}\n",
        "L_1\\left(\\mathbf{w}\\right) = \\frac{1}{N}\\sum_{i=1}^{N} | y_i^\\mathrm{true} -  y_i^\\mathrm{predicted} |\n",
        "\\end{equation}\n",
        "\n",
        "Lastly, we also need to define the optimizer. In this case we choose the \"Adam\" optimizer, which is a variant of gradient-descent optimization.\n",
        "*  Kingma & Ba (2014) *arXiv*, https://arxiv.org/abs/1412.6980"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OKy5Wn_QHTrI",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# L2 loss function (Mean Squared Error)\n",
        "# loss_fn = nn.MSELoss() \n",
        "\n",
        "# L1 loss (Mean Absoulte Error)\n",
        "loss_fn = nn.L1Loss()\n",
        "\n",
        "# Optimizer, lr = learning rate\n",
        "optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oy7b0RadOPsV",
        "colab_type": "text"
      },
      "source": [
        "Now, we are write the optimization procedure\n",
        "\n",
        "1.   Make prediction\n",
        "2.   Calculate the loss\n",
        "3.   Calculate the gradient of the loss function\n",
        "4.   Take a step in the direction suggested by the gradient.\n",
        "\n",
        "For example 30000 epoch (about 30-60 sec on Google Colab):"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "R16-NVtwwZDL",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "for epoch in range(30000):\n",
        "\n",
        "  # Make a prediction with the model\n",
        "  y_predicted = model.forward(x)\n",
        " \n",
        "  # Calculate the Loss-function\n",
        "  loss = loss_fn(y_predicted, y)\n",
        "\n",
        "  # Set gradient to zero\n",
        "  optimizer.zero_grad()\n",
        "\n",
        "  # Take the gradient of the loss function\n",
        "  loss.backward()\n",
        "\n",
        "  # Take a step with the optimizer\n",
        "  optimizer.step()\n",
        "\n",
        "  \n",
        "  # Print additional aoutput\n",
        "  if (epoch % 1000 == 0):\n",
        "    mse =(y - y_predicted).pow(2).sum()\n",
        "    rmsd = torch.sqrt(mse/n_train) * abs(y_mean)\n",
        "    mae = (y - y_predicted).abs().sum() * abs(y_mean) / n_train\n",
        "    \n",
        "    yvs = model.forward(xv)\n",
        "    rmsd_v = torch.sqrt((yvs - yv).pow(2).sum()/n_valid) * abs(y_mean)\n",
        "    mae_v = (yvs - yv).abs().sum() * abs(y_mean) / n_valid\n",
        "    \n",
        "    mae = mae.to(\"cpu\").detach().numpy()\n",
        "    rmsd = rmsd.to(\"cpu\").detach().numpy()\n",
        "    mae_v = mae_v.to(\"cpu\").detach().numpy()\n",
        "    rmsd_v = rmsd_v.to(\"cpu\").detach().numpy()\n",
        "    \n",
        "    print(\"EPOCH %7i    MAE(train) = %7.2f   RMSE(train) = %7.2f   MAE(valid) = %7.2f   RMSE(valid) = %7.2f    [kcal/mol]\" %\n",
        "        (epoch, mae, rmsd, mae_v, rmsd_v))\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BC12gYMJTM47",
        "colab_type": "text"
      },
      "source": [
        "### Calcluate test errors:\n",
        "With the optimized neural network, we can now predict the error on the test set. \n",
        "\n",
        "Because everything is stored on the GPU, we also have to copy back to the CPU and convert to Numpy array for convenience:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nVEfj60Kysdd",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Calculate test, copy back to CPU and convert to numpy \n",
        "y_pred = model(xs).to(\"cpu\").detach().numpy().flatten() * y_mean\n",
        "\n",
        "# Copy true test values back to CPU and convert to numpy\n",
        "y_true = ys.to(\"cpu\").detach().numpy().flatten() * y_mean\n",
        "\n",
        "print(y_pred)\n",
        "print(y_true)\n",
        "\n",
        "mae = np.mean(np.abs(y_pred - y_true))\n",
        "rmse = np.sqrt(np.mean(np.square(y_pred - y_true)))\n",
        "\n",
        "print(\"MAE:\", mae)\n",
        "print(\"RMSE:\", rmse)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BBUSzecdTgul",
        "colab_type": "text"
      },
      "source": [
        "Lastly, lets plot a correlation plot between the true and predicted atomization energies:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "EiTA-xI-TctO",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "\n",
        "plt.scatter(y_true,y_pred)\n",
        "plt.xlabel(\"True Atomization Energy [kcal/mol]\")\n",
        "plt.ylabel(\"Predicted Atomization Energy [kcal/mol]\")\n",
        "plt.show()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "X72ohLzIT4Zw",
        "colab_type": "text"
      },
      "source": [
        "### Note on hyper parameters and optimization"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "720A8DK6TBn9",
        "colab_type": "text"
      },
      "source": [
        "In order to have the best performing neural network, we could have changed a lot of things:\n",
        "\n",
        "*   Number of hidden layers\n",
        "*   Hidden layer sizes\n",
        "*   Type of loss function\n",
        "*   Optimizer + setting\n",
        "*   Regularization of parameters\n",
        "\n",
        "\n",
        "If you are wondering, how we could have added regularization to the loss function, here is an example for example L2 regularization:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "J3pIdFVtKFiX",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "l2reg = 1e-4\n",
        "\n",
        "for param in model.parameters():\n",
        "    loss += torch.norm(param) * l2reg"
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}