tsujimotterのノートブック

日曜数学者 tsujimotter の「趣味で数学」実践ノート

日曜化学(2.5):メトロポリス・ヘイスティングス法を用いた電子雲の可視化(Python/matplotlib)

前回の記事で、水素原子の電子雲の可視化方法して、3つの方法を紹介しました。
tsujimotter.hatenablog.com


特に、方法2「散布図としてプロット」については、棄却サンプリング法を使った方法を紹介しました。

f:id:tsujimotter:20210704031955j:plain:w400

今回、電子の確率分布を表す  f(x, y, z) = |\Psi(x, y, z)|^2 なる3次元の確率密度関数にしたがう乱数を得る必要がありました。

棄却サンプリング法は、 f(x, y, z) とは直接関係ない(たとえば一様乱数のような)確率分布  g(x, y, z) にしたがう乱数から、 f g の確率密度関数としての差を考慮しつつ、適切な確率で乱数を「棄却」することによって、目的の  f(x, y, z) にしたがう乱数を得る手法のことでした。

ところがこの方法、非常に時間がかかることが知られています。実際、上記の散布図を描くだけでかなりの時間を要しました。特に、 f g の差が大きければ大きいほど、棄却率が高くなり、乱数を生成するのに時間がかかるというわけです。


今回紹介する メトロポリス・ヘイスティングス法 を用いれば、上記の問題の多くは解決します。

マルコフ連鎖モンテカルロ法(MCMC)という一般的な枠組みの一手法であり、棄却サンプリングのような素朴な手法と比べて、効率的に乱数生成を行うことができます。


これまでこの手の手法(マルコフ連鎖モンテカルロ法)については、話には聞いてはいたものの、なんとなく手が出せないでいました。しかしながら、今回の目的に向けて調べてみると、案外簡単に実装できることがわかりました。(やっぱりモチベーションって大事ですね。)

f:id:tsujimotter:20210711111649p:plain:w360

仕組みはちゃんと理解できているわけではないですが、アルゴリズムの流れと実装方法は理解できたので、電子雲の可視化に応用してみたというのが今回の記事です。

当初の予定では、日曜化学シリーズの記事として書く予定はなかったので、番外編として日曜化学(2.5)としています。よろしければご覧になってください!



メトロポリス・ヘイスティングス法

簡単にアルゴリズムの流れを紹介します。

このアルゴリズムによって生成したいのは、波動関数を  \Psi(x, y, z) として、確率密度関数  P(x, y, z) = |\Psi(x, y, z)|^2 にしたがう  N 個のランダムな3次元ベクトルの列

 {\bf x}_1, \; {\bf x}_2, \; {\bf x}_3, \; \ldots, \; {\bf x}_N

です。


まず、3次元ベクトル  {\bf x} = (x, y, z) および  {\bf u} = (u, v, w) を変数に持つ、条件付き確率密度関数  Q({\bf x} \mid {\bf u}) を用意します。この  Q提案分布といいます。

この  Q については、 {\bf u} に対して  Q({\bf x} \mid {\bf u}) にしたがうランダムなベクトル  {\bf x} を生成できるような関数(乱数が作れる確率密度関数)に決める必要があります。また、あとに示す「条件」を満たすようにすると話が簡単になります。


初期値となる3次元ベクトル  {\bf x}_0 = (x_0, y_0, z_0) を用意します。

 {\bf x}_0 からスタートし、 {\bf x}_0 の情報から {\bf x}_{1} を生成します。続けて  {\bf x}_1 の情報から  {\bf x}_{2} を生成します。こんな風に  {\bf x}_t の情報から  {\bf x}_{t+1} を生成するというのを  t = 0, 1, 2, \ldots, N-1 に対して行います。
(時刻  t の情報に基づいて、時刻  t+1 における状態の確率が定まる系のことを「マルコフ連鎖」といいます。)

より具体的な生成方法として、先ほど用意した  Q を用いて、 Q({\bf x}_{t+1}^* \mid {\bf x}_{t+1}) によって、 {\bf x}_{t+1}候補となる乱数  {\bf x}_{t+1}^* を生成します。以下の確率  \alpha によって、候補  {\bf x}_{t+1}^* {\bf x}_{t+1} として採用します。

 \displaystyle \alpha = \frac{P({\bf x}_{t+1}^*) \cdot Q({\bf x}_{t} \mid {\bf x}_{t+1}^*)}{P({\bf x}_{t}) \cdot Q({\bf x}_{t+1}^* \mid {\bf x}_{t})} \tag{1}

