Introduction to QML.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Introduction to Quantum Machine Learning (QML)\n",
    "## Installation\n",
    "We will mostly use qml (http://www.qmlcode.org/) to generate the molecular representations and use some custom atomic/local kernels for kernel ridge regression (KRR).\n",
    "\n",
    "First you install qml (I recommend the develop branch which is more up to date, and is needed for e.g. the FCHL19 representation) as detailed here http://www.qmlcode.org/installation.html. Just check the install works:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import qml"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Theory\n",
    "In QML we always use KRR to regress a (quantum) property $y$ for a query molecule represented by a vector $\\mathbf{X'}$. We expand the property of a query molecule in a basis of kernel functions placed on N training molecules represented by vectors $\\mathbf{X}$: \n",
    "\n",
    "$$\n",
    "y(\\mathbf{X'}) = \\sum_i^N \\alpha_i K(\\mathbf{X'}, \\mathbf{X}_i)\n",
    "$$\n",
    "\n",
    "\n",
    "In matrix form this reads: \n",
    "\n",
    "$$\n",
    "\\mathbf{y} = \\mathbf{K} \\mathbf{\\alpha}\n",
    "$$\n",
    "\n",
    "\n",
    "where $\\mathbf{y}$ is a vector containing predicted properties, $\\mathbf{K}$ is the kernel matrix containing the pairwise kernel elements between molecules in the training and prediction set and $\\mathbf{\\alpha}$ is the vector containing the regression coefficients, which can be obtained by minimizing the cost function:\n",
    "\n",
    "$$\n",
    "J(\\mathbf{\\alpha} | \\{\\mathbf{y}^{\\mathrm{ref}}\\}) = \\frac{1}{2} ||\\mathbf{K} \\alpha - \\mathbf{y}^{\\mathrm{ref}} ||^2_2 + \\frac{\\lambda}{2} \\alpha^T \\mathbf{K} \\alpha\n",
    "$$\n",
    "\n",
    "\n",
    "where $\\lambda$ is the regularization coefficient. The minimum of the cost function has the following closed-form solution:\n",
    "\n",
    "$$\n",
    "\\alpha = (\\mathbf{K} + \\mathbf{I}\\lambda)^{-1} \\mathbf{y}^{\\mathrm{ref}}\n",
    "$$\n",
    "\n",
    "\n",
    "In QML we distinguish between global and local kernels. Global kernels use a single vector to represent the molecule $\\mathbf{X}$ and a corresponding global kernel function:\n",
    "\n",
    "$$\n",
    "K_{ij} = K(\\mathbf{X}_i, \\mathbf{X}_j) = \\exp\\Big(-\\frac{||\\mathbf{X}_i - \\mathbf{X}_j||_2^2}{2 \\sigma^2}\\Big)\n",
    "$$\n",
    "\n",
    "which is here a gaussian function with width $\\sigma$ with Euclidean norm but could be any other smooth and differentiable function. \n",
    "\n",
    "\n",
    "A local kernel on the other hand is to be used in combination with local representations for atoms in molecules $\\mathbf{q}$:\n",
    "\n",
    "$$\n",
    "K_{ij} = \\sum_{I\\in i}\\sum_{J\\in j} \\delta_{Z_I Z_J} \\exp\\Bigg(-\\frac{||\\mathbf{q}_I - \\mathbf{q}_J||^2_2}{2\\sigma^2} \\Bigg)\n",
    "$$\n",
    "\n",
    "where here $\\mathbf{q}_I$ and $\\mathbf{q}_J$ are the representations of the Ith and Jth atoms in molecules i and j respectively. The Kronecker delta $\\delta$ is not included in all representations but is for example in the kernel for FCHL. $Z_I$ and $Z_J$ are the atomic numbers (nuclear charges) of each atom, respectively."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dataset \n",
    "Here we will use the QM7 dataset, a dataset of atomization energies and relaxed geometries at PBE0/def2-TZVP level of theory for approx. 7k small organic molecules (up to 7 heavy atoms).\n",
    "\n",
    "## Using QML\n",
    "QML has `Compound` objects that read the nuclear charges and coordinates from xyz files:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "mol = qml.Compound(xyz=\"qm7/0001.xyz\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([6, 1, 1, 1, 1])"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mol.nuclear_charges"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[ 1.041682, -0.0562  , -0.071481],\n",
       "       [ 2.130894, -0.056202, -0.071496],\n",
       "       [ 0.678598,  0.174941, -1.072044],\n",
       "       [ 0.678613,  0.694746,  0.62898 ],\n",
       "       [ 0.678614, -1.038285,  0.228641]])"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mol.coordinates"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generating global representations with QML \n",
    "### Coulomb matrices\n",
    "The Coulomb matrix $\\mathbf{X}$ (Rupp et al., Phys. Rev. Lett., 108, 058301, 2012) is defined as: \n",
    "\n",
    "$$\n",
    "M_{ij} = \n",
    "\\begin{cases}\n",
    "  0.5 Z_i^{2.4}  & \\mathrm{for} \\ i = j\\\\\n",
    "  \\frac{Z_i Z_j}{R_{ij}}  & \\mathrm{for} \\ i \\neq j\n",
    "\\end{cases}\n",
    "$$\n",
    "\n",
    "where the diagonal elements are the interactions of atoms with themselves (effectively a polynomial fit of the atomic energies to the nuclear charge $Z_i$) and the off-diagonal elements represent the Coulomb repulsion between nuclei $i$ and $j$. \n",
    "\n",
    "Normally the Coulomb matrix rows and columns are then sorted by their L2-norm to handle permutation invariance.\n",
    "\n",
    "We can generate the (sorted) coulomb matrix for a qml Compound object as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from qml import representations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "m = representations.generate_coulomb_matrix(\n",
    "mol.nuclear_charges, mol.coordinates, size=len(mol.nuclear_charges),\n",
    "sorting='row-norm')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([36.8581052 ,  5.50857022,  0.5       ,  5.50857007,  0.56221501,\n",
       "        0.5       ,  5.5085695 ,  0.56221605,  0.56221611,  0.5       ,\n",
       "        5.50856526,  0.56221669,  0.56221777,  0.56221405,  0.5       ])"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Bag of Bonds "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The Bag of Bonds representation (Hansen et al., J. Phys. Chem. Lett., 2015, 6, 2326-2331) expands on the CM representation. For each element a bag (vector) is constructed for a self interaction and a pairwise interaction (e.g. C-C, C-O, C-N, etc). \n",
    "\n",
    "The self interaction is as in the CM:\n",
    "\n",
    "$$\n",
    "\\frac{1}{2} Z_I^{2.4}\n",
    "$$\n",
    "\n",
    "\n",
    "The pairwise interacton between atom $i$ of element $I$ and atom $j$ of element $J$ is given by:\n",
    "\n",
    "$$\n",
    "\\frac{Z_I Z_J}{|\\mathbf{R}_i - \\mathbf{R}_j|}\n",
    "$$\n",
    "\n",
    "where $\\mathbf{R}_i$ is the euclidean coordinates of an atom $i$. \n",
    "\n",
    "The sorted bags are concatenated to a 1D vector representation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([6, 1, 1, 1, 1])"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mol.nuclear_charges"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([1, 6])"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "atomtypes = np.unique(mol.nuclear_charges)\n",
    "atomtypes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "5"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "size = len(mol.nuclear_charges)\n",
    "size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "bob = representations.generate_bob(mol.nuclear_charges,\n",
    "                                  mol.coordinates,\n",
    "                                  atomtypes,\n",
    "                                  size=size,\n",
    "                                  asize={\n",
    "                                      \"C\":1,\n",
    "                                      \"O\":0,\n",
    "                                      \"N\":0,\n",
    "                                      \"H\":4\n",
    "                                  })"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([36.8581052 ,  5.50857022,  5.50857007,  5.5085695 ,  5.50856526,\n",
       "        0.5       ,  0.5       ,  0.5       ,  0.5       ,  0.56221777,\n",
       "        0.56221669,  0.56221611,  0.56221605,  0.56221501,  0.56221405])"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bob"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(15,)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bob.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training a KRR model from global features\n",
    "For any global representation, taking as an example the CM, they can be used to train a KRR model like: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "from glob import glob"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "xyzs = sorted(glob(\"qm7/*.xyz\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "mols = [qml.Compound(x) for x in xyzs]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "max_natoms = max([len(mol.nuclear_charges) for mol in mols])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "reps = [representations.generate_coulomb_matrix(\n",
    "mol.nuclear_charges, mol.coordinates, size=max_natoms,\n",
    "sorting='row-norm') for mol in mols]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "# read PBE0 (DFT) data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/puckvg/miniconda3/lib/python3.7/site-packages/ipykernel_launcher.py:1: ParserWarning: Falling back to the 'python' engine because the 'c' engine does not support regex separators (separators > 1 char and different from '\\s+' are interpreted as regex); you can avoid this warning by specifying engine='python'.\n",
      "  \"\"\"Entry point for launching an IPython kernel.\n"
     ]
    }
   ],
   "source": [
    "energies = pd.read_csv('hof_qm7.txt', header=None, sep='\\s')\n",
    "energies.columns = ['filename', 'PBE0', 'DFTB']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "energies['filename'].to_list() == xyzs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>filename</th>\n",
       "      <th>PBE0</th>\n",
       "      <th>DFTB</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>qm7/0001.xyz</td>\n",
       "      <td>-417.031</td>\n",
       "      <td>-431.787272</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>qm7/0002.xyz</td>\n",
       "      <td>-711.117</td>\n",
       "      <td>-730.835895</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>qm7/0003.xyz</td>\n",
       "      <td>-563.084</td>\n",
       "      <td>-573.104249</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>qm7/0004.xyz</td>\n",
       "      <td>-403.695</td>\n",
       "      <td>-412.902289</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>qm7/0005.xyz</td>\n",
       "      <td>-858.499</td>\n",
       "      <td>-869.887692</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7096</th>\n",
       "      <td>qm7/7168.xyz</td>\n",
       "      <td>-1224.850</td>\n",
       "      <td>-1339.264266</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7097</th>\n",
       "      <td>qm7/7169.xyz</td>\n",
       "      <td>-1065.110</td>\n",
       "      <td>-1180.593348</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7098</th>\n",
       "      <td>qm7/7170.xyz</td>\n",
       "      <td>-1309.130</td>\n",
       "      <td>-1446.772902</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7099</th>\n",
       "      <td>qm7/7171.xyz</td>\n",
       "      <td>-1296.090</td>\n",
       "      <td>-1430.446486</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7100</th>\n",
       "      <td>qm7/7172.xyz</td>\n",
       "      <td>-1320.970</td>\n",
       "      <td>-1465.410635</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>7101 rows × 3 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "          filename      PBE0         DFTB\n",
       "0     qm7/0001.xyz  -417.031  -431.787272\n",
       "1     qm7/0002.xyz  -711.117  -730.835895\n",
       "2     qm7/0003.xyz  -563.084  -573.104249\n",
       "3     qm7/0004.xyz  -403.695  -412.902289\n",
       "4     qm7/0005.xyz  -858.499  -869.887692\n",
       "...            ...       ...          ...\n",
       "7096  qm7/7168.xyz -1224.850 -1339.264266\n",
       "7097  qm7/7169.xyz -1065.110 -1180.593348\n",
       "7098  qm7/7170.xyz -1309.130 -1446.772902\n",
       "7099  qm7/7171.xyz -1296.090 -1430.446486\n",
       "7100  qm7/7172.xyz -1320.970 -1465.410635\n",
       "\n",
       "[7101 rows x 3 columns]"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "energies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "y = energies.PBE0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.model_selection import train_test_split"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.metrics.pairwise import rbf_kernel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train, X_test, y_train, y_test = train_test_split(reps, y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "# hyperparams to opt (fix for now)\n",
    "sigma = 100\n",
    "l2reg = 1e-8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "gamma = 1. / (2 * sigma ** 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "K = rbf_kernel(X_train, X_train, gamma=gamma)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "K[np.diag_indices_from(K)] += l2reg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "alpha = np.dot(np.linalg.inv(K), y_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "K_test = rbf_kernel(X_test, X_train, gamma=gamma)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_pred = np.dot(K_test, alpha)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "mae = np.mean(np.abs(y_test - y_pred))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "8.33677468604611"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mae"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generating local representations with QML"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### aSLATM and SLATM"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The aSLATM and SLATM representations (Huang and von Lilienfeld, Nat. Chem. 12, 945-951, 2020) is constructed by considering first the environments around single atoms (one body), then a potential resulting from pairs of atomic environments (two body) and another potential resulting from combinations of three atomic environments (three body).\n",
    "\n",
    "The one-body term for an atom $I$ is just the nuclear charge $Z_I$. \n",
    "\n",
    "The two-body term is given by: \n",
    "\n",
    "$$\n",
    "\\frac{1}{2}\\sum_{J \\neq I} Z_J \\delta(\\mathbf{r} - R_{IJ})g(\\mathbf{r})\n",
    "$$\n",
    "\n",
    "where $g(r)$ is the London potential $g(r) = 1/r^6$.\n",
    "\n",
    "\n",
    "And the three-body term is given by: \n",
    "\n",
    "$$\n",
    "\\frac{1}{3} \\sum_{J \\neq K \\neq I} Z_J Z_K \\delta(\\theta - \\theta_{IJK}) h(\\theta, \\mathbf{R}_{IJ}, \\mathbf{R}_{IK})\n",
    "$$\n",
    "\n",
    "where $\\theta$ is the angle spanned by the vector $\\mathbf{R}_{IJ}$ and $\\mathbf{R}_{IK}$ (i.e. $\\theta_{IJK}$). The three-body contribution $h(\\theta, \\mathbf{R}_{IJ}, \\mathbf{R}_{IK})$ is based on the Axilrod-Teller-Muto (ATM) potential:\n",
    "\n",
    "$$\n",
    "h(\\theta, \\mathbf{R}_{IJ}, \\mathbf{R}_{JK} = \\frac{1 + \\cos\\theta \\cos\\theta_{JKI} \\cos\\theta_{KIJ}}{(\\mathbf{R}_{IJ} \\mathbf{R}_{IK} \\mathbf{R}_{KJ})^3}\n",
    "$$\n",
    "\n",
    "where:\n",
    "\n",
    "$$\n",
    "\\theta_{JKI} = f_1(\\theta, \\mathbf{R}_{IJ}, \\mathbf{R}_{IK}), \\theta_{KIJ} = f_2(\\theta, \\mathbf{R}_{IJ}, \\mathbf{R}_{IK}), \\mathbf{R}_{KJ} = f_3(\\theta, \\mathbf{R}_{IJ}, \\mathbf{R}_{IK})\n",
    "$$\n",
    "\n",
    "\n",
    "The one, two and three body terms are concatenated for each atomtype to get the aSLATM representation. \n",
    "\n",
    "A global version, SLATM, can be obtained by summing over all atomic contributions.\n",
    "\n",
    "Both can be generated using the `qml` package, where the combinations of nuclear charges (atomtypes) to be considered (in a one, two, three body way) is first specified before constructing the representation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "mbtypes = qml.representations.get_slatm_mbtypes([mol.nuclear_charges])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "aslatm = np.array(qml.representations.generate_slatm(mol.coordinates, \n",
    "                                          mol.nuclear_charges,\n",
    "                                          mbtypes,\n",
    "                                          local=True))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(5, 857)"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "aslatm.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can see there is one representation per nuclear charge. This local representation has a corresponding local kernel which also has an argument for the number of atoms $N$:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note we cut the training set size for efficiency."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/puckvg/miniconda3/lib/python3.7/site-packages/ipykernel_launcher.py:5: VisibleDeprecationWarning: Creating an ndarray from ragged nested sequences (which is a list-or-tuple of lists-or-tuples-or ndarrays with different lengths or shapes) is deprecated. If you meant to do this, you must specify 'dtype=object' when creating the ndarray\n",
      "  \"\"\"\n"
     ]
    }
   ],
   "source": [
    "reps = np.array([np.array(qml.representations.generate_slatm(mol.coordinates, \n",
    "                                          mol.nuclear_charges,\n",
    "                                          mbtypes,\n",
    "                                          local=True))\n",
    "       for mol in mols[:500]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "y = y[:500]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [],
   "source": [
    "natoms = np.array([len(mol.nuclear_charges) for mol in mols[:500]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train, X_test, N_train, N_test, y_train, y_test = train_test_split(\n",
    "reps, natoms, y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train = np.concatenate(X_train)\n",
    "X_test = np.concatenate(X_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [],
   "source": [
    "K = qml.kernels.get_local_kernels_gaussian(X_train, X_train,\n",
    "                                          N_train, N_train,\n",
    "                                          [sigma])[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [],
   "source": [
    "K[np.diag_indices_from(K)] += l2reg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [],
   "source": [
    "alpha = np.dot(np.linalg.inv(K), y_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [],
   "source": [
    "K_test = qml.kernels.get_local_kernels_gaussian(X_test, X_train, \n",
    "                                                N_test, N_train,\n",
    "                                                [sigma])[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_pred = np.dot(K_test, alpha)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [],
   "source": [
    "mae = np.mean(np.abs(y_test - y_pred))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "3.212686569335936"
      ]
     },
     "execution_count": 52,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mae"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The global version can be generated using the `local=False` flag:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [],
   "source": [
    "slatm = np.array(qml.representations.generate_slatm(mol.coordinates,\n",
    "                                                   mol.nuclear_charges,\n",
    "                                                   mbtypes,\n",
    "                                                   local=False))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(857,)"
      ]
     },
     "execution_count": 54,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "slatm.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Which can be used with a regular RBF kernel like the global CM."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## FCHL"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There are two variants of the FCHL representation: the original, dubbed FCHL18 (Faber et al, J. Chem. Phys., 148, 241717, 2018), and a newer variant dubbed FCHL19 (Christensen et al, J. Chem. Phys., 152, 044107, 2020). \n",
    "\n",
    "### FCHL18\n",
    "FCHL18 is based on Gaussian distribution functions scaled by power laws which explicitly counts for structural and elemental degrees of freedom. Like SLATM, there are first, second and third order terms. The first order expansion $A_1(I)$ accounts for chemical composition and is modelled by a Gaussian function placed on period $P_I$ and group $G_I$ in the periodic table of element $I$:\n",
    "\n",
    "$$\n",
    "A_1(I) = \\mathcal{N}(\\mathbf{x}_I^{(1)}) = \\exp\\Big(-\\frac{(P_I - \\chi_1)^2}{2\\sigma_P^2} - \\frac{(G_I - \\chi_2)^2}{2\\sigma_G^2} \\Big)\n",
    "$$\n",
    "\n",
    "where $\\mathbf{x}_I^{(1)} = \\{P_I, \\sigma_P; G_I, \\sigma_G\\}$ with respective widths $\\sigma_P$ and $\\sigma_G$, which can be seen as elemental smearing parameters which control the nearsightedness of elements in the periodic table. $\\chi_1$ and $\\chi_2$ represent dummy variables for period and group which are integrated out when evaluating the Euclidean distance. \n",
    "\n",
    "The second order expansion $A_2(I)$ is a product of $A_1(I)$ and a sum that runs over neighbouring atoms $i$: \n",
    "\n",
    "$$\n",
    "A_2(I) = \\mathcal{N}(\\mathbf{x}_I^{(1)}) \\sum_{i \\neq I} \\mathcal{N}(\\mathbf{x}_{iI}^{(2)}) \\xi_2(d_{iI})\n",
    "$$\n",
    "\n",
    "where: \n",
    "\n",
    "$$\n",
    "\\mathbf{x}_{iI}^{(2)} = \\{d_{iI}, \\sigma_d; P_i, \\sigma_P; G_i, \\sigma_G\\}\n",
    "$$\n",
    "\n",
    "where $d_{iI}$ and $\\sigma_d$ correspond to the interatomic distance where a Gaussian is placed and its width. $\\xi_2$ corresponds to the 2-body scaling function which takes the form of a power law.\n",
    "\n",
    "\n",
    "Finally the third order expansion $A_3(I)$ runs over all neighbouring atoms: \n",
    "\n",
    "$$\n",
    "A_3(I) = \\mathcal{N}(\\mathbf{x}^{(1)}) \\sum_{i \\neq I} \\mathcal{N}(\\mathbf{x}^{(2)}_{iI}) \\sum_{j \\neq i, I} \\mathcal{N}(\\mathbf{x}_{ijI}^{(3)}) \\xi_3(d_{iI}, d_{jI}, \\theta_{ij}^I)\n",
    "$$\n",
    "\n",
    "where: \n",
    "\n",
    "$$\n",
    "\\mathbf{x}_{ijI}^{(3)} = \\{\\theta_{ij}^I, \\sigma_{\\theta}; P_j, \\sigma_P; G_j, \\sigma_G \\}\n",
    "$$\n",
    "\n",
    "where $P_j$ and $G_j$ correspond to the period and group of atom $j$. $\\xi_3$ is the three-body scaling function, $\\theta_{ij}^I$ is the principle angle between the two distance vectors which span from $I$ to $i$ and $I$ to $j$ respectively. $\\sigma_{\\theta}$ is the width of the Gaussian placed at $\\theta_{ij}^I$.\n",
    "\n",
    "For the forms of the scaling functions please refer to the paper. \n",
    "FCHL18 is constructed to be used with its own kernel with specific dimensions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [],
   "source": [
    "rep = np.array(qml.fchl.generate_representation(\n",
    "mol.coordinates, mol.nuclear_charges, max_size=max_natoms))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(23, 5, 23)"
      ]
     },
     "execution_count": 56,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rep.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As for the Coulomb matrix example above, a `max_size` flag is used to pad representations to a maximum dimension.\n",
    "\n",
    "The FCHL18 representation is 3D and has a specific kernel to handle these dimensions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [],
   "source": [
    "mols = [qml.Compound(x) for x in xyzs]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [],
   "source": [
    "max_natoms = max([len(mol.nuclear_charges) for mol in mols])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [],
   "source": [
    "reps = np.array([qml.fchl.generate_representation(\n",
    "mol.coordinates, mol.nuclear_charges, max_size=max_natoms,\n",
    ") for mol in mols[:500]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train, X_test, y_train, y_test = train_test_split(reps, y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [],
   "source": [
    "# hyperparams to opt (fix for now)\n",
    "sigma = 5\n",
    "l2reg = 1e-8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [],
   "source": [
    "K = qml.fchl.get_local_symmetric_kernels(X_train, sigma)[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [],
   "source": [
    "K[np.diag_indices_from(K)] += l2reg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {},
   "outputs": [],
   "source": [
    "alpha = np.dot(np.linalg.inv(K), y_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {},
   "outputs": [],
   "source": [
    "K_test = qml.fchl.get_local_kernels(X_test, X_train, sigma)[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(375, 375)"
      ]
     },
     "execution_count": 66,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "K.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_pred = np.dot(K_test, alpha)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {},
   "outputs": [],
   "source": [
    "mae = np.mean(np.abs(y_test - y_pred))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.8068252896528206"
      ]
     },
     "execution_count": 69,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mae"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### FCHL19\n",
    "Whereas FCHL18 solves an analytical integral to compare atomic environments, FCHL19 uses a discretised representation to enable greater computational efficiency. The representation is a vector per atomic environment in an atom, like aSLATM.\n",
    "\n",
    "The two body term encodes radial distributions between the central atoms and neighbouring atoms of a given element type. The three body term encodes the mean distances and angles between an atom and neighbouring pairs of atoms of given element types. It does not include an explicit one body term.\n",
    "\n",
    "The two body term has the form:\n",
    "\n",
    "$$\n",
    "G^{\\mathrm{2-body}} = \\xi_2(r_{IJ})f_{\\mathrm{cut}}(r_{IJ}) \\exp\\Bigg(-\\frac{(\\ln R_s - \\mu(r_{ij}))^2}{2\\sigma(r_{ij})^2}\\Bigg)\n",
    "$$\n",
    "\n",
    "where $R_s$ is the distance location of the grid point and $\\mu(r_{ij})$ and $\\sigma(r_{ij})$ are parameters of the log-normal distribution which in turn depend on the interatomic distance $r_{IJ}$ and a hyperparameter $w$:\n",
    "\n",
    "$$\n",
    "\\mu(r_{ij}) = \\ln \\Big( \\frac{r_{IJ}}{\\sqrt{1 + w/r^2_{IJ}}} \\Big)\n",
    "$$\n",
    "\n",
    "and \n",
    "\n",
    "$$\n",
    "\\sigma(r_{ij})^2 = \\ln\\Bigg(1 + \\frac{w}{r_{IJ}^2} \\Bigg)\n",
    "$$\n",
    "\n",
    "with the two body scaling function $\\xi_2(r_{IJ})$: \n",
    "\n",
    "$$\n",
    "\\xi_2(r_{IJ}) = \\frac{1}{r_{IJ}^{N_2}}\n",
    "$$\n",
    "\n",
    "where the exponent $N_2$ is a hyperparameter of the representation. A soft cut-off function is used."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "metadata": {},
   "outputs": [],
   "source": [
    "elements = np.unique(np.concatenate([(mol.nuclear_charges) for mol in mols]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {},
   "outputs": [],
   "source": [
    "rep = qml.representations.generate_fchl_acsf(mol.nuclear_charges, \n",
    "                                            mol.coordinates,\n",
    "                                            elements=elements,\n",
    "                                            gradients=False,\n",
    "                                            pad=max_natoms)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(23, 720)"
      ]
     },
     "execution_count": 72,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rep.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The representation is 2D like aSLATM but it is recommended to use a kernel which includes a Dirac delta function over nuclear charges, implemented in `qml`. For this kernel the nuclear charges of the train and test molecules need to be provided."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "metadata": {},
   "outputs": [],
   "source": [
    "reps = np.array([qml.representations.generate_fchl_acsf(mol.nuclear_charges, \n",
    "                                            mol.coordinates,\n",
    "                                            elements=elements,\n",
    "                                            gradients=False,\n",
    "                                            pad=max_natoms)\n",
    "        for mol in mols[:500]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/puckvg/miniconda3/lib/python3.7/site-packages/ipykernel_launcher.py:1: VisibleDeprecationWarning: Creating an ndarray from ragged nested sequences (which is a list-or-tuple of lists-or-tuples-or ndarrays with different lengths or shapes) is deprecated. If you meant to do this, you must specify 'dtype=object' when creating the ndarray\n",
      "  \"\"\"Entry point for launching an IPython kernel.\n"
     ]
    }
   ],
   "source": [
    "nuclear_charges = np.array([mol.nuclear_charges for mol in mols[:500]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train, X_test, Q_train, Q_test, y_train, y_test = train_test_split(\n",
    "    reps, nuclear_charges, y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "metadata": {},
   "outputs": [],
   "source": [
    "K = qml.kernels.get_local_symmetric_kernel(X_train, Q_train, sigma)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "metadata": {},
   "outputs": [],
   "source": [
    "K[np.diag_indices_from(K)] += l2reg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "metadata": {},
   "outputs": [],
   "source": [
    "alpha = np.dot(np.linalg.inv(K), y_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "metadata": {},
   "outputs": [],
   "source": [
    "K_test = qml.kernels.get_local_kernel(X_train, X_test, Q_train, Q_test,\n",
    "                                     sigma)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_pred = np.dot(K_test, alpha)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "metadata": {},
   "outputs": [],
   "source": [
    "mae = np.mean(np.abs(y_test - y_pred))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1.0355328826656074"
      ]
     },
     "execution_count": 82,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mae"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Other local representations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There are other popular local representations, perhaps the most notable being SOAP and ACSF. SOAP in particular is not very easy to understand to non physicists / quantum chemists so I won't mention it here. The most important thing to understand about these representations is:\n",
    "- they are constructed directly from nuclear charges and coordinates read in from xyz files \n",
    "- they are constructed in order to ensure certain properties: uniqueness, invariance to changes in atom indexing,  translation, rotation and atom reflection\n",
    "- this is because the property we want to learn (i.e. an energy) will not change if we re-order the atom indexing, translate the molecules, rotate them, or reflect them in a mirror plane \n",
    "\n",
    "For a nice review on this, please refer to this review: \n",
    "https://arxiv.org/pdf/2003.12081.pdf\n",
    "\n",
    "Others are now trying to incorporate these symmetries / physical laws into neural networks to predict quantum properties."
   ]
  }
 ],
 "metadata": {
  "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.7.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}