【第１話】ルンゲクッタで行こう！①【Pythonコピペで古典力学完全攻略マニュアル】

############################################################################
#　単振動運動シミュレーション
############################################################################
import numpy as np
import matplotlib.pyplot as plt

#図全体
fig = plt.figure(figsize=(12, 8))
#全体設定
plt.rcParams['font.family'] = 'Times New Roman' #フォント
plt.rcParams['font.size'] = 24 #フォントサイズ
plt.rcParams["mathtext.fontset"] = 'cm' #数式用フォント
#カラーリストの取得
colors = plt.rcParams['axes.prop_cycle'].by_key()['color']

#################################################################
## 物理定数
#################################################################
#質量
m = 1.0
#ばね定数
k = 1.0

######################################
#　４次のルンゲ・クッタ　クラス
######################################
class RK4:
    #コンストラクタ
    def __init__(self, dt = 0.01, r0 = np.array([0.0, 0.0, 0.0]), v0 = np.array([0.0, 0.0, 0.0]) ):
        self.r = r0.copy()
        self.v = v0.copy()
        self.dr = np.array([0.0, 0.0, 0.0])
        self.dv = np.array([0.0, 0.0, 0.0])
        self.dt = dt
        self.__v1 = np.array([0.0, 0.0, 0.0])
        self.__v2 = np.array([0.0, 0.0, 0.0])
        self.__v3 = np.array([0.0, 0.0, 0.0])
        self.__v4 = np.array([0.0, 0.0, 0.0])
        self.__a1 = np.array([0.0, 0.0, 0.0])
        self.__a2 = np.array([0.0, 0.0, 0.0])
        self.__a3 = np.array([0.0, 0.0, 0.0])
        self.__a4 = np.array([0.0, 0.0, 0.0])

    #速度を与えるメソッド
    def V(self, t, r, v):
        return v.copy()
    ##########################################################
    #加速度を与えるメソッド
    def A(self, t, r, v):
        return -r.copy() * k / m
    ##########################################################

    #時間発展を計算するメソッド
    def timeEvolution(self, t):
        #１段目
        self.__v1 = self.V(
            t, 
            self.r, 
            self.v
        )
        self.__a1 = self.A(
            t, 
            self.r, 
            self.v
        );
        #２段目
        self.__v2 = self.V(
            t + self.dt / 2.0, 
            self.r + self.__v1 * self.dt / 2.0, 
            self.v + self.__a1 * self.dt / 2.0
        )
        self.__a2 = self.A(
            t + self.dt / 2.0, 
            self.r + self.__v1 * self.dt / 2.0, 
            self.v + self.__a1 * self.dt / 2.0            
        )
        #３段目
        self.__v3 = self.V(
            t + self.dt / 2.0, 
            self.r + self.__v2 * self.dt / 2.0, 
            self.v + self.__a2 * self.dt / 2.0                    
        )
        self.__a3 = self.A(
            t + self.dt / 2.0, 
            self.r + self.__v2 * self.dt / 2.0, 
            self.v + self.__a2 * self.dt / 2.0            
        )
        #４段目
        self.__v4 = self.V(
            t + self.dt, 
            self.r + self.__v3 * self.dt, 
            self.v + self.__a3 * self.dt                    
        );
        self.__a4 = self.A(
            t + self.dt, 
            self.r + self.__v3 * self.dt, 
            self.v + self.__a3 * self.dt            
        );        
        #差分の計算
        self.dr = self.dt * ( self.__v1 + 2.0 * self.__v2 + 2.0 * self.__v3 + self.__v4 ) / 6.0 
        self.dv = self.dt * ( self.__a1 + 2.0 * self.__a2 + 2.0 * self.__a3 + self.__a4 ) / 6.0 

#################################################################
## 物理系の設定
#################################################################

#初期位置と速度
r0 = np.array([0.0, 0.0, 10.0])
v0 = np.array([0.0, 0.0, 0.0])
#時間区間
t_min = 0
t_max = 20.0

#時間区間数
NT = 1000
skip = 10
#時間間隔
dt = (t_max - t_min) / NT
#座標点配列の生成
ts = np.linspace(t_min, t_max, NT + 1)
#インスタンス
rk4 = RK4(dt, r0, v0)
#グラフ描画用位置データ
tgs = []
rgs = []
ags = []

#解析解：初期位置のみ指定
def AnalyticalSolution( t ):
    omega = np.sqrt( k / m )
    return r0[2] *np.cos(omega * t)

#各時刻における運動方程式の計算
for tn in range(len(ts)):
    t = ts[tn]
    if( tn == 0 ):
        tgs.append( t ) 
        rgs.append( rk4.r[2] )
        ags.append( AnalyticalSolution( t ) )
        continue
    #ルンゲ・クッタ法による時間発展
    rk4.timeEvolution( t )
    #位置と速度を更新
    rk4.r += rk4.dr
    rk4.v += rk4.dv
    if( tn%skip == 0  ): 
        tgs.append( t ) 
        rgs.append( rk4.r[2] )
        ags.append( AnalyticalSolution( t ) )

plt.title( u"単振動運動（数値解と解析解）", fontsize=20, fontname="Yu Gothic", fontweight=1000)
plt.xlabel(u"時刻" + r"$\,[{\rm s}]$" , fontsize=20, fontname="Yu Gothic", fontweight=1000)
plt.ylabel(u"位置" + r"$\,[{\rm m}]$" , fontsize=20, fontname="Yu Gothic", fontweight=1000)
#罫線の描画
plt.grid(which = "major", axis = "x", alpha = 0.8, linestyle = "-", linewidth = 1)
plt.grid(which = "major", axis = "y", alpha = 0.8, linestyle = "-", linewidth = 1)
#スケールの指定
#ax = plt.gca()
#ax.set_yscale('log')
#ax.set_xscale('log')
#描画範囲を設定
plt.xlim([0, 20])
plt.ylim([-10.5, 10.5])
#プロット
plt.plot(tgs, ags, linestyle='solid', linewidth = 1 )
plt.plot(tgs, rgs, marker = 'o', markersize = 10, linestyle='solid', linewidth = 0 )
#グラフの表示
plt.show()