MD.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import parmed\n",
    "from parmed.charmm import CharmmParameterSet\n",
    "from parmed.formats.registry import load_file\n",
    "\n",
    "%matplotlib inline\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "psf_file = 'inp_files/da.psf'\n",
    "pdb_file = 'inp_files/smd_ini.pdb'\n",
    "prm_file = 'inp_files/par_all27_prot_lipid_cmap.prm'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "HERE\n",
      "READING\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/shaunak/anaconda3/envs/MD/lib/python3.6/site-packages/parmed/charmm/parameters.py:754: ParameterWarning: Atom type C not present in AtomType list\n",
      "  warnings.warn('Atom type %s not present in AtomType list' % key, ParameterWarning)\n"
     ]
    }
   ],
   "source": [
    "prm = CharmmParameterSet.load_set(pfile = prm_file)\n",
    "psf = load_file(psf_file)\n",
    "pdb = load_file(pdb_file)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Energy Calculation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Energies:\n",
    "    def __init__(self, psf, pdb, prm):\n",
    "        self.psf = psf\n",
    "        self.pdb = pdb\n",
    "        self.prm = prm\n",
    "        \n",
    "        ## initializing bond types and coordinates\n",
    "        coords = pdb.coordinates\n",
    "        bonds = psf.bonds\n",
    "        angles = psf.angles\n",
    "        dihedrals = psf.dihedrals\n",
    "        impropers = psf.impropers\n",
    "        for bond in bonds:\n",
    "            ## bond types\n",
    "            tid1 = bond.atom1.type\n",
    "            tid2 = bond.atom2.type\n",
    "            \n",
    "            t = (tid1, tid2)\n",
    "            bond.type = prm.bond_types[t]\n",
    "    \n",
    "            ## coordinates\n",
    "            b1 = bond.atom1\n",
    "            b1.xx, b1.xy, b1.xz = coords[b1.idx]\n",
    "            \n",
    "            b2 = bond.atom2\n",
    "            b2.xx, b2.xy, b2.xz = coords[b2.idx]\n",
    "            \n",
    "            ## non bonded parameters eps and rmin\n",
    "            name = b1.type\n",
    "            eps, rmin, _, _ = prm.nonbonded_types[name]\n",
    "            b1.rmin = rmin\n",
    "            b1.epsilon = eps\n",
    "            \n",
    "            name = b2.type\n",
    "            eps, rmin, _, _ = prm.nonbonded_types[name]\n",
    "            b2.rmin = rmin\n",
    "            b2.epsilon = eps\n",
    "        \n",
    "            \n",
    "            \n",
    "        ## initializing angle types\n",
    "        for angle in angles:\n",
    "            a1 = angle.atom1.type\n",
    "            a2 = angle.atom2.type\n",
    "            a3 = angle.atom3.type\n",
    "            \n",
    "            t = (a1, a2, a3)\n",
    "            \n",
    "            angle.type = prm.angle_types[t]\n",
    "            a = angle.measure()\n",
    "            a0 = angle.type\n",
    "            \n",
    "        for i, dihedral in enumerate(dihedrals):\n",
    "            a1 = dihedral.atom1.type\n",
    "            a2 = dihedral.atom2.type\n",
    "            a3 = dihedral.atom3.type\n",
    "            a4 = dihedral.atom4.type\n",
    "            \n",
    "            t = (a1, a2, a3, a4)\n",
    "            g = self._get_bond_type(t, prm.dihedral_types)\n",
    "            dihedrals[i].type = prm.dihedral_types[g]\n",
    "            \n",
    "        for i, improper in enumerate(impropers):\n",
    "            a1 = improper.atom1.type\n",
    "            a2 = improper.atom2.type\n",
    "            a3 = improper.atom3.type\n",
    "            a4 = improper.atom4.type\n",
    "            \n",
    "            t = (a1, a2, a3, a4)\n",
    "            g = self._get_bond_type(t, prm.improper_types)\n",
    "            impropers[i].type = prm.improper_types[g]\n",
    "        \n",
    "    def bond(self):\n",
    "        bonds = psf.bonds\n",
    "        s = 0\n",
    "        for bond in bonds:\n",
    "            r1 = pdb.coordinates[bond.atom1.idx]\n",
    "            r2 = pdb.coordinates[bond.atom2.idx]\n",
    "            k = bond.type.k\n",
    "            dist = np.linalg.norm(r1 - r2)\n",
    "            \n",
    "            r0 = bond.type.req\n",
    "            cust_energy = k * ((dist - r0)**2)\n",
    "            s += cust_energy\n",
    "            \n",
    "        return s    \n",
    "    \n",
    "    def angle(self):\n",
    "        s = 0\n",
    "        \n",
    "        for angle in psf.angles:\n",
    "            theta = angle.measure()\n",
    "            theta0 = angle.type.theteq\n",
    "            k = angle.type.k\n",
    "            theta_diff = (np.pi * (theta - theta0)) / 180\n",
    "            en = k * (theta_diff ** 2)\n",
    "#             print(en, angle.energy())\n",
    "            s += en\n",
    "        return s\n",
    "    \n",
    "    ## DIHEDRAL AND IMPROPER HELPER FUNCTION\n",
    "    def _get_bond_type(self, t, search_list):\n",
    "        '''\n",
    "        Input :\n",
    "        t : a 4-atom tuple, \n",
    "        search_list: usually prm.dihedral_types or prm.improper_types\n",
    "\n",
    "        Return type: BondType\n",
    "        '''\n",
    "        for dtype in search_list:\n",
    "            val = True\n",
    "            for i in range(4):\n",
    "                val = val and (dtype[i] == t[i] or dtype[i] == 'X')\n",
    "\n",
    "            if val:\n",
    "                return dtype\n",
    "\n",
    "            ## Attempt reverse search\n",
    "            val = True\n",
    "            for i in range(4):\n",
    "                val = val and (dtype[i] == t[3 - i] or dtype[i] == 'X')\n",
    "\n",
    "            if val:\n",
    "                return dtype\n",
    "        return None\n",
    "    \n",
    "    def dihedral(self):\n",
    "        dihedrals = psf.dihedrals\n",
    "        k_sum = 0\n",
    "        for dihed in dihedrals:\n",
    "            dtype = dihed.type\n",
    "            \n",
    "            total_energy = 0\n",
    "            chi = dihed.measure() * (np.pi / 180)\n",
    "            if chi > np.pi or chi < -np.pi:\n",
    "                print(\"Houston, we have a problem\")\n",
    "#             print(dihed.measure())\n",
    "            if isinstance(dtype, parmed.DihedralType):\n",
    "                n = dtype.per\n",
    "                kchi = dtype.phi_k\n",
    "                phase = dtype.phase * (np.pi / 180)\n",
    "                en = kchi * (1 + np.cos(n * chi + phase))\n",
    "                total_energy += en\n",
    "                \n",
    "            elif isinstance(dtype, parmed.DihedralTypeList):\n",
    "                e = 0\n",
    "                for term in dtype:\n",
    "                    n = term.per\n",
    "                    kchi = term.phi_k\n",
    "                    phase = term.phase * (np.pi / 180)\n",
    "#                     print(n)\n",
    "#                     print(kchi, n, chi, phase)\n",
    "                    if n > 0:\n",
    "                        en = kchi * (1 + np.cos(n * chi - phase))\n",
    "#                         print(np.cos(n * chi - phase), n * chi - phase)\n",
    "                    else:\n",
    "                        en = kchi * (chi - phase)**2\n",
    "                    k_sum += kchi\n",
    "                    print(en)\n",
    "                    e += np.abs(en)\n",
    "                total_energy += np.abs(e)\n",
    "                \n",
    "            else:\n",
    "                printf(\"Cant determine parameters for {}\".format(t))\n",
    "                return 0\n",
    "        print(\"2k : \", 2 * k_sum)\n",
    "        return total_energy\n",
    "    \n",
    "    def improper(self):\n",
    "        impropers = psf.impropers\n",
    "        total_energy = 0\n",
    "        for imp in impropers:\n",
    "            psi = imp.measure() * (np.pi / 180)\n",
    "            dtype = imp.type\n",
    "            psi_0 = dtype.psi_eq * (np.pi / 180)\n",
    "            k = dtype.psi_k \n",
    "            en = k *(psi - psi_0)**2\n",
    "            \n",
    "            total_energy += en\n",
    "            \n",
    "        return total_energy\n",
    "    \n",
    "    # NON BONDED\n",
    "    def LJ(self):\n",
    "        energy = 0\n",
    "        coords = psf.coordinates\n",
    "        for i, a in enumerate(psf.atoms):\n",
    "            for j, b in enumerate(psf.atoms):\n",
    "                if j <= i: \n",
    "                    continue\n",
    "                eps_ij = np.sqrt(a.epsilon * b.epsilon)\n",
    "                r_ij = (a.rmin + b.rmin) / 2\n",
    "    \n",
    "                r1 = coords[a.idx]\n",
    "                r2 = coords[b.idx]\n",
    "                r = np.linalg.norm(r1 - r2)\n",
    "                ratio = r_ij / r\n",
    "                en = 4 * eps_ij * (ratio**12 - (ratio**6))\n",
    "                energy += en\n",
    "                \n",
    "        return energy\n",
    "    \n",
    "    def Elec(self):\n",
    "        energy = 0\n",
    "        coords = psf.coordinates\n",
    "        for i, a in enumerate(psf.atoms):\n",
    "            for j, b in enumerate(psf.atoms):\n",
    "                if j <= i: \n",
    "                    continue\n",
    "                k = 8.9 * 1e9\n",
    "                r1 = coords[a.idx]\n",
    "                r2 = coords[b.idx]\n",
    "                r = np.linalg.norm(r1 - r2)\n",
    "                q1 = a.charge\n",
    "                q2 = b.charge\n",
    "                en = k * q1 * q2 / (r * 1e-9)\n",
    "                energy += en\n",
    "        return energy\n",
    "    \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "h = Energies(psf, pdb, prm)\n",
    "# h.dihedral()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Force calculation : Resources used in `resources` folder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from parmed.geometry import angle\n",
    "\n",
    "class Forces:\n",
    "    def __init__(self, psf, pdb, prm):\n",
    "        self.psf = psf\n",
    "        self.pdb = pdb\n",
    "        self.prm = prm\n",
    "        \n",
    "        ## initializing bond types and coordinates\n",
    "        coords = pdb.coordinates\n",
    "        bonds = psf.bonds\n",
    "        angles = psf.angles\n",
    "        dihedrals = psf.dihedrals\n",
    "        impropers = psf.impropers\n",
    "        for bond in bonds:\n",
    "            ## bond types\n",
    "            tid1 = bond.atom1.type\n",
    "            tid2 = bond.atom2.type\n",
    "            \n",
    "            t = (tid1, tid2)\n",
    "            bond.type = prm.bond_types[t]\n",
    "    \n",
    "            ## coordinates\n",
    "            b1 = bond.atom1\n",
    "            b1.xx, b1.xy, b1.xz = coords[b1.idx]\n",
    "            \n",
    "            b2 = bond.atom2\n",
    "            b2.xx, b2.xy, b2.xz = coords[b2.idx]\n",
    "            \n",
    "            ## non bonded parameters eps and rmin\n",
    "            name = b1.type\n",
    "            eps, rmin, _, _ = prm.nonbonded_types[name]\n",
    "            b1.rmin = rmin\n",
    "            b1.epsilon = eps\n",
    "            \n",
    "            name = b2.type\n",
    "            eps, rmin, _, _ = prm.nonbonded_types[name]\n",
    "            b2.rmin = rmin\n",
    "            b2.epsilon = eps\n",
    "        \n",
    "            \n",
    "            \n",
    "        ## initializing angle types\n",
    "        for angle in angles:\n",
    "            a1 = angle.atom1.type\n",
    "            a2 = angle.atom2.type\n",
    "            a3 = angle.atom3.type\n",
    "            \n",
    "            t = (a1, a2, a3)\n",
    "            \n",
    "            angle.type = prm.angle_types[t]\n",
    "            a = angle.measure()\n",
    "            a0 = angle.type\n",
    "            \n",
    "        for i, dihedral in enumerate(dihedrals):\n",
    "            a1 = dihedral.atom1.type\n",
    "            a2 = dihedral.atom2.type\n",
    "            a3 = dihedral.atom3.type\n",
    "            a4 = dihedral.atom4.type\n",
    "            \n",
    "            t = (a1, a2, a3, a4)\n",
    "            g = self._get_bond_type(t, prm.dihedral_types)\n",
    "            dihedrals[i].type = prm.dihedral_types[g]\n",
    "            \n",
    "        for i, improper in enumerate(impropers):\n",
    "            a1 = improper.atom1.type\n",
    "            a2 = improper.atom2.type\n",
    "            a3 = improper.atom3.type\n",
    "            a4 = improper.atom4.type\n",
    "            \n",
    "            t = (a1, a2, a3, a4)\n",
    "            g = self._get_bond_type(t, prm.improper_types)\n",
    "            impropers[i].type = prm.improper_types[g]\n",
    "\n",
    "    @classmethod\n",
    "    def bond_force(self, bond):\n",
    "        b1 = bond.atom1\n",
    "        b2 = bond.atom2\n",
    "        coords = psf.coordinates\n",
    "\n",
    "        r1 = coords[b1.idx]\n",
    "        r2 = coords[b2.idx]\n",
    "        r = r1 - r2\n",
    "        r_val = np.linalg.norm(r)\n",
    "        r_hat = r / r_val\n",
    "        k = bond.type.k\n",
    "        r0 = bond.type.req\n",
    "\n",
    "        f_1 = -2 * k * (r_val - r0) * r_hat\n",
    "        f_2 = -f_1\n",
    "        return f_1, f_2\n",
    "    \n",
    "    @classmethod\n",
    "    def angular_force(self, angle):\n",
    "        coords = psf.coordinates\n",
    "        theta = angle.measure()\n",
    "        theta0 = angle.type.theteq\n",
    "        k = angle.type.k\n",
    "\n",
    "        a1 = angle.atom1\n",
    "        a2 = angle.atom2\n",
    "        a3 = angle.atom3\n",
    "\n",
    "        a = coords[a1.idx]\n",
    "        b = coords[a2.idx]\n",
    "        c = coords[a3.idx]\n",
    "\n",
    "        ba = (b - a)\n",
    "        bc = (b - c)\n",
    "        p_a = np.cross(ba, np.cross(ba, bc))\n",
    "        p_a_hat = p_a / np.linalg.norm(p_a)\n",
    "        f_a = - 2 * k * (theta - theta0) * p_a_hat / (np.linalg.norm(ba))\n",
    "\n",
    "        p_c = np.cross(bc, np.cross(bc, ba))\n",
    "        p_c_hat = p_c / np.linalg.norm(p_c)\n",
    "        f_b = - 2 * k * (theta - theta0) * p_c_hat / (np.linalg.norm(bc))\n",
    "\n",
    "        f_c = - f_a - f_b\n",
    "        return f_a, f_b, f_c\n",
    "        \n",
    "    \n",
    "    \n",
    "    @classmethod\n",
    "    def dihedral_force(self, dihed):\n",
    "        coords = psf.coordinates\n",
    "        dtype = dihed.type\n",
    "        chi = dihed.measure() * (np.pi / 180)\n",
    "        f_a, f_b, f_c, f_d = 0, 0, 0, 0\n",
    "        \n",
    "        if isinstance(dtype, parmed.DihedralType):\n",
    "            n = dtype.per\n",
    "            kchi = dtype.phi_k\n",
    "            phase = dtype.phase * (np.pi / 180)\n",
    "\n",
    "            a1 = dihed.atom1\n",
    "            a2 = dihed.atom2\n",
    "            a3 = dihed.atom3\n",
    "            a4 = dihed.atom4\n",
    "\n",
    "            a = coords[a1.idx]\n",
    "            b = coords[a2.idx]\n",
    "            c = coords[a3.idx]\n",
    "            d = coords[a4.idx]\n",
    "\n",
    "            ba = (b - a)\n",
    "            bc = (b - c)\n",
    "            p_1 = np.cross(ba, bc)\n",
    "            p_1 /= np.linalg.norm(p_1)\n",
    "\n",
    "            theta_1 = angle(a1, a2, a3) * (np.pi / 180)\n",
    "            f_a = - k * np.sin(n * chi - phase) * n * p_1 / (np.linalg.norm(ab) * np.sin(theta_1))\n",
    "\n",
    "            cd = (c - d)\n",
    "            cb = -bc\n",
    "            p_2 = np.cross(cd, cb)\n",
    "            p_2 /= np.linalg.norm(p_2)\n",
    "\n",
    "            theta_2 = angle(a2, a3, a4) * (np.pi / 180)\n",
    "            f_d = - k * np.sin(n * chi - phase) * n * p_2 / (np.linalg.norm(cd) * np.sin(theta_2))\n",
    "\n",
    "            o = (b + c) / 2\n",
    "            oc = (o - c)\n",
    "            t_c = - np.cross(oc, f_d) + 0.5 * np.cross(cd, f_d) + 0.5 * np.cross(ba, f_a)\n",
    "            f_c = np.cross(t_c, oc) / (np.linalg.norm(cd)**2)\n",
    "\n",
    "\n",
    "            f_b =  - f_a - f_c - f_d\n",
    "            \n",
    "        elif isinstance(dtype, parmed.DihedralTypeList):\n",
    "            for term in dtype:\n",
    "                n = term.per\n",
    "                kchi = term.phi_k\n",
    "                phase = term.phase * (np.pi / 180)\n",
    "\n",
    "                a1 = dihed.atom1\n",
    "                a2 = dihed.atom2\n",
    "                a3 = dihed.atom3\n",
    "                a4 = dihed.atom4\n",
    "\n",
    "                a = coords[a1.idx]\n",
    "                b = coords[a2.idx]\n",
    "                c = coords[a3.idx]\n",
    "                d = coords[a4.idx]\n",
    "\n",
    "                ba = (b - a)\n",
    "                bc = (b - c)\n",
    "                p_1 = np.cross(ba, bc)\n",
    "                p_1 /= np.linalg.norm(p_1)\n",
    "\n",
    "                theta_1 = angle(a1, a2, a3) * (np.pi / 180)\n",
    "                f_a = (-1)*(kchi * np.sin(n * chi - phase) * n * p_1) / (np.linalg.norm(ba) * np.sin(theta_1))\n",
    "\n",
    "\n",
    "                cd = (c - d)\n",
    "                cb = -bc\n",
    "                p_2 = np.cross(cd, cb)\n",
    "                p_2 /= np.linalg.norm(p_2)\n",
    "\n",
    "                theta_2 = angle(a2, a3, a4) * (np.pi / 180)\n",
    "                f_d = (-1)* kchi * np.sin(n * chi - phase) * n * p_2 / (np.linalg.norm(cd) * np.sin(theta_2))\n",
    "\n",
    "\n",
    "                o = (b + c) / 2\n",
    "                oc = (o - c)\n",
    "                t_c = (-1)* np.cross(oc, f_d) + 0.5 * np.cross(cd, f_d) + 0.5 * np.cross(ba, f_a)\n",
    "                f_c = np.cross(t_c, oc) / (np.linalg.norm(cd)**2)\n",
    "\n",
    "\n",
    "                f_b =  - f_a - f_c - f_d\n",
    "        else:\n",
    "            print(\"Cant determine parameters for {}\".format(dtype))\n",
    "        return f_a, f_b, f_c, f_d\n",
    "    \n",
    "    @classmethod\n",
    "    def improper_force(self, imp):\n",
    "        coords = psf.coordinates\n",
    "        psi = imp.measure() * (np.pi / 180)\n",
    "        dtype = imp.type\n",
    "        psi_0 = dtype.psi_eq * (np.pi / 180)\n",
    "        kpsi = dtype.psi_k \n",
    "\n",
    "        a1 = imp.atom1\n",
    "        a2 = imp.atom2\n",
    "        a3 = imp.atom3\n",
    "        a4 = imp.atom4\n",
    "        \n",
    "        a = coords[a1.idx]\n",
    "        b = coords[a2.idx]\n",
    "        c = coords[a3.idx]\n",
    "        d = coords[a4.idx]\n",
    "\n",
    "        ba = (b - a)\n",
    "        bc = (b - c)\n",
    "        p_1 = np.cross(ba, bc)\n",
    "        p_1 /= np.linalg.norm(p_1)\n",
    "\n",
    "        theta_1 = angle(a1, a2, a3) * (np.pi / 180)\n",
    "        f_a = (-2)* kpsi * (psi - psi_0) * p_1 / (np.linalg.norm(ba) * np.sin(theta_1))\n",
    "\n",
    "\n",
    "        cd = (c - d)\n",
    "        cb = -bc\n",
    "        p_2 = np.cross(cd, cb)\n",
    "        p_2 /= np.linalg.norm(p_2)\n",
    "\n",
    "        theta_2 = angle(a2, a3, a4) * (np.pi / 180)\n",
    "        f_d = (-2)* kpsi * (psi - psi_0) * p_2 / (np.linalg.norm(cd) * np.sin(theta_2))\n",
    "\n",
    "\n",
    "        o = (b + c) / 2\n",
    "        oc = (o - c)\n",
    "        t_c = (-1)* np.cross(oc, f_d) + 0.5 * np.cross(cd, f_d) + 0.5 * np.cross(ba, f_a)\n",
    "        f_c = np.cross(t_c, oc) / (np.linalg.norm(cd)**2)\n",
    "\n",
    "\n",
    "        f_b =  - f_a - f_c - f_d\n",
    "        return f_a, f_b, f_c, f_d\n",
    "    \n",
    "                     \n",
    "    ## Total forces calculation\n",
    "    def bond_forces(self):\n",
    "        bonds = psf.bonds\n",
    "        coords = psf.coordinates\n",
    "        force = np.zeros(coords.shape)\n",
    "        for bond in bonds:\n",
    "            f_1, f_2 = Forces.bond_force(bond)\n",
    "            \n",
    "            force[bond.atom1.idx] += f_1\n",
    "            force[bond.atom2.idx] += f_2\n",
    "        return force\n",
    "    \n",
    "    def angular_forces(self):\n",
    "        coords = psf.coordinates\n",
    "        force = np.zeros(coords.shape)\n",
    "        angles = psf.angles\n",
    "        \n",
    "        for angle in angles:\n",
    "            theta = angle.measure()\n",
    "            theta0 = angle.type.theteq\n",
    "            k = angle.type.k\n",
    "            \n",
    "            f_a, f_b, f_c = Forces.angular_force(angle)\n",
    "            \n",
    "            force[angle.atom1.idx] += f_a\n",
    "            force[angle.atom2.idx] += f_b\n",
    "            force[angle.atom3.idx] += f_c\n",
    "        return force\n",
    "    \n",
    "    def dihedral_forces(self):\n",
    "        dihedrals = psf.dihedrals\n",
    "        force = np.zeros(psf.coordinates.shape)\n",
    "        coords = psf.coordinates\n",
    "        \n",
    "        for dihed in dihedrals:\n",
    "            dtype = dihed.type\n",
    "            \n",
    "            chi = dihed.measure() * (np.pi / 180)\n",
    "\n",
    "            f_a, f_b, f_c, f_d = Forces.dihedral_force(dihed)\n",
    "            force[dihed.atom1.idx] += f_a\n",
    "            force[dihed.atom2.idx] += f_b\n",
    "            force[dihed.atom3.idx] += f_c\n",
    "            force[dihed.atom4.idx] += f_d\n",
    "            \n",
    "        return force\n",
    "    \n",
    "    def improper_forces(self):\n",
    "        impropers = psf.impropers\n",
    "        coords = psf.coordinates\n",
    "        force = np.zeros(coords.shape)\n",
    "        for imp in impropers:\n",
    "            f_a, f_b, f_c, f_d = Forces.improper_force(imp)\n",
    "            force[imp.atom1.idx] += f_a\n",
    "            force[imp.atom2.idx] += f_b\n",
    "            force[imp.atom3.idx] += f_c\n",
    "            force[imp.atom4.idx] += f_d\n",
    "            \n",
    "        return force\n",
    "    \n",
    "    def LJ_forces(self):\n",
    "        coords = psf.coordinates\n",
    "        force = np.zeros(coords.shape)\n",
    "        for i, a in enumerate(psf.atoms):\n",
    "            for j, b in enumerate(psf.atoms):\n",
    "                if j <= i: \n",
    "                    continue\n",
    "                eps_ij = np.sqrt(a.epsilon * b.epsilon)\n",
    "                r_ij = (a.rmin + b.rmin) / 2\n",
    "    \n",
    "                r1 = coords[a.idx]\n",
    "                r2 = coords[b.idx]\n",
    "                delta_r = r1 - r2\n",
    "                r = np.linalg.norm(delta_r)\n",
    "                ratio = r_ij / r\n",
    "                \n",
    "                f_1 = -24 * eps_ij * (2 * (ratio ** 12) - (ratio**6)) * delta_r / r\n",
    "                f_2 = -f_1\n",
    "                \n",
    "                force[a.idx] += f_1\n",
    "                force[b.idx] += f_2\n",
    "        return force\n",
    "        \n",
    "    \n",
    "#     def Elec(self):\n",
    "#         energy = 0\n",
    "#         coords = psf.coordinates\n",
    "#         for i, a in enumerate(psf.atoms):\n",
    "#             for j, b in enumerate(psf.atoms):\n",
    "#                 if j <= i: \n",
    "#                     continue\n",
    "#                 k = 8.9 * 1e9\n",
    "#                 r1 = coords[a.idx]\n",
    "#                 r2 = coords[b.idx]\n",
    "#                 r = np.linalg.norm(r1 - r2)\n",
    "#                 q1 = a.charge\n",
    "#                 q2 = b.charge\n",
    "#                 en = k * q1 * q2 / (r * 1e-9)\n",
    "#         return energy\n",
    "            \n",
    "    ## DIHEDRAL AND IMPROPER HELPER FUNCTION\n",
    "    def _get_bond_type(self, t, search_list):\n",
    "        '''\n",
    "        Input :\n",
    "        t : a 4-atom tuple, \n",
    "        search_list: usually prm.dihedral_types or prm.improper_types\n",
    "\n",
    "        Return type: BondType\n",
    "        '''\n",
    "        for dtype in search_list:\n",
    "            val = True\n",
    "            for i in range(4):\n",
    "                val = val and (dtype[i] == t[i] or dtype[i] == 'X')\n",
    "\n",
    "            if val:\n",
    "                return dtype\n",
    "\n",
    "            ## Attempt reverse search\n",
    "            val = True\n",
    "            for i in range(4):\n",
    "                val = val and (dtype[i] == t[3 - i] or dtype[i] == 'X')\n",
    "\n",
    "            if val:\n",
    "                return dtype\n",
    "        return None\n",
    "    \n",
    "    def bonded(self):\n",
    "        return self.bond_forces() + self.angular_forces() + self.dihedral_forces() + self.improper_forces()\n",
    "\n",
    "    def nonbonded(self):\n",
    "        return self.LJ_forces()\n",
    "    \n",
    "    def all_forces(self):\n",
    "        return self.bonded() + self.nonbonded()\n",
    "    \n",
    "    def update(self):\n",
    "        coords = psf.coordinates\n",
    "        \n",
    "        for atom in psf.atoms:\n",
    "            atom.xx, atom.xy, atom.xz = coords[atom.idx]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# forces = Forces(psf, pdb, prm)\n",
    "# f_bond = forces.bond_forces()\n",
    "# f_angle = forces.angular_forces()\n",
    "# f_dihedral = forces.dihedral_forces()\n",
    "# f_impropers = forces.improper_forces()\n",
    "# f_LJ = forces.LJ_forces()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Coding the velocity verlet integrator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = []\n",
    "dt = 1e-3\n",
    "\n",
    "v = np.random.rand(*psf.coordinates.shape)\n",
    "\n",
    "m = 0\n",
    "for atom in psf.atoms:\n",
    "    m += atom.mass\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "x_orig = psf.coordinates.copy()\n",
    "\n",
    "psf.coordinates = x_orig\n",
    "\n",
    "x = psf.coordinates\n",
    "\n",
    "F_k = forces.all_forces()\n",
    "print(F_k.sum())\n",
    "x_new = x + dt * v + (F_k * (dt**2) / (2 * m))\n",
    "psf.coordinates = x_new\n",
    "F_k1 = forces.all_forces()\n",
    "print(np.linalg.norm(F_k1 - F_k))\n",
    "print(psf.coordinates[0])\n",
    "v = v + (dt * (F_k + F_k1) / (2 * m))\n",
    "\n",
    "print(np.linalg.norm(x_new - x))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "forces.update()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plotting forces for validation - Incomplete"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from mpl_toolkits.mplot3d import Axes3D"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_bond(bond):\n",
    "    a1 = bond.atom1\n",
    "    a2 = bond.atom2\n",
    "    \n",
    "    r1 = np.array([a1.xx, a1.xy, a1.xz])\n",
    "    r2 = np.array([a2.xx, a2.xy, a2.xz])\n",
    "    \n",
    "    f_1, f_2 = Forces.bond_force(bond)\n",
    "\n",
    "    f_1 /= (np.linalg.norm(f_1))\n",
    "    f_2 /= (np.linalg.norm(f_2))\n",
    "\n",
    "    fig = plt.figure(figsize = (8, 8))\n",
    "    ax = fig.add_subplot(111,projection='3d')\n",
    "    ax.scatter(a1.xx, a1.xy, a1.xz,  color='green', s = 100)\n",
    "    ax.scatter(a2.xx, a2.xy, a2.xz,  color='blue', s = 100)\n",
    "    \n",
    "    ax.quiver(*r1, f_1[0], f_1[1], f_1[2], color = ['r'], arrow_length_ratio = 0.1)\n",
    "    ax.quiver(*r2, f_2[0], f_2[1], f_2[2], color = ['g'], arrow_length_ratio = 0.1)\n",
    "    ax.view_init(70, -10)\n",
    "    r1 = np.array([a1.xx, a1.xy, a1.xz])\n",
    "    r2 = np.array([a2.xx, a2.xy, a2.xz])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_bond(psf.bonds[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_angle(angle):\n",
    "    theta = angle.measure() * (np.pi / 180)\n",
    "    theta0 = angle.type.theteq * (np.pi / 180)\n",
    "    k = angle.type.k\n",
    "    \n",
    "    coords = psf.coordinates\n",
    "    \n",
    "\n",
    "    a1 = angle.atom1\n",
    "    a2 = angle.atom2\n",
    "    a3 = angle.atom3\n",
    "\n",
    "    a = coords[a1.idx]\n",
    "    b = coords[a2.idx]\n",
    "    c = coords[a3.idx]\n",
    "\n",
    "    ba = (b - a)\n",
    "    bc = (b - c)\n",
    "\n",
    "    p_c = np.cross(bc, np.cross(bc, ba))\n",
    "\n",
    "    f_a, f_b, f_c = Forces.angular_force(angle)\n",
    "    \n",
    "    fig = plt.figure(figsize = (8, 8))\n",
    "    ax = fig.add_subplot(111,projection='3d')\n",
    "    ax.scatter(*a,  color='green', s = 100)\n",
    "    ax.scatter(*b,  color='blue', s = 100)\n",
    "    ax.scatter(*c,  color='red', s = 100)\n",
    "\n",
    "    \n",
    "    ax.quiver(*b, *(-1 * ba), color = ['g'], arrow_length_ratio = 0.05)\n",
    "    ax.quiver(*b, *(-1 * bc), color = ['r'], arrow_length_ratio = 0.05)\n",
    "    \n",
    "    ax.quiver(*a, *f_a, color = ['y'], length = 0.01, arrow_length_ratio = 0.05)\n",
    "    ax.quiver(*b, *f_b, color = ['y'], length = 0.005, arrow_length_ratio = 0.05)\n",
    "    ax.quiver(*c, *f_c, color = ['y'], length = 0.005, arrow_length_ratio = 0.05)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_angle(psf.angles[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "MD",
   "language": "python",
   "name": "md"
  },
  "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.6.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}