NeuralNetwork - AtomicCharges.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Atomic Charge Prediction\n",
    "\n",
    "## Introduction\n",
    "\n",
    "In this notebook we will machine-learn the relationship between an atomic descriptor and its electron density using neural networks.\n",
    "\n",
    "The atomic descriptor is a numerical representation of the chemical environment of the atom. Several choices are available for testing.\n",
    "Reference Mulliken charges were calculated for 134k molecules at the CCSD level: each took 1-4 hours. \n",
    "\n",
    "\n",
    "The problem is not trivial, even for humans.\n",
    "<table><tr>\n",
    "    <td>\n",
    "        <img src=\"./images/complex-CX.png\" width=\"350pt\">\n",
    "    </td><td>\n",
    "        <img src=\"./images/complex-CH3-X.png\" width=\"350pt\">\n",
    "    </td></tr>\n",
    "</table>\n",
    "\n",
    "On the left we see the distribution of s electron density on C atoms in the database. Different chemical environments are shown with different colours. The stacked histograms on the right show the details of charge density for CH$_3$-CX depending on the environment of CX. The total amount of possible environments for C up to the second order exceeds 100, and the figure suggests the presence of third order effects. This is too complex to treat accurately with human intuition. \n",
    "\n",
    "Let's see if we can train neural networks to give accurate predictions in milliseconds!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "Here we use the ANN to model the relationship between the descriptors of atoms in molecules and the partial atomic charge density."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# --- INITIAL DEFINITIONS ---\n",
    "from sklearn.neural_network import MLPRegressor\n",
    "import numpy, math, random\n",
    "from scipy.sparse import load_npz\n",
    "import matplotlib.pyplot as plt\n",
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "from ase import Atoms\n",
    "from visualise import view"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's pick a descriptor and an atomic type.\n",
    "\n",
    "Available descriptors are:\n",
    "* boba: bag of bonds (per atom)\n",
    "* boba2: bag of bonds - advanced (per atom)\n",
    "* acsf: atom centered symmetry functions - 40k atoms max for each type\n",
    "* gnn: graph based fingerprint from NanoLayers\n",
    "* soap: smooth overlap of atomic positions (per atom)\n",
    "* mbtr: manybody tensor representation (per atom)\n",
    "\n",
    "Possible atom types are:\n",
    "* 1 = Hydrogen\n",
    "* 6 = Carbon\n",
    "* 7 = Nitrogen\n",
    "* 8 = Oxygen\n",
    "* 9 = Fluorine\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Z is the atom type: allowed values are 1, 6, 7, 8, or 9\n",
    "Z = 6\n",
    "\n",
    "# TYPE is the descriptor type\n",
    "TYPE = \"mbtr\"\n",
    "\n",
    "#show descriptor details\n",
    "print(\"\\nDescriptor details\")\n",
    "desc = open(\"./data/descriptor.\"+TYPE+\".txt\",\"r\").readlines()\n",
    "for l in desc: print(l.strip())\n",
    "print(\" \")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "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/charge.\"+str(Z)+\".input.\"+TYPE+\".npz\").toarray()\n",
    "trainOut = numpy.load(\"./data/charge.\"+str(Z)+\".output.npy\")\n",
    "trainOut = trainOut[0:trainIn.shape[0]]\n",
    "\n",
    "# decide how many samples to take from the database\n",
    "samples  = min(trainIn.shape[0], 9000) # change the number here if needed!\n",
    "\n",
    "print(\"training samples:   \"+str(samples))\n",
    "print(\"validation samples: \"+str(trainIn.shape[0]-samples))\n",
    "print(\"number of features: {}\".format(trainIn.shape[1]))\n",
    "\n",
    "# 70-30 split between training and validation\n",
    "validIn = trainIn[samples:]\n",
    "validOut = trainOut[samples:]\n",
    "\n",
    "trainIn  = trainIn[0:samples]\n",
    "trainOut = trainOut[0:samples]\n",
    "\n",
    "# shift and scale the inputs - OPTIONAL\n",
    "train_mean = numpy.mean(trainIn, axis=0)\n",
    "train_std = numpy.std(trainIn, axis=0)\n",
    "train_std[train_std == 0] = 1\n",
    "for a in range(trainIn.shape[1]):\n",
    "    trainIn[:,a] -= train_mean[a]\n",
    "    trainIn[:,a] /= train_std[a]\n",
    "\n",
    "# also for validation set\n",
    "for a in range(validIn.shape[1]):\n",
    "    validIn[:,a] -= train_mean[a]\n",
    "    validIn[:,a] /= train_std[a]\n",
    "\n",
    "\n",
    "# show the first few descriptors\n",
    "print(\"\\nDescriptors for the first 5 atoms:\")\n",
    "print(trainIn[0:5])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we setup a multilayer perceptron of suitable size. Out package of choice is scikit-learn, but more efficient ones are available.<br>\n",
    "Check the scikit-learn <a href=\"http://scikit-learn.org/stable/modules/generated/sklearn.neural_network.MLPRegressor.html\">documentation</a> for a list of parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# setup the neural network\n",
    "# alpha is a regularisation parameter, explained later\n",
    "\n",
    "nn = MLPRegressor(hidden_layer_sizes=(10),  activation='tanh', solver='adam', alpha=0.01, learning_rate='adaptive')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training\n",
    "Now comes the tough part! The idea of training is to evaluate the ANN with the training inputs and measure its error (since we know the correct outputs). It is then possible to compute the derivative (gradient) of the error w.r.t. each parameter (connections and biases). By shifting the parameters in the opposite direction of the gradient, we obtain a better set of parameters, that should give smaller error.\n",
    "This procedure can be repeated until the error is minimised.<br><br>\n",
    "It may take a while..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# use this to change some parameters during training if needed\n",
    "nn.set_params(solver='lbfgs')\n",
    "\n",
    "nn.fit(trainIn, trainOut);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Check the ANN quality with a regression plot, showing the mismatch between the exact and NN predicted outputs for the validation set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# evaluate the training and validation error\n",
    "trainMLOut = nn.predict(trainIn)\n",
    "validMLOut = nn.predict(validIn)\n",
    "\n",
    "print (\"Mean Abs Error (training)  : \", (numpy.abs(trainMLOut-trainOut)).mean())\n",
    "print (\"Mean Abs Error (validation): \", (numpy.abs(validMLOut-validOut)).mean())\n",
    "\n",
    "plt.plot(validOut,validMLOut,'o')\n",
    "plt.plot([Z-1,Z+1],[Z-1,Z+1]) # perfect fit line\n",
    "plt.xlabel('correct output')\n",
    "plt.ylabel('NN output')\n",
    "plt.show()\n",
    "\n",
    "# error histogram\n",
    "plt.hist(validMLOut-validOut,50)\n",
    "plt.xlabel(\"Error\")\n",
    "plt.ylabel(\"Occurrences\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Exercises"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1. Compare descriptors\n",
    "Test the accuracy of different descriptors with the same NN size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# DIY code here...\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2. Optimal NN\n",
    "Find the smallest NN that gives good accuracy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# DIY code here...\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3. Training sample size issues\n",
    "Check whether the descriptor fails because it does not contain enough information, or because there was not enough training data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# DIY code here..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4. Combine with Principal Component Analysis - Advanced\n",
    "\n",
    "Reduce the descriptor size with PCA (check the PCA.ipynb notebook) and train again. Can you get similar accuracy with much smaller networks?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# DIY code here..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 5. Putting it all together\n",
    "After training NNs for each atomic species (MBTR), combine them into one program that predicts charges for all atoms in a molecule.\n",
    "\n",
    "* Compute local MBTR for the molecule below.\n",
    "* Compute all atomic charges with the NNs.\n",
    "* Is the total charge zero? If not, normalise it.\n",
    "\n",
    "**Note: Careful about the training: if the training data was transformed, the MBTR here should be as well.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# atomic positions as matrix\n",
    "molxyz = numpy.load(\"./data/molecule.coords.npy\")\n",
    "# atom types\n",
    "moltyp = numpy.load(\"./data/molecule.types.npy\")\n",
    "\n",
    "atoms_sys = Atoms(positions=molxyz, numbers=moltyp)\n",
    "view(atoms_sys)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compute MBTR descriptor for the molecule\n",
    "# ...\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compute all atomic charghes using previously trained NNs\n",
    "# ...\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 5. Analyse the chemical environments\n",
    "\n",
    "Try to plot the ACSF/SOAP chemical environemnts of C using t-SNE. Can you identify clusters of similar C atoms? What about their partial charge?"
   ]
  },
  {
   "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
}