２種類の気体分子が円形の容器内で動く様子のシミュレーションを行い、その様子をアニメーションとして保存する。詳細は https://ushitora.net/archives/2197 を参照

molecule.py

# ２種類の気体分子が円形の容器内で動く様子のシミュレーションを行い、
# その様子をアニメーションとして保存する。
# 詳細は https://ushitora.net/archives/2197 を参照

import numpy as np
from matplotlib import pyplot as plt
from matplotlib import animation as animation

""" データの準備 """
# 乱数で２種類の気体の極座標生成
r = np.random.random((2, 100, 1))
theta = np.random.random((2, 100, 1)) * 2 * np.pi
# 極座標を直交座標に変換
x = np.concatenate((r * np.cos(theta), r * np.sin(theta)), axis=2)
# 乱数で気体の速度方向（直交座標）を生成
v = np.random.random((2, 100, 2)) - 0.5

""" 直交座標から、x軸との角度を計算 """
def atan_pi(x, y):
    alpha = np.arctan(y / x)
    # numpy.arctan は [-pi/2, pi/2] で値を返す
    alpha[x < 0] += np.pi
    return alpha

""" 動画用の関数 """
# FuncAnimation は、この関数を実行してはフレームに保存、というプロセスを繰り返す。
# その際、変数 i の値は 1 ずつ増えていく
def _update_plot(i, fig, im):
    # 先に定義した x, v を使用
    global x, v
    # 各気体のプロットを消す。
    # im[0] には円形の枠線がプロットされているため、これは残す
    for j in range(len(im) - 1):
        im[-1].remove() # 最後のプロットを削除
        im.pop()        # 最後のプロットの内容（情報）も削除

    """ 以下、数学パート """
    """ Matplotlib でアニメーションを作りたいだけなら無視する（ここから） """
    
    # x を移動させる
    x = x + 0.05 * v
    # x の原点からの距離
    r = np.square(x).sum(axis=2)
    # 移動の結果、半径 1 の円の外に出てしまった気体分子
    out = r > 1
    # 外に出た気体分子は演習上に配置（ r = 1 にする）
    x[out] /= np.sqrt(r[out]).reshape(-1, 1)
    # 外に出た気体分子について、速度を極座標化したときの半径
    v_r = np.sqrt(np.square(v[out]).sum(axis=1))
    # 外に出た気体分子について、速度を極座標化したときの角度
    v_theta = atan_pi(v[out][:, 0], v[out][:, 1])
    # 外に出た気体分子について、位置を極座標化したときの角度
    x_theta = atan_pi(x[out][:, 0], x[out][:, 1])
    # 壁との衝突した後の、速度の角度（極座標）
    v_theta_new = np.pi + 2 * x_theta - v_theta
    # 外に出た気体分子の速度を更新。速さ（ v_r ）は変えず、向きだけ変える（ v_theta_new ）
    v[out] = np.concatenate(((v_r * np.cos(v_theta_new)).reshape(-1, 1), (v_r * np.sin(v_theta_new)).reshape(-1, 1)), axis=1)
    
    """ （ここまで） """

    # 気体 A の位置をプロット
    im.append(plt.scatter(x[0, :, 0], x[0, :, 1], alpha=0.5, c="C1", label="Gas A"))
    # 気体 B の位置をプロット
    im.append(plt.scatter(x[1, :, 0], x[1, :, 1], alpha=0.5, c="C2", label="Gas B"))
    # 凡例を表示（位置を右上に固定する）
    plt.legend(bbox_to_anchor=(1, 1), loc='upper right', borderaxespad=1)
    # タイトルに経過時間を表示
    plt.title("Time: {}".format(i))


""" アニメーションを作成 """
fig = plt.figure(figsize=(6, 6)) # 正方形の動画にする
plt.xlim(-1, 1) # x 軸の範囲を指定
plt.ylim(-1, 1) # y 軸の範囲を指定

# ここに、各プロットの内容が保存される。
# im[i].remove() とすれば、そのプロットを消すことができる
im = []

# 外枠の円の極座標（角度）を生成
circle_theta = np.array([np.pi / 50 * i for i in range(101)])
# 円を描画し、このプロットは im[0] に保存する
im.append(plt.plot(np.cos(circle_theta), np.sin(circle_theta)))

# _update_plot を実行しながら、結果をアニメーション（長さ 100 フレーム）として保存。
# その際、変数 fig, im を利用する（ i は書かずとも勝手に増えていく）
ani = animation.FuncAnimation(fig, _update_plot, fargs=(fig, im), frames=100, interval=1)
# アニメーションを gas.gif として保存
ani.save("gas.gif", writer="pillow")