#!/usr/bin/env python
# -*- coding:utf-8 -*-
import numpy as np
import sympy as sp
import matplotlib.pyplot as plt

#定義四個力系數(shù)矩陣
K1=np.array([[36.50,0,0],[0,24.50,0],[0,0,9.82]])*14000
K2=np.array([[8.80,0,0],[0,-3.23,0],[0,0,-0.40]])*14000
K3=np.array([[3.00,0,0],[0,-5.25,0],[0,0,0.15]])*14000
K4=np.array([[-1.92,0,0],[0,2.29,0],[0,0,-0.58]])*14000
a=1.42*10**(-10)

#定義旋轉(zhuǎn)矩陣：
I=np.array([[1,0,0],[0,1,0],[0,0,1]])
I120=np.array([[np.cos(2/3*np.pi),np.sin(2/3*np.pi),0],[-np.sin(2/3*np.pi),np.cos(2/3*np.pi),0],[0,0,1]])
I60=np.array([[np.cos(1/3*np.pi),np.sin(1/3*np.pi),0],[-np.sin(1/3*np.pi),np.cos(1/3*np.pi),0],[0,0,1]])
I30=np.array([[np.cos(1/6*np.pi),np.sin(1/6*np.pi),0],[-np.sin(1/6*np.pi),np.cos(1/6*np.pi),0],[0,0,1]])
I30f=np.array([[np.cos(-1/6*np.pi),np.sin(-1/6*np.pi),0],[-np.sin(-1/6*np.pi),np.cos(-1/6*np.pi),0],[0,0,1]])
I4=np.array([[5/(2*np.sqrt(7)),np.sqrt(3)/(2*np.sqrt(7)),0],[-np.sqrt(3)/(2*np.sqrt(7)),5/(2*np.sqrt(7)),0],[0,0,1]])
I4f=np.array([[5/(2*np.sqrt(7)),-np.sqrt(3)/(2*np.sqrt(7)),0],[np.sqrt(3)/(2*np.sqrt(7)),5/(2*np.sqrt(7)),0],[0,0,1]])
I4t=np.array([[2/np.sqrt(7),np.sqrt(3)/np.sqrt(7),0],[-np.sqrt(3)/np.sqrt(7),2/np.sqrt(7),0],[0,0,1]])
I4tf=np.array([[2/np.sqrt(7),-np.sqrt(3)/np.sqrt(7),0],[np.sqrt(3)/np.sqrt(7),2/np.sqrt(7),0],[0,0,1]])
#原子A的方向矩陣
U11=I
U12=I120
U13=I120@I120

U21=I30
U26=I30f
U22=np.matmul(I120,U26)
U23=np.matmul(I120,U21)
U24=np.matmul(I120,U22)
U25=np.matmul(I120,U23)

U31=I60
U32=np.matmul(I120,U31)
U33=np.matmul(I120,U32)

U41=I4
U46=I4f
U42=np.matmul(I120,U46)
U43=np.matmul(I120,U41)
U44=np.matmul(I120,U42)
U45=np.matmul(I120,U43)

#原子A的力系數(shù)矩陣。
A11=K1
A12=np.linalg.inv(U12)@K1@U12
A13=np.linalg.inv(U13)@K1@U13

A21=np.linalg.inv(U21)@K2@U21
A22=np.linalg.inv(U22)@K2@U22
A23=np.linalg.inv(U23)@K2@U23
A24=np.linalg.inv(U24)@K2@U24
A25=np.linalg.inv(U25)@K2@U25
A26=np.linalg.inv(U26)@K2@U26

A31=np.linalg.inv(U31)@K3@U31
A32=np.linalg.inv(U32)@K3@U32
A33=np.linalg.inv(U33)@K3@U33

A41=np.linalg.inv(U41)@K4@U41
A42=np.linalg.inv(U42)@K4@U42
A43=np.linalg.inv(U43)@K4@U43
A44=np.linalg.inv(U44)@K4@U44
A45=np.linalg.inv(U45)@K4@U45
A46=np.linalg.inv(U46)@K4@U46

A=A11+A12+A13+A21+A22+A23+A24+A25+A26+A31+A32+A33+A41+A42+A43+A44+A45+A46

#原子B的位置矩陣
X11=I60
X12=np.matmul(I120,X11)
X13=np.matmul(I120,X12)

