{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# SMArt\n",
    "## Welcome to Simulation & Modelling Art\n",
    "### importing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import os\n",
    "\"\"\"\n",
    "SMArt_fd = os.path.abspath(os.getcwd())\n",
    "for _ in range(2):\n",
    "    SMArt_fd = os.path.split(smart_fd)[0]\n",
    "sys.path.insert(0, smart_fd)\n",
    "#\"\"\"\n",
    "# if SMArt not in the path, uncomment the code above\n",
    "\n",
    "import SMArt\n",
    "data_fd = os.path.abspath(os.path.join(os.path.split(SMArt.__file__)[0], '..', 'doc', 'some_data'))\n",
    "\n",
    "from SMArt.md import parse_top, parse_ff # functions to parse different files (e.g. topology or force field)\n",
    "from SMArt.md import FF, Topology, Configuration # examples of some of the most important classes\n",
    "\n",
    "#help(SMArt) # to get some general info on the package"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## parsing / writing"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### GROMOS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# just defining some paths to the data\n",
    "out_fd = os.path.abspath(os.path.join(os.getcwd(), 'out_data'))\n",
    "if not os.path.isdir(out_fd):os.mkdir(out_fd)\n",
    "fd_gromos = os.path.join(data_fd, 'gromos')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Force field"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# parsing\n",
    "ifp_file = os.path.join(fd_gromos, '54a8.ifp')\n",
    "ff_gr = parse_ff(ifp_file)\n",
    "# same as\n",
    "#ff_gr = FF(ifp_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['sys_title',\n",
       " 'm_type',\n",
       " 'gr_bonds',\n",
       " 'gr_angles',\n",
       " 'gr_impropers',\n",
       " 'gr_dihedrals',\n",
       " 'a_type',\n",
       " 'vdw']"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# data of any object is stored in _containers\n",
    "ff_gr._containers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "OrderedDict([('1', 1 1.008),\n",
       "             ('3', 3 13.019),\n",
       "             ('4', 4 14.027),\n",
       "             ('5', 5 15.035),\n",
       "             ('6', 6 16.043),\n",
       "             ('12', 12 12.011),\n",
       "             ('14', 14 14.0067),\n",
       "             ('16', 16 15.9994),\n",
       "             ('19', 19 18.9984),\n",
       "             ('23', 23 22.9898),\n",
       "             ('24', 24 24.305),\n",
       "             ('28', 28 28.08),\n",
       "             ('31', 31 30.9738),\n",
       "             ('32', 32 32.06),\n",
       "             ('35', 35 35.453),\n",
       "             ('39', 39 39.948),\n",
       "             ('40', 40 40.08),\n",
       "             ('56', 56 55.847),\n",
       "             ('63', 63 63.546),\n",
       "             ('65', 65 65.37),\n",
       "             ('80', 80 79.904)])"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# mass types\n",
    "ff_gr.m_type"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1 1.008"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ff_gr.m_type['1'] # mass type 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2 OM"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ff_gr.a_type['2'] # atom type 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "BondType (15700000.0, 314000.0, 0.1)"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ff_gr.gr_bonds['1'] # bond type 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Topology"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "top_file = os.path.join(fd_gromos, 'villin', 'wt.top')\n",
    "#top_gr = parse_top(top_file)\n",
    "# same as\n",
    "top_gr = Topology(top_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['a_type', 'gr_bonds', 'gr_angles', 'gr_impropers', 'gr_dihedrals', 'vdw']"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "top_gr.ff._containers # the topology object contains the force field data as well"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['sys_title',\n",
       " 'ff',\n",
       " 'residues',\n",
       " 'cg',\n",
       " 'excl_pair',\n",
       " 'bonds',\n",
       " 'angles',\n",
       " 'impropers',\n",
       " 'dihedrals']"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "top_gr._containers # data stored in:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# write out FF ifp file\n",
    "out_file = os.path.join(out_fd, 'new_ff.ifp')\n",
    "ff_gr.write_ff(out_file, format_type = 'gr')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "# write out top file\n",
    "out_file = os.path.join(out_fd, 'new_top.top')\n",
    "top_gr.write_top(out_file)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### some operations with the topology object"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "389"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(top_gr.atoms)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[1 H1, 2 H2, 3 N, 4 H3, 5 CA, 6 CB, 7 CG, 8 SD, 9 CE, 10 C, 11 O]"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "top_gr.residues['1'].atoms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "# reduce topolgy to first 11 atoms (first residue)\n",
    "red_top = top_gr.reduce_top([str(i) for i in range(1,12)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "OrderedDict([('1', <SMArt.md.data_st.Residue at 0x7fdaa40b1fd0>)])"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "red_top.residues"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "OrderedDict([('1', 1 H1),\n",
       "             ('2', 2 H2),\n",
       "             ('3', 3 N),\n",
       "             ('4', 4 H3),\n",
       "             ('5', 5 CA),\n",
       "             ('6', 6 CB),\n",
       "             ('7', 7 CG),\n",
       "             ('8', 8 SD),\n",
       "             ('9', 9 CE),\n",
       "             ('10', 10 C),\n",
       "             ('11', 11 O)])"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "red_top.atoms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "out_file = os.path.join(out_fd, 'new_red_top.top')\n",
    "red_top.write_top(out_file)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### GROMACS"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Force field"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "ff_gm = parse_ff('./gromos54a7.ff/forcefield.itp', format_type = 'gm') # default for format_type is 'gr' for GROMOS\n",
    "# note that the force-field files can be automatically found from the gmx (if installed) - otherwise, provide the full path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[('_FF_GROMOS96', []),\n",
       " ('_FF_GROMOS54A7', []),\n",
       " ('gb_1', ['0.1000', '1.5700e+07']),\n",
       " ('gb_2', ['0.1000', '1.8700e+07']),\n",
       " ('gb_3', ['0.1090', '1.2300e+07']),\n",
       " ('gb_4', ['0.112', '3.7000e+07']),\n",
       " ('gb_5', ['0.1230', '1.6600e+07']),\n",
       " ('gb_6', ['0.1250', '1.3400e+07']),\n",
       " ('gb_7', ['0.1320', '1.2000e+07']),\n",
       " ('gb_8', ['0.1330', '8.8700e+06'])]"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "list(ff_gm.defines.items())[:10] # list of definitions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['_segments',\n",
       " 'sys_title',\n",
       " 'a_type',\n",
       " 'gm_vdw_normal',\n",
       " 'gm_vdw_pairs',\n",
       " 'gm_constraints',\n",
       " 'gm_bonds',\n",
       " 'gm_angles',\n",
       " 'gm_dihedrals']"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ff_gm._containers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[BondType (0.204, 5030000.0), BondType (0.198, 640000.0)]"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ff_gm.gm_bonds # bonde types"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Topology"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "fd_gromacs = fd_gromos = os.path.join(data_fd, 'gromacs')\n",
    "top_file = os.path.join(fd_gromacs, 'gromos', 'topol.top') # GROMACS topology using the GROMOS force field\n",
    "top_gm = parse_top(top_file, format_type = 'gm')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['_segments', 'sys_title', 'ff', 'molecule_types', 'molecules']"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "top_gm._containers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Protein_chain_A Protein_chain_A 1]"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "top_gm.molecules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "OrderedDict([('Protein_chain_A', Protein_chain_A Protein_chain_A),\n",
       "             ('SOL', SOL SOL),\n",
       "             ('CU1', CU1 CU1),\n",
       "             ('CU', CU CU),\n",
       "             ('ZN', ZN ZN),\n",
       "             ('MG', MG MG),\n",
       "             ('CA', CA CA),\n",
       "             ('NA', NA NA),\n",
       "             ('CL', CL CL)])"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "top_gm.molecule_types"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['residues',\n",
       " 'cg',\n",
       " 'bonds',\n",
       " 'vdw_pairs',\n",
       " 'angles',\n",
       " 'dihedrals',\n",
       " 'impropers',\n",
       " 'excl_pair']"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mol_type = top_gm.molecule_types['Protein_chain_A']\n",
    "mol_type._containers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.chdir(out_fd)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "ff_gm.write_ff('gmx_ff.itp', from_str = 1, flag_segment_order=1, flag_split_non_bonded = 1, flag_if=1, format_type = 'gm')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "top_gm.write_top('gm_top.top', format_type = 'gm')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## g2g (converting between GROMOS and GROMACS formats)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "# gromos to gromacs\n",
    "ff_gr.gr2gm()\n",
    "top_gr.gr2gm()\n",
    "# gromacs to gromos\n",
    "ff_gm.gm2gr()\n",
    "top_gm.gm2gr()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "# write out into files\n",
    "ff_gr.write_ff('gr2gm_ff.itp', format_type = 'gm')\n",
    "top_gr.write_top('gr2gm_top.top', from_str = 1, flag_segment_order=1, flag_if=1, flag_include = 1, sep_ff2itp='top_ff.itp', flag_defines_first_include = 1, sep_mol2itp=1, format_type='gm')\n",
    "\n",
    "ff_gm.write_ff('gm2gr_ff.ifp', format_type = 'gr')\n",
    "top_gm.write_top('gm2gr_top.top', format_type = 'gr')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# some additional features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "([177 CD1, 181 CE1, 185 CZ, 183 CE2, 179 CD2, 176 CG],\n",
       " [242 CE3,\n",
       "  246 CZ3,\n",
       "  248 CH2,\n",
       "  244 CZ2,\n",
       "  241 CE2,\n",
       "  238 CD2,\n",
       "  236 CD1,\n",
       "  239 NE1,\n",
       "  235 CG],\n",
       " [377 CD1, 381 CE1, 385 CZ, 383 CE2, 379 CD2, 376 CG],\n",
       " [216 CA, 217 CB, 218 CG, 219 CD, 215 N],\n",
       " [106 CD1, 110 CE1, 114 CZ, 112 CE2, 108 CD2, 105 CG],\n",
       " [62 CD1, 66 CE1, 70 CZ, 68 CE2, 64 CD2, 61 CG])"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# find rings\n",
    "mol_type.find_rings()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[12 N, 11 O, 7 CG]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[7 CG, 11 O, 12 N]"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "at = mol_type.atoms[1]\n",
    "print(list(mol_type.BFS_l(at, 3))) # 3rd neighbours\n",
    "el, pl = mol_type._generate_atom_excl_pair_list()\n",
    "pl[at] # list of 1-4 pairs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[3 H2, 5 CA, 2 H1, 4 H3]\n"
     ]
    }
   ],
   "source": [
    "print(list(mol_type.BFS_l(at, 1))) # first neighbours"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1 N, 3 H2, 5 CA, 2 H1, 4 H3, 10 C, 6 CB, 12 N, 11 O, 7 CG]\n",
      "[1 N, 2 H1, 3 H2, 4 H3, 5 CA, 6 CB, 10 C, 7 CG, 11 O, 12 N]\n"
     ]
    }
   ],
   "source": [
    "print(list(mol_type.BFS_d(at, 3))) # all up to 3rd neighbours (including itself)\n",
    "print([at] + el[at] + pl[at])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "#\"\"\"\n",
    "# remove all backed-up files\n",
    "import glob\n",
    "for f_path in glob.glob('#bck*'):\n",
    "    os.remove(f_path)\n",
    "#\"\"\""
   ]
  }
 ],
 "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}