최단강하곡선 파이썬 구현

from numpy import *
import matplotlib.pyplot as plt
from scipy.integrate import quad
from scipy.optimize import newton
from scipy import interpolate

g = 9.81    #중력가속도
x_0, x_end = 0, pi #그래프 x 범위
y_0, y_end = 0, 2  #그래프 y 범위

# ---------- functional ----------

def func(x, f, fp, offset):
    return sqrt((1+fp(x)**2) / (2 * g * f(x, offset)))

# ---------- cycloid ----------

R = 1

def f(x, offset=0):
    def p(th) :
        return x - R*(th - sin(th))
    th = newton(p, pi)
    # newton 함수를 사용하여 th(세타값) 찾기.
    return R*(1 - cos(th)) - offset
    #p(th)= x − R(th−sin(th)) 방정식을 만족하는 th를 찾은 뒤에 R * (1-cos(th)) - offset 을 반환


def fp(x):
    def p(th) :
        return x - R*(th - sin(th))
    th = newton(p, pi)
    return sin(th)/(1-cos(th))

T = pi * sqrt(R/g)
print('T_cycloid  = {:.3f}'.format(T))

# ---------- cal ----------

N = 1000
plt.figure('equal')

for x_s, y_s in [[x_0, y_0],[pi/2, f(pi/2)]]:   #x=0 -> pi/2 y=0 -> f(pi/2)
    xs = linspace(x_s, x_end, N)    #xs 배열 선언
    ys = array([-y_s])              #ys배열 선언, y_s값을 뒤집어서 선언
    h = xs[1] - xs[0]               #h선언,

    t = 0
    sol = array([0, x_s])

    for i in range(N-1):
        t += quad(func, xs[i], xs[i+1], args=(f, fp, y_s))[0]
        sol = append(sol, [t, xs[i+1]])
    sol = sol.reshape((N,2))

    ts = arange(0, t, 0.01)
    interp = interpolate.interp1d(sol[:,0],sol[:,1])
    for x in interp(ts)[1:] :
        ys =  append(ys, -f(x))
    new = array([ts, interp(ts), ys]).transpose()

    if x_s == x_0 and y_s == y_0 :
        line = plt.plot(new[:,1], new[:,2], 'k')
    print(t)

    plt.plot(new[range(0,len(new),10), 1], new[range(0,len(new),10), 2], '.', markersize=10)

plt.show(block=False)