X21=I30
X26=I30f
X22=np.matmul(I120,X26)
X23=np.matmul(I120,X21)
X24=np.matmul(I120,X22)
X25=np.matmul(I120,X23)

X31=I
X32=I120
X33=I120@I120

X41=I4t
X46=I4tf
X42=np.matmul(I120,X46)
X43=np.matmul(I120,X41)
X44=np.matmul(I120,X42)
X45=np.matmul(I120,X43)

#原子B的力系數(shù)矩陣

B11=np.linalg.inv(X11)@K1@X11
B12=np.linalg.inv(X12)@K1@X12
B13=np.linalg.inv(X13)@K1@X13

B21=np.linalg.inv(X21)@K2@X21
B22=np.linalg.inv(X22)@K2@X22
B23=np.linalg.inv(X23)@K2@X23
B24=np.linalg.inv(X24)@K2@X24
B25=np.linalg.inv(X25)@K2@X25
B26=np.linalg.inv(X26)@K2@X26

B31=np.linalg.inv(X31)@K3@X31
B32=np.linalg.inv(X32)@K3@X32
B33=np.linalg.inv(X33)@K3@X33

B41=np.linalg.inv(X41)@K4@X41
B42=np.linalg.inv(X42)@K4@X42
B43=np.linalg.inv(X43)@K4@X43
B44=np.linalg.inv(X44)@K4@X44
B45=np.linalg.inv(X45)@K4@X45
B46=np.linalg.inv(X46)@K4@X46

B=B11+B12+B13+B21+B22+B23+B24+B25+B26+B31+B32+B33+B41+B42+B43+B44+B45+B46

#按照公式
#第一段
w1=np.zeros([850,6])
for i in range(850):
? ?kx=i/(850/(2/3*np.pi))/a
? ?ky=0
? ?# OME=np.array([[w**2,0,0],[0,w**2,0],[0,0,w**2]],dtype=complex)

? ?DAA=A-A21*np.exp(np.sqrt(3)*a*1j*(kx*np.cos(1/6*np.pi)+ky*np.sin(1/6*np.pi)))\
? ? ? ? -A22*np.exp(np.sqrt(3)*a*1j*(kx*np.cos(3/6*np.pi)+ky*np.sin(3/6*np.pi)))\
? ? ? ? -A23*np.exp(np.sqrt(3)*a*1j*(kx*np.cos(5/6*np.pi)+ky*np.sin(5/6*np.pi)))\
? ? ? ? -A24*np.exp(np.sqrt(3)*a*1j*(kx*np.cos(7/6*np.pi)+ky*np.sin(7/6*np.pi)))\
? ? ? ? -A25*np.exp(np.sqrt(3)*a*1j*(kx*np.cos(9/6*np.pi)+ky*np.sin(9/6*np.pi)))\
? ? ? ? -A26*np.exp(np.sqrt(3)*a*1j*(kx*np.cos(11/6*np.pi)+ky*np.sin(11/6*np.pi)))\

? ?DAB=-A11*np.exp(a*1j*(kx*np.cos(0*np.pi)+ky*np.sin(0*np.pi)))\
? ? ? ?-A12*np.exp(a*1j*(kx*np.cos(2/3*np.pi)+ky*np.sin(2/3*np.pi)))\
? ? ? ?-A13*np.exp(a*1j*(kx*np.cos(4/3*np.pi)+ky*np.sin(4/3*np.pi)))\
? ? ? ?-A31*np.exp(2*a*1j*(kx*np.cos(1/3*np.pi)+ky*np.sin(1/3*np.pi)))\
? ? ? ?-A32*np.exp(2*a*1j*(kx*np.cos(1*np.pi)+ky*np.sin(1*np.pi)))\
? ? ? ?-A33*np.exp(2*a*1j*(kx*np.cos(5/3*np.pi)+ky*np.sin(5/3*np.pi)))\
? ? ? ?-A41*np.exp(a*1j*(kx*5/2+ky*np.sqrt(3)/2))\
? ? ? ?-A42*np.exp(a*1j*(kx*(-1/2)+ky*3*np.sqrt(3)/2))\
? ? ? ?-A43*np.exp(a*1j*(kx*(-2)+ky*np.sqrt(3)))\
? ? ? ?-A44*np.exp(a*1j*(kx*(-2)-ky*np.sqrt(3)))\
? ? ? ?-A45*np.exp(a*1j*(kx*(-1/2)-ky*3*np.sqrt(3)/2))\
? ? ? ?-A46*np.exp(a*1j*(kx*5/2-ky*np.sqrt(3)/2))\

