【第３話】拘束条件の改良！【ルンゲクッタで行こう！③】

20210.08.27 アルゴリズムを改良しました。破綻しません。

############################################################################
#　単振り子運動シミュレーション（補正力あり）
############################################################################
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation

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

#################################################################
## 物理系の設定
#################################################################
#質量
m = 1.0
#ひもの長さ
r_abs0 = 1.0
#重力加速度
g = np.array([0.0, 0.0, -10.0])
#補正パラメータ
compensationK = 100.0
compensationGamma = 0.0

#張力ベクトル
def S ( t, r_t, v_t ):
	r_abs = np.sqrt( np.sum(r_t**2) )
	return -m / r_abs ** 2 * ( np.sum(v_t**2) + np.dot(g, r_t) ) * r_t

def F ( t, r_t, v_t ):
	r = np.sqrt( np.sum( r_t**2 ) )
	v = np.sqrt( np.sum( v_t**2 ) )
	r_hat = r_t.copy() / r
	v_hat = v_t.copy() / v
	fc  = - compensationK * ( r - r_abs0 ) * r_hat
	fc += - compensationGamma * np.dot(v, r_hat) * r_hat
	return S( t, r_t, v_t ) + m * g + fc


#初期位置と速度
r0 = np.array([0.0 , 0.0, -r_abs0])
v0 = np.array([6.0, 0.0, 0.0])
#時間区間
t_min = 0
t_max = 10.0
#時間区間数
NT = 1000
skip =  1
#時間間隔
dt = (t_max - t_min) / NT
#座標点配列の生成
ts = np.linspace(t_min, t_max, NT + 1)

######################################
#　４次のルンゲ・クッタ　クラス
######################################
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_t, v_t):
		return F( t, r_t, v_t ) / 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 

#################################################################
## 計算開始

#インスタンス
rk4 = RK4(dt, r0, v0)

#アニメーション作成用
ims=[]
def plot( t, r ):
	print( "{:.2f}".format(t),  np.sqrt(r[0]**2 + r[2]**2) - 1  )
	#各コマを描画
	img  = plt.plot([0], [0], colors[0], marker = 'o', markersize = 10, linestyle='solid', linewidth = 0 )
	img  += plt.plot([r[0]], [r[2]], colors[1], marker = 'o', markersize = 20, linestyle='solid', linewidth = 0 )
	img  += plt.plot([0, r[0]], [0, r[2]], colors[1], linestyle='solid', linewidth = 1 )
	time = plt.text(0.8, 1.25, "t = " +  str("{:.2f}".format(t)) + r"$[{\rm s}]$" , ha='left', va='center', fontsize=18)
	#テキストをグラフに追加
	img.append( time )
	ims.append( img )

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



plt.title( u"単振り子運動", fontsize=20, fontname="Yu Gothic", fontweight=1000)
plt.xlabel(r"$x\,[{\rm m}]$" , fontsize=20, fontname="Yu Gothic", fontweight=1000)
plt.ylabel(r"$z\,[{\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)

#描画範囲を設定
plt.xlim([-1.2, 1.2])
plt.ylim([-1.2, 1.2])

#アニメーションの生成
#ani = animation.ArtistAnimation(fig, ims, interval=10)
#ani.save("output.html", writer=animation.HTMLWriter())
#ani.save("output.gif", writer="imagemagick")
#ani.save("output.mp4", writer="ffmpeg", dpi=300)

#グラフの表示
#plt.show()