{
 "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 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": 2,
   "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": 3,
   "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": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compamatrix1D(N):\n",
    "    C=np.zeros((N,N))\n",
    "    for p in range(N):\n",
    "        if p < N-1:\n",
    "            C[p][p]=1\n",
    "            C[p][p+1]=-1\n",
    "        if p == N-1:\n",
    "            C[p][p]=1\n",
    "            C[p][0]=-1\n",
    "    return C"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compamatrix(bondlist,x0,y0):\n",
    "    N=np.max(bondlist)\n",
    "    M=len(bondlist)\n",
    "    C=np.zeros((M,2*N))\n",
    "    for p in range(M):\n",
    "        a1,a2=bondlist[p]\n",
    "        xa1,xa2=x0[a1-1],x0[a2-1]\n",
    "        ya1,ya2=y0[a1-1],y0[a2-1]\n",
    "        if a1 < a2:\n",
    "            C[p][2*a1-2]=xa1-xa2\n",
    "            C[p][2*a2-2]=xa2-xa1\n",
    "        elif a1 > a2:\n",
    "            C[p][2*a1-2]=xa2-xa1\n",
    "            C[p][2*a2-2]=xa1-xa2            \n",
    "        C[p][2*a1-1]=ya1-ya2\n",
    "        C[p][2*a2-1]=ya2-ya1\n",
    "    return C"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def wilsoncurva(Vk,kk,eigv):\n",
    "    def mmatrix(Vk,kx,eigv):\n",
    "        M1x=np.zeros((eigv,eigv), dtype=complex)\n",
    "        for e in range(eigv):\n",
    "            for p in range(eigv):\n",
    "                Vec0=Vk[kx,:,e]\n",
    "                Vecx=Vk[kx+1,:,p]\n",
    "                M1x[e,p]=np.dot(Vec0,Vecx)#/np.abs(np.dot(Vec0,Vecx))\n",
    "        return M1x\n",
    "    chernn=[]\n",
    "    berryphase=np.zeros((eigv,len(kk)))\n",
    "    for kx in range(-1,len(kk)-1):\n",
    "        Wx=np.identity(eigv, dtype=complex)\n",
    "        for ky in range(-1,len(kk)-1):\n",
    "            M1x=mmatrix(Vk,ky,eigv)\n",
    "            Wx=np.matmul(Wx,M1x)\n",
    "        Ex,Vx=np.linalg.eig(Wx)\n",
    "        for e in range(len(Ex)):\n",
    "            berryphase[e,kx]=-1*np.imag(np.log(Ex[e]))\n",
    "    chernumb=np.sum(berryphase,axis=1)/(2*np.pi)\n",
    "    return berryphase,chernumb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def wilsoncurva1D(Vk,kk,eigv):\n",
    "    def mmatrix(Vk,kx,eigv):\n",
    "        M1x=np.zeros((eigv,eigv), dtype=complex)\n",
    "        for e in range(eigv):\n",
    "            for p in range(eigv):\n",
    "                Vec0=Vk[kx,:,e]\n",
    "                Vecx=Vk[kx+1,:,p]\n",
    "                ux=np.dot(Vec0,Vecx)#/np.abs(np.dot(Vec0,Vecx))\n",
    "                M1x[e,p]=ux\n",
    "        return M1x\n",
    "    chernn=[]\n",
    "    berryphase=np.zeros((eigv))\n",
    "    #for kx in range(-1,len(kk)-1):\n",
    "    Wx=np.identity(eigv, dtype=complex)\n",
    "    for kx in range(-1,len(kk)-1):\n",
    "        M1x=mmatrix(Vk,kx,eigv)\n",
    "        Wx=np.matmul(Wx,M1x)\n",
    "    Ex,Vx=scipy.linalg.eig(Wx)\n",
    "    for e in range(len(Ex)):\n",
    "        berryphase[e]=-1*np.imag(np.log(Ex[e]))\n",
    "    chernumb=np.sum(berryphase)/(2*np.pi)\n",
    "    return berryphase,chernumb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 113,
   "metadata": {},
   "outputs": [],
   "source": [
    "def berrycurvature(Vk,kk,eigv=0):\n",
    "    berry=[]\n",
    "    for kx in range(-1,len(kk)-1):\n",
    "        Vec0=Vk[kx,:,eigv]\n",
    "        Vec1=Vk[kx+1,:,eigv]\n",
    "        u1=np.dot(Vec0,Vec1)/np.absolute(np.dot(Vec0,Vec1)) # problem!\n",
    "        bcur=np.log(u1)\n",
    "        berry.append(bcur)\n",
    "    berryc=np.asarray(berry, dtype=complex)\n",
    "    chernn=np.imag(np.sum(berryc)/(2*np.pi))\n",
    "    return berryc,chernn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 107,
   "metadata": {},
   "outputs": [],
   "source": [
    "def berrycurvature2(Vk,kk,eigv=0):\n",
    "    k0=len(kk)//2\n",
    "    berry=[]\n",
    "    for kx in range(-1,len(kk)-1):\n",
    "        Vec0=Vk[k0,:,eigv]\n",
    "        Vec1=Vk[kx,:,eigv]\n",
    "        u1=np.dot(Vec0,Vec1)/np.absolute(np.dot(Vec0,Vec1))\n",
    "        bcur=np.log(u1)\n",
    "        berry.append(bcur)\n",
    "    berryc=np.asarray(berry, dtype=complex)\n",
    "    chernn=np.imag(np.sum(berryc)/(2*np.pi))\n",
    "    return berryc,chernn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 210,
   "metadata": {},
   "outputs": [],
   "source": [
    "def berrycurvature3(Vk,kk,eigv=0):\n",
    "    berry=[]\n",
    "    for kx in range(-1,len(kk)-1):\n",
    "        Vec0=Vk[kx,:,eigv]\n",
    "        Vec1=Vk[kx+1,:,eigv]\n",
    "        u1=np.dot(Vec0,Vec1)/np.absolute(np.dot(Vec0,Vec1))\n",
    "       # u2=-1*np.dot(Vec0,Vec1)/np.absolute(np.dot(Vec0,Vec1))\n",
    "        bcur=np.log(u1)\n",
    "       # bcur2=np.log(u2)\n",
    "        if len(berry) > 1:\n",
    "            if np.absolute(bcur-berry[-1]) > np.pi/4 and np.absolute(bcur-berry[-1]) < 5*np.pi/4:\n",
    "                bcur=np.log(-1*u1)\n",
    "                berry.append(bcur)\n",
    "            else:\n",
    "                berry.append(bcur)\n",
    "        else:\n",
    "            berry.append(bcur)\n",
    "    berryc=np.asarray(berry, dtype=complex)\n",
    "    chernn=np.sum(np.imag(berryc)/(2*np.pi))\n",
    "    return berryc,chernn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 211,
   "metadata": {},
   "outputs": [],
   "source": [
    "def chernnumb(Vk,kk,eig=0):\n",
    "    bpha=0\n",
    "    for kx in range(-1,len(kk)-1):\n",
    "        for ky in range(-1,len(kk)-1):\n",
    "            Vec0=Vk[kx,ky,:,eig]\n",
    "            Vec1=Vk[kx+1,ky,:,eig]\n",
    "            Vec12=Vk[kx+1,ky+1,:,eig]\n",
    "            Vec2=Vk[kx,ky+1,:,eig]\n",
    "            u1=np.dot(Vec0,Vec1)/np.abs(np.dot(Vec0,Vec1))\n",
    "            u2=np.dot(Vec1,Vec12)/np.abs(np.dot(Vec1,Vec12))\n",
    "            u3=np.dot(Vec2,Vec12)/np.abs(np.dot(Vec2,Vec12))\n",
    "            u4=np.dot(Vec0,Vec2)/np.abs(np.dot(Vec0,Vec2))\n",
    "            bcur=np.log(u1*u2/(u3*u4))\n",
    "            bpha+=bcur\n",
    "    chen=bpha/(2j*np.pi)\n",
    "    return chen"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 212,
   "metadata": {},
   "outputs": [],
   "source": [
    "def kvalues(N,k1,k2,order=\"one\"):\n",
    "    if order == \"one\":\n",
    "        return np.ones((N))\n",
    "    elif order == \"random\":\n",
    "        kvals=np.random.randint(2, size=N)\n",
    "    elif order == \"alternate\":\n",
    "        kvals=np.ones((N))\n",
    "        for p in range(N):\n",
    "            if p%2 == 0:\n",
    "                kvals[p]=0\n",
    "    for p in range(N):\n",
    "        if kvals[p] == 0:\n",
    "            kvals[p]=k1\n",
    "        elif kvals[p] == 1:\n",
    "            kvals[p]=k2\n",
    "    return kvals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 213,
   "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": 214,
   "metadata": {},
   "outputs": [],
   "source": [
    "def bearlponds(N):\n",
    "    pbonds=[]\n",
    "    for p in range(N):\n",
    "        if p == N-1:\n",
    "            pbonds.append([p,0])\n",
    "        elif p < N-1:\n",
    "            pbonds.append([p,p+1])\n",
    "    return np.asarray(pbonds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "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,:,eg],Vec[ka,:,eg]))\n",
    "    return phaseV"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "def corrematrix(Vec):\n",
    "    #phaseV=np.zeros(np.shape(Vec[:,2,0]))\n",
    "    npearl=len(Vec.T)\n",
    "    phaseM=np.zeros((len(Vec),npearl,npearl,npearl), dtype = complex)\n",
    "    for ka in range(len(Vec)): # k-vector\n",
    "        for eg in range(npearl):# eigenmode\n",
    "            for p in range(npearl): # pearl\n",
    "                for q in range(npearl): # other pearl\n",
    "                    phaseM[ka,eg,p,q]=np.vdot(Vec[ka,p,eg],Vec[ka,q,eg])\n",
    "#                    phaseM[ka,p,q]=np.real(Vec[ka,p,eg]*Vec[ka,q,eg])\n",
    "    return phaseM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "def JJphase(Vec):\n",
    "    #phaseV=np.zeros(np.shape(Vec[:,2,0]))\n",
    "    npearl=len(Vec.T)\n",
    "    phaseM=np.zeros((len(Vec),npearl), dtype = complex)\n",
    "    for ka in range(len(Vec)): # k-vector\n",
    "        for eg in range(npearl):# eigenmode\n",
    "            phase=0\n",
    "            for p in range(npearl-1): # pearl\n",
    "                phase+=np.vdot(Vec[ka,p,eg],Vec[ka,p+1,eg])\n",
    "            phase+=np.vdot(Vec[ka,0,eg],Vec[ka,-1,eg])\n",
    "            phaseM[ka,eg]=phase\n",
    "    return np.real(phaseM)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "def ordereig(E,V):\n",
    "    V2=np.zeros(np.shape(V), dtype = complex)\n",
    "    E2=np.zeros(np.shape(E), dtype = complex)\n",
    "    for p in range(len(E)):\n",
    "        Eval=np.real(E[p])\n",
    "        sortindex=sorted(range(len(Eval)), key=lambda k: s[k])\n",
    "        for q in range(len(Eval)):\n",
    "            E2[p][q]=E[p][sortindex[q]]\n",
    "            V2[p,:,q]=V[p,:,sortindex[q]]\n",
    "    return E2,V2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### One dimensional pbc chain of pearls with alternating spring constant"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 229,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Spring constant values\n",
    "#k1,k2=3,1 # trivial case\n",
    "k1,k2=3,1 # topological case\n",
    "#k1,k2=5,5 # equal case"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 230,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([3., 1.]),\n",
       " array([1., 1.]),\n",
       " array([[ 1., -1.],\n",
       "        [-1.,  1.]]))"
      ]
     },
     "execution_count": 230,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "a=np.pi # spacing between pearls\n",
    "Np=2# number of pearls\n",
    "C=compamatrix1D(Np)\n",
    "segm=100\n",
    "kk=np.array([1*np.pi/(a*segm)*e for e in range(-segm,segm+1)]) #k points\n",
    "Ckk=np.zeros((len(kk),len(C),len(C.T)), dtype = complex)\n",
    "for e in range(len(kk)):\n",
    "    Ckk[e]=C\n",
    "    Ckk[e][Np-1][0]=C[Np-1][0]*np.exp(1j * (kk[e] * a))\n",
    "#k1,k2=1,5\n",
    "k=kvalues(Np,k1,k2,\"alternate\")# spring constants could be all \"one\", \"random\", \"alternate\" (k1,k2,k1,...)\n",
    "bondlist=bearlponds(Np)\n",
    "mA,mB=1,2\n",
    "masses=massvalues(Np,mA,mB,\"one\")# mass values could be all \"one\", \"random\", \"alternate\" (k1,k2,k1,...)\n",
    "#masses=np.array([1, 1, 1, 1, 100, 1, 1, 1, 1, 1])\n",
    "k,masses,C"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 231,
   "metadata": {},
   "outputs": [],
   "source": [
    "D=dynamicmatrixC(C,k,masses)\n",
    "Dk=np.zeros((len(kk),len(D),len(D.T)), dtype = complex)\n",
    "for e in range(len(kk)):\n",
    "    Dk[e]=dynamicmatrixC(Ckk[e],k,masses)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 232,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Diagonalization of Mwh\n",
    "E,V = np.linalg.eigh(D)\n",
    "#for i in range(len(E)):\n",
    "#    print(E[i])\n",
    "#    if abs(np.log10(abs(E[i]))) > 10:\n",
    "#        zerom=i+1\n",
    "#zerom,len(E)\n",
    "Ek=np.zeros((len(kk),len(E)), dtype = complex)\n",
    "Vk=np.zeros((len(kk),len(V),len(V.T)), dtype = complex)\n",
    "Ek2=np.zeros((len(kk),len(E)), dtype = complex)\n",
    "Vk2=np.zeros((len(kk),len(V),len(V.T)), dtype = complex)\n",
    "for e in range(len(kk)):\n",
    "    Ek[e],Vk[e]=scipy.linalg.eigh(Dk[e])\n",
    "#    Ek2[e],Vk2[e]=np.linalg.eig(Dk[e])\n",
    "#Ek3,Vk3=ordereig(Ek2,Vk2)\n",
    "phaseVk=phasevec(Vk,0,1)\n",
    "#np.real(corrematrix(Vk)[100])\n",
    "jjphase=JJphase(Vk)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 233,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Calculate Berry curvature and chern number using method from Fukui, Hatsugai and Suzuki J.Phys.Soc.Jpn. 2005\n",
    "#chernnumb(Vk,kk,eig=0) # Chern number per eigenvalue, the real part must be an integer\n",
    "berrc=np.zeros((len(Ek.T),len(kk)), dtype=complex)\n",
    "chernn=np.zeros((len(Ek.T)))\n",
    "for eig in range(len(berrc)):\n",
    "    bc,cn=berrycurvature3(Vk,kk,eig)\n",
    "    berrc[eig]=bc\n",
    "    chernn[eig]=cn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 234,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "6.83980106397914"
      ]
     },
     "execution_count": 234,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.sum(np.abs(np.imag(berrc[0,:segm])))/np.pi#/(len(kk)//2)/np.pi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 235,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "201"
      ]
     },
     "execution_count": 235,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(kk)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 236,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([8.8817842e-16, 0.0000000e+00])"
      ]
     },
     "execution_count": 236,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chernn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 237,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 36412 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",
    "Ylim=np.sqrt(np.abs(np.real(Ek.T)))\n",
    "#plt.title('Eigenmodes for {:d} Pearls'.format(Np))\n",
    "#plt.title('Eigenmodes for {:d} Pearls: -{:d}-[-{:d}-]-{:d}- '.format(Np,k2,k1,k2))\n",
    "Z=np.sort(jjphase.reshape(len(jjphase)*len(jjphase.T)))\n",
    "X=kk\n",
    "for egv in range(len(Ek.T)):\n",
    "    #X=kk\n",
    "    Y=np.sqrt(np.abs(np.real(Ek.T[egv])))\n",
    "    Z2=jjphase[:,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('$\\lambda (k)$', fontsize=20)\n",
    "plt.xlabel('k ($\\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": 97,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 32104 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",
    "Ylim=np.sqrt(np.abs(np.real(Ek.T)))\n",
    "#plt.title('Eigenmodes for {:d} Pearls'.format(Np))\n",
    "#plt.title('Eigenmodes for {:d} Pearls: -{:d}-[-{:d}-]-{:d}- '.format(Np,k2,k1,k2))\n",
    "Z=np.sort(phaseVk.reshape(len(phaseVk)*len(phaseVk.T)))\n",
    "for egv in range(len(Ek.T)):\n",
    "    X=kk\n",
    "    Y=np.sqrt(np.abs(np.real(Ek.T[egv])))\n",
    "    Z2=phaseVk[:,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",
    "fig.colorbar(line, ax=ax)\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)$', fontsize=20)\n",
    "plt.xlabel('k ($\\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": 136,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 35992 characters]\n",
      "text/plain": [
       "<Figure size 504x504 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Overview surface and atomic indexes\n",
    "#eig,dire=2,0 # dire=0/1 for dependence on x/y\n",
    "dire=0 # dire=0/1 for dependence on x/y\n",
    "if dire == 0:\n",
    "    s='k'\n",
    "elif dire ==1:\n",
    "    s='k$_y$'    \n",
    "fig,ax = plt.subplots(figsize=(7,7)) # Figure size\n",
    "X=kk#/np.max(kk)\n",
    "#plt.title('Berry phase for {:d} Pearls: -{:d}-[-{:d}-]-{:d}- '.format(Np,k2,k1,k2))\n",
    "plt.title('Berry phase for {:d} Pearls: -{:d}-[-{:d}-]-{:d}- '.format(Np,k2,k1,k2))\n",
    "#colors=['yo','ro','co','go','mo','bo','ko','yo','ro','co'] #mbcgyrk\n",
    "colors = iter(cm.jet(np.linspace(0, 1, Np)))\n",
    "for eig in range(Np):\n",
    "    col=next(colors)\n",
    "    #X=np.real(berrc[eig])\n",
    "    Y=np.imag(berrc[eig])#%(np.pi)\n",
    "    ax.plot(X,Y/np.pi+0*eig,c=col,marker='o',linestyle=\"\",markersize=5,label=\"C$_{:.0f}$ = {:.4f}\".format(eig+1,chernn[eig]))\n",
    "#    ax.plot([-1,1],[-np.pi,-np.pi],c=\"red\")\n",
    "#    ax.plot([-1,1],[np.pi,np.pi],c=\"red\")\n",
    "#    ax.plot([-1,1],[np.pi/2,np.pi/2],c=\"red\")\n",
    "    plt.grid()\n",
    "plt.xticks(fontsize=25)\n",
    "plt.yticks(fontsize=25)\n",
    "plt.ylabel('$\\\\theta (\\pi)$', fontsize=25)\n",
    "plt.xlabel(s+' ($\\pi/a$)', fontsize=25)\n",
    "#plt.xlim(-1,1)\n",
    "plt.ylim(-1.1,1.1)\n",
    "#plt.legend(fontsize=12)#,loc=(1.02,0))\n",
    "ax.xaxis.set_major_locator(ticker.MultipleLocator(0.5))\n",
    "ax.yaxis.set_major_locator(ticker.MultipleLocator(0.5))\n",
    "plt.tight_layout()\n",
    "plt.grid()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 106,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 44828 characters]\n",
      "text/plain": [
       "<Figure size 576x576 with 1 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=(8,8))\n",
    "colors = iter(cm.jet(np.linspace(0, 1, Np)))\n",
    "for egv in range(len(jjphase.T)):\n",
    "    col=next(colors)\n",
    "    X=kk\n",
    "    Y=jjphase[:,egv]\n",
    "    cjj=np.sum(Y)/segm\n",
    "    ax.plot(X,Y,marker='o',c=col,linestyle=\"\",markersize=5,label=\"Cj$_{:.0f}$ = {:.4f}\".format(egv+1,cjj))\n",
    "#plt.xlim(-1,-0.5)\n",
    "#plt.ylim(0.9,1.1)\n",
    "plt.grid()\n",
    "plt.legend(loc=1)\n",
    "plt.xticks(fontsize=16)\n",
    "plt.yticks(fontsize=16)\n",
    "plt.ylabel('JJPhase', fontsize=20)\n",
    "plt.xlabel('k ($\\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": 141,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Text(0.5, 0, 'k$_x$ ($\\\\pi/a$)')"
      ]
     },
     "execution_count": 141,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 34032 characters]\n",
      "text/plain": [
       "<Figure size 576x576 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Overview surface and atomic indexes\n",
    "fig,ax = plt.subplots(figsize=(8,8)) # Figure size\n",
    "X=kk#/np.max(kk)*1\n",
    "Y=Ek.T\n",
    "plt.title('Eigenmodes for {:d} Pearls: -{:d}-[-{:d}-]-{:d}- '.format(Np,k2,k1,k2))\n",
    "#for e in range(0,5):\n",
    "#plt.plot(X,np.real(Y[e]),'k--')\n",
    "#colors=['y-','r-','c-','g-','m-','b-','k-','y-','r-','c-','g-','m-'] #mbcgyrk\n",
    "colors = iter(cm.jet(np.linspace(0, 1, Np)))\n",
    "for e in range(0,Np):\n",
    "    col = next(colors)\n",
    "    ax.plot(X,np.sqrt(np.abs(np.real(Y[e]))),c=col,label=\"Eigenmode \"+str(e+1))\n",
    "#   plt.plot(X,np.real(Y[e]),colors[e],label=\"Eigenvalue \"+str(e+1))\n",
    "#   plt.title('k$_y$ = {:.2f} '.format(wnum))\n",
    "#plt.xlim(-1,1)\n",
    "plt.grid()\n",
    "#plt.legend()\n",
    "#plt.ylim(0,30)\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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Fixed boundary conditions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([1., 2., 1., 2., 1., 2., 1., 2., 1., 2.]),\n",
       " array([[-1.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],\n",
       "        [ 1., -1.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],\n",
       "        [ 0.,  1., -1.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],\n",
       "        [ 0.,  0.,  1., -1.,  0.,  0.,  0.,  0.,  0.,  0.],\n",
       "        [ 0.,  0.,  0.,  1., -1.,  0.,  0.,  0.,  0.,  0.],\n",
       "        [ 0.,  0.,  0.,  0.,  1., -1.,  0.,  0.,  0.,  0.],\n",
       "        [ 0.,  0.,  0.,  0.,  0.,  1., -1.,  0.,  0.,  0.],\n",
       "        [ 0.,  0.,  0.,  0.,  0.,  0.,  1., -1.,  0.,  0.],\n",
       "        [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  1., -1.,  0.],\n",
       "        [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  1., -1.],\n",
       "        [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  1.]]))"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "a=np.pi # spacing between pearls\n",
    "Np=10# number of pearls\n",
    "C=compamatrix1D(Np+2)[:-1,1:-1]\n",
    "#bondlist=bearlponds(Np)\n",
    "mA,mB=1,2\n",
    "masses=massvalues(Np,mA,mB,\"alternate\")# mass values could be all \"one\", \"random\", \"alternate\" (k1,k2,k1,...)\n",
    "#masses=np.array([1, 1, 1, 1, 100, 1, 1, 1, 1, 1])\n",
    "masses,C"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a=np.pi # spacing between pearls\n",
    "Np=3# number of pearls\n",
    "C=compamatrix1D(Np)\n",
    "segm=100\n",
    "kk=np.array([1*np.pi/(a*segm)*e for e in range(-segm,segm+1)]) #k points\n",
    "Ckk=np.zeros((len(kk),len(C),len(C.T)), dtype = complex)\n",
    "for e in range(len(kk)):\n",
    "    Ckk[e]=C\n",
    "    Ckk[e][Np-1][0]=C[Np-1][0]*np.exp(1j * (kk[e] * a))\n",
    "#k1,k2=1,5\n",
    "k=kvalues(Np,k1,k2,\"alternate\")# spring constants could be all \"one\", \"random\", \"alternate\" (k1,k2,k1,...)\n",
    "bondlist=bearlponds(Np)\n",
    "mA,mB=1,2\n",
    "masses=massvalues(Np,mA,mB,\"one\")# mass values could be all \"one\", \"random\", \"alternate\" (k1,k2,k1,...)\n",
    "#masses=np.array([1, 1, 1, 1, 100, 1, 1, 1, 1, 1])\n",
    "k,masses,C"
   ]
  }
 ],
 "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
}
