krr_homo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# HOMO energy prediction with kernel ridge regression\n",
    "\n",
    "\n",
    "In this notebook we will machine-learn the relationship between molecular structure (represented by the Coulomb matrix CM) and their HOMO energy using kernel regression (KRR).\n",
    "\n",
    "KRR is a machine learning method that performs regression (fitting). This tutorial shows step by step how to load the data, visualize them, select the hyperparameters, train the model and validate it. We use the QM7 dataset of 7k small organic molecules. The HOMO energies of all molecules were pre-computed with first principles quantum mechanical methods (DFT) to obtain the target data that our model can be trained on. Detailed descriptions and results for a similar dataset (QM9) can be found in [A. Stuke, et al. \"Chemical diversity in molecular orbital energy predictions with kernel ridge regression.\" J. Chem. Phys. 150. 204121 (2019)](https://aip.scitation.org/doi/10.1063/1.5086105).\n",
    "\n",
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# initial imports\n",
    "\n",
    "import numpy as np\n",
    "import math, random\n",
    "import matplotlib.pyplot as plt\n",
    "import pandas as pd\n",
    "import json\n",
    "import seaborn as sns\n",
    "from scipy.sparse import load_npz\n",
    "from matplotlib.colors import LinearSegmentedColormap\n",
    "from sklearn.model_selection import GridSearchCV\n",
    "from sklearn.kernel_ridge import KernelRidge\n",
    "from sklearn.metrics import r2_score"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load and visualize data \n",
    "\n",
    "At first, we load the data. The input data x is an array that contains all 7k molecules of the QM7 dataset, represented by their Coulomb matrices. The output data y is a list that contains the corresponding (pre-computed) HOMO energies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = load_npz(\"./data/qm7/cm.npz\").toarray()\n",
    "y = np.genfromtxt(\"./data/qm7/HOMO.txt\")\n",
    "\n",
    "print(\"Number of molecules:\", len(y))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Print the Coulomb matrix of a random molecule in the dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rand_mol = random.randint(0, len(y))\n",
    "\n",
    "print(x[rand_mol])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Visualize the Coulomb matrix of the random molecule."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "shape = (23, 23)\n",
    "mat = x[rand_mol].reshape(shape)\n",
    "\n",
    "plt.figure()\n",
    "plt.figure(figsize = (6,6))\n",
    "plt.imshow(mat, origin=\"upper\", cmap='rainbow', vmin=-15, vmax=90, interpolation='nearest')\n",
    "plt.colorbar(fraction=0.046, pad=0.04).ax.tick_params(labelsize=20)\n",
    "plt.axis('off')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Visualize the target data by plotting the distribution of HOMO energies in the dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.hist(y, bins=20, density=False, facecolor='blue')\n",
    "plt.xlabel(\"Energy [eV]\")\n",
    "plt.ylabel(\"Number of molecules\")\n",
    "plt.title(\"Distribution of HOMO energies\")\n",
    "plt.show()\n",
    "\n",
    "## mean value of distribution\n",
    "print(\"Mean value of HOMO energies in QM9 dataset: %0.2f eV\" %np.mean(y))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before dividing the dataset into training and test set, we shuffle the data. Data are often stored in a certain order, and simply taking the first part for training and the second for testing would not result in a well trained model, since the training set would not represent the test data well (and vice versa)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## shuffle the data\n",
    "\n",
    "c = list(zip(x, y))\n",
    "random.shuffle(c)\n",
    "\n",
    "x, y = zip(*c)\n",
    "\n",
    "x = np.array(x)\n",
    "y = np.array(y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we divide the data into training and test set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# decide how many samples to take from the database for training and testing\n",
    "n_train = 2000\n",
    "n_test = 1000\n",
    "\n",
    "# split data in training and test\n",
    "# take first n_train molecules for training\n",
    "x_train  = x[0:n_train] \n",
    "y_train = y[0:n_train]\n",
    "\n",
    "# take the next n_test data for testing\n",
    "x_test = x[n_train:n_train + n_test]\n",
    "y_test = y[n_train:n_train + n_test]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Check that the training data resemble the test data well by plotting the distribution of HOMO energies for both sets. The distributions should be centered around the same mean value and have the same shape."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.hist(y_test, bins=20, density=False, alpha=0.5, facecolor='red', label='test set')\n",
    "plt.hist(y_train, bins=20, density=False, alpha=0.5, facecolor='gray', label='training set')\n",
    "plt.xlabel(\"Energy [eV]\")\n",
    "plt.ylabel(\"Number of molecules\")\n",
    "plt.legend()\n",
    "plt.show()\n",
    "\n",
    "## mean value of distributions\n",
    "print(\"Mean value of HOMO energies in training set: %0.2f eV\" %np.mean(y_train))\n",
    "print(\"Mean value of HOMO energies in test set: %0.2f eV\" %np.mean(y_test))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training\n",
    "\n",
    "In the training phase we use a kernel function to measure the distance between all pairs of molecules (represented by their Coulomb matrices) in the training set. We here employ one of two kernels, the Gaussian kernel or the Laplacian kernel. The Gaussian kernel is given by\n",
    "\n",
    "\\begin{equation}\n",
    "k_{Gaussian}(\\boldsymbol{x},\\boldsymbol{x}')=e^{-\\frac{||{\\boldsymbol{x}-\\boldsymbol{x}'}||_2^2}{2\\gamma^2}},\n",
    "\\end{equation}\n",
    "\n",
    "which employs the Euclidean distance as similarity measure. The parameter $\\gamma$ is defined as $\\frac{1}{2\\sigma^2}$, where $\\sigma$ is the standard deviation of the Gaussian kernel (kernel width). The Laplacian kernel is given by\n",
    "\n",
    "\\begin{equation}\n",
    "    k_{Laplacian}(\\boldsymbol{x},\\boldsymbol{x}')=e^{-\\frac{||{\\boldsymbol{x}-\\boldsymbol{x}'}||_1}{\\gamma}},\n",
    "\\end{equation}\n",
    "\n",
    "which uses the 1-norm as similarity measure. Here, $\\gamma$ is defined as $\\frac{1}{\\sigma}$, where $\\sigma$ is the kernel width of the Laplacian kernel.\n",
    "\n",
    "In the KRR training phase with $N$ training molecules, the machine learns the relationship between the molecules (represented by their Coulomb matrix) and their corresponding (pre-computed) HOMO energies. It does so by employing a function $f(\\boldsymbol{x})$ that maps a training molecule $\\boldsymbol{x}$ to its reference HOMO energy:\n",
    "\n",
    "\\begin{equation}\n",
    "f(\\boldsymbol{x}) = \\sum_{i=1}^N \\omega_i k(\\boldsymbol{x}, \\boldsymbol{x}_i) = HOMO^{ref},\n",
    "\\end{equation}\n",
    "\n",
    "For a given training molecule $\\boldsymbol{x}$, the distance to each molecule in the training set is computed by employing the kernel function $k$ (either Gaussian or Laplacian). Each kernel contribution (distance) is then weighted by a regression weight $\\omega_i$. The above function is thus given by the weighted sum of kernel contributions (sum over $N$ training molecules). The purpose of training is to fit the regression weight $\\omega_i$ so that HOMO$_{ref}$ is matched for each training molecule. In practice, the machine solves the minimization problem\n",
    "\n",
    "\n",
    "\\begin{equation}\n",
    " \\underset{\\omega}{min} \\sum_{i=1}^N (f(\\boldsymbol{x}_i) - HOMO^{ref}_i)^2 + \\alpha \\boldsymbol{\\omega}^T \\mathbf{K} \\boldsymbol{\\omega}.\n",
    "\\end{equation}\n",
    "\n",
    "for a vector $\\boldsymbol{\\omega} \\in \\mathbb{R}^N = (\\omega_1, \\omega_2, ..., \\omega_N)$ of regression weights. In KRR, the penalty term $ \\alpha \\boldsymbol{\\omega}^T \\mathbf{K} \\boldsymbol{\\omega}$ is added to the minimization problem in order to avoid over- and underfitting. Overfitting occurs when the model learns the training data too well, even the noise and other unimportant details. The model is unable to generalize on unseen data and therefore yields high prediction errors on the test data. Underfitting occurs when the model is too simple and does not learn the training data at all, and therefore is not able to predict test data well either. Both behaviours can be avoided by tuning the parameter $\\alpha \\in \\left[0,1\\right]$ to a reasonable value. This has do be done separately from training. Both the regularization parameter $\\alpha$ and the kernel width $\\gamma$ are so called hyperparameters. Hyperparameters cannot be learned during training and have to be selected beforehand. However, it is not always obvious how to choose these hyperparameters and it often requires intuition or rules of thumb. We here employ a cross-validated grid search in order to find the best values for these two hyperparameters. \n",
    "\n",
    "In grid search, a part of the training set is split off as validation set. We set up a grid of pre-defined hyperparameter values and train the machine on the remaining training set, for each possible combination of $\\alpha$ and $\\gamma$ values. We validate each possible combination by making predictions on the validation set. The two hyperparameter values that yield the best performance (lowest error) are then selected for the final model to make predictions on the test set.\n",
    "\n",
    "In cross-validation, the roles of training and validation sets alternate. As described above, a part from the training set is split off as validation set. After training one combination of hyperparameters on the remaining training set and validating on the validation set, the validation set becomes the training set and vice versa, and the model is trained on the new training set and validated on the new validation set for the same combination of hyperparameters. The ratio can be varied, for example in 5-fold cross-validation, the training set is split in 5 equal parts. For each combination of hyperparameters, the model is trained on 80% of the data and validated on the other 20%. Then the roles of training and validation set rotate until each part has served as validation set exactly once. The final validation error for one particular combination of hyperparameters is computed as the mean from all 5 errors on the 5 validation sets. The combination with lowest average error is chosen for the final model.\n",
    "\n",
    "The cross-validated grid search routine is implemented in scikit-learn."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set up grids for alpha and gamma hyperparameters. \n",
    "# first value: lower bound; second value: upper bound; \n",
    "# third value: number of points to evaluate (here set to '3' --> '-2', '-1' and '0' are evaluated)\n",
    "# --> make sure to change third value as well when changing the bounds!\n",
    "alpha = np.logspace(-5, -2, 4)\n",
    "gamma = np.logspace(-5, -2, 4)\n",
    "\n",
    "cv_number = 5 ## choose into how many parts training set is divided for cross-validation\n",
    "kernel = 'laplacian' # select kernel function here ('rbf': Gaussian kernel, 'laplacian': Laplacian kernel)\n",
    "scoring_function = 'neg_mean_absolute_error' # it is called \"negative\" because scikit-learn interprets\n",
    "                                             # highest scoring value as best, but we want small errors\n",
    "\n",
    "## define settings for grid search routine in scikit-learn with above defined grids as input\n",
    "\n",
    "grid_search = GridSearchCV(KernelRidge(),  #machine learning method (KRR here)\n",
    "                           [{'kernel':[kernel],'alpha': alpha, 'gamma': gamma}], \n",
    "                           cv = cv_number, \n",
    "                           scoring = scoring_function,\n",
    "                           verbose=1000)  ## produces detailed output statements of grid search \n",
    "                                          # routine so we can see what is computed\n",
    "    \n",
    "# call the fit function in scikit-learn which fits the Coulomb matrices in the training set \n",
    "# to their corresponding HOMO energies.\n",
    "grid_search.fit(x_train, y_train)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Grid search results\n",
    "\n",
    "Print out the average validation errors and corresponding hyperparameter combinations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "means = grid_search.cv_results_['mean_test_score']\n",
    "stds = grid_search.cv_results_['std_test_score']\n",
    "for mean, std, params in zip(-means, stds, grid_search.cv_results_['params']):\n",
    "    print(\"%0.3f (+/-%0.03f) for %r\" % (mean, std * 2, params))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we visualize the grid search results by plotting a heatmap."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results = pd.DataFrame(grid_search.cv_results_)\n",
    "#pd.DataFrame(grid_search.cv_results_)\n",
    "\n",
    "pvt = pd.pivot_table(results, values='mean_test_score', \n",
    "                     index='param_gamma', columns='param_alpha')\n",
    "heatmap = sns.heatmap(-pvt, annot=True, cmap='viridis', cbar_kws={'label': \"Mean absolute error [eV]\"})\n",
    "figure = heatmap.get_figure()\n",
    "plt.show()\n",
    "\n",
    "\n",
    "print(\"The best combinations of parameters are %s with a score of %0.3f eV on the validation set.\"\n",
    "      % (grid_search.best_params_, -grid_search.best_score_))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Testing\n",
    "\n",
    "With the best combination of hyperparameters, the model is once again trained on the entire training set (this is done automatically in scikit-learn). Then, with the best combination of hyperparameters, predictions are made on the test set to evaluate the final model. With the fitted regressions weights $\\omega_i$ and the selected hyperparameter $\\gamma$ (kernel width), the final model is used to predict the energies of the test molecules. The energy of a particular test molecule $\\boldsymbol{x}$ is predicted by computing the weighted sum of kernel contributions $k(\\boldsymbol{x}, \\boldsymbol{x}_i)$ between the test molecule $\\boldsymbol{x}$ and each of the $N$ molecules $\\boldsymbol{x}_i$ in the training set (sum over $N$):\n",
    "\n",
    "\\begin{equation}\n",
    "f(x) = \\sum_{i=1}^N \\omega_i k(\\boldsymbol{x}, \\boldsymbol{x}_i) = HOMO^{pred},\n",
    "\\end{equation}\n",
    "\n",
    "The deviation of the predicted HOMO energies to the true reference HOMO energies yields the final error of the model. We compute the mean absolute error between predicted and reference HOMO energies for all $M$ test molecules (sum over $M$):\n",
    "\n",
    "\\begin{equation}\n",
    "\\sum_{i=1}^M \\frac{1}{M} \\big|HOMO^{pred} - HOMO^{ref}\\big|\n",
    "\\end{equation}\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# predicted HOMO energies for all test molecules\n",
    "\n",
    "y_pred = grid_search.predict(x_test) # scikit-learn automatically takes the best combination\n",
    "                                     # of hyperparameters from grid search\n",
    "\n",
    "print(\"Mean absolute error on test set: %0.3f eV\" %(np.abs(y_pred-y_test)).mean())\n",
    "\n",
    "# do the regression plot\n",
    "plt.plot(y_test, y_pred, 'o')\n",
    "plt.plot([np.min(y_test),np.max(y_test)], [np.min(y_test),np.max(y_test)], '-')\n",
    "plt.xlabel('reference HOMO energy [eV]')\n",
    "plt.ylabel('predicted HOMO energy [eV]')\n",
    "plt.show()\n",
    "print(\"R^2 score on test set: %.3f\" % r2_score(y_test, y_pred))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The $R^2$ score indicates how close the predicted energies in the test set are to the reference energies. The closer the points in the above figure are located to the diagonal, the better the predictions. Points on the diagonal (\"predicted energy\"=\"reference energy\") correspond to $R^2=1$. Therefore, $R^2$ values close to 1 indicate good model performance."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "## Exercises\n",
    "\n",
    "### 1. Grid search\n",
    "\n",
    "Increase the number and range of grid points used for grid search. Which combination of $\\alpha$ and $\\gamma$ works best? How does the computational time increase? Choose a reasonable number of grid points that don't take too long to evaluate."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2. Cross-validation\n",
    "\n",
    "Increase the number of folds used for cross-validation. Does the quality of the model increase? Take note as well of the increasing computational time and choose a number of folds that does not require too much computational time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3. Kernel function\n",
    "\n",
    "Use the Laplacian kernel instead of the Gaussian kernel. Which kernel leads to better model performance?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4. Training set size\n",
    "\n",
    "Using the better performing kernel, let's investigate the effects of training set size on the final model.\n",
    "\n",
    "Increase the size of the training set and plot the mean absolute error on the test set as a function of training set size (e.g. use 1000, 2000, 3000 etc. as training set size). How does the computational time increase? What is the largest training set size that allows training within a reasonable time, given computational resources?\n",
    "\n",
    "The optimal hyperparameter values for $\\alpha$ and $\\gamma$ can change throughout varying training set sizes. Therefore, when increasing the training set size, it is recommended to perform a cross-validated grid search for each training set size. Use a reasonable amount of grid points as well as a reasonable number of folds for cross-validation. Take note of the optimal hyperparameters for each training set size. Do they change? "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}