{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from numpy.fft import fft, ifft, fftn, ifftn\n",
    "import scipy\n",
    "import glob\n",
    "import sys\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.collections import LineCollection\n",
    "from matplotlib import cm\n",
    "from IPython.display import Image\n",
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "#from functools import partial\n",
    "#import pathos.multiprocessing as mp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def dynamicmatrixQ(Q,k,masses):\n",
    "    kf=np.diag(k)\n",
    "    #C=np.matrix.transpose(Q)# C is the compatibilita matrix\n",
    "    C=np.conj(Q.T)\n",
    "    m1=np.matmul(kf,C)\n",
    "    K=np.matmul(Q,m1)# the dynamical matrix D = transpose(C)*kf*C = Q*kf*C\n",
    "    M=np.diag(1/np.asarray(masses))#M=massmatrix(masses)\n",
    "    D=np.matmul(M,K)\n",
    "    return D"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def dynamicmatrixC(C,k,masses):\n",
    "    kf=np.diag(k)\n",
    "    m1=np.matmul(kf,C)\n",
    "    Q=np.conj(C.T)\n",
    "    K=np.matmul(Q,m1)# the dynamical matrix D = transpose(C)*kf*C = Q*kf*C\n",
    "    M=np.diag(1/np.asarray(masses))#M=massmatrix(masses)\n",
    "    D=np.matmul(M,K)\n",
    "    return D"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def massvalues(N,m1,m2,order=\"one\"):\n",
    "    if order == \"one\":\n",
    "        return np.ones((N))\n",
    "    elif order == \"random\":\n",
    "        mvals=np.random.randint(2, size=N)\n",
    "    elif order == \"alternate\":\n",
    "        mvals=np.ones((N))\n",
    "        for p in range(N):\n",
    "            if p%2 == 0:\n",
    "                mvals[p]=0\n",
    "    for p in range(N):\n",
    "        if mvals[p] == 0:\n",
    "            mvals[p]=m1\n",
    "        elif mvals[p] == 1:\n",
    "            mvals[p]=m2\n",
    "    return mvals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def phasevec(Vec,pearl1=0,pearl2=1):\n",
    "    #phaseV=np.zeros(np.shape(Vec[:,2,0]))\n",
    "    phaseV=np.zeros((len(Vec),len(Vec.T)))\n",
    "    for ka in range(len(Vec)):\n",
    "        for eg in range(len(Vec.T)):\n",
    "            #phaseV[ka,eg]=np.real(Vec[ka,pearl1,eg]*Vec[ka,pearl2,eg])\n",
    "            phaseV[ka,eg]=np.real(np.dot(Vec[ka,pearl1,eg],Vec[ka,pearl2,eg])*2*np.pi)\n",
    "            #phaseV[ka,eg]=(np.dot(Vec[ka,pearl1,eg],Vec[ka,pearl2,eg])/np.abs(np.dot(Vec[ka,pearl1,eg],Vec[ka,pearl2,eg])))/(2j*np.pi)\n",
    "    return phaseV"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 315,
   "metadata": {},
   "outputs": [],
   "source": [
    "def Cmatrix(tdeg,adeg):\n",
    "    #substrate resonator in z\n",
    "    theta=tdeg*np.pi/180 # in radians\n",
    "    alpha=adeg*np.pi/180 # in radians\n",
    "    cost=np.cos(theta)\n",
    "    sint=np.sin(theta)\n",
    "    cosa=np.cos(alpha)\n",
    "    sina=np.sin(alpha)\n",
    "    C=np.array([[cost,-1*cost,sint,-1*sint,0.,0.],[-1*cosa,cosa,sina,-1*sina,0.,0.],[1.,0.,0.,0.,0.,0.],\n",
    "            [0.,1.,0.,0.,0.,0.],[0.,0.,0.,0.,1.,0.],[0.,0.,0.,0.,0.,1.]])\n",
    "    return C"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 334,
   "metadata": {},
   "outputs": [],
   "source": [
    "def Cmatrixy(tdeg,adeg):\n",
    "    #substrate resonator in y\n",
    "    theta=tdeg*np.pi/180 # in radians\n",
    "    alpha=adeg*np.pi/180 # in radians\n",
    "    cost=np.cos(theta)\n",
    "    sint=np.sin(theta)\n",
    "    cosa=np.cos(alpha)\n",
    "    sina=np.sin(alpha)\n",
    "    C=np.array([[cost,-1*cost,sint,-1*sint],[-1*cosa,cosa,sina,-1*sina],[1.,0.,0.,0.],\n",
    "            [0.,1.,0.,0.],[0.,0.,1.,0.],[0.,0.,0.,1.]])\n",
    "    return C"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Two dimensional pbc triangular mesh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 355,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Spring constant values\n",
    "k1,k2,k3,k4,k5,k6=3,1,0.5,0.5,3,3 # case1\n",
    "#k1,k3,k2,k4=1,3,0.01,0.01 # case2\n",
    "### how many betapoint and angles\n",
    "betapoints=200\n",
    "#beta=np.array([e/betapoints for e in range(-betapoints//2,betapoints//2+1)])\n",
    "#angles=np.array([120*e/betapoints for e in range(-betapoints//2,betapoints//2+1)])\n",
    "angles=np.array([120*e/betapoints for e in range(-betapoints//2,betapoints//2+1,10)])\n",
    "##### Alpha angle\n",
    "alpangle=60 #90"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 356,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([-60., -54., -48., -42., -36., -30., -24., -18., -12.,  -6.,   0.,\n",
       "         6.,  12.,  18.,  24.,  30.,  36.,  42.,  48.,  54.,  60.])"
      ]
     },
     "execution_count": 356,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "angles"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 357,
   "metadata": {},
   "outputs": [],
   "source": [
    "a=np.pi # spacing between pearls\n",
    "#Np=6\n",
    "segm=1000\n",
    "kk=np.array([1*np.pi/(a*segm)*e for e in range(-segm,segm+1)]) #k points\n",
    "C=Cmatrix(0,0)\n",
    "Np=np.shape(C)[1]# degrees of freedom\n",
    "Ckk=np.zeros((len(angles),len(kk),len(C),len(C.T)), dtype = complex)\n",
    "for e in range(len(angles)):\n",
    "    C=Cmatrix(angles[e],alpangle)\n",
    "    for p in range(len(kk)):\n",
    "        Ckk[e][p]=C\n",
    "        Ckk[e][p][2][0]=1-C[2][0]*np.exp(1j * (kk[p] * a))\n",
    "        Ckk[e][p][1][0]=C[1][0]*np.exp(1j * (kk[p] * a))\n",
    "        Ckk[e][p][1][2]=C[1][2]*np.exp(1j * (kk[p] * a))\n",
    "        Ckk[e][p][3][1]=1-C[3][1]*np.exp(1j * (kk[p] * a))\n",
    "#C=np.conj(Q.T)\n",
    "#bondlist=np.array([[0,1],[1,2],[0,2],[1,0],[2,1],[2,0]])\n",
    "#Nbnopbc=len(bondlist)\n",
    "#N=6\n",
    "k=np.array([k1,k2,k3,k4,k5,k6]) \n",
    "mA,mB=1,1\n",
    "masses=massvalues(Np,mA,mB,\"one\")# mass values could be all \"one\", \"random\", \"alternate\" (k1,k2,k1,...)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 358,
   "metadata": {},
   "outputs": [],
   "source": [
    "D=dynamicmatrixC(C,k,masses)\n",
    "Dk=np.zeros((len(angles),len(kk),len(D),len(D.T)), dtype = complex)\n",
    "for a in range(len(angles)):\n",
    "    for e in range(len(kk)):\n",
    "        Dk[a][e]=dynamicmatrixC(Ckk[a][e],k,masses)\n",
    "#########\n",
    "E,V = np.linalg.eigh(D)\n",
    "Ek=np.zeros((len(angles),len(kk),len(E)), dtype = complex)\n",
    "Vk=np.zeros((len(angles),len(kk),len(V),len(V.T)), dtype = complex)\n",
    "for e in range(len(angles)):\n",
    "    for p in range(len(kk)):\n",
    "        Ek[e][p],Vk[e][p]=scipy.linalg.eigh(Dk[e][p])\n",
    "#######\n",
    "phaseVk=np.zeros((len(angles),len(kk),len(E)))\n",
    "for e in range(len(angles)):\n",
    "    phaseVk[e]=phasevec(Vk[e],0,1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 360,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 44688 characters]\n",
      "text/plain": [
       "<Figure size 648x576 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Create a continuous norm to map from data points to colors\n",
    "fig,ax = plt.subplots(figsize=(9,8))\n",
    "angl=20\n",
    "phaseVkang=phaseVk[angl]\n",
    "Ekang=Ek[angl]\n",
    "Ylim=np.sqrt(np.abs(np.real(Ekang.T)))\n",
    "#plt.title('Colored-by-phase eigenmodes for $\\\\theta$={:.1f}deg and $\\\\alpha$={:.1f}deg'.format(angles[angl],alpangle))\n",
    "#plt.title('Eigenmodes for {:d} Pearls: -{:d}-[-{:d}-]-{:d}- '.format(Np,k2,k1,k2))\n",
    "Z=np.sort(phaseVkang.reshape(len(phaseVkang)*len(phaseVkang.T)))\n",
    "for egv in range(len(Ekang.T)):\n",
    "    X=kk\n",
    "    Y=np.sqrt(np.abs(np.real(Ekang.T[egv])))\n",
    "    Z2=phaseVkang[:,egv] #values used for colormapping\n",
    "    points = np.array([X, Y]).T.reshape(-1, 1, 2)\n",
    "    segments = np.concatenate([points[:-1], points[1:]], axis=1)\n",
    "    norm = plt.Normalize(Z.min(), Z.max())\n",
    "    lc = LineCollection(segments, cmap='rainbow', norm=norm)\n",
    "    lc.set_array(Z2) # Set the values used for colormapping\n",
    "    lc.set_linewidth(2)\n",
    "    line = ax.add_collection(lc)\n",
    "cbar=fig.colorbar(line, ax=ax)\n",
    "cbar.ax.tick_params(labelsize=15)\n",
    "plt.xlim(X.min(),X.max())\n",
    "plt.ylim(0,Ylim.max()*1.02)\n",
    "plt.grid()\n",
    "#plt.legend()\n",
    "plt.xticks(fontsize=16)\n",
    "plt.yticks(fontsize=16)\n",
    "plt.ylabel('$\\omega (k_x)$', fontsize=20)\n",
    "plt.xlabel('k$_x$ ($\\pi/a$)', fontsize=20)\n",
    "ax.xaxis.set_major_locator(ticker.MultipleLocator(0.5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 251,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Text(0.5, 0, 'k ($\\\\pi/a$)')"
      ]
     },
     "execution_count": 251,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 168696 characters]\n",
      "text/plain": [
       "<Figure size 792x792 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Overview surface and atomic indexes\n",
    "fig,ax = plt.subplots(figsize=(11,11)) # Figure size\n",
    "X=range(np.shape(Ek)[1])\n",
    "#eig=0\n",
    "#colors=['y-','r-','c-','g-','m-','b-','k-'] #mbcgyrk\n",
    "#colors = iter(cm.jet(np.linspace(0, 1, len(Ek.T))))\n",
    "#bb=50\n",
    "#plt.title(\"$\\\\beta$ = {:.2f} along $\\\\Gamma$-$K$-$M$\".format(beta[bb]), fontsize=14)\n",
    "for e in range(len(Ek.T)):\n",
    "    colors = iter(cm.jet(np.linspace(0, 1, 5)))\n",
    "    #colors = iter(cm.jet(np.linspace(0, 1, 26)))\n",
    "    #colors = iter(cm.jet(np.linspace(0, 1, len(Ek))))\n",
    "    #col = next(colors)\n",
    "    #for a in range(len(Ek)-10):\n",
    "    leg=[]\n",
    "    for a in range(0,14,3):\n",
    "    #for a in range(25,76,2):\n",
    "    #for a in range(len(Ek)):\n",
    "        leg.append(str(angles[a])+\"°\")\n",
    "        col = next(colors)\n",
    "        Y=Ek[a,:,e].T\n",
    "        ax.plot(X,np.sqrt(np.abs(np.real(Y))),c=col)#, label=\"Ang = {:.1f}\".format(angles[a]))#,colors[e])\n",
    "plt.grid()\n",
    "plt.legend(leg,bbox_to_anchor=(1, 1),fontsize=15)\n",
    "#plt.xticks([0,133,199,314],['$\\Gamma$','K','M','$\\Gamma$'],fontsize=25)#[0,40,60,94]\n",
    "plt.yticks(fontsize=25)\n",
    "plt.ylabel('$\\omega (k)$', fontsize=25)\n",
    "plt.xlabel('k ($\\pi/a$)', fontsize=25)\n",
    "#plt.xlim(0,314)\n",
    "#plt.ylim(1.20,1.24)\n",
    "#ax.xaxis.set_major_locator(ticker.MultipleLocator(0.5))\n",
    "#ax.yaxis.set_major_locator(ticker.MultipleLocator(4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "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
}
