CUDA サンプルにもある nbody (多体問題) を cupy で解いてみた。
うまくはまれば namba を使うよりも少ない行数で記述できる。
実行速度は namba とほぼ変わらなかった。
初速は銀河面上を円運動するように与えているため、シミュレーション
時間を長くしてもあまり面白い絵になっていませんので、
適宜工夫してみて下さい。
ソースコードは以下の通り。
# -*- coding: utf-8 -*-
import datetime
import cupy as cp
import numpy as np
import matplotlib.pyplot as plt
# 質点数
N = 1024
REP = 5000
AREAX = 1e3
AREAY = 1e3
FLIP = 0
# 物理定数
cp.cuda.set_allocator(cp.cuda.MemoryPool().malloc)
MASS = cp.float32(1.0)
M = cp.ones(N, dtype=cp.float32) * MASS
G = cp.float32(1.0)
EPSILON = cp.float32(1e-6)
DT = cp.float32(0.01)
# 分布、初速を決定する
DAX = cp.zeros(N, dtype=cp.float32)
DAY = cp.zeros(N, dtype=cp.float32)
def advance(x0, y0, xx1, xx2, yy1, yy2, M, EPSILON):
diffx = x0[xx2] - x0[xx1]
diffy = y0[yy2] - y0[yy1]
r = cp.sqrt(diffx * diffx + diffy * diffy)
r3 = r * r * r + EPSILON
dax = cp.sum(M * diffx / r3, axis = 1)
day = cp.sum(M * diffy / r3, axis = 1)
return dax, day
def update(x0, y0, vx0, vy0, DAX, DAY, G, DT):
vx1 = vx0 + DAX * G * DT
vy1 = vy0 + DAY * G * DT
x1 = x0 + vx1 * DT
y1 = y0 + vy1 * DT
return x1, y1, vx1, vy1
def eccentric(x):
return (cp.exp(2 - 8 * x) - cp.exp(2)) / (cp.exp(2 - 8) - cp.exp(2))
def init_dist():
global AREAX, AREAY, MASS, G
RR = eccentric(cp.random.rand(N)) + 1e-7
THETA = cp.random.rand(N)
XS = RR * cp.cos(THETA * 2 * cp.pi)
YS = RR * cp.sin(THETA * 2 * cp.pi)
VAREA = cp.ones(N, dtype=cp.float32) * cp.sqrt(G/N) / RR
RR = cp.random.rand(N) / 4
THETA = cp.random.rand(N)
VSX = RR * cp.cos(THETA * 2 * cp.pi)
VSY = RR * cp.sin(THETA * 2 * cp.pi)
r = cp.sqrt(XS * XS + YS * YS)
VSX = VAREA * (-YS + VSX * (1 - r))
VSY = VAREA * (XS + VSY * (1 - r))
XS *= AREAX
YS *= AREAY
# 重心を求める
x0 = cp.sum(XS) / N
y0 = cp.sum(YS) / N
vx0 = cp.sum(VSX) / N
vy0 = cp.sum(VSY) / N
# 重心座標系に移行する
XS -= x0
YS -= y0
VSX -= vx0
VSY -= vy0
return XS, YS, VSX, VSY
# 分布、初速を決定する
XS, YS, VSX, VSY = init_dist()
X = [XS, cp.zeros(N, dtype=cp.float32)]
Y = [YS, cp.zeros(N, dtype=cp.float32)]
VX = [VSX, cp.zeros(N, dtype=cp.float32)]
VY = [VSY, cp.zeros(N, dtype=cp.float32)]
print("intialized")
# REP 時間繰り返す
x0s = range(N)
y0s = range(N)
xx1, xx2 = np.meshgrid(x0s, x0s)
yy1, yy2 = np.meshgrid(y0s, y0s)
st = datetime.datetime.now()
for t in range(REP):
DAX, DAY = advance(X[FLIP], Y[FLIP], xx1, xx2, yy1, yy2, M, EPSILON)
X[1-FLIP], Y[1-FLIP], VX[1-FLIP], VY[1-FLIP] = \
update(X[FLIP], Y[FLIP], VX[FLIP], VY[FLIP], \
DAX, DAY, G, DT)
if t % (REP // 50) == 0:
print('*', end = '', flush = True)
# 結果をグラフで表示
if t == 0:
plt.plot(cp.asnumpy(X[FLIP]), cp.asnumpy(Y[FLIP]), 'x', color='red')
FLIP = 1 - FLIP
ed = datetime.datetime.now()
print("\nelapsed time:", ed - st)
x = cp.asnumpy(X[FLIP])
y = cp.asnumpy(Y[FLIP])
plt.plot(x, y, '.', color='green')
plt.show()