Python光学仿真学习衍射算法初步理解

#基尔霍夫衍射,衍射屏坐标范围-dGrid:dGrid,光源坐标(0,0)
#简单的矩孔衍射,dSource为光源到小孔的距离；dScreen为衍射屏到小孔距离
#dHole为矩孔网格尺寸；dGrid为衍射屏网格尺寸；nGrid为网格数目
def squareDiff(dSource=1,dScreen=1.4,dWave=1e-6,
            dHole=3e-5,nGrid=100,dGrid=1e-5):
    nX,nY = nGrid*np.array([1,1])
    axisX = np.arange(-nX,nX+1)*dGrid       
    axisY = np.arange(-nY,nY+1)*dGrid
    xAxis,yAxis = np.meshgrid(axisX,axisY)  #此为衍射屏的x坐标
    axisX = np.arange(-nX,nX+1)*dHole       
    axisY = np.arange(-nY,nY+1)*dHole
    xHole,yHole = np.meshgrid(axisX,axisY)  #此为矩孔的x坐标
    dArrS = np.sqrt(xHole**2+yHole**2+dSource**2)   #孔平面到光源的距离
    nSide = int(nGrid*2+1)              #格点个数
    pane = np.zeros([nSide,nSide])      #衍射屏强度
    for m in range(nSide):
        for n in range(nSide):
            dArr = np.sqrt((xHole-xAxis[m,n])**2+(yHole-yAxis[m,n])**2+dScreen**2)
            pane[m,n] = np.sum(
                np.cos(np.pi*2*(dArr+dArrS)/dWave)/dArr/dArrS)
    pane = np.abs(pane)
    pane = pane/np.max(np.array(pane))
    fig = plt.figure()
    ax = axd(fig)
    ax.plot_surface(xAxis,yAxis,pane)
    plt.show()
    return pane, xAxis, yAxis

#输入对于M*M矩阵的第一个值到N*N矩阵的距离，返回M(m,n)的距离矩阵
def getDisMat(dMat,N,m,n):
    dMat = np.mat(dMat)
    A = dMat[1:m,1:n]
    B = dMat[1:m,0:N-n+1]
    C = dMat[0:N-m+1,1:n]
    D = dMat[0:N-m+1,0:N-n+1]
    return np.vstack((np.hstack((np.flip(A),np.flipud(B))),np.hstack((np.fliplr(C),D))))    #stack矩阵拼接，flip翻转矩阵

    dArrS = np.sqrt(xHole**2+yHole**2+dSource**2) #孔平面到光源的距离
    dScreen = np.sqrt(                             #衍射平上第(0,0)个点的距离矩阵
        (xHole-xAxis[0,0])**2+(yHole-yAxis[0,0])**2+dScreen**2)
    nSide = int(nGrid*2+1)              #格点个数
    pane = np.zeros([nSide,nSide])      #衍射屏强度
    for m in range(nSide):
        for n in range(nSide):
            dArr = getDisMat(dScreen,nSide,m,n)