{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from numpy.fft import fft, ifft, fftn, ifftn\n",
    "import glob\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": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "### crystal\n",
    "fftkx98cryy=np.load(\"../gromacs/pdp98c/fftkx_100_nvt.pdp.98c.yy.10fs.DAB.1470at.npy\")\n",
    "fftky98cryy=np.load(\"../gromacs/pdp98c/fftky_100_nvt.pdp.98c.yy.10fs.DAB.1470at.npy\")\n",
    "fftkm98cryy=np.load(\"../gromacs/pdp98c/fftkm_100_nvt.pdp.98c.yy.10fs.DAB.1470at.npy\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 195,
   "metadata": {},
   "outputs": [],
   "source": [
    "### ribbon\n",
    "fftkx98nyy=np.load(\"../gromacs/pdp98n/fftkx_100_nvt.pdp.98n.yy.20fs.1470at.npy\")\n",
    "fftky98nyy=np.load(\"../gromacs/pdp98n/fftky_100_nvt.pdp.98n.yy.20fs.1470at.npy\")\n",
    "### ribbon with 3-fold masses at one edge\n",
    "fftkx98nyy1M=np.load(\"../gromacs/pdp98n/fftkx_100_nvt.pdp.98n.yy.1M.1470at.npy\")\n",
    "fftky98nyy1M=np.load(\"../gromacs/pdp98n/fftky_100_nvt.pdp.98n.yy.1M.1470at.npy\")\n",
    "### ribbon with 3-fold masses at both edges\n",
    "fftkx98nyy2M=np.load(\"../gromacs/pdp98n/fftkx_100_nvt.pdp.98n.yy.2M.1470at.npy\")\n",
    "fftky98nyy2M=np.load(\"../gromacs/pdp98n/fftky_100_nvt.pdp.98n.yy.2M.1470at.npy\")\n",
    "### ribbon with 2 rows of 3-fold masses at both edges\n",
    "fftkx98nyy3M=np.load(\"../gromacs/pdp98n/fftkx_100_nvt.pdp.98n.yy.3M.1470at.npy\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 220,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 66484 characters]\n",
      "text/plain": [
       "<Figure size 720x432 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "dataplot=fftky98nyy\n",
    "X=dataplot[0]*33.35641\n",
    "Y=dataplot[1:]\n",
    "X1=np.real(X)\n",
    "Y1=np.real(Y)\n",
    "stp=50 # spectra to plot in PSD\n",
    "deltaf=(X1[1]-X1[0])\n",
    "initfreq=10\n",
    "finfreq=30\n",
    "ifq=int(initfreq/deltaf)\n",
    "ffq=int(finfreq/deltaf)\n",
    "#y=np.real(X1[ifq:])\n",
    "y=np.real(X1[ifq:ffq])\n",
    "X2,Y2=np.meshgrid(range(-len(Y1)//2,len(Y1)//2),y)\n",
    "Z=np.transpose(np.real(Y1[:,ifq:ffq]))\n",
    "\n",
    "plt.figure(figsize=(10,6))#figsize=(16,8))\n",
    "#plt.suptitle(\"Dispersion relation\", fontsize=14)\n",
    "\n",
    "ax=plt.subplot(1, 5, (1,3))\n",
    "s = \"Heat map - PSD of Rv\"\n",
    "#plt.title(s,fontsize=14)\n",
    "#plt.pcolormesh(X2/150,Y2,Z, cmap='jet',vmin=np.min(Y1), vmax=np.max(Y1[stp,ifq+int(0.3/deltaf):ffq-int(0.3/deltaf)])/1.3)\n",
    "plt.pcolormesh(X2/120,Y2,Z, cmap='jet', norm=colors.LogNorm(vmin=10, vmax=Z.max()*0.5),shading='auto')\n",
    "plt.colorbar()\n",
    "#plt.xlim(0,40)\n",
    "#plt.zlim(0,30)\n",
    "plt.ylabel('Frequency (cm$^{-1}$)', fontsize=20)\n",
    "plt.xlabel('k$_x$ ($\\pi/a$)', fontsize=20)\n",
    "plt.xticks(fontsize=16)#,color=\"w\")\n",
    "#plt.xticks([0,133,199,314],['$\\Gamma$','K','M','$\\Gamma$'],fontsize=25)#[0,40,60,94]\n",
    "plt.yticks(fontsize=16)\n",
    "ax.xaxis.set_major_locator(ticker.MultipleLocator(0.5))\n",
    "\n",
    "bx=plt.subplot(1, 5, (4,5))\n",
    "#s2='PSD'\n",
    "#plt.title(s2)\n",
    "#plt.legend()\n",
    "plt.xlabel('PSD (Intensity)', fontsize=18, labelpad=10)\n",
    "#plt.ylabel('Frequency (cm$^{-1}$)',fontsize=16,rotation=90)\n",
    "plt.plot(Y1[stp]/1+0, X1,'b-')\n",
    "#This X limit acts like the filter\n",
    "plt.ylim(initfreq,finfreq)\n",
    "plt.xlim(np.min(Y1),np.max(Y1[stp,ifq+int(0.3/deltaf):ffq-int(0.3/deltaf)])/0.8)\n",
    "#plt.xlim(np.min(Y1),np.max(Y1[stp,ifq+int(0.3/deltaf):ffq-int(0.3/deltaf)])/1.2)\n",
    "bx.set_xscale(\"log\")\n",
    "plt.xticks(fontsize=16)\n",
    "plt.yticks(fontsize=16)\n",
    "plt.grid()\n",
    "plt.tight_layout(pad=1.0)\n",
    "#bx.yaxis.set_major_locator(ticker.MultipleLocator(0.1))\n",
    "bx.xaxis.set_major_locator(ticker.MultipleLocator(20))\n",
    "#plt.savefig(\"dispM_pdp.112n.CT3.10K.\"+str(initfreq)+\"-\"+str(finfreq)+\"cm-1.png\") # The graphs can be as PNG, SVG, BMP or PDF"
   ]
  },
  {
   "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
}