? ?DBA=-B11*np.exp(a*1j*(kx*np.cos(1/3*np.pi)+ky*np.sin(1/3*np.pi)))\
? ? ? ?-B12*np.exp(a*1j*(kx*np.cos(1*np.pi)+ky*np.sin(1*np.pi)))\
? ? ? ?-B13*np.exp(a*1j*(kx*np.cos(5/3*np.pi)+ky*np.sin(5/3*np.pi)))\
? ? ? ?-B31*np.exp(2*a*1j*(kx*np.cos(0*np.pi)+ky*np.sin(0*np.pi)))\
? ? ? ?-B32*np.exp(2*a*1j*(kx*np.cos(2/3*np.pi)+ky*np.sin(2/3*np.pi)))\
? ? ? ?-B33*np.exp(2*a*1j*(kx*np.cos(4/3*np.pi)+ky*np.sin(4/3*np.pi)))\
? ? ? ?-B41*np.exp(a*1j*(kx*2+ky*np.sqrt(3)))\
? ? ? ?-B42*np.exp(a*1j*(kx*1/2+ky*3*np.sqrt(3)/2))\
? ? ? ?-B43*np.exp(a*1j*(kx*(-5/2)+ky*np.sqrt(3)/2))\
? ? ? ?-B44*np.exp(a*1j*(kx*(-5/2)-ky*np.sqrt(3)/2))\
? ? ? ?-B45*np.exp(a*1j*(kx*1/2-ky*3*np.sqrt(3))/2)\
? ? ? ?-B46*np.exp(a*1j*(kx*2-ky*np.sqrt(3)))\

? ?DBB=B-B21*np.exp(np.sqrt(3)*a*1j*(kx*np.cos(1/6*np.pi)+ky*np.sin(1/6*np.pi)))\
? ? ? ? -B22*np.exp(np.sqrt(3)*a*1j*(kx*np.cos(3/6*np.pi)+ky*np.sin(3/6*np.pi)))\
? ? ? ? -B23*np.exp(np.sqrt(3)*a*1j*(kx*np.cos(5/6*np.pi)+ky*np.sin(5/6*np.pi)))\
? ? ? ? -B24*np.exp(np.sqrt(3)*a*1j*(kx*np.cos(7/6*np.pi)+ky*np.sin(7/6*np.pi)))\
? ? ? ? -B25*np.exp(np.sqrt(3)*a*1j*(kx*np.cos(9/6*np.pi)+ky*np.sin(9/6*np.pi)))\
? ? ? ? -B26*np.exp(np.sqrt(3)*a*1j*(kx*np.cos(11/6*np.pi)+ky*np.sin(11/6*np.pi)))\


? ?#組合成總矩陣D：
? ?D=np.zeros([6,6],dtype=complex)
? ?D[0:3,0:3]=DAA
? ?D[0:3,3:6]=DAB
? ?D[3:6,0:3]=DBA
? ?D[3:6,3:6]=DBB


? ?c=np.linalg.eig(D)
? ?# e=np.zeros([6,6],dtype=complex)
? ?# for p in range(6):
? ?# ? ? e[p,p]=c[0][p]

? ?for p in range(6):
? ? ?w1[i,p]=np.real(np.sqrt(c[0][p]))

#第二段
w2=np.zeros([500,6])
for i in range(500):
? ?kx=2/3*np.pi/a
? ?ky=i/(500/(2/9*np.sqrt(3)*np.pi))/a
? ?# OME=np.array([[w**2,0,0],[0,w**2,0],[0,0,w**2]],dtype=complex)

