简单的地震射线追踪

#问题一
import math
from turtle import *
from matplotlib.pyplot import *
# print(radians(45))
def Ray(theta1, h_down, v,x_hypocenter):
    theta1=math.radians(theta1)#将角度转化为弧度
    n=len(h_down)#提取列表个数
    sum_t=0#总时间
    sum_x=0#最后位置
    p=math.sin(theta1)/v[0]#snell比例系数
    h_up=list(reversed(h_down))#向上厚度列表
    h=h_down+h_up#总厚度列表
    x_step=list(range(n))#水平方向步长
    sin_i=list(range(n))#每一层的sin值
    angle_i = list(range(n))  # 每一层的角度
    x=list(range(2*n+1))#横坐标
    y=list(range(2*n+1))#纵坐标
    for i in range(n-1):
        sin_i[i]=p*v[i]
        angle_i[i]=math.asinh(sin_i[i])
        if sin_i[i]>1 or sin_i[i]<0:#入射角过大、下层速度过大等
            print('参数设置错误')
            exit(1)
        sum_t=2*(h_down[i]/(v[i]*(1-sin_i[i]**2)**0.5))+sum_t
        sum_x = 2 * h_down[i] * sin_i[i] / ((1 - sin_i[i] ** 2) ** 0.5)+sum_x
        x_step[i]=h_down[i]*math.tan(angle_i[i])
    x_step=x_step+list(reversed(x_step))
    #坐标
    for j in range(2*n):
        x[0]=x_hypocenter#初始坐标,震源位置
        y[0]=0
        x[j + 1] = x[j] + x_step[j]
        if j<n:
            y[j+1]=y[j]+h[j]
        else:
            y[j+1]=y[j]-h[j]
    # #画图
    # plot(x,y,color='b',alpha=1)#射线
    # for i in range(6):
    #     axhline(y[i])#地层
    #print('总时间：', sum_t)
    # print('位置：', sum_x)
    # show()
    result=x_hypocenter+sum_x
    return result
# #问题二
def T(h_down,v,x_hypocenter,x_detector):#已知震源位置与检波器位置求解总时间
    theta1=list(range(1,90,1))#打靶范围，不应取到0和90度
    x_diff=list(range(89))#差值
    #求解最小差值时角度
    for i in range(89):
        x_diff[i]=abs(Ray(theta1[i],h_down,v,x_hypocenter)-x_detector)
    j=x_diff.index(min(x_diff))#最小值的索引
    Ray(theta1[j],h_down,v,x_hypocenter)
#以五层为例
theat1=45#出射角
h_down=[10,20,30,20,10]#每层厚度
v=[30,20,10,20,30]#每层速度
S=0#震源位置
#Ray(theat1,h_down,v,S)#问题一调用
#T(h_down,v,S,77)#问题二调用