added runge-kutta

maths/runge_kutta.py

@@ -0,0 +1,44 @@
+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()