from util.FuncVec2 import *
# ############## #
# Core algorithm #
# ############## #
def _runge_kutta_(t, f: FuncVec2, vec: Vec2, h: float):
r1 = f.apply(t, vec)
r2 = f.apply(t + h / 2.0, vec + r1 * h / 2.0)
r3 = f.apply(t + h / 2.0, vec + r2 * h / 2.0)
r4 = f.apply(t + h, vec + r3 * h)
r_result = (r1 + r2 * 2.0 + r3 * 2.0 + r4) * h / 6.0
return r_result.x, r_result.y
def runge_kutta(start: Vec2, func: FuncVec2, h: float):
t = 0
xs = [0]
ys = [0]
vec = start
while True:
t += h
x = xs[len(xs) - 1] + vec.x
y = ys[len(ys) - 1] + vec.y
xs.append(x)
ys.append(y)
vx, vy = _runge_kutta_(t, func, vec, h)
vec = Vec2(vx, vy)
if y <= 0:
break
return xs, ys