{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Neighbour list\n",
    "\n",
    "## Basic usage\n",
    "\n",
    "`matscipy` neighbour lists are stored in a format similar to the coordinate (`COO`) format of sparse matrices. The basic neighbor list consists of two array that each contain the indices of the atoms that constitute the pair."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(['O', 'H', 'H'],\n",
       " array([0, 0, 1, 2], dtype=int32),\n",
       " array([2, 1, 0, 0], dtype=int32))"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import numpy as np\n",
    "\n",
    "from ase.build import molecule\n",
    "from matscipy.neighbours import neighbour_list\n",
    "\n",
    "# Single water in a box with vacuum\n",
    "a = molecule('H2O')\n",
    "a.center(vacuum=5)\n",
    "\n",
    "# Create neighbor list\n",
    "i, j = neighbour_list('ij', a, cutoff=1.2)\n",
    "\n",
    "# Return list of neighbor pairs\n",
    "a.get_chemical_symbols(), i, j"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The water molecule has four pairs at a cutoff of 1.2, which are the O-H bonds. Each of the bonds occurs twice in the neighbor list.\n",
    "\n",
    "This list format allows simple analysis. For example, coordination numbers can be computed by counting the number of entries in the index arrays."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([2, 1, 1])"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Bincount counts the number of times a specific entry shows up in the array\n",
    "np.bincount(i)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The oxygen atom has a coordination of 2 (both hydrogens) while each of the hydrogens has a coordination of 1 (since only the oxygen is the neighbor)."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Per-atom properties\n",
    "\n",
    "The neighbour list can also compute per atom properties, in particular distances and distance vectors. The first argument to the `neighbour_list` function is a string that identifies the members of the return tuple. If we want distances between atoms, we additionally specific a 'd' in this string. The return tuple then has three members."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0.96856502, 0.96856502, 0.96856502, 0.96856502])"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "i, j, d = neighbour_list('ijd', a, cutoff=1.2)\n",
    "\n",
    "d"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This the O-H bond length. If we increase the cutoff to 2 Å, we also capture the H-H distance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0.96856502, 0.96856502, 0.96856502, 1.526478  , 0.96856502,\n",
       "       1.526478  ])"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "neighbour_list('d', a, cutoff=2.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Interatomic potential\n",
    "\n",
    "As an advances usage of the neighbor list, consider the implementation of a pair potential on top of the neighbour list data structure. (This is actually how it is done in the calculators that ship with `matscipy`.) The following code example implements an attractive {math}`\\propto r^{-6}` potential, i.e. a London dispersion force with a prefactor {math}:`C`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[  0.        ,   0.        ,  17.33444449],\n",
       "       [  0.        ,  12.54138009,  -8.66722225],\n",
       "       [  0.        , -12.54138009,  -8.66722225]])"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from matscipy.numpy_tricks import mabincount\n",
    "\n",
    "C = 1.0  # just some number\n",
    "\n",
    "i, j, d, D = neighbour_list('ijdD', a, 5.0)\n",
    "energy = (-C/d**6).sum()\n",
    "pair_forces = (6*C/d**5  * D.T/d).T\n",
    "forces = mabincount(j, pair_forces, len(a)) - mabincount(i, pair_forces, len(a))\n",
    "\n",
    "forces"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The code computes the energy as a sum over all pair contributions. The variable `pair_forces` contains the force vectors between pairs of atoms, which are then summed onto the respective components of the force. The utility function `mabincount` works like `np.bincount` but can handle multidimensional arrays."
   ]
  }
 ],
 "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.11.4"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}