cms.diffeq_rungekutta4 のソースコード

import numpy as np
from math import exp, sqrt, sin, cos, pi

"""
  Solve first order diffrential equation by four-stage fourth-order Runge-kutta method
"""


#===================
# parameters
#===================
x0 = 1.0
dt = 0.1
nt = 10000001
iprint_interval = 1000000


# dx/dt = dxdt(x,t)
# define function to be integrated


def dxdt(t, x):
    return -x*x


# solution: x = 1 / (C + t), C = 1 for x(0) = 1.0


def fsolution(t):
    return 1.0 / (1.0 + t)




def diffeq_euler(diff1func, t0, x0, dt):
   k1 = dt * diff1func(t0, x0)
   x1 = x0 + k1
   return x1




def diffeq_heun(diff1func, t0, x0, dt):
   k0 = dt * diff1func(t0, x0)
   k1 = dt * diff1func(t0+dt, x0+k0)
   x1 = x0 + (k0 + k1) / 2.0
   return x1



#===================
# main routine
#===================


def main(x0, dt, nt):
    print("Solve first order diffrential equation by four-stage fourth-order Runge-kutta method")

    print("{:^10}  {:^16}  {:^16}".format('t', 'x(cal)', 'x(exact)'))
    t0 = 0.0
    f0 = dxdt(t0, x0)
# The next x (x1) must be predicted by Euler or Heum method
#    x1 = diffeq_euler(dxdt, t0, x0, dt)
    x1 = diffeq_heun(dxdt, t0, x0, dt)
    xexact = fsolution(t0)

    print("t={:10.2f}  {:16.10e}  {:16.10e}".format(t0, x0, xexact))
    for i in range(1, nt):
        t1 = i * dt
        k0 = dt * dxdt(t0,        x0)
        k1 = dt * dxdt(t0+dt,     x0+k0)
        k2 = dt * dxdt(t0+dt,     x0+k1)
        k3 = dt * dxdt(t0+2.0*dt, x0+2.0*k2)
        x2 = x0 + 1.0 / 3.0 * (k0 + 2.0 * k1 + 2.0 * k2 + k3)
        xexact = fsolution(t1)

        if i % iprint_interval == 0:
            print("t={:10.2f}  {:16.10e}  {:16.10e}".format(t1, x1, xexact))

        t0 = t1
        x0 = x1
        x1 = x2


if __name__ == '__main__':
    main(x0, dt, nt)