KRR - TotalEnergy.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Total Energy Prediction - Kernel Ridge Regression\n",
    "\n",
    "## Introduction\n",
    "\n",
    "In this notebook we will machine-learn the relationship between descriptors of molecules and their total energy using kernel regression.\n",
    "\n",
    "KRR is a one of the ML methods to perform regression (fitting).\n",
    "\n",
    "In practice, we use kernel function $d$ that measures the distance between two data points.\n",
    "With the function we build the kernel matrix:\n",
    "\n",
    "\\begin{equation}\n",
    "K_{ij} =  d(M_{i}, M_{j}) - \\alpha I_{ij}\n",
    "\\end{equation}\n",
    "\n",
    "for all pairs of samples $M_i$ and $M_j$ in the training set. $I_{ij}$ is the identity matrix and $\\alpha$ is a hyperparameter chosen between 0 and 1.\n",
    "\n",
    "When we want to predict the output of a new sample or set of samples that were not in the training set, we need to first evaluate the distance between them and the training samples:\n",
    "\n",
    "\\begin{equation}\n",
    "D_{ij} =  d(T_{i}, M_{j})\n",
    "\\end{equation}\n",
    "\n",
    "where $T_i$ is one of the new samples. The output prediction $\\mathbf{y_T}$ for the test set $T$ is obtained by:\n",
    "\n",
    "\\begin{equation}\n",
    "\\mathbf{y_T} = \\mathbf{y_M} \\cdot K^{-1} \\cdot D^T\n",
    "\\end{equation}\n",
    "\n",
    "where $\\mathbf{y_M}$ is the vector of known outputs for the training set $M$.\n",
    "\n",
    "The energy of ~134k molecules was calculated at the CCSD level."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# INITIAL DEFINITIONS\n",
    "from pyKRR import KRRsolver  # import our KRR solver object\n",
    "import numpy, random\n",
    "import numpy, math, random\n",
    "import matplotlib.pyplot as plt\n",
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "from scipy.sparse import load_npz"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's pick a descriptor. Allowed types are:\n",
    "\n",
    "1. cnt: atom counts\n",
    "2. bob: bag of bonds\n",
    "4. soap: smooth overlap of atomic positions; choose from:\n",
    "    1. soap.sum - all atoms summed together\n",
    "    4. soap.mean - mean of all atom SOAP\n",
    "    4. soap.centre - computed at the central point\n",
    "5. mbtr: many-body tensor representation\n",
    "6. cm: Coulomb matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TYPE is the descriptor type\n",
    "TYPE = \"cnt\"\n",
    "\n",
    "#show descriptor details\n",
    "print(\"\\nDescriptor details\")\n",
    "desc = open(\"./data/descriptor.\"+TYPE.split('.')[0]+\".txt\",\"r\").readlines()\n",
    "for l in desc: print(l.strip())\n",
    "print(\" \")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "and load the databases with the descriptors (input) and the correct charge densities (output). Databases are quite big, so we can decide how many samples to use for training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load input/output data\n",
    "trainIn = load_npz(\"./data/energy.input.\"+TYPE+\".npz\").toarray()\n",
    "trainOut = numpy.load(\"./data/energy.output.npy\")\n",
    "trainIn = trainIn.astype(dtype=numpy.float64, casting='safe')\n",
    "\n",
    "# decide how many samples to take from the database\n",
    "samples  = min(trainIn.shape[0], 1000) # change the number here!\n",
    "vsamples = min(trainIn.shape[0]-samples,1000)\n",
    "print(\"training samples:   \"+str(samples))\n",
    "print(\"validation samples: \"+str(vsamples))\n",
    "\n",
    "# split between training and validation\n",
    "validIn = trainIn[samples:samples+vsamples]\n",
    "validOut = trainOut[samples:samples+vsamples]\n",
    "\n",
    "trainIn  = trainIn[0:samples]\n",
    "trainOut = trainOut[0:samples]\n",
    "\n",
    "# show the first few descriptors\n",
    "print(\"\\nDescriptors for the first 5 molecules:\")\n",
    "print(trainIn[0:5])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training\n",
    "\n",
    "In the training phase we use the kernel function $d$ to measure the distance between all pairs of molecules in the dataset. The results for the kernel matrix:\n",
    "\n",
    "\\begin{equation}\n",
    "K_{ij} =  d\\left(M_{i}, M_{j}\\right) - \\alpha I_{ij}\n",
    "\\end{equation}\n",
    "\n",
    "$I_{ij}$ is the identity matrix and $\\alpha$ is a regularisation hyperparameter $\\in \\left[0,1\\right]$.\n",
    "The matrix is then inverted."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create a new solver\n",
    "solver = KRRsolver()\n",
    "\n",
    "# set the regularisation hyperparameter\n",
    "solver.alpha = 0.1\n",
    "\n",
    "# call its training function with the training inputs and outputs\n",
    "# WARNING: building the kernel matrix is O(N^2)\n",
    "# WARNING: inverting the kernel matris is O(N^3)\n",
    "# Keep the training set small\n",
    "solver.Train(trainIn, trainOut)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The evaluation function computes the distance between each validation sample $T_i$ and the training ones $M_i$:\n",
    "\n",
    "$$ D_{ij} =  d(T_{i}, M_{j}) $$\n",
    "\n",
    "The predicted energies $E\\left(T\\right)$ are then given by:\n",
    "\n",
    "$$ E\\left( T \\right) = E\\left( M \\right) \\cdot K^{-1} \\cdot D^T $$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get a list of predicted outputs for the validation inputs\n",
    "predict = solver.Evaluate(validIn)\n",
    "\n",
    "print(\"Mean Abs Error (validation): \" + str((numpy.abs(predict-validOut)).mean()))\n",
    "\n",
    "# do the regression plot\n",
    "plt.plot(validOut, predict, 'o')\n",
    "plt.plot([numpy.min(validOut),numpy.max(validOut)], [numpy.min(validOut),numpy.max(validOut)], '-')\n",
    "plt.xlabel('correct output')\n",
    "plt.ylabel('KRR output')\n",
    "plt.show()\n",
    "\n",
    "# check the distribution of energies in the training set\n",
    "plt.hist(validOut, bins=20, density=True, alpha=0.5, facecolor='gray')\n",
    "plt.hist(trainOut, bins=20, density=True, alpha=0.2, facecolor='red')\n",
    "plt.xlabel(\"Energy [H]\")\n",
    "plt.ylabel(\"Probability\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Exercises"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1. Compare descriptors\n",
    "\n",
    "Using the same training and validation set size, find which descriptor gives the smallest error."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2. Training set size\n",
    "\n",
    "Choose a descriptor and investigate the effects of training set size on the quality of the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3. Combine with Principal Component Analysis - Advanced\n",
    "Reduce the descriptor size with PCA (check the PCA.ipynb notebook) and train again. Can you get similar accuracy with a reduced input size?"
   ]
  },
  {
   "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.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}