Runge–Kutta method Sample by using simple problem

ode_rk_raw.py

# -*- coding: utf-8 -*-
 
import numpy as np
import matplotlib.pyplot as plt



def main():

    x_ini = 0.0 # x(0) = 0
    v_ini = 0.0 # v(0) = d/dt * x(0) = 0
    g = 9.8

    def eq_F(t,x,v):
        '''
        dx/dt = F(t,x,v) = v
        '''
        return v

    def eq_G(t,x,v):
        '''
        dv/dt = G(t,x,v) = -g
        '''
        return -g

    # ここから
    h = 0.1 # ステップ幅(s)
    T = 1000 # ステップ回数
    t = 0.0 # 現在時刻
    X = [x_ini] # 位置を記録する配列
    V = [v_ini] # 速度を記録する配列

    # 理論値(ここは比較のため)
    X_ideal =[x_ini]
    for i in range(T):
        tt = h*(i+1)
        X_ideal.append(-(g*tt**2)/2)

    for i in range(T):

        # 必要な変数を算出
        k1 = h*eq_F(t,X[i],V[i])
        l1 = h*eq_G(t,X[i],V[i])
        k2 = h*eq_F(t+h/2,X[i]+k1/2,V[i]+l1/2)
        l2 = h*eq_G(t+h/2,X[i]+k1/2,V[i]+l1/2)
        k3 = h*eq_F(t+h/2,X[i]+k2/2,V[i]+l2/2)
        l3 = h*eq_G(t+h/2,X[i]+k2/2,V[i]+l2/2)
        k4 = h*eq_F(t,X[i]+k3,V[i]+l3)
        l4 = h*eq_G(t,X[i]+k3,V[i]+l3)

        # 値を求める
        t = t + h
        X.append(X[i]+(k1+2.0*k2+2.0*k3+k4)/6)
        V.append(V[i]+(l1+2*l2+2*l3+l4)/6)

    # 描画
    t = np.arange(0,T*h+h/10,h)
    a = np.array(X)
    b = np.array(X_ideal)
    plt.plot(t,b,color='green',label='X_correct')
    plt.plot(t,a,color='black',label='X_calculate')
    plt.legend(bbox_to_anchor=(0.01, 0.01), loc='lower left', borderaxespad=0)
    plt.xlabel("Time [s]")
    plt.show()

if __name__ == '__main__':
    main()