pythonでブラウン運動(ランダムウォーク)のシミュレーション(1次元・3次元)

1次元

matplotlib.animation.FuncAnimation の機能を使って、Jupyter Notebook上でブラウン運動をシミュレーションすることが出来る（注意：Jupyter lab, Google Colabなどで動かない可能性があります。）

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from IPython.display import HTML

# parameters
n = 100  # タイムステップ数
trend = 0.0001  # トレンド成分
variance = 2  # 分散

fig, ax = plt.subplots()
x, y = [], []
g, = plt.plot([], [], '-', animated=True)
plt.close()

movement = np.cumsum(np.random.randn(n + 1)*variance + np.cumsum(np.linspace(0, (trend)*n, n+1)))
movement = np.append(np.array([0]), movement[:n])


def init_func():
    ax.set_xlim(0, n)
    ax.set_ylim(movement.min()*1.2, movement.max()*1.2)
    return g,

def func(frame):
    x.append(frame)
    y.append(movement[int(frame)])
    g.set_data(x, y)
    return g,

a = FuncAnimation(fig, func, frames=np.linspace(0, n, n),
                    init_func=init_func, blit=True)

HTML(a.to_jshtml())

ちなみに、動かないバージョンは以下の簡単なコードで実装できる

import matplotlib.pyplot as plt
variance = 1
plt.plot(np.cumsum(variance*np.random.randn(100)))
plt.show()

3次元

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from IPython.display import HTML

n = 100

fig = plt.figure(figsize = (8, 8))
ax = fig.add_subplot(111, projection='3d')
ax.set_xlabel("x", size = 14)
ax.set_ylabel("y", size = 14)
ax.set_zlabel("z", size = 14)


x, y, z = [], [], []
g, = ax.plot([], [], [], '-', animated=True)
plt.close()

xmovement = np.cumsum(np.random.randn(n + 1)*0.1)
xmovement = np.append(np.array([0]), xmovement[:n])
ymovement = np.cumsum(np.random.randn(n + 1)*0.1)
ymovement = np.append(np.array([0]), ymovement[:n])
zmovement = np.cumsum(np.random.randn(n + 1)*0.1)
zmovement = np.append(np.array([0]), zmovement[:n])

ax.set_xlim(min(xmovement), max(xmovement))
ax.set_ylim(min(ymovement), max(ymovement))
ax.set_zlim(min(zmovement), max(zmovement))

def init_func():
    return g,

def func(frame):
    x.append(xmovement[int(frame)])
    y.append(ymovement[int(frame)])
    z.append(zmovement[int(frame)])
    g.set_data(x, y)
    g.set_3d_properties(z)
    return g,

a = FuncAnimation(fig, func, frames=np.linspace(0, n, n),
                    init_func=init_func, blit=True)

HTML(a.to_jshtml())