古典的ルンゲ＝クッタ法のPythonプログラムをWikiPediaの記述と同じ変数名で書いてみた

教育用の覚書
Python だと汎用の Runge-Kutta 法のルーチンが６行で書けてしまう*1。以下のコードではルンゲクッタ法ルーチン rk4() を書き換えることなく、1階, 2階, 3階の常微分方程式(ordinary differential equation(s), ODE)を解いている*2 *3。関数 rk4() で使っている変数名はRunge–Kutta methods - Wikipediaに揃えた。C言語 *4で数値計算をコーディング/教育する気力が一気に失われてしまったｗｗｗｗｗｗ
次のコンプリートコードは次の微分方程式を積分して、あらかじめ求めておいた解の関数の値と比較している*5：

class getExpt():
[1階] $\displaystyle\frac{dy}{dt}=y$ 　(比較する解： $y=e^t$ );
class getSint():
[2階] $\displaystyle\frac{d^2y}{dt^2}=-y$ 　(比較する解： $\displaystyle\left(y,\frac{dy}{dt}\right)=(\sin t,\cos t)$ );
[1階連立ODE] $\displaystyle\frac{d}{dt}\Big(\underbrace{y}_{y_0},\underbrace{\frac{dy}{dt}}_{y_1}\Big)=\underbrace{\left(\frac{dy}{dt},-y\right)}_{f(t,(y_0,y_1))}$
class getExptSint():
[3階] $\displaystyle\frac{d^3y}{dt^3}+\frac{d^2y}{dt^2}+\frac{dy}{dt}+y=0$
(比較する解： $\displaystyle\left(y,\frac{dy}{dt},\frac{d^2y}{dt^2}\right)=(e^{-t}+\sin t,-e^{-t}+\cos t,e^{-t}-\sin t)$ );
[1階連立ODE] $\displaystyle\frac{d}{dt}\bigg(\underbrace{y}_{y_0},\underbrace{\frac{dy}{dt}}_{y_1},\underbrace{\frac{d^2y}{dt^2}}_{y_2}\bigg)=\underbrace{\left(\frac{dy}{dt},\frac{d^2y}{dt^2},-y-\frac{dy}{dt}-\frac{d^2y}{dt^2}\right)}_{f(t,(y_0,y_1,y_2))}$
class spring():
[2階]単振動(調和振動子)の問題 $\displaystyle m a=- k x$ , すなわち
位置： $\displaystyle x(t)$ , 速度： $\displaystyle v(t)=\frac{dx}{dt}$ , 加速度： $\displaystyle a(t)=\frac{d^2x}{dt^2}$ より $\displaystyle \frac{d^2x}{dt^2}=-\frac{k}{m}x$
(比較する解： $(x(t),v(t))=(\sin\omega t,\omega\cos\omega t)$ , $\displaystyle\omega=\sqrt{\frac{k}{m}}$ );
[1階連立ODE] $\displaystyle\frac{d}{dt}\Big(\underbrace{x}_{y_0},\underbrace{v}_{y_1}\Big)=\underbrace{\left(v,-\frac{k}{m}x\right)}_{f(t,(y_0,y_1))}$

初期条件は比較する解に t=0 を代入して作っている。

import numpy

# classical 4-th order Runge-Kutta method subroutine
# variables letters are the same as those in
# https://en.wikipedia.org/wiki/Runge%E2%80%93Kutta_methods
def rk4 ( y, t, h, f ):
    k1 = h * f( t, y )
    k2 = h * f( t + h/2, y + k1/2 )
    k3 = h * f( t + h/2, y + k2/2 )
    k4 = h * f( t + h, y + k3 )
    return y + ( k1 + 2*k2 + 2*k3 + k4 )/6, t + h

# each of the following four classes has
#   __init__(): constructor,
#   rhsODE(t,y): f(t,y) of the 1st order ODE dy/dx=f(t,y),
#   ref(t): a reference solution to the ODE.
class getExpt():
    def __init__ (self):
        print(self.__class__,' y(t)=exp(t)')
    def rhsODE ( self, t, y ):
        return y
    def ref ( self, t ):
        return numpy.exp(t)

class getSint():
    def __init__ (self):
        print(self.__class__,' y(t)=sin(t)')
    def rhsODE ( self, t, y ):
        return numpy.array([y[1],-y[0]])
    def ref ( self, t ):
        return numpy.array([numpy.sin(t),numpy.cos(t)])

class getExptSint():
    def __init__ (self):
        print(self.__class__,' y(t)=exp(-t)+sin(t)')
    def rhsODE ( self, t, y ):
        return numpy.array([y[1],y[2],-y[0]-y[1]-y[2]])
    def ref ( self, t ):
        et=numpy.exp(-t); st=numpy.sin(t);
        return numpy.array([et+st,-et+numpy.cos(t),et-st])

class spring():
    def __init__ (self,mass=1,springConstant=1):
        self.mass= mass
        self.springConstant= springConstant
        m= mass
        k= springConstant
        self.angularFrequency= (k/m)**0.5
        print(self.__class__,' linear spring motion m=',m,' k=',k)
    def rhsODE ( self, t, y ):
        k= self.springConstant
        m= self.mass
        x= y[0]
        v= y[1]
        a= -(k/m)*x
        return numpy.array([v,a])
    def ref ( self, t ):
        w= self.angularFrequency
        x= numpy.sin(w*t)
        v= w*numpy.cos(w*t)
        return numpy.array([x,v])

if __name__=="__main__":

    # choose class to e[x]ecute
    #x= getExpt();
    #x= getSint();
    #x= getExptSint();
    x= spring(mass=2,springConstant=3);

    # set [ref]erence [sol]ution function
    #   and [f]unction appears in the [r]ight [h]and [s]ide of
    #   [O]rdinary [D]ifferential [E]quation
    sol= x.ref; f= x.rhsODE;

    # set parameters
    h = 0.01;   # step-size
    t = 0;      # initial time
    y = sol(t); # initial condition (y[,dy/dt[,d^2y/dt^2]])_t=0
    tEnd = 0.1; # terminal time
    print(t,y,sol(t),y-sol(t))

    # integration and comparison to reference solution
    while t < tEnd:
        y, t = rk4( y, t, h, f )
        print(t,y,sol(t),y-sol(t))

*1:ルンゲクッタ python - Google 検索で検索すると上位に上がるページに意外と汎用性のあるコードが無い。

*2:ルンゲクッタ法は1階の連立常微分方程式を解く計算スキームなので、n階の微分方程式はn個の関数 $\boldsymbol{y}=(y,y',...,y^{(n-1)})$ に対する連立方程式に書き換える。