用python模拟高度多模非线性光学系统中的热化

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转成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)     #必须设定类型为复数
   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已经有了
   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()