{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "7bbccb87",
   "metadata": {},
   "outputs": [],
   "source": [
    "# importing and preparing some variables\n",
    "import sys\n",
    "import os\n",
    "import glob\n",
    "import re\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",
    "out_fd = os.path.abspath(os.path.join(os.getcwd(), 'PTP_out_data'))\n",
    "if not os.path.isdir(out_fd):os.mkdir(out_fd)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d735109c",
   "metadata": {},
   "source": [
    "## Perturbation topologies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c2a1250a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from SMArt.md import parse_top\n",
    "from SMArt import alchemy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "fdf34469",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "            top        \n",
       "              0       1\n",
       "atom 0     6 CA    6 CA\n",
       "     1     7 CB    7 CB\n",
       "     2     8 CG    8 CG\n",
       "     3     9 CD    9 CD\n",
       "     4     15 C    15 C\n",
       "     5      4 N     4 N\n",
       "     6    10 CE   10 CE\n",
       "     7    11 NZ   11 NZ\n",
       "     8   17 NTE  17 NTE\n",
       "     9    2 CN1   2 CN1\n",
       "     10    16 O    16 O\n",
       "     11     5 H     5 H\n",
       "     12  12 HZ1  14 CH3\n",
       "     13  14 HZ3  12 CH1\n",
       "     14  13 HZ2  13 CH2\n",
       "     15   19 H2   19 H2\n",
       "     16   18 H1   18 H1\n",
       "     17   3 ON2   3 ON2\n",
       "     18   1 CN2   1 CN2"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# GROMACS\n",
    "# parse the topologies\n",
    "top_wt_file = os.path.join(data_fd, 'gromacs', 'LYSH.top')\n",
    "top_wt = parse_top(top_wt_file, format_type='gm')\n",
    "top_PTM_file = os.path.join(data_fd, 'gromacs', 'K3C.top')\n",
    "top_PTM = parse_top(top_PTM_file, format_type='gm')\n",
    "mt_wt, mt_PTP = top_wt.molecule_types['mol_1'], top_PTM.molecule_types['mol_1']\n",
    "\n",
    "# run the MCS algorithm\n",
    "mcs = alchemy.point_mutation(mt_wt, mt_PTP)\n",
    "mcs.solutions[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "52a87131-4191-43ba-8361-9365e96da4fe",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "            top        \n",
       "              0       1\n",
       "atom 0     6 CA    6 CA\n",
       "     1     7 CB    7 CB\n",
       "     2     8 CG    8 CG\n",
       "     3     9 CD    9 CD\n",
       "     4     15 C    15 C\n",
       "     5      4 N     4 N\n",
       "     6    10 CE   10 CE\n",
       "     7    11 NZ   11 NZ\n",
       "     8   17 NTE  17 NTE\n",
       "     9    2 CN1   2 CN1\n",
       "     10    16 O    16 O\n",
       "     11     5 H     5 H\n",
       "     12  12 HZ1  14 CH3\n",
       "     13  14 HZ3  13 CH2\n",
       "     14  13 HZ2  12 CH1\n",
       "     15   19 H2   19 H2\n",
       "     16   18 H1   18 H1\n",
       "     17   3 ON2   3 ON2\n",
       "     18   1 CN2   1 CN2"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mcs.solutions[1] # an alternative solution (6 effectively identical solutions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "41698b52-ac79-4f1e-a87c-a455e858c2ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "sol = alchemy.generate_2state_top(mcs, ff_dumm = top_wt.get_DUM_type) # solution = 0 by default\n",
    "out_itp = os.path.join(out_fd, 'toptp.itp')\n",
    "sol.toptp.write_itp(out_itp, flag_generate_excl_pairs = True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "287780b0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "            top        \n",
       "              0       1\n",
       "atom 0     6 CA    6 CA\n",
       "     1     7 CB    7 CB\n",
       "     2     8 CG    8 CG\n",
       "     3     9 CD    9 CD\n",
       "     4      4 N     4 N\n",
       "     5     15 C    15 C\n",
       "     6    10 CE   10 CE\n",
       "     7    11 NZ   11 NZ\n",
       "     8    2 CN1   2 CN1\n",
       "     9   17 NTE  17 NTE\n",
       "     10     5 H     5 H\n",
       "     11    16 O    16 O\n",
       "     12  13 HZ2  14 CH3\n",
       "     13  14 HZ3  13 CH2\n",
       "     14  12 HZ1  12 CH1\n",
       "     15   3 ON2   3 ON2\n",
       "     16   1 CN2   1 CN2\n",
       "     17   18 H1   19 H2\n",
       "     18   19 H2   18 H1"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# GROMOS\n",
    "# parse topologies\n",
    "top_wt_file = os.path.join(data_fd, 'gromos', 'LYS_K3C_KAC', 'single_AA', 'LYS.top')\n",
    "top_wt = parse_top(top_wt_file)\n",
    "top_PTM_file = os.path.join(data_fd, 'gromos', 'LYS_K3C_KAC', 'single_AA', 'K3C.top')\n",
    "top_PTM = parse_top(top_PTM_file)\n",
    "\n",
    "# run the MCS algorithm\n",
    "mcs = alchemy.point_mutation(top_wt, top_PTM)\n",
    "sol = alchemy.generate_2state_top(mcs) # solution = 0 by default\n",
    "sol.toptp.write_top(os.path.join(out_fd, 'toptp.top'))\n",
    "sol.toptp.write_ptp(os.path.join(out_fd, 'toptp.ptp'))\n",
    "\n",
    "#check_names(sol, at_map)\n",
    "sol"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5033a771-123a-4faa-b9f9-6374ee18f254",
   "metadata": {},
   "source": [
    "## Alternative MCS search options"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd964f6b-e627-4334-a65e-35cb44b43780",
   "metadata": {},
   "source": [
    "### no bond perturbations allowed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "25251f56-e8ff-4500-87ed-ee4604a8be56",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "            top        \n",
       "              0       1\n",
       "atom 0     8 CG    8 CG\n",
       "     1     9 CD    9 CD\n",
       "     2    10 CE   10 CE\n",
       "     3    11 NZ   11 NZ\n",
       "     4     7 CB    7 CB\n",
       "     5   13 HZ2     DUM\n",
       "     6   12 HZ1     DUM\n",
       "     7   14 HZ3     DUM\n",
       "     8     6 CA    6 CA\n",
       "     9      DUM  12 CH1\n",
       "     10     DUM  14 CH3\n",
       "     11     DUM  13 CH2"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# for simplicity, let's reduce the topologies to the sidechains\n",
    "top_wt = top_wt.reduce_top([str(i) for i in range(6,15)])\n",
    "top_PTM = top_PTM.reduce_top([str(i) for i in range(6,15)])\n",
    "\n",
    "\n",
    "mcs = alchemy.point_mutation(top_wt, top_PTM, flag_top_prune = 'bond') # makes sure no bonds are perturbed\n",
    "mcs.solutions[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c23cc69d-0858-4fa8-a015-27623037bed9",
   "metadata": {},
   "source": [
    "### add RMSD"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "2c7845de-a970-4dfb-b527-3090de6a75c5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "           top        \n",
       "             0       1\n",
       "atom 0    8 CG    8 CG\n",
       "     1    9 CD    9 CD\n",
       "     2   10 CE   10 CE\n",
       "     3   11 NZ   11 NZ\n",
       "     4    7 CB    7 CB\n",
       "     5  13 HZ2  13 CH2\n",
       "     6  12 HZ1  12 CH1\n",
       "     7  14 HZ3  14 CH3\n",
       "     8    6 CA    6 CA"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import numpy as np\n",
    "coord = np.random.rand(len(top_wt.atoms) * 3).reshape(len(top_wt.atoms), 3) # random coordinates\n",
    "coord2 = coord.copy()\n",
    "\n",
    "mcs = alchemy.point_mutation(top_wt, top_PTM, add_RMSD='simple', coords = [coord, coord2])\n",
    "mcs.solutions[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "d08cd464-554e-4457-a0d1-55c6d35bc52b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[0.40479395 0.89759386 0.03049702]\n",
      " [0.15758507 0.15801852 0.00853602]]\n",
      "after swapping\n",
      "[[0.15758507 0.15801852 0.00853602]\n",
      " [0.40479395 0.89759386 0.03049702]]\n"
     ]
    }
   ],
   "source": [
    "# let's swapp the the coordinates of last 2 atoms\n",
    "print(coord2[-2:])\n",
    "coord2[-2:] = coord2[-2:][::-1]\n",
    "print('after swapping')\n",
    "print(coord2[-2:])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "5ce6aaf1-127f-4342-8e10-8cc7c4e55242",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "           top        \n",
       "             0       1\n",
       "atom 0    8 CG    8 CG\n",
       "     1    9 CD    9 CD\n",
       "     2   10 CE   10 CE\n",
       "     3   11 NZ   11 NZ\n",
       "     4    7 CB    7 CB\n",
       "     5  13 HZ2  14 CH3\n",
       "     6  12 HZ1  12 CH1\n",
       "     7  14 HZ3  13 CH2\n",
       "     8    6 CA    6 CA"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mcs = alchemy.point_mutation(top_wt, top_PTM, add_RMSD='simple', coords = [coord, coord2])\n",
    "mcs.solutions[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55ed8302-cdb4-4b6c-a39b-7a71cdd2c174",
   "metadata": {},
   "source": [
    "other options include:\n",
    "\n",
    "- flag_partial_ring - defines if partial ring match is allowed or not (only for fused rings) - True by default\n",
    "- max_partial_ring_match - defines how many ring to non-ring atoms are allowed to be matched (2 by default)\n",
    "- dihedral_match_v - procedure to match proper dihedrals - 1 (default) based on the middle 2 atoms of a dihedral, or 2 match based on all atoms\n",
    "- use_ptp_flags - defines if the multiplicity of proper dihedrals is allowed to be perturbed (for GROMACS)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a75d6545",
   "metadata": {},
   "source": [
    "## EDS topology"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "06d93e52-34d7-45a3-9e6b-42a429188447",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "            top                                \n",
       "              0       1       2       3       4\n",
       "atom 0     3 C4    3 C4    3 C4    3 C4    3 C4\n",
       "     1   20 C25  20 C25  20 C25  20 C25  20 C25\n",
       "     2     1 C1    1 C1    1 C1    1 C1    1 C1\n",
       "     3    9 C10   9 C10   9 C10   9 C10   9 C10\n",
       "     4     5 C6    5 C6    5 C6    5 C6    5 C6\n",
       "     5     7 C8    7 C8    7 C8    7 C8    7 C8\n",
       "     6   17 S22  17 S22  17 S22  17 S22  17 S22\n",
       "     7   15 N20  15 N20  15 N20  15 N20  15 N20\n",
       "     8   10 N11  10 N11  10 N11  10 N11  10 N11\n",
       "     9   12 C13  12 C13  12 C13  12 C13  12 C13\n",
       "     10  11 C12  11 C12  11 C12  11 C12  11 C12\n",
       "     11  13 C14  14 C19  13 C14  13 C14  14 C19\n",
       "     12    4 H5    4 H5     DUM    4 H5    4 H5\n",
       "     13    2 H2     DUM    2 H2    2 H2    2 H2\n",
       "     14    6 H7    6 H7    6 H7     DUM    6 H7\n",
       "     15    8 H9    8 H9    8 H9    8 H9     DUM\n",
       "     16  19 O24  18 O23  18 O23  18 O23  19 O24\n",
       "     17  18 O23  19 O24  19 O24  19 O24  18 O23\n",
       "     18  16 H21  16 H21  16 H21  16 H21  16 H21\n",
       "     19  14 C19  13 C14  14 C19  14 C19  13 C14\n",
       "     20     DUM    2 F2     DUM     DUM     DUM\n",
       "     21     DUM     DUM    4 F5     DUM     DUM\n",
       "     22     DUM     DUM     DUM    6 F7     DUM\n",
       "     23     DUM     DUM     DUM     DUM    8 F9"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "TOPs_path = os.path.join(data_fd, 'gromos', 'GRA2', 'tops', '*top')\n",
    "TOPs = glob.glob(TOPs_path)\n",
    "\n",
    "tops = [parse_top(fpath) for fpath in TOPs]\n",
    "\n",
    "mcs = alchemy.get_EDS(*tops[:5]) # first 5 tops; faster than all 16 topologies (line below)\n",
    "#mcs = alchemy.get_EDS(*tops)\n",
    "\n",
    "sol_EDS, state_names = alchemy.generate_EDS_top(mcs)\n",
    "sol_EDS.toptp.write_top(os.path.join(out_fd, 'EDS.top'))\n",
    "sol_EDS.toptp.write_EDS(os.path.join(out_fd, 'EDS.ptp'), state_names=state_names)\n",
    "sol_EDS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c611d8c-7859-48cb-a3ca-4a6f929540aa",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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": 5
}