? ?DAA=A-A21*np.exp(np.sqrt(3)*a*1j*(kx*np.cos(1/6*np.pi)+ky*np.sin(1/6*np.pi)))\
? ? ? ? ? ? -A22*np.exp(np.sqrt(3)*a*1j*(kx*np.cos(3/6*np.pi)+ky*np.sin(3/6*np.pi)))\
? ? ? ? ? ? -A23*np.exp(np.sqrt(3)*a*1j*(kx*np.cos(5/6*np.pi)+ky*np.sin(5/6*np.pi)))\
? ? ? ? ? ? -A24*np.exp(np.sqrt(3)*a*1j*(kx*np.cos(7/6*np.pi)+ky*np.sin(7/6*np.pi)))\
? ? ? ? ? ? -A25*np.exp(np.sqrt(3)*a*1j*(kx*np.cos(9/6*np.pi)+ky*np.sin(9/6*np.pi)))\
? ? ? ? ? ? -A26*np.exp(np.sqrt(3)*a*1j*(kx*np.cos(11/6*np.pi)+ky*np.sin(11/6*np.pi)))\

? ?DAB=-A11*np.exp(a*1j*(kx*np.cos(0*np.pi)+ky*np.sin(0*np.pi)))\
? ? ? ?-A12*np.exp(a*1j*(kx*np.cos(2/3*np.pi)+ky*np.sin(2/3*np.pi)))\
? ? ? ?-A13*np.exp(a*1j*(kx*np.cos(4/3*np.pi)+ky*np.sin(4/3*np.pi)))\
? ? ? ?-A31*np.exp(2*a*1j*(kx*np.cos(1/3*np.pi)+ky*np.sin(1/3*np.pi)))\
? ? ? ?-A32*np.exp(2*a*1j*(kx*np.cos(1*np.pi)+ky*np.sin(1*np.pi)))\
? ? ? ?-A33*np.exp(2*a*1j*(kx*np.cos(5/3*np.pi)+ky*np.sin(5/3*np.pi)))\
? ? ? ?-A41*np.exp(a*1j*(kx*5/2+ky*np.sqrt(3)/2))\
? ? ? ?-A42*np.exp(a*1j*(kx*(-1/2)+ky*3*np.sqrt(3)/2))\
? ? ? ?-A43*np.exp(a*1j*(kx*(-2)+ky*np.sqrt(3)))\
? ? ? ?-A44*np.exp(a*1j*(kx*(-2)-ky*np.sqrt(3)))\
? ? ? ?-A45*np.exp(a*1j*(kx*(-1/2)-ky*3*np.sqrt(3)/2))\
? ? ? ?-A46*np.exp(a*1j*(kx*5/2-ky*np.sqrt(3)/2))\

? ?DBA=-B11*np.exp(a*1j*(kx*np.cos(1/3*np.pi)+ky*np.sin(1/3*np.pi)))\
? ? ? ?-B12*np.exp(a*1j*(kx*np.cos(1*np.pi)+ky*np.sin(1*np.pi)))\
? ? ? ?-B13*np.exp(a*1j*(kx*np.cos(5/3*np.pi)+ky*np.sin(5/3*np.pi)))\
? ? ? ?-B31*np.exp(2*a*1j*(kx*np.cos(0*np.pi)+ky*np.sin(0*np.pi)))\
? ? ? ?-B32*np.exp(2*a*1j*(kx*np.cos(2/3*np.pi)+ky*np.sin(2/3*np.pi)))\
? ? ? ?-B33*np.exp(2*a*1j*(kx*np.cos(4/3*np.pi)+ky*np.sin(4/3*np.pi)))\
? ? ? ?-B41*np.exp(a*1j*(kx*2+ky*np.sqrt(3)))\
? ? ? ?-B42*np.exp(a*1j*(kx*1/2+ky*3*np.sqrt(3)/2))\
? ? ? ?-B43*np.exp(a*1j*(kx*(-5/2)+ky*np.sqrt(3)/2))\
? ? ? ?-B44*np.exp(a*1j*(kx*(-5/2)-ky*np.sqrt(3)/2))\
? ? ? ?-B45*np.exp(a*1j*(kx*1/2-ky*3*np.sqrt(3))/2)\
? ? ? ?-B46*np.exp(a*1j*(kx*2-ky*np.sqrt(3)))\

