{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import glob\n",
    "import MDAnalysis as mda\n",
    "from MDAnalysis.analysis import lineardensity as lin\n",
    "#import nglview as nv\n",
    "#import ipywidgets\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.colors as colors\n",
    "import matplotlib.ticker as ticker\n",
    "from matplotlib import cm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2D 3D-Mapping of nomal modes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def dcdtovel(u,sel,nm=1):\n",
    "    \"\"\" u and sel are MDAnalysis\n",
    "    universe and atomic selection.\n",
    "    returns the velocity vector\n",
    "    of the nm mode for the selection\"\"\"\n",
    "    posit=np.zeros((12, len(sel), 3))\n",
    "    k=0\n",
    "    for ts in u.trajectory[(nm-1)*12:nm*12]:\n",
    "        posit[k]=sel.positions\n",
    "        k+=1\n",
    "    velo=np.zeros((12,len(sel),3))\n",
    "    for at in range(len(sel)):        \n",
    "        velo[:,at]=np.gradient(posit[:,at], axis=0)\n",
    "    return velo[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def dcd2omega(u,sel,nm=1):\n",
    "    \"\"\" u and sel are MDAnalysis\n",
    "    universe and atomic selection.\n",
    "    returns the velocity vector\n",
    "    of the nm mode for the selection\"\"\"\n",
    "    enedab=u.select_atoms(\"type 9 and resname DAB\")\n",
    "    posMres=np.zeros((12,len(sel),3))\n",
    "    enevecs=np.zeros((enedab.n_residues,3))\n",
    "    posit=np.zeros((12, len(sel), 3))\n",
    "    atmperres=np.shape(posit)[1]//enedab.n_residues\n",
    "    k=0\n",
    "    for ts in u.trajectory[(nm-1)*12:nm*12]:\n",
    "        posit[k]=sel.positions\n",
    "        k+=1\n",
    "    for e in range(enedab.n_residues):\n",
    "        enevecs[e]=sel[atmperres*e:atmperres*(e+1)].center_of_mass()\n",
    "        posMres[:,atmperres*e:atmperres*(e+1)]=posit[:,atmperres*e:atmperres*(e+1)]-enevecs[e]\n",
    "    velo=np.zeros((12,len(sel),3))\n",
    "    for at in range(len(sel)):        \n",
    "        velo[:,at]=np.gradient(posit[:,at], axis=0)\n",
    "    rsqr=np.linalg.norm(posMres,axis=2)**(-2)\n",
    "    invpos=np.zeros((np.shape(posMres)))\n",
    "    for trj in range(np.shape(posMres)[0]):\n",
    "        for at in range(np.shape(posMres)[1]):\n",
    "            invpos[trj,at]=posMres[trj,at]*rsqr[trj,at]\n",
    "    ome=np.cross(invpos,velo)\n",
    "    return ome[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mergeeigfrq(entry_list,dec=1):\n",
    "    # groups the the eigenvalues/NM within a certain deltaFreq\n",
    "    # return a list of list with initial and final limits for each mode\n",
    "    incr=np.power(10.,dec)\n",
    "    round_list=np.around(entry_list*incr)\n",
    "    nestM=[]\n",
    "    count=0\n",
    "    while count < len(round_list):\n",
    "        nest1=[]\n",
    "        nest1.append(count+1)\n",
    "        ref_value=round_list[count]\n",
    "        try:\n",
    "            next_value=round_list[count+1] # next value\n",
    "        except IndexError: # account for last value\n",
    "            nest1.append(count+1)\n",
    "            nestM.append(nest1)\n",
    "            break\n",
    "        while ref_value == next_value:\n",
    "            count+=1\n",
    "            try:\n",
    "                next_value=round_list[count+1] ### problem!!!!\n",
    "            except IndexError:\n",
    "                break#nextf+=deltf\n",
    "        nest1.append(count+1) # final index\n",
    "        nestM.append(nest1)\n",
    "        count+=1\n",
    "    return nestM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def intensnm(NM,minf):\n",
    "    # return a list with the NM with intensities above minf\n",
    "    intensnm=NM[~np.any(NM <= 0, axis=1)].T\n",
    "    intensnm1=np.vstack((np.arange(1,len(intensnm.T)+1),intensnm[0],intensnm[1]))\n",
    "    intensnm2=intensnm1.T[~np.any(intensnm1.T <= minf, axis=1)].T\n",
    "    return intensnm2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mergebigfrq(entry_list,NM,minf,dec=1):\n",
    "    intpeaks=np.asarray(intensnm(NMspec,minf)[0]).astype(int)\n",
    "    bigfreq=[]\n",
    "    for p in range(len(intpeaks)):\n",
    "        ix=intpeaks[p]\n",
    "        bigfreq.append(eigfreq[ix])\n",
    "    bigfreq=np.asarray(bigfreq)\n",
    "    frecum=mergeeigfrq(bigfreq,dec)\n",
    "    return intpeaks,frecum"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1. Input data for mapping\n",
    "Choose one to run before going to point 2. Check which trajectory nmode file (nmtraj) are you loading#### 1. Input data for mapping"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 149,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load NM trajectories, coordinates, NMA spectra and eigenfreq\n",
    "# pdp 2*14 PBC in x\n",
    "nmtraj=mda.Universe(\"../charmm/inp/inpcharmm.pdp.28x.yy.psf\",\"../charmm/normalmodestraj/nmodes.pdp.28px.yy.dcd\")\n",
    "coord=mda.Universe(\"../charmm/inp/inpcharmm.pdp.28x.yy.psf\",\"../charmm/inp/miniabnr.pdp.28x.yy.nma.pdb\")\n",
    "NMspec=np.loadtxt(\"../charmm/NMAspectra_nma.pdp.28px.yy.DAB.dat\")\n",
    "eigfreq=np.loadtxt(\"../charmm/eigenfrq_nma.pdp.28px.yy.DAB.dat\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 142,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load NM trajectories, coordinates, NMA spectra and eigenfreq\n",
    "# pdp 2*14 island\n",
    "nmtraj=mda.Universe(\"../charmm/inp/inpcharmm.pdp.28x.yy.psf\",\"../charmm/normalmodestraj/nmodes.pdp.28x.yy.dcd\")\n",
    "coord=mda.Universe(\"../charmm/inp/inpcharmm.pdp.28x.yy.psf\",\"../charmm/inp/miniabnr.pdp.28x.yy.nma.pdb\")\n",
    "NMspec=np.loadtxt(\"../charmm/NMAspectra_nma.pdp.28x.yy.DAB.dat\")\n",
    "eigfreq=np.loadtxt(\"../charmm/eigenfrq_nma.pdp.28x.yy.DAB.dat\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load NM trajectories, coordinates, NMA spectra and eigenfreq\n",
    "# pdp 2*14 island with heavy boundary\n",
    "nmtraj=mda.Universe(\"../charmm/inp/inpcharmm.pdp.28x.yy.psf\",\"../charmm/normalmodestraj/nmodes.pdp.28x.yy.mass.dcd\")\n",
    "coord=mda.Universe(\"../charmm/inp/inpcharmm.pdp.28x.yy.psf\",\"../charmm/inp/miniabnr.pdp.28x.yy.mass.nma.pdb\")\n",
    "NMspec=np.loadtxt(\"../charmm/NMAspectra_nma.pdp.28x.yy.mass.DAB.dat\")\n",
    "eigfreq=np.loadtxt(\"../charmm/eigenfrq_nma.pdp.28x.yy.mass.DAB.dat\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2. Plot NMA overview\n",
    "The data ploted would be the last one run from point 1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "#mergefrecu=mergeeigfrq(eigfreq,-0.1)\n",
    "bigfreq,mergefrecu=mergebigfrq(eigfreq,NMspec,0.05,0.5)\n",
    "lentraj=len(nmtraj.trajectory)//12\n",
    "leneigfrq=len(eigfreq)\n",
    "#assert (lentraj == leneigfrq), \"Trajectory movie and eigenfrequency count are not the same length\"\n",
    "strsel=\"resname DAB\" # atomic system selection for the plot\n",
    "atmsel=nmtraj.select_atoms(strsel)\n",
    "initsel=coord.select_atoms(strsel)\n",
    "initpos=initsel.positions\n",
    "X=initpos[:,0]\n",
    "Y=initpos[:,1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 24408 characters]\n",
      "text/plain": [
       "<Figure size 576x576 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "###### Plot the spectrum/spectra\n",
    "fig, ax = plt.subplots(figsize=(8,8)) # Figure size\n",
    "#plt.title(\"NMA crystal and Nanoribbon\")\n",
    "#plt.ylabel('Intensity',fontsize=30,labelpad=10)\n",
    "plt.xlabel('Frequency (cm$^{-1}$)',fontsize=20,labelpad=15)\n",
    "\n",
    "#ax.plot(NMspec[:,0], NMspec[:,1], color='tab:gren' , label='Nanoribbon')\n",
    "ax.plot(NMspec[:,0], NMspec[:,1], color='black' )#, label='Hex 73') \n",
    "#ax.plot(NMspec2[:,0], NMspec2[:,1]*-1+2.5, color='tab:red' , label='Crystal')\n",
    "\n",
    "plt.xlim(30,38) #Set the x axis limits\n",
    "#plt.xscale(\"log\")\n",
    "plt.ylim(0,0.8) # Set the y axis limits\n",
    "plt.xticks(fontsize=20)\n",
    "plt.yticks(fontsize=10)\n",
    "#ax.yaxis.set_major_formatter(ticker.FormatStrFormatter(''))\n",
    "#ax.xaxis.set_major_locator(ticker.MultipleLocator(20))\n",
    "ax.xaxis.set_minor_locator(ticker.MultipleLocator(0.5))\n",
    "ax.xaxis.set_major_locator(ticker.MultipleLocator(2))\n",
    "for axis in ['top','bottom','left','right']:\n",
    "    ax.spines[axis].set_linewidth(2.0)\n",
    "ax.xaxis.set_tick_params(width=2.0,length=10,which='major')\n",
    "ax.xaxis.set_tick_params(width=1.5,length=5,which='minor')\n",
    "ax.yaxis.set_tick_params(width=0.0,length=10,which='major')\n",
    "plt.tight_layout()\n",
    "#plt.legend(fontsize=15,loc=0)\n",
    "plt.grid(which='both')\n",
    "#plt.savefig(\"NMA.pdp.112.aa.80-110.png\") # The last spectra can be saved as an image file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3. Quiver plot based on initial and final normal mode velocities\n",
    "Remember to check the nmtraj file and modify the entry \"velonm[nm-fini]=dcdtovel(nmtraj,atmsel,nm)\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 53564 characters]\n",
      "text/plain": [
       "<Figure size 1033.2x180 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "#mm=43\n",
    "#fini,ffin=bigfreq[mergefrecu[mm-1][:2][0]-1],bigfreq[mergefrecu[mm-1][:2][1]-1]\n",
    "fck=90\n",
    "fini,ffin=fck,fck+0 # initial normal mode, final normal mode\n",
    "plt.figure(figsize=(14.35,2.5)) #2x14, \n",
    "#plt.figure(figsize=(4,20))\n",
    "ff=np.average(eigfreq[fini-1:ffin])\n",
    "#plt.title('ROC Curve; Enrichment factor 1% = {:.2f}'.format(ff))\n",
    "plt.title('eigfrequency: {:.2f} cm-1\\n Modes: {:d} - {:d}'.format(ff,fini,ffin))\n",
    "velonm=np.zeros((ffin+1-fini,len(atmsel),3))\n",
    "for nm in range(fini,ffin+1):\n",
    "    #velonm[nm-fini]=dcdtovel(nmtraj,atmsel,nm)\n",
    "    velonm[nm-fini]=dcdtovel(nmtraj,atmsel,nm) # if not the first nmode traj file\n",
    "velonma=np.sum(np.abs(velonm),axis=0) #np.average(velonm,axis=0)\n",
    "veloabs=np.linalg.norm(velonma,axis=1)\n",
    "#velabsn=veloabs/np.power(np.max(veloabs),2)\n",
    "veloz=np.abs(velonma[:,2])\n",
    "ax=plt.quiver(X,Y, velonma[:,0], velonma[:,1], veloz, scale_units='xy', scale=np.max(velonma)/9, cmap=\"jet\",linewidth=1)\n",
    "#    ax=plt.scatter(X,Y, c=veloabs, s=velabsn, cmap=\"rainbow\")\n",
    "#plt.xlim(-60,70) # for 7*14\n",
    "#plt.ylim(-100,110)  # for 7*14\n",
    "plt.xlim(-60,170) # for 2*14\n",
    "plt.ylim(-95,-45) # for 2*14\n",
    "cb = plt.colorbar(ax)\n",
    "cb.ax.tick_params(labelsize=16)\n",
    "#plt.savefig(\"nmode\"+str(fini)+\"-\"+str(ffin)+\".pdbr.96.nr.ave.png\",transparent=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4. Quiver plot a range of normal modes\n",
    "Remember to check the nmtraj file and modify the entry \"velonm[nm-fini]=dcdtovel(nmtraj,atmsel,nm)\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 148,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 52232 characters]\n",
      "text/plain": [
       "<Figure size 1033.2x180 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 82464 characters]\n",
      "text/plain": [
       "<Figure size 1033.2x180 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 81732 characters]\n",
      "text/plain": [
       "<Figure size 1033.2x180 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 70736 characters]\n",
      "text/plain": [
       "<Figure size 1033.2x180 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 69904 characters]\n",
      "text/plain": [
       "<Figure size 1033.2x180 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 59344 characters]\n",
      "text/plain": [
       "<Figure size 1033.2x180 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 48596 characters]\n",
      "text/plain": [
       "<Figure size 1033.2x180 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 53912 characters]\n",
      "text/plain": [
       "<Figure size 1033.2x180 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 50484 characters]\n",
      "text/plain": [
       "<Figure size 1033.2x180 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 51824 characters]\n",
      "text/plain": [
       "<Figure size 1033.2x180 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "for i in range(70,80):\n",
    "    fini,ffin=i,i # initial normal mode, final normal mode\n",
    "    plt.figure(figsize=(14.35,2.5))#,(8.75,7.5)\n",
    "    #plt.figure(figsize=(4,20))\n",
    "    ff=np.average(eigfreq[fini-1:ffin])\n",
    "    #plt.title('ROC Curve; Enrichment factor 1% = {:.2f}'.format(ff))\n",
    "    plt.title('eigfrequency: {:.2f} cm-1\\n Modes: {:d} - {:d}'.format(ff,fini,ffin))\n",
    "    velonm=np.zeros((ffin+1-fini,len(atmsel),3))\n",
    "    for nm in range(fini,ffin+1):\n",
    "        #velonm[nm-fini]=dcdtovel(nmtraj,atmsel,nm)\n",
    "        velonm[nm-fini]=dcdtovel(nmtraj,atmsel,nm) # if not the first nmode traj file\n",
    "    velonma=np.sum(velonm,axis=0)\n",
    "    veloabs=np.linalg.norm(velonma,axis=1)\n",
    "    velabsn=veloabs/np.power(np.max(veloabs),2)\n",
    "    veloz=np.abs(velonma[:,2])\n",
    "    ax=plt.quiver(X,Y, velonma[:,0], velonma[:,1], veloz, scale_units='xy', scale=np.max(velonma)/7, cmap=\"jet\",linewidth=1)\n",
    "    #    ax=plt.scatter(X,Y, c=veloabs, s=velabsn, cmap=\"rainbow\")\n",
    "    #plt.xlim(-60,70) # for 7*14\n",
    "    #plt.ylim(-100,110)  # for 7*14\n",
    "    plt.xlim(-60,170) # for 2*14\n",
    "    plt.ylim(-95,-45) # for 2*14\n",
    "    cb = plt.colorbar(ax)\n",
    "    cb.ax.tick_params(labelsize=15)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 5. Plot NMA comparision"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 150,
   "metadata": {},
   "outputs": [],
   "source": [
    "NMspec1=np.loadtxt(\"../charmm/NMAspectra_nma.pdp.28px.yy.DAB.dat\")\n",
    "NMspec2=np.loadtxt(\"../charmm/NMAspectra_nma.pdp.28x.yy.DAB.dat\")\n",
    "NMspec3=np.loadtxt(\"../charmm/NMAspectra_nma.pdp.28x.yy.mass.DAB.dat\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 151,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 28504 characters]\n",
      "text/plain": [
       "<Figure size 432x576 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "###### Plot the spectrum/spectra\n",
    "fig, ax = plt.subplots(figsize=(6,8)) # Figure size\n",
    "#plt.title(\"NMA crystal and Nanoribbon\")\n",
    "#plt.ylabel('Intensity',fontsize=30,labelpad=10)\n",
    "plt.xlabel('Frequency (cm$^{-1}$)',fontsize=20,labelpad=15)\n",
    "\n",
    "#ax.plot(NMspec[:,0], NMspec[:,1], color='tab:gren' , label='Nanoribbon')\n",
    "ax.plot(NMspec3[:,0], NMspec3[:,1], color='tab:blue' , label='Heavy Edge') \n",
    "#ax.plot(NMspec2[:,0], NMspec2[:,1]*-1+1., color='tab:red' , label='Island')\n",
    "ax.plot(NMspec1[:,0], NMspec1[:,1]*-1+1., color='tab:red' , label='PBC in x')\n",
    "\n",
    "plt.xlim(28,32) #Set the x axis limits\n",
    "#plt.xscale(\"log\")\n",
    "plt.ylim(0,1) # Set the y axis limits\n",
    "plt.xticks(fontsize=20)\n",
    "plt.yticks(fontsize=10)\n",
    "ax.yaxis.set_major_formatter(ticker.FormatStrFormatter(''))\n",
    "ax.xaxis.set_major_locator(ticker.MultipleLocator(1))\n",
    "ax.xaxis.set_minor_locator(ticker.MultipleLocator(0.5))\n",
    "#ax.yaxis.set_major_locator(ticker.MultipleLocator(5))\n",
    "for axis in ['top','bottom','left','right']:\n",
    "    ax.spines[axis].set_linewidth(2.0)\n",
    "ax.xaxis.set_tick_params(width=2.0,length=10,which='major')\n",
    "ax.xaxis.set_tick_params(width=1.5,length=5,which='minor')\n",
    "ax.yaxis.set_tick_params(width=0.0,length=10,which='major')\n",
    "plt.tight_layout()\n",
    "plt.legend(fontsize=15,loc=0)\n",
    "plt.grid(which='both')\n",
    "#plt.savefig(\"NMA.pdp.112.aa.80-110.png\") # The last spectra can be saved as an image file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