ただし、 \alpha > 1 のときは、 \alpha = 1 とします。

逆に、確率  1-\alpha によって、 {\bf x}_{t+1}^* を棄却します。この場合、再度同じ手順で  {\bf x}_{t+1}^* を得て、同様に上記の確率で採用するか棄却するかを判定します。

f:id:tsujimotter:20210711105032p:plain:w500


以上を繰り返して、採用されたランダムなベクトルの列

 {\bf x}_1, \; {\bf x}_2, \; {\bf x}_3, \; \ldots, \; {\bf x}_N

を考えると、これが目的の分布  P(x, y, z) にしたがうランダムなベクトルの列となっているというわけです。


たったこれだけなので、実装も非常に簡単です。理屈を理解していないので、こんな方法でなぜうまくいくのかはよく分かりませんが、とにかくこれでうまくいくのだそうです。


また、 Q({\bf x} \mid {\bf u}) を決める必要がありますが、これについては多次元正規分布を用いることにします。

 \displaystyle Q({\bf x} \mid {\bf u}) = \frac{\exp\left( -\frac{1}{2}({\bf x} - {\bf u})^T \Sigma^{-1}({\bf x} - {\bf u}) \right)}{\sqrt{(2\pi)^k |\Sigma|}} \tag{2}

ここで、 \Sigma は分散共分散行列です。変数間の相関を考える必要がないので、単に対角行列を考えれば十分です。


この  Q は「①  Q({\bf x}_{t+1}^* \mid {\bf x}_{t+1}) にしたがうランダムなベクトル  {\bf x}_{t+1}^* を生成する」と「② 式  (1) を計算するのに用いる」という2つの役割があります。

①については、numpyに多次元正規分布にしたがう乱数を生成する関数

np.random.multivariate_normal

がありますので、これを使えば簡単です。具体的な使い方は、こちらのページ を参照。


②については、そもそも計算する必要がないという話があります。というのも、式  (2) の形を見ると、明らかに  {\bf x} {\bf u} で対称的な形になっていますね。すなわち

 Q({\bf u} \mid {\bf x}) = Q({\bf x} \mid {\bf u}) \tag{3}

ということです。

この性質を使うと、式  (1) は次のように劇的になります:

 \require{cancel} \displaystyle \alpha = \frac{P({\bf x}_{t+1}^*) \cdot \cancel{Q({\bf x}_{t} \mid {\bf x}_{t+1}^*)}}{P({\bf x}_{t}) \cdot \cancel{Q({\bf x}_{t+1}^* \mid {\bf x}_{t}})}  = \frac{P({\bf x}_{t+1}^*)}{P({\bf x}_{t})} \tag{4}

なるほど、これなら  Q の値は計算する必要がなくなりますね。むしろこうなるように、多次元正規分布を選んだということでしょう。


さらにいうと、棄却率を計算する  P({\bf x}_{t+1}^*) P({\bf x}_{t}) という式は、その比がわかれば十分です。つまり、波動関数の規格化定数はわからなくても全く問題ないのです。波動関数の規格化定数の決定は大変面倒だったので、この性質はありがたいですね。



電子雲を可視化するPythonプログラム

そんなわけで、メトロポリス・ヘイスティングス法を用いた可視化のプログラムを紹介します。


例によってPythonのmatplotlibを使っていますので、以下をインストールしておいてください。

  • numpy
  • matplotlib

また、前回紹介したPlotlyというライブラリを使った可視化もできるようになっています。使いたい人は、Plotlyをインストールした上で、適切にコード内のコメントアウトを外してください。

# メトロポリス・ヘイスティング法による電子分布の可視化

import numpy as np
import matplotlib.pyplot as plt
import mpl_toolkits.mplot3d.axes3d as axes3d

import math


