Python/maths/runge_kutta.py
Ankur Chattopadhyay 7592cba417 psf/black code formatting (#1421)
* added sol3.py for problem_20

* added sol4.py for problem_06

* ran `black .` on `\Python`
2019-10-22 19:13:48 +02:00

45 lines
1014 B
Python

import numpy as np
def runge_kutta(f, y0, x0, h, x_end):
"""
Calculate the numeric solution at each step to the ODE f(x, y) using RK4
https://en.wikipedia.org/wiki/Runge-Kutta_methods
Arguments:
f -- The ode as a function of x and y
y0 -- the initial value for y
x0 -- the initial value for x
h -- the stepsize
x_end -- the end value for x
>>> # the exact solution is math.exp(x)
>>> def f(x, y):
... return y
>>> y0 = 1
>>> y = runge_kutta(f, y0, 0.0, 0.01, 5)
>>> y[-1]
148.41315904125113
"""
N = int(np.ceil((x_end - x0) / h))
y = np.zeros((N + 1,))
y[0] = y0
x = x0
for k in range(N):
k1 = f(x, y[k])
k2 = f(x + 0.5 * h, y[k] + 0.5 * h * k1)
k3 = f(x + 0.5 * h, y[k] + 0.5 * h * k2)
k4 = f(x + h, y[k] + h * k3)
y[k + 1] = y[k] + (1 / 6) * h * (k1 + 2 * k2 + 2 * k3 + k4)
x += h
return y
if __name__ == "__main__":
import doctest
doctest.testmod()