iframe-proxy | Sunbelt Computer Software

44 lines (34 loc) · 1014 Bytes

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()

Sunbelt Computer Software