import numpy as np
import math
import matplotlib.pyplot as plt
import random

Z=200
N=20*Z
dz=Z/N
z=np.linspace(0,Z,N+1)

L=20
N=L**2
k1=1
k2=1.5

H=np.diag(k2*np.ones((N-1)), 1)+k2*np.diag(np.ones((N-1)), -1)+k1*np.diag(np.ones((N-L)),L)+k1*np.diag(np.ones((N-L)),-L)? #哈密頓矩陣

w, v = np.linalg.eig(H)
idx_sorted = np.argsort(w)
V=v[:, idx_sorted] ? #特征向量
W=w[idx_sorted]


a=np.zeros([L+2,L+2,len(z)],dtype=complex) ?#a_mn
psi=np.zeros([L**2,len(z)],dtype=complex) ? #光矢量
c2=np.zeros([L**2,len(z)]) ?#算出ci0的平方

xuan=(W>=1)*(W<=2)
c2[xuan,0]=1
for i in range(L**2):
? ?psi[:,0]+=np.sqrt(c2[i,0])*V[:,i]*np.exp(1j*2*np.pi*random.random()) ? ?#初始條件


#把psi轉(zhuǎn)成amn
a[1:(L+1),1:(L+1),0]=psi[:,0].reshape(L,L)


for i in range(len(z)-1):
? ?s1=np.zeros([L+2,L+2],dtype=complex) ? ? #必須設(shè)定類型為復(fù)數(shù)
? ?s2=np.zeros([L+2,L+2],dtype=complex)
? ?s3=np.zeros([L+2,L+2],dtype=complex)
? ?s4=np.zeros([L+2,L+2],dtype=complex)
? ?for m9 in range(L):
? ? ? ?m=m9+1
? ? ? ?for n9 in range(L):
? ? ? ? ? ?n=n9+1
? ? ? ? ? ?s1[m,n]=1j*(k1*(a[m-1,n,i]+a[m+1,n,i])+k2*(a[m,n-1,i]+a[m,n+1,i])+abs(a[m,n,i])**2*a[m,n,i])

? ?for m9 in range(L):
? ? ? ?m=m9+1
? ? ? ?for n9 in range(L):
? ? ? ? ? ?n=n9+1
? ? ? ? ? ?s2[m,n]=1j*(k1*(a[m-1,n,i]+s1[m-1,n]*dz/2+a[m+1,n,i]+s1[m+1,n]*dz/2)\
? ? ? ? ? ? ? ? ? ? ? ?+k2*(a[m,n-1,i]+s1[m,n-1]*dz/2+a[m,n+1,i]+s1[m,n+1]*dz/2)\
? ? ? ? ? ? ? ? ? ? ? ?+abs(a[m,n,i]+s1[m,n]*dz/2)**2*(a[m,n,i]+s1[m,n]*dz/2))

? ?for m9 in range(L):
? ? ? ?m=m9+1
? ? ? ?for n9 in range(L):
? ? ? ? ? ?n=n9+1
? ? ? ? ? ?s3[m,n]=1j*(k1*(a[m-1,n,i]+s2[m-1,n]*dz/2+a[m+1,n,i]+s2[m+1,n]*dz/2)\
? ? ? ? ? ? ? ? ? ? ? ?+k2*(a[m,n-1,i]+s2[m,n-1]*dz/2+a[m,n+1,i]+s2[m,n+1]*dz/2)\
? ? ? ? ? ? ? ? ? ? ? ?+abs(a[m,n,i]+s2[m,n]*dz/2)**2*(a[m,n,i]+s2[m,n]*dz/2))

? ?for m9 in range(L):
? ? ? ?m=m9+1
? ? ? ?for n9 in range(L):
? ? ? ? ? ?n=n9+1
? ? ? ? ? ?s4[m,n]=1j*(k1*(a[m-1,n,i]+s3[m-1,n]*dz+a[m+1,n,i]+s3[m+1,n]*dz)\
? ? ? ? ? ? ? ? ? ? ? ?+k2*(a[m,n-1,i]+s3[m,n-1]*dz+a[m,n+1,i]+s3[m,n+1]*dz)\
? ? ? ? ? ? ? ? ? ? ? ?+abs(a[m,n,i]+s3[m,n]*dz)**2*(a[m,n,i]+s3[m,n]*dz))

? ?for m9 in range(L):
? ? ? ?m=m9+1
? ? ? ?for n9 in range(L):
? ? ? ? ? ?n=n9+1
? ? ? ? ? ?a[m,n,i+1]=a[m,n,i]+1/6*(s1[m,n]+2*s2[m,n]+2*s3[m,n]+s4[m,n])*dz

? ? ? ? ? ?# a[m,n,i+1]=a[m,n,i]+1j*(k1*(a[m-1,n,i]+a[m+1,n,i])+k2*(a[m,n-1,i]+a[m,n+1,i])+abs(a[m,n,i])**2*a[m,n,i])*dz



for i in range(len(z)-1):
? ?psi[:,i+1]=a[1:(L+1),1:(L+1),i+1].ravel() ? ?#從1開始賦值，0已經(jīng)有了
? ?for j in range(L**2):
? ? ? ?c2[j,i+1]=abs(np.dot(V[:,j],psi[:,i+1]))**2




plt.figure()
plt.subplot(321)
plt.plot(W,c2[:,0])
plt.subplot(322)
plt.plot(W,c2[:,round(40/dz)])
plt.subplot(323)
plt.plot(W,c2[:,round(80/dz)])
plt.subplot(324)
plt.plot(W,c2[:,round(120/dz)])
plt.subplot(325)
plt.plot(W,c2[:,round(160/dz)])
plt.subplot(326)
plt.plot(W,c2[:,round(200/dz)])

plt.show()