# 動径波動関数(ボーア半径を a_0 = 1 に正規化)
def radial_wave_function(r, n, l, Z_eff):
    rho = 2.0 * Z_eff * r / n
    if n == 1 and l == 0:
        return (Z_eff ** 1.5) * np.exp(-rho/2.0) * 2.0
    elif n == 2 and l == 0:
        return (Z_eff ** 1.5) * np.exp(-rho/2.0) * (2.0 - rho) / (2.0*np.sqrt(2.0))
    elif n == 2 and l == 1:
        return (Z_eff ** 1.5) * np.exp(-rho/2.0) * rho / (2.0*np.sqrt(6.0))
    elif n == 3 and l == 0:
        return (Z_eff ** 1.5) * np.exp(-rho/2.0) * (6.0 - 6.0*rho + rho**2) / (9.0*np.sqrt(3.0))
    elif n == 3 and l == 1:
        return (Z_eff ** 1.5) * np.exp(-rho/2.0) * (4.0 - rho) * rho / (9.0*np.sqrt(6.0))
    elif n == 3 and l == 2:
        return (Z_eff ** 1.5) * np.exp(-rho/2.0) * (rho**2) / (9.0*np.sqrt(30.0))


# 球面調和関数(ただし、実関数表示したもの)
def spherical_harmonics(theta, phi, l, m):
    if l == 0:
        if m == 0:
            # l=0, m=0
            return np.sqrt(1.0/(4*np.pi))
    if l == 1:
        if m == 0:
            # l=1, m=0
            return np.sqrt(3.0/(4.0*np.pi)) * np.cos(theta)
        if m == 1:
            # l=1, m=+1
            return np.sqrt(3.0/(4.0*np.pi)) * np.sin(theta) * np.cos(phi)
        if m == -1:
            # l=1, m=+1
            return np.sqrt(3.0/(4.0*np.pi)) * np.sin(theta) * np.sin(phi)
    if l == 2:
        if m == 0:
            return np.sqrt(5.0/(16.0*np.pi)) * (3.0*(np.cos(theta)**2) - 1.0)
        if m == 1:
            return np.sqrt(15.0/(4.0*np.pi)) * np.cos(theta) * np.sin(theta) * np.cos(phi)
        if m == -1:
            return np.sqrt(15.0/(4.0*np.pi)) * np.cos(theta) * np.sin(theta) * np.sin(phi)
        if m == 2:
            return np.sqrt(15.0/(16.0*np.pi)) * (np.sin(theta)**2) * np.cos(2*phi)
        if m == -2:
            return np.sqrt(15.0/(16.0*np.pi)) * (np.sin(theta)**2) * np.sin(2*phi)

        
        
# 水素様原子の波動関数 Ψ_{n,l,m}(x, y, z)
#
# (n,l,m): 量子数
# Z_eff:   有効核電荷
# 
def f(x, y, z, n, l, m, Z_eff):
    # 座標系を (x,y,z) -> (r,θ,φ) に変換
    r = np.sqrt(x*x + y*y + z*z)
    
    theta = 0.0
    phi = 0.0
    if r > 0:
        theta = np.arccos( z / r )
    if y == 0:
        if x < 0:
            phi = np.pi
    elif x*x+y*y > 0:
        phi = np.sign(y) * np.arccos(x/np.sqrt(x*x+y*y))
        
    # 動径波動関数と球面調和関数の積を計算して出力
    return (radial_wave_function(r, n, l, Z_eff) * spherical_harmonics(theta, phi, l, m))


# 量子数
n = 3
l = 2
m = 0

# 有効核電荷
Z = 1        # 水素原子を想定しているので Z = 1
Z_eff = Z


# サンプリング数の設定
N = 5000                  # サンプリング目標個数
x_range = 18              # 座標の範囲 (-x_range, x_range)


# メトロポリス・ヘイスティング法を使った電子雲のサンプリング
num_of_points = 0  # 採択されたデータ数

x = 1
y = 1
z = 1


# 波動関数が正のデータを格納する変数
x1_list = []
y1_list = []
z1_list = []

# 波動関数が負のデータを格納する変数
x2_list = []
y2_list = []
z2_list = []



while num_of_points < N:
    if num_of_points % 10 == 0:
        print(num_of_points)
    # 確率分布 Q(x*|x_t) から次の状態の候補 x* の乱数生成(今回は Q として3次元ガウス分布を選択)
    values = np.random.multivariate_normal([x, y, z], [[10,0,0],[0,10,0],[0,0,10]], 1)  # 乱数生成
    x_star = values[0,0]
    y_star = values[0,1]
    z_star = values[0,2]

    f_t = f(x, y, z, n, l, m, Z_eff)                    # xの波動関数
    f_star = f(x_star, y_star, z_star, n, l, m, Z_eff)  # x* の波動関数

    # 採択率を計算  α = p(x*) / p(x_t)
    alpha = (f_star ** 2) / (f_t ** 2)

    # 採択率により決定
    r = np.random.rand(1)[0]   # 0 <= r <= 1 の一様乱数 r を生成
    if r <= alpha:
        x = x_star
        y = y_star
        z = z_star
        if f_star > 0:
            x1_list.append(x)
            y1_list.append(y)
            z1_list.append(z)
        else:
            x2_list.append(x)
            y2_list.append(y)
            z2_list.append(z)            
        num_of_points += 1



