{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "8bcfb4a8",
   "metadata": {},
   "source": [
    "# Topology building for LAMMPS non-reactive MD simulations of amorphous carbon\n",
    "\n",
    "This is an example of how to generate LAMMPS input files for simulations of water-lubricated amorphous carbon (a-C) interfaces with the non-reactive interatomic potential described in [Reichenbach et al.](https://doi.org/10.48550/arXiv.2310.12619) starting from a set of atomic coordinates. The parameter file can be found in `matscipy`'s `docs/topology/` directory. The example structure file `aC_H2O.xyz` (provided in the same directory) consists of two H/OH-terminated a-C slabs that are serparated by a couple of water molecules. Periodic boundary conditions are used in the interface plane."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e012c1a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from nglview import show_ase, ASEStructure\n",
    "from ase.io import read\n",
    "\n",
    "atoms = read('aC_H2O.xyz')\n",
    "atoms.rotate(90, 'x')\n",
    "\n",
    "view = show_ase(atoms)\n",
    "component = view.add_component(ASEStructure(atoms), default_representation=False)\n",
    "component.add_ball_and_stick(cylinderOnly=True, radiusType='covalent', radiusScale=0.6, aspectRatio=0.1)\n",
    "component.add_ball_and_stick(radiusType='covalent', radiusScale=0.6)\n",
    "view.control.zoom(0.8)\n",
    "view"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a4de5112",
   "metadata": {},
   "source": [
    "## Assignment of the different atom types\n",
    "\n",
    "The following script assigns an atom type for each atom based on its neighbours and generates a new structure file `struct.extxyz` that includes these atom types, the atomic coordinates and molecular ids, which can at a later stage faciliate group assignments in LAMMPS. The script only assigns atom types for which parameters are available in `parameters_Reichenbach2023.in` and gives an error otherwise. Please modify at your convenience if you use other parameter sets.  Note that the assignment is based on the cutoffs specified when building the neighbour list: `('C', 'C'): 1.85, ('C', 'H'): 1.15, ('C', 'O'): 1.55, ('O', 'H'): 1.3`. In most cases, it is desirable to choose these cutoffs in line with the cutoffs that are specified in the parameter file, which is used to generate the bonding topology later (here  `parameters_Reichenbach2023.in`). The naming convention of the atom types follows Table 1 in [Reichenbach et al.](https://doi.org/10.48550/arXiv.2310.12619). Note that the script is rather strict and raises an error as soon as it encounters atoms that are in an unphysical or chemically not very stable environment such as C atoms with 5 neighbours or reactive sp-hybridised C on the surfaces. In these situations you may want to passivate the reactive atoms on the surfaces e.g. with H. Regarding C atoms with 5 C neighbours in the bulk, it may be acceptable to simply assign the `CD` atom type to them as long as you ensure that the bulk's elastic response is not affected by this. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "36d067fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "from ase.io import write, read\n",
    "import numpy as np\n",
    "from matscipy.neighbours import neighbour_list\n",
    "from ase import Atoms\n",
    "\n",
    "# read initial structure and initialize neighbour list\n",
    "a = read('aC_H2O.xyz', index=0)\n",
    "\n",
    "types = [0] * len(a)\n",
    "\n",
    "# create neighbour list, note that the cutoffs here should be\n",
    "# consistent with the ones in the parameter file.\n",
    "i,  j = neighbour_list('ij', a, {('C', 'C'): 1.85, ('C', 'H'): 1.15,\n",
    "                                 ('C', 'O'): 1.55, ('O', 'H'): 1.3})\n",
    "\n",
    "# define surface region to\n",
    "# assign CA atom types to sp2-C surface atoms in the contact region\n",
    "h_surface_region = 2.0  # arbitrarily defined depth of surface region\n",
    "zmid_C = 65  # z coordinate separating upper and lower slab\n",
    "\n",
    "mask_C = a.numbers == 6\n",
    "\n",
    "mask_C_lower_slab = np.logical_and(mask_C, a.get_positions()[:, 2] < zmid_C)\n",
    "mask_C_upper_slab = np.logical_and(mask_C, a.get_positions()[:, 2] > zmid_C)\n",
    "\n",
    "zmax_C_lower_slab = np.max(a[mask_C_lower_slab].get_positions()[:, 2])\n",
    "zmin_C_upper_slab = np.min(a[mask_C_upper_slab].get_positions()[:, 2])\n",
    "\n",
    "CA_region_lower_slab = (zmax_C_lower_slab - h_surface_region,\n",
    "                        zmax_C_lower_slab)\n",
    "CA_region_upper_slab = (zmin_C_upper_slab,\n",
    "                        zmin_C_upper_slab + h_surface_region)\n",
    "\n",
    "def atom_in_CA_region(atom):\n",
    "    # returns True if the atom is in the surface region\n",
    "    zpos = atom.position[2]\n",
    "    return ((CA_region_lower_slab[0] < zpos <= CA_region_lower_slab[1]) or\n",
    "            (CA_region_upper_slab[0] <= zpos < CA_region_upper_slab[1]))\n",
    "\n",
    "# assign C atom types\n",
    "for k, atom in enumerate(a):\n",
    "    neighs = j[i == k]  # all neighbor indices of atom k\n",
    "\n",
    "    if atom.symbol == 'C':\n",
    "        assert len(neighs) <= 4\n",
    "        # count neighbors by element\n",
    "        nC = 0\n",
    "        nO = 0\n",
    "        nH = 0\n",
    "        for neigh in neighs:\n",
    "            if a[neigh].symbol == 'C':\n",
    "                nC += 1\n",
    "            elif a[neigh].symbol == 'O':\n",
    "                nO += 1\n",
    "            elif a[neigh].symbol == 'H':\n",
    "                nH += 1\n",
    "            else:\n",
    "                raise NotImplementedError('unknown atom type.')\n",
    "\n",
    "        if len(neighs) == 4:  # sp3-C atom types\n",
    "            assert nO <= 1  # only allow at most 1 OH group per C\n",
    "            if nC == 4:\n",
    "                types[k] = 'CD'\n",
    "            elif nC == 3:\n",
    "                if nO == 1:\n",
    "                    types[k] = 'C2'\n",
    "                elif nH == 1:\n",
    "                    types[k] = 'C1'\n",
    "            elif nC == 2:\n",
    "                assert nO + nH == 2\n",
    "                if nO == 1:\n",
    "                    assert nH == 1\n",
    "                    types[k] = 'C4'\n",
    "                elif nO == 0:\n",
    "                    types[k] = 'C3'\n",
    "            elif nC == 1:\n",
    "                assert nO + nH == 3\n",
    "                if nO == 1:\n",
    "                    assert nH == 2\n",
    "                    types[k] = 'C6'\n",
    "                elif nO == 0:\n",
    "                    types[k] = 'C5'\n",
    "            else:\n",
    "                raise NotImplementedError('unknown atom type.')\n",
    "        elif len(neighs) == 3:  # sp2-C atom types\n",
    "            assert nC + nO + nH == 3\n",
    "            if nC == 3:\n",
    "                N_sp3neigh = 0  # count number of sp3 neighbors\n",
    "                for neigh in neighs:\n",
    "                    neighs2 = j[i == neigh]\n",
    "                    if len(neighs2) == 4:\n",
    "                        N_sp3neigh += 1\n",
    "                if atom_in_CA_region(atom):\n",
    "                    types[k] = 'CA'\n",
    "                else:\n",
    "                    if N_sp3neigh == 0:\n",
    "                        types[k] = 'CG'  # bulk sp2 C with LJ interaction\n",
    "                    else:\n",
    "                        types[k] = 'CB'  # bulk sp2 C, no LJ interaction\n",
    "\n",
    "            elif nC == 2:\n",
    "                if nO == 1:\n",
    "                    types[k] = 'CY'\n",
    "                elif nH == 1:\n",
    "                    types[k] = 'CZ'\n",
    "                else:\n",
    "                    raise NotImplementedError('unknown atom type.')\n",
    "\n",
    "            elif nC == 1:\n",
    "                assert nH == 2\n",
    "                types[k] = 'CX'\n",
    "                # only allow bonds to CA, CB, CG, CY, CZ with 3 bonds\n",
    "                for neigh in neighs:\n",
    "                    if a[neigh].symbol == 'C':\n",
    "                        neighs2 = j[i == neigh]\n",
    "                        if len(neighs2) != 3:\n",
    "                            raise NotImplementedError('CX bonds to CA, CB, CG,'\n",
    "                                                      ' CY, CZ with 3 bonds')\n",
    "\n",
    "        elif len(neighs) == 2: # sp atoms\n",
    "            assert nC == 2\n",
    "\n",
    "            N_sp3neigh = 0\n",
    "            for neigh in neighs:\n",
    "                neighs2 = j[i == neigh]\n",
    "                if len(neighs2) == 4:\n",
    "                    N_sp3neigh += 1\n",
    "\n",
    "            if atom_in_CA_region(atom):\n",
    "                print(k)\n",
    "                raise NotImplementedError('Sp-C atom in surface region')\n",
    "            else:\n",
    "                if N_sp3neigh == 0:\n",
    "                    types[k] = 'CG'  # bulk sp C with LJ interaction\n",
    "                else:\n",
    "                    types[k] = 'CB'  # bulk sp C, no LJ interaction\n",
    "\n",
    "# assign all O atom types\n",
    "for k, atom in enumerate(a):\n",
    "    neighs = j[i == k]\n",
    "\n",
    "    if atom.symbol == 'O':\n",
    "        assert len(neighs) == 2\n",
    "        nH = 0\n",
    "        for neigh in neighs:\n",
    "            sy = a[neigh].symbol\n",
    "            if sy == 'C':\n",
    "                if types[neigh] == 'CY':\n",
    "                    types[k] = 'O2'\n",
    "                elif types[neigh] in ['C2', 'C4', 'C6']:\n",
    "                    types[k] = 'O1'\n",
    "            elif sy == 'H':\n",
    "                nH += 1\n",
    "        if nH == 2:\n",
    "            types[k] = 'OW'\n",
    "\n",
    "# assign all H atom types\n",
    "for k, atom in enumerate(a):\n",
    "    neighs = j[i == k]\n",
    "\n",
    "    if atom.symbol == 'H':\n",
    "        assert len(neighs) == 1\n",
    "        for neigh in neighs:\n",
    "            if a[neigh].symbol == 'C':\n",
    "                if types[neigh] in ['CZ', 'CX']:\n",
    "                    types[k] = 'H2'\n",
    "                elif types[neigh] in ['C1', 'C3', 'C4', 'C5', 'C6']:\n",
    "                    types[k] = 'H1'\n",
    "            elif a[neigh].symbol == 'O':\n",
    "                if types[neigh] == 'O1':\n",
    "                    types[k] = 'H4'\n",
    "                elif types[neigh] == 'O2':\n",
    "                    types[k] = 'H5'\n",
    "                elif types[neigh] == 'OW':\n",
    "                    types[k] = 'HW'\n",
    "            else:\n",
    "                raise NotImplementedError('unknown atom type.')\n",
    "\n",
    "a.set_array('type', np.array(types))\n",
    "\n",
    "# assign molecule ids\n",
    "molid1 = np.ones(np.sum(a.get_positions()[:, 2] < 63))  # id=1 for lower slab\n",
    "molid2 = np.ones(np.sum(a.get_positions()[:, 2] > 67)) + 1  # id=2 for upper slab\n",
    "molid_H2O = np.arange(180) // 3 + 3  # H2O ids from 3 to 63\n",
    "molid = np.append(molid1, molid2)\n",
    "molid = np.append(molid, molid_H2O)\n",
    "a.set_array('molid', molid.astype(int))\n",
    "\n",
    "# save structure to .extxyz file\n",
    "write('struct.extxyz', a)\n",
    "\n",
    "# check for missed atoms\n",
    "for k, t in enumerate(types):\n",
    "    if t == 0:\n",
    "        neighs = j[i == k]\n",
    "        print(k, a[k].symbol, [(neigh, types[neigh]) for neigh in neighs])\n",
    "assert 0 not in types"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "64605110",
   "metadata": {},
   "source": [
    "## Generation of the bonding topology and LAMMPS input files\n",
    "\n",
    "In the next step, we use `matscipy`'s routines to generate the bonding topology and create LAMMPS input files based on the parameters specified in `parameters_Reichenbach2023.in`.\n",
    "This script generates three files:\n",
    "- `struct.lammps.atoms`, which contains the atomic structure and the bonding topology in a LAMMPS-readable way.\n",
    "- `struct.lammps.opls`, which contains the pair_styles and force field parameters in a LAMMPS-readable way.\n",
    "- `struct.lammps.in`, which is an example LAMMPS input script how to run a geometry optimisation.\n",
    "Note that the script gives warnings of missing Dihedrals that are not part of the force field described in [Reichenbach et al.](https://doi.org/10.48550/arXiv.2310.12619) (e.g., C-C-C-C Dihedrals)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "af6aeaef",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "import matscipy\n",
    "import matscipy.opls\n",
    "import matscipy.io.opls\n",
    "\n",
    "parameter_file = 'parameters_Reichenbach2023.in'\n",
    "\n",
    "s = matscipy.io.opls.read_extended_xyz('struct.extxyz')\n",
    "\n",
    "parameter_data = matscipy.io.opls.read_parameter_file(parameter_file)\n",
    "cutoffs, atom_data, bond_data, angle_data, dihed_data = parameter_data\n",
    "\n",
    "s.set_cutoffs(cutoffs)\n",
    "s.set_atom_data(atom_data)\n",
    "\n",
    "assert abs(sum(s.get_charges())) < 0.01  # ensure charge neutrality\n",
    "\n",
    "s.get_bonds(bond_data)\n",
    "s.get_angles(angle_data)\n",
    "s.get_dihedrals(dihed_data)\n",
    "\n",
    "matscipy.io.opls.write_lammps('struct.lammps', s);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c431f15",
   "metadata": {},
   "source": [
    "To remove all files created during this tutorial, you can use the following code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "f2cc9cca",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "for tmp_file in ['struct.extxyz', 'struct.lammps.atoms',\n",
    "                 'struct.lammps.opls', 'struct.lammps.in']:\n",
    "    try:\n",
    "        os.remove(tmp_file)\n",
    "    except FileNotFoundError:\n",
    "        continue"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}