Python模拟光的折射（未优化）

import matplotlib.pyplot as plt
import math


def refraction(startpoint=(0,0),n=1,n1=1.5,y=2,d=2):
    beta0 = 60/180*math.pi
    beta1 = 45/180*math.pi
    beta2 = 30/180*math.pi

    x0 = (y-startpoint[1])*(math.tan(beta0))
    x1 = (y-startpoint[1])*(math.tan(beta1))
    x2 = (y-startpoint[1])*(math.tan(beta2))

    beta00 = math.asin(math.sin(beta0)*n/n1)
    beta11 = math.asin(math.sin(beta1)*n/n1)
    beta22 = math.asin(math.sin(beta2)*n/n1)  
    
    x00 = d * math.sin(beta00) + x0 
    x11 = d * math.sin(beta11) + x1 
    x22 = d * math.sin(beta22) + x2 

    y00 = d * math.cos(beta00) + y
    y11 = d * math.cos(beta11) + y
    y22 = d * math.cos(beta22) + y
    
    x000 = -2*d * math.sin(beta00) + x0 
    x111 = -2*d * math.sin(beta11) + x1 
    x222 = -2*d * math.sin(beta22) + x2 

    y000 = -2*d * math.cos(beta00) + y
    y111 = -2*d * math.cos(beta11) + y
    y222 = -2*d * math.cos(beta22) + y   
    
    
    #光源，交界点，折射点，反向延长点  的x,y坐标
    return [startpoint[0],x0,x00,x000],[startpoint[1],y,y00,y000],\
           [startpoint[0],x1,x11,x111],[startpoint[1],y,y11,y111],\
           [startpoint[0],x2,x22,x222],[startpoint[1],y,y22,y222]

a,b,c,d,e,f = refraction()
plt.axhline(2)
plt.plot([a[0],a[1]],[b[0],b[1]])
plt.plot([a[1],a[2]],[b[1],b[2]])
plt.plot([a[2],a[3]],[b[2],b[3]])

plt.plot([c[0],c[1]],[d[0],d[1]])
plt.plot([c[1],c[2]],[d[1],d[2]])
plt.plot([c[2],c[3]],[d[2],d[3]])

plt.plot([e[0],e[1]],[f[0],f[1]])
plt.plot([e[1],e[2]],[f[1],f[2]])
plt.plot([e[2],e[3]],[f[2],f[3]])
plt.savefig("fig.jpg")

效果图