? ?DBB=B-B21*np.exp(np.sqrt(3)*a*1j*(kx*np.cos(1/6*np.pi)+ky*np.sin(1/6*np.pi)))\
? ? ? ? ? ? -B22*np.exp(np.sqrt(3)*a*1j*(kx*np.cos(3/6*np.pi)+ky*np.sin(3/6*np.pi)))\
? ? ? ? ? ? -B23*np.exp(np.sqrt(3)*a*1j*(kx*np.cos(5/6*np.pi)+ky*np.sin(5/6*np.pi)))\
? ? ? ? ? ? -B24*np.exp(np.sqrt(3)*a*1j*(kx*np.cos(7/6*np.pi)+ky*np.sin(7/6*np.pi)))\
? ? ? ? ? ? -B25*np.exp(np.sqrt(3)*a*1j*(kx*np.cos(9/6*np.pi)+ky*np.sin(9/6*np.pi)))\
? ? ? ? ? ? -B26*np.exp(np.sqrt(3)*a*1j*(kx*np.cos(11/6*np.pi)+ky*np.sin(11/6*np.pi)))\


? ?#組合成總矩陣D：
? ?D=np.zeros([6,6],dtype=complex)
? ?D[0:3,0:3]=DAA
? ?D[0:3,3:6]=DAB
? ?D[3:6,0:3]=DBA
? ?D[3:6,3:6]=DBB


? ?c=np.linalg.eig(D)
? ?e=np.zeros([6,6],dtype=complex)
? ?for p in range(6):
? ? ? ?e[p,p]=c[0][p]

? ?for p in range(6):
? ? ?w2[i,p]=np.real(np.sqrt(e[p,p]))



#第三段
w3=np.zeros([1000,6])
for i in range(1000):
? ?kx=(2/3*np.pi-2/3000*np.pi*i)/a
? ?ky=(2/9*np.sqrt(3)*np.pi-2/9000*np.sqrt(3)*np.pi*i)/a
? ?# OME=np.array([[w**2,0,0],[0,w**2,0],[0,0,w**2]],dtype=complex)

? ?DAA=A-A21*np.exp(np.sqrt(3)*a*1j*(kx*np.cos(1/6*np.pi)+ky*np.sin(1/6*np.pi)))\
? ? ? ? ? ? -A22*np.exp(np.sqrt(3)*a*1j*(kx*np.cos(3/6*np.pi)+ky*np.sin(3/6*np.pi)))\
? ? ? ? ? ? -A23*np.exp(np.sqrt(3)*a*1j*(kx*np.cos(5/6*np.pi)+ky*np.sin(5/6*np.pi)))\
? ? ? ? ? ? -A24*np.exp(np.sqrt(3)*a*1j*(kx*np.cos(7/6*np.pi)+ky*np.sin(7/6*np.pi)))\
? ? ? ? ? ? -A25*np.exp(np.sqrt(3)*a*1j*(kx*np.cos(9/6*np.pi)+ky*np.sin(9/6*np.pi)))\
? ? ? ? ? ? -A26*np.exp(np.sqrt(3)*a*1j*(kx*np.cos(11/6*np.pi)+ky*np.sin(11/6*np.pi)))\

? ?DAB=-A11*np.exp(a*1j*(kx*np.cos(0*np.pi)+ky*np.sin(0*np.pi)))\
? ? ? ?-A12*np.exp(a*1j*(kx*np.cos(2/3*np.pi)+ky*np.sin(2/3*np.pi)))\
? ? ? ?-A13*np.exp(a*1j*(kx*np.cos(4/3*np.pi)+ky*np.sin(4/3*np.pi)))\
? ? ? ?-A31*np.exp(2*a*1j*(kx*np.cos(1/3*np.pi)+ky*np.sin(1/3*np.pi)))\
? ? ? ?-A32*np.exp(2*a*1j*(kx*np.cos(1*np.pi)+ky*np.sin(1*np.pi)))\
? ? ? ?-A33*np.exp(2*a*1j*(kx*np.cos(5/3*np.pi)+ky*np.sin(5/3*np.pi)))\
? ? ? ?-A41*np.exp(a*1j*(kx*5/2+ky*np.sqrt(3)/2))\
? ? ? ?-A42*np.exp(a*1j*(kx*(-1/2)+ky*3*np.sqrt(3)/2))\
? ? ? ?-A43*np.exp(a*1j*(kx*(-2)+ky*np.sqrt(3)))\
? ? ? ?-A44*np.exp(a*1j*(kx*(-2)-ky*np.sqrt(3)))\
? ? ? ?-A45*np.exp(a*1j*(kx*(-1/2)-ky*3*np.sqrt(3)/2))\
? ? ? ?-A46*np.exp(a*1j*(kx*5/2-ky*np.sqrt(3)/2))\

? ?DBA=-B11*np.exp(a*1j*(kx*np.cos(1/3*np.pi)+ky*np.sin(1/3*np.pi)))\
? ? ? ?-B12*np.exp(a*1j*(kx*np.cos(1*np.pi)+ky*np.sin(1*np.pi)))\
? ? ? ?-B13*np.exp(a*1j*(kx*np.cos(5/3*np.pi)+ky*np.sin(5/3*np.pi)))\
? ? ? ?-B31*np.exp(2*a*1j*(kx*np.cos(0*np.pi)+ky*np.sin(0*np.pi)))\
? ? ? ?-B32*np.exp(2*a*1j*(kx*np.cos(2/3*np.pi)+ky*np.sin(2/3*np.pi)))\
? ? ? ?-B33*np.exp(2*a*1j*(kx*np.cos(4/3*np.pi)+ky*np.sin(4/3*np.pi)))\
? ? ? ?-B41*np.exp(a*1j*(kx*2+ky*np.sqrt(3)))\
? ? ? ?-B42*np.exp(a*1j*(kx*1/2+ky*3*np.sqrt(3)/2))\
? ? ? ?-B43*np.exp(a*1j*(kx*(-5/2)+ky*np.sqrt(3)/2))\
? ? ? ?-B44*np.exp(a*1j*(kx*(-5/2)-ky*np.sqrt(3)/2))\
? ? ? ?-B45*np.exp(a*1j*(kx*1/2-ky*3*np.sqrt(3))/2)\
? ? ? ?-B46*np.exp(a*1j*(kx*2-ky*np.sqrt(3)))\

? ?DBB=B-B21*np.exp(np.sqrt(3)*a*1j*(kx*np.cos(1/6*np.pi)+ky*np.sin(1/6*np.pi)))\
? ? ? ? ? ? -B22*np.exp(np.sqrt(3)*a*1j*(kx*np.cos(3/6*np.pi)+ky*np.sin(3/6*np.pi)))\
? ? ? ? ? ? -B23*np.exp(np.sqrt(3)*a*1j*(kx*np.cos(5/6*np.pi)+ky*np.sin(5/6*np.pi)))\
? ? ? ? ? ? -B24*np.exp(np.sqrt(3)*a*1j*(kx*np.cos(7/6*np.pi)+ky*np.sin(7/6*np.pi)))\
? ? ? ? ? ? -B25*np.exp(np.sqrt(3)*a*1j*(kx*np.cos(9/6*np.pi)+ky*np.sin(9/6*np.pi)))\
? ? ? ? ? ? -B26*np.exp(np.sqrt(3)*a*1j*(kx*np.cos(11/6*np.pi)+ky*np.sin(11/6*np.pi)))\


? ?#組合成總矩陣D：
? ?D=np.zeros([6,6],dtype=complex)
? ?D[0:3,0:3]=DAA
? ?D[0:3,3:6]=DAB
? ?D[3:6,0:3]=DBA
? ?D[3:6,3:6]=DBB


? ?c=np.linalg.eig(D)
? ?e=np.zeros([6,6],dtype=complex)
? ?for p in range(6):
? ? ? ?e[p,p]=c[0][p]

? ?for p in range(6):
? ? ?w3[i,p]=np.real(np.sqrt(e[p,p]))



k=np.linspace(0,2/3*(np.sqrt(3)+1)*np.pi/a,2350)
w=np.zeros([2350,6])
w[0:850,:]=w1[:,:]
w[850:1350,:]=w2[:,:]
w[1350:2350,:]=w3[:,:]
# w=w/(3*10**10)



plt.rcParams['font.family']='SimHei'
plt.rcParams['axes.unicode_minus']=False

plt.plot(k,w[:,0],'.')
plt.plot(k,w[:,1],'.')
plt.plot(k,w[:,2],'.')
plt.plot(k,w[:,3],'.')
plt.plot(k,w[:,4],'.')
plt.plot(k,w[:,5],'.')
plt.ylabel('$\omega$')
plt.title(u'石墨烯在特定方向上的聲子的色散關(guān)系')
plt.show()