# figureを生成する
X1 = np.array(x1_list)
Y1 = np.array(y1_list)
Z1 = np.array(z1_list)

X2 = np.array(x2_list)
Y2 = np.array(y2_list)
Z2 = np.array(z2_list)



# matplotlibによる可視化
fig = plt.figure(figsize=(8.0, 8.0))
ax = fig.add_subplot(1,1,1, projection='3d')

ax.set_xlim3d(-x_range, x_range)
ax.set_ylim3d(-x_range, x_range)
ax.set_zlim3d(-x_range, x_range)

# 軸の設定
ax.set_xlabel("x", size = 14)
ax.set_ylabel("y", size = 14)
ax.set_zlabel("z", size = 14)

ax.plot(X2,Y2,Z2,color='b',marker=".",linestyle='None')
ax.plot(X1,Y1,Z1,color='r',marker=".",linestyle='None')

plt.show()


'''
# Plotlyによる可視化
import plotly.graph_objects as go

fig = go.Figure(data=[go.Scatter3d(
    x = X1,
    y = Y1,
    z = Z1,
    mode='markers',
    marker = dict(
        size = 1,
        color = 'red'
    )
), go.Scatter3d(
    x = X2,
    y = Y2,
    z = Z2,
    mode='markers',
    marker = dict(
        size = 1,
        color = 'blue'
    )
)])

fig.update_layout(
    scene = dict(
        camera=dict(eye=dict(x=1.7, y=-1.7, z=1.7)),  #eye=dict(x=1.25, y=1.25, z=1.25)), 
        xaxis = dict(range=[-x_range, x_range]),
        yaxis = dict(range=[-x_range, x_range]),
        zaxis = dict(range=[-x_range, x_range]),
        aspectmode='cube', #this string can be 'data', 'cube', 'auto', 'manual'
        #a custom aspectratio is defined as follows:
        aspectratio=dict(x=1, y=1, z=0.95)
    ),
    width=700,
    height=700,
)

fig.show()
'''


上記のコードは、量子数  (n, l, m) = (3, 2, 0) に対応する水素原子の波動関数(すなわち、d軌道)を図示するものです。量子数については、以下の部分を変更すれば、違うものをプロットすることができます。

# 量子数
n = 3
l = 2
m = 0

また、デフォルトではサンプル数5000としていますが、これも変更可能です。


実行してみると、冒頭に紹介した以下の図が得られます。

f:id:tsujimotter:20210711111649p:plain:w400
図はPlotlyによるもの

たしかにd軌道の形に電子が分布していますね!!すごい!!


実際実行してみると分かりますが、5000個のサンプリングであっても、ものの数秒で実行完了します。これは棄却サンプリングの結果とは大違いです。


超楽しい!!!

f:id:tsujimotter:20210711112401g:plain:w500


次回の記事はこちら

tsujimotter.hatenablog.com


追記:2021.07.12

@dc1394さんという方が開発されている「SchracVisualize_Direct3D_11」をご紹介させてください。水素原子の可視化ツールです。実行ファイルによりローカルで実行でき、とても高速に、綺麗な絵を出力してくれます。(3次元をマウスでグリグリ動かして見ることができて便利です。)

このツールは、元々棄却サンプリングを使用されていたそうなのですが、こちらの記事を見てメトロポリス・ヘイスティングス法を使った実装に書き換え、高速化に成功されたそうです。


この方に教えていただいたのですが、メトロポリス・ヘイスティングス法に使っている正規分布について、私は適当に分散  10.0 とおいていたのですが、どうやら波動関数の動径分布関数、すなわち

 \Psi = R(r) Y(\theta, \phi)

と変数分離した  R(r) を用いて  D(r) = r^2 |R(r)|^2 によって定義された関数を考えて、そのピーク値をとるような  r を分散に採用すればよいとのことです。

また、3次元正規分布によって乱数生成していた部分も、1次元正規分布を独立に  x, y, z について適用すれば問題ないということもご指摘いただきました。(言われてみれば確かにそうですね。)

教えていただいてありがとうございます!