python光学仿真通过菲涅耳公式实现波动模型

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
z = np.arange(15,200)*10    #单位为nm
x = np.arange(15,200)*10
x,z = np.meshgrid(x,z)      #创建坐标系
E = 1/np.sqrt(x**2+z**2)*np.cos(2*np.pi*np.sqrt(x**2+z**2)/(532*1e-9))
fig = plt.figure()
ax = Axes3D(fig)
ax.plot_surface(x,z,E)
plt.show()

import matplotlib.pyplot as plt
import numpy as np
def fresnel(theta, n1, n2):
    theta = theta*np.pi/180
    xTheta = np.cos(theta)
    mid = np.sqrt(1-(n1/n2*np.sin(theta))**2)          #中间变量
    rp = (n2*xTheta-n1*mid)/(n2*xTheta+n1*mid)  #p分量振幅反射率
    rs = (n1*xTheta-n2*mid)/(n1*xTheta+n2*mid)
    tp = 2*n1*xTheta/(n2*xTheta+n1*mid)
    ts = 2*n1*xTheta/(n1*xTheta+n2*mid)
    return rp, rs, tp, ts
def testFres(n1=1,n2=1.45):         #默认n2为1.45
    theta = np.arange(0,90,0.1)+0j
    a = theta*np.pi/180
    rp,rs,tp,ts = fresnel(theta,n1,n2)
    fig = plt.figure(1)
    plt.subplot(1,2,1)
    plt.plot(theta,rp,'-',label='rp')
    plt.plot(theta,rs,'-.',label='rs')
    plt.plot(theta,np.abs(rp),'--',label='|rp|')
    plt.plot(theta,np.abs(rs),':',label='|rs|')
    plt.legend()
    plt.subplot(1,2,2)
    plt.plot(theta,tp,'-',label='tp')
    plt.plot(theta,ts,'-.',label='ts')
    plt.plot(theta,np.abs(tp),'--',label='|tp|')
    plt.plot(theta,np.abs(ts),':',label='|ts|')
    plt.legend()
    plt.show()
if __init__=="__main__":
    testFres()

def testFres(n1=1,n2=1.45):
    theta = np.arange(0,90,0.1)+0j
    a = theta*np.pi/180
    rp,rs,tp,ts = fml.fresnel(theta,n1,n2)
    Rp = np.abs(rp)**2
    Rs = np.abs(rs)**2
    Rn = (Rp+Rs)/2
    Tp = n2*np.sqrt(1-(n1/n2*np.sin(a))**2)/(n1*np.cos(a))*np.abs(tp)**2
    Ts = n2*np.sqrt(1-(n1/n2*np.sin(a))**2)/(n1*np.cos(a))*np.abs(ts)**2
    Tn = (Tp+Ts)/2
    fig = plt.figure(2)
    plt.subplot(1,2,1)
    plt.plot(theta,Rp,'-',label='R_p')
    plt.plot(theta,Rs,'-.',label='R_s')
    plt.plot(theta,Rn,'-',label='R_n')
    plt.legend()
    plt.subplot(1,2,2)
    plt.plot(theta,Tp,'-',label='T_p')
    plt.plot(theta,Ts,'-.',label='T_s')
    plt.plot(theta,Tn,'--',label='T_n')
    plt.legend()
    plt.show()