{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 09. Computing coarse-grained molecular features\n",
    "\n",
    "This notebook shows how to compute pairwise ditances, angles and dihedrals between CG beads given a CG mapping. The CG mapping used in this example is generated from [DSGPM](https://github.com/rochesterxugroup/DSGPM) model. \n",
    "You must need MDAnalysis and NetworkX in your working environment to run this example. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2 Physical GPUs, 2 Logical GPUs\n"
     ]
    }
   ],
   "source": [
    "import hoomd\n",
    "import hoomd.htf as htf\n",
    "import tensorflow as tf\n",
    "import MDAnalysis as mda\n",
    "import numpy as np\n",
    "\n",
    "import os\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = '-1'\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this example we read from a trajectory file to compute coarse-grained (CG) bond distances, angles and dihedrals. Here, we use two MDAnalysis universes with and without hydrogens as the mapping model used to compute the CG mappings (DSGPM model) only maps heavy atoms of a given molecule. Hence, we have to add the missing hydrogen atoms to the corresponding CG beads. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Read frames from the trajectory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "universe = mda.Universe('ex9_segA.pdb','ex9_segA.trr')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [

      "100%|██████████| 801/801 [01:30<00:00,  8.83it/s] \n"
     ]
    }
   ],
   "source": [
    "avg_cgr = tf.keras.metrics.MeanTensor()\n",
    "avg_cga = tf.keras.metrics.MeanTensor()\n",
    "avg_cgd = tf.keras.metrics.MeanTensor()\n",
    "directory = os.getcwd()\n",
    "jfile = os.path.join(directory,'ex9_cgmap_segA.json')\n",
    "\n",
    "#mda universe without H's\n",
    "uxh = mda.Universe(os.path.join(directory,'ex9_segA_xH.pdb')) \n",
    "#mda universe with H's\n",
    "uh = mda.Universe(os.path.join(directory,'ex9_segA.pdb')) \n",
    "\n",
    "for inputs, ts in htf.iter_from_trajectory(16, universe, r_cut=10, start=400, end=700):\n",
    "    cg_fts = []\n",
    "    r_tensor = []\n",
    "    a_tensor = []\n",
    "    d_tensor = []\n",
    "    \n",
    "    box = inputs[2]\n",
    "    \n",
    "    #get CG bead indices that make bonds, angles, dihedrals and\n",
    "    #CG coordinates\n",
    "    cg_fts = htf.compute_cg_graph(DSGPM=True,infile=jfile,group_atoms=True,\n",
    "                                  u_no_H=uxh, u_H=uh)\n",
    "\n",
    "    for i in range(len(cg_fts[0])):\n",
    "        cg_r = htf.mol_bond_distance(CG = True, cg_positions = cg_fts[-1],\n",
    "                                     b1=cg_fts[0][i][0],b2=cg_fts[0][i][1],box=box)\n",
    "        r_tensor.append(cg_r)\n",
    "\n",
    "    for j in range(len(cg_fts[1])):       \n",
    "        cg_a = htf.mol_angle(CG= True, cg_positions=cg_fts[-1],\n",
    "                             b1=cg_fts[1][j][0],b2=cg_fts[1][j][1],b3=cg_fts[1][j][2],box=box)\n",
    "        a_tensor.append(cg_a)\n",
    "\n",
    "\n",
    "    for k in range(len(cg_fts[2])):\n",
    "        cg_d = htf.mol_dihedral(CG=True,cg_positions=cg_fts[-1],b1=cg_fts[2][k][0],\n",
    "                                b2=cg_fts[2][k][1],b3=cg_fts[2][k][2],b4=cg_fts[2][k][3],box=box)\n",
    "        d_tensor.append(cg_d)\n",
    "\n",
    "    avg_cgr.update_state(r_tensor)\n",
    "    avg_cga.update_state(a_tensor)\n",
    "    avg_cgd.update_state(d_tensor)\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CG pairwise distances: [ 5.4447026  1.1312671  6.8375177  2.9382892  2.4656532  4.4416947\n",
      "  3.199694   4.2150507  3.5845404  2.153627   7.9029765  3.8829455\n",
      "  6.7589035  6.4774413  2.255304   4.924929  15.143286 ] \n",
      "\n",
      "CG angles: [ 57.06865   75.22357   83.657074 113.90926   30.8918    61.174572\n",
      "  40.556293  27.594091  50.535973 149.74725   46.7441    91.21376\n",
      "  44.42922  157.15317   45.61479  121.53312  140.93109   90.67879\n",
      "  51.733078 156.72841 ] \n",
      "\n",
      "CG dihedral angles: [ 61.196575  177.25443     4.7860584 111.41965   176.07312   133.15497\n",
      "  84.99461   135.7665    147.13869     4.834345  168.7402    124.28182\n",
      " 175.61597    21.146255  163.78894    32.634514    9.021241  175.17809\n",
      "  10.565324    7.1954145]\n"
     ]
    }
   ],
   "source": [
    "cgR = avg_cgr.result().numpy()\n",
    "cgD = avg_cgd.result().numpy()*180./np.pi\n",
    "cgA = avg_cga.result().numpy()*180./np.pi\n",
    "print('CG pairwise distances:',cgR,'\\n')\n",
    "print('CG angles:',cgA,'\\n')\n",
    "print('CG dihedral angles:',cgD)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Application to multiple molecules\n",
    "\n",
    "Note that the above computation is only applied to one molecule in the system. If a user has multiple copies of a molecule, the calculation of indices of CG beads making bonds, angles and dihedrals must be applied to all molecules. \n",
    "\n",
    "Let's assume there are 2 of the above molecules are available in the system. Each molecule has 18 CG beads. We can obtain the indices as follows. Here we use the outputs from `compute_cg_graph` to both molecules."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "angles in molecule 1:  [[ 0  1  2]\n",
      " [ 1  2  3]\n",
      " [ 2  3  5]\n",
      " [ 3  5  4]\n",
      " [ 3  5  6]\n",
      " [ 4  5  6]\n",
      " [ 5  6  7]\n",
      " [ 5  6  8]\n",
      " [ 6  8 10]\n",
      " [ 7  6  8]\n",
      " [ 8 10  9]\n",
      " [ 8 10 11]\n",
      " [ 9 10 11]\n",
      " [10 11 12]\n",
      " [11 12 13]\n",
      " [12 13 15]\n",
      " [13 15 14]\n",
      " [13 15 16]\n",
      " [14 15 16]\n",
      " [15 16 17]]\n",
      "\n",
      " angles in molecule 2: [[18 19 20]\n",
      " [19 20 21]\n",
      " [20 21 23]\n",
      " [21 23 22]\n",
      " [21 23 24]\n",
      " [22 23 24]\n",
      " [23 24 25]\n",
      " [23 24 26]\n",
      " [24 26 28]\n",
      " [25 24 26]\n",
      " [26 28 27]\n",
      " [26 28 29]\n",
      " [27 28 29]\n",
      " [28 29 30]\n",
      " [29 30 31]\n",
      " [30 31 33]\n",
      " [31 33 32]\n",
      " [31 33 34]\n",
      " [32 33 34]\n",
      " [33 34 35]]\n"
     ]
    }
   ],
   "source": [
    "r_ids,a_ids,d_ids = htf.mol_features_multiple(bnd_indices=cg_fts[0],ang_indices=cg_fts[1],\n",
    "                                              dih_indices=cg_fts[2],molecules=2,beads=18)\n",
    "\n",
    "# For example here are the CG bead indices involved in making angles\n",
    "print('angles in molecule 1: ',a_ids[:20])\n",
    "print('\\n angles in molecule 2:',a_ids[20:])\n",
    "\n",
    "# Now the same calculation with mol_bond_distance,mol_angle and mol_dihedral can be used  \n",
    "# to calculate CG bond distances, angles and dihedrals "
   ]
  }
 ],
 "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}