投稿日: 投稿者: simu-kun

【統計力学001】理想気体で状態方程式シミュレーション！【Pythonで分子動力学入門①】

############################################################################
#　シリンダーの中のN個の粒子
############################################################################
import os
import numpy as np
import random 
import matplotlib.pyplot as plt
import matplotlib.animation as animation
random.seed(0)

#図全体
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
M = 1000.0
#粒子数
N = 100
#箱の長さ
L = 1.0
#シリンダーに加わる外力：重力
Fex = np.array([0.0, 0.0, -10.0 * M ])
#シリンダーの粘性係数
beta = 0
#粒子に加わる力
fex = np.array([0.0, 0.0, 0.0 ])
#初期位置と速度
v_max = 10.0

#粒子に加わる力
def F (self, t, r_t, v_t ):
	return self.fex + self.fc - self.m * self.beta * v_t

#時間区間
t_min = 0
t_max = 2.0
#時間区間数
NT =   2000
skip =    5
#時間間隔
dt = (t_max - t_min) / NT
#座標点配列の生成
ts = np.linspace(t_min, t_max, NT + 1)
#
lw = 0.01
######################################
#　４次のルンゲ・クッタ　クラス
######################################
class RK4:
	#コンストラクタ
	def __init__(self, m, dt = 0.01, r0 = np.array([0.0, 0.0, 0.0]), v0 = np.array([0.0, 0.0, 0.0]) ):
		self.m = m
		#現時刻の位置と速度
		self.r = r0.copy()
		self.v = v0.copy()
		#１ステップ前の時刻の位置と速度
		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.fex = np.array([0.0, 0.0, 0.0])
		#衝突力
		self.fc = np.array([0.0, 0.0, 0.0])
		#粘性係数
		self.beta = 0
		#内部
		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(self, t, r_t, v_t ) / self.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 

#################################################################
## 時間発展の計算
#################################################################

#衝突判定（戻値：法線ベクトル）
def CollisionDetection( r, R ):
	#衝突リスト
	collisionlist = []
	if(r[0] > L):
		collisionlist.append({ "type": 1, "normal" : np.array([-1.0, 0.0, 0.0]) })
	if(r[0] < 0):
		collisionlist.append({ "type": 1, "normal" : np.array([1.0, 0.0, 0.0]) })
	if(r[1] > L):
		collisionlist.append({ "type": 1, "normal" : np.array([0.0, -1.0, 0.0]) })
	if(r[1] < 0):
		collisionlist.append({ "type": 1, "normal" : np.array([0.0, 1.0, 0.0]) })
	if(r[2] > R[2]):
		collisionlist.append({ "type": 2, "normal" : np.array([0.0, 0.0, -1.0]) })
	if(r[2] < 0):
		collisionlist.append({ "type": 1, "normal" : np.array([0.0, 0.0, 1.0]) })
	return collisionlist

#運動エネルギー
def T():
	_T = 1.0/2.0 * M * np.sum( rk4_M.v**2 )
	for n in range(N):
		_T += 1.0/2.0 * m * np.sum( rk4_ms[n].v**2 )
	return  _T
#位置エネルギー
def U():
	_R = - np.dot(rk4_M.fex, rk4_M.r)
	for n in range(N):
		_R += - np.dot(rk4_ms[n].fex, rk4_ms[n].r)
	return _R

#粒子の運動エネルギー
def T_particle():
	_T = 0
	for n in range(N):
		_T += 1.0/2.0 * m * rk4_ms[n].v[2]**2
	return  _T
#シリンダーの位置エネルギー（底を基準）
def U_cylinder():
	_R = - np.dot(rk4_M.fex, rk4_M.r) + M * L/2 
	return _R

#系の時間発展の計算
def timeEvolution( t, timeScale ):
	#時間スケールの指定
	rk4_M.dt = dt / timeScale

	rk4_M.fc = np.array([0.0, 0.0, 0.0])
	rk4_M.timeEvolution( t )
	rk4_M.r += rk4_M.dr
	rk4_M.v += rk4_M.dv

	#複数衝突
	cylinderCollisionList = []

	for mn in updateList:
		#時間スケールの指定
		rk4_ms[mn].dt = dt / timeScale

		#ルンゲ・クッタ法による時間発展
		rk4_ms[mn].fc = np.array([0.0, 0.0, 0.0])
		rk4_ms[mn].timeEvolution( t )
		#位置と速度を更新
		rk4_ms[mn].r += rk4_ms[mn].dr
		rk4_ms[mn].v += rk4_ms[mn].dv
		#衝突リスト
		collisionlist = CollisionDetection( rk4_ms[mn].r, rk4_M.r )
		#衝突あり
		if( len(collisionlist) > 0 ):

			#２つ以上の壁orシリンダーとの衝突がある場合　→　高精度計算モードへ
			if( len(collisionlist) > 1 ):
				return False

			#以下、衝突が１箇所
			normal = collisionlist[ 0 ]["normal"]

			#いったん巻き戻しておく
			rk4_ms[mn].r -= rk4_ms[mn].dr
			rk4_ms[mn].v -= rk4_ms[mn].dv

			#衝突力
			if( collisionlist[ 0 ]["type"] == 1 ): #壁との衝突
				#相対速度
				barV = - rk4_ms[mn].v
				rk4_ms[mn].fc = (2.0 * m * np.dot(barV, normal) * normal ) / rk4_ms[mn].dt - np.dot(rk4_ms[mn].fex, normal) * normal
				#位置と速度を更新
				rk4_ms[mn].timeEvolution( t )
				rk4_ms[mn].r += rk4_ms[mn].dr
				rk4_ms[mn].v += rk4_ms[mn].dv

			if( collisionlist[ 0 ]["type"] == 2 ): #ピストンとの衝突
				cylinderCollisionList.append( { "mn": mn, "normal":normal}  )


	#ピストンとの衝突ありの場合
	if( len(cylinderCollisionList) > 0 ):
		#ピストンをいったん巻き戻して
		rk4_M.r -= rk4_M.dr
		rk4_M.v -= rk4_M.dv
		#衝突数が１個の場合
		if( len(cylinderCollisionList) == 1 ):
			mn = cylinderCollisionList[0]["mn"]
			normal = cylinderCollisionList[0]["normal"]
			#衝突力の計算
			barV = rk4_M.v - rk4_ms[mn].v
			barFex =  rk4_M.fex / M - rk4_ms[mn].fex / m 
			rk4_ms[mn].fc = 2.0 * m * M / ( m + M ) * np.dot(barV, normal) * normal / rk4_ms[mn].dt
			rk4_ms[mn].fc += m * M / ( m + M ) * np.dot(barFex, normal) * normal

			rk4_M.fc = - rk4_ms[mn].fc
			#時間発展を計算
			rk4_ms[mn].timeEvolution( t )
			rk4_M.timeEvolution( t )
			#位置と速度を更新
			rk4_M.r += rk4_M.dr
			rk4_M.v += rk4_M.dv
			rk4_ms[mn].r += rk4_ms[mn].dr
			rk4_ms[mn].v += rk4_ms[mn].dv

		#衝突数が２個以上の場合
		if(len(cylinderCollisionList) > 1):
			#時間発展：失敗
			return False

	#時間発展：正常
	return True

#高精度計算モード
def highPrecisionMode(t, level):
	print ( "レベル" , level, "t=", round(t, level+3 ))

	timeScale = 10**level
	#時間ステップ計算前に巻き戻す
	rk4_M.r = rk4_M._r.copy()
	rk4_M.v = rk4_M._v.copy()
	for mn in range(N):
		rk4_ms[mn].r = rk4_ms[mn]._r.copy()
		rk4_ms[mn].v = rk4_ms[mn]._v.copy()

	for _n in range(0, 10):
		t_level = t + dt / timeScale * _n

		result_level = timeEvolution( t_level, timeScale )

		if( result_level == False ): 
			#もっと高精度が必要
			level += 1
			highPrecisionMode(t_level, level)
			level -= 1
			
		else:
			#時間発展前の位置と速度を保持
			rk4_M._r = rk4_M.r.copy()
			rk4_M._v = rk4_M.v.copy()
			for mn in updateList:
				rk4_ms[mn]._r = rk4_ms[mn].r.copy()
				rk4_ms[mn]._v = rk4_ms[mn].v.copy()		


#################################################################
## 計算開始
#################################################################
#初速度配列
v0s = []
#初期運動エネルギーの計算（z成分のみ）
Tz0 = 0
for n in range(N):
	theta_v = np.pi * random.random()
	phi_v = 2.0 * np.pi * random.random()
	_v_max = v_max * random.random()
	v0s.append( np.array([
		_v_max * np.sin( theta_v ) * np.cos( phi_v ), 
		_v_max * np.sin( theta_v ) * np.sin( phi_v ), 
		_v_max * np.cos( phi_v )
	]))
	Tz0 += 1.0/2.0 * m * v0s[len(v0s)-1][2]**2

#ピストンの初期位置
h0 = L #2*Tz0/(M*10)

rk4_ms = []
for n in range(N):
	r0 = np.array([
		L * random.random(), 
		L * random.random(), 
		h0 * random.random()
	])
	rk4_ms.append( RK4(m, dt, r0, v0s[n]) )
#粒子外力の設定
for n in range(N):
	rk4_ms[n].fex = fex


#シリンダーの初期値

R0 = np.array([0.0, 0.0, h0  ])
V0 = np.array([0.0, 0.0, 0.0])
rk4_M = RK4(M, dt, R0, V0)
rk4_M.beta = beta
rk4_M.fex = Fex

####################################
## 描画用関数
####################################
#アニメーション作成用
ims=[]
#ピストン位置の配列
zs = []

def plot( t ):
	#各コマを描画
	img  = plt.plot([0-lw, L+lw], [rk4_M.r[2] + lw, rk4_M.r[2] + lw], colors[1], linestyle='solid', linewidth = 10 )
	for n in range(N):
		img += plt.plot([rk4_ms[n].r[0]], [rk4_ms[n].r[2]], colors[0], marker = 'o', markersize = 10, linestyle='solid', linewidth = 0 )
	time = plt.text(0.9, 1.12, "t = " +  str("{:.2f}".format(t)) + r"$[{\rm s}]$" , ha='left', va='center', fontsize=18)
	#テキストをグラフに追加
	img.append( time )
	ims.append( img )
	E = T() + U()
	print( "{:.2f}".format(t), E )
	zs.append( [t, rk4_M.r[2]] )
	#print( "{:.2f}".format(t), U_cylinder(), T_particle(), 3.0 * U_cylinder()/T_particle() )


####################################
## 各時刻における運動方程式の計算
####################################
for tn in range(len(ts)):
	t = ts[tn]
	if( tn%skip == 0 ): 
		plot(t)

	#計算対象粒子リスト（ここでは全部）
	updateList = list(range(0, N))

	#時間発展前の位置と速度を保持
	rk4_M._r = rk4_M.r.copy()
	rk4_M._v = rk4_M.v.copy()
	for mn in updateList:
		rk4_ms[mn]._r = rk4_ms[mn].r.copy()
		rk4_ms[mn]._v = rk4_ms[mn].v.copy()		

	level = 0

	#時間発展の計算
	result = timeEvolution( t, 1 )

	#複数衝突を検知した場合
	if( result == False ): 
		level += 1
		highPrecisionMode(t, level )

#フォルダ指定
dir_path = "output/"
#フォルダ生成
os.makedirs(dir_path, exist_ok = True)
filepath = dir_path + "z.txt"
#ファイルへ保存
np.savetxt( filepath, zs)

#################################################################
## 分子運動アニメーション
#################################################################
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([-0.1, 1.1])
plt.ylim([-0.1, 1.1])


#箱の概形
plt.vlines([0-lw], -2*lw, L, "black", linestyles='solid', linewidth = 10 )
plt.vlines([L+lw], -2*lw, L, "black", linestyles='solid', linewidth = 10 )
plt.hlines([-lw], 0, L, "black", linestyles='solid', linewidth = 10 )

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

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

#########################################################################
# グラフの描画（エネルギー固有値）
#########################################################################
fig1 = plt.figure(figsize=(15, 8))
plt.title( u"ピストンの位置", fontsize=20, fontname="Yu Gothic", fontweight=1000)
plt.xlabel( u"時間" + r"$\,[{\rm s}]$", fontsize=30, fontname="Yu Gothic", fontweight=1000  )
plt.ylabel( u"高さ" + r"$\,[{\rm m}]$", fontsize=30, fontname="Yu Gothic", fontweight=1000 )
#余白の調整
plt.subplots_adjust(left = 0.1, right = 0.98, bottom=0.12, top = 0.95)
#罫線の描画
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([0, t_max])
plt.ylim([0, L])

xs = []
ys = []
for zt in zs:
	xs.append( zt[0] )
	ys.append( zt[1] )

#グラフを描画	
plt.plot(xs, ys, linewidth = 5)

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

コメントを残す

メールアドレスが公開されることはありません。 ※ が付いている欄は必須項目です