scipyの特殊関数

量子力学の教科書を読んでいて、球Bessel関数とかLaguerreの陪多項式がある微分方程式の解だと言われてもいまいちピンと来なかったので適当に実装してmatplotlibでグラフを書いてみようと思っていたところ、scipyに特殊関数のmoduleがあることを知った。

Special functions (scipy.special) — SciPy v0.15.1 Reference Guide

有名所の特殊関数はだいたい実装されており、折角なので特殊関数の概形を見てみるだけでなく水素原子のSchrödinger方程式の解もプロットしてみることにした。

ここで問題なのは3次元上の各点での波動関数の振幅の2乗をカラースケールで表示しても、意味の分からないグラフとなってしまうことだ。これを回避するために、各座標に点がプロットされる確率を波動関数の振幅の２乗と対応するようにした。n=1,2ではこれでうまくいくが、n>=3だと空間的に広がるために毎回パラメターをいじる必要がある。

下図は(n,l,m)=(2,1,0)のときの結果。ボーア半径を1としてある。

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import scipy.special
import numpy as np
import random
import math

def freeParticle(n,l,m,r,theta,phi):
    res = (-1)**(m+abs(m))
    res *= scipy.special.sph_harm(m,l,phi,theta)
    res *= scipy.special.sph_jn(0,r)[0][0]
    return res

def hydrogen(n,l,m,r,theta,phi):
    res = (-1)**(m+abs(m))
    res *= scipy.special.sph_harm(m,l,phi,theta)
    
    rho = 2.0*r/n
    res *= scipy.special.assoc_laguerre(rho,n+l,2*l+1)
    tmp = np.sqrt(math.factorial(n-l-1)/2.0/n*math.factorial(n+l)**-3*(2.0/n)**3)
    res *= -tmp
    res *= np.exp(-rho/2.0)*(rho**l)
    
    return res

def plotGraph(n,l,m,waveFunction):
    L = 2.0
    dl = 0.1

    fig = plt.figure()
    ax = Axes3D(fig)
    
    ax.set_xlabel("X-axis")
    ax.set_ylabel("Y-axis")
    ax.set_zlabel("Z-axis")
    
    ax.set_xlim(-L,L)
    ax.set_ylim(-L,L)
    ax.set_zlim(-L,L)

    for z in np.arange(-L,L,dl):
        for y in np.arange(-L,L,dl):
            for x in np.arange(-L,L,dl):
                r = np.sqrt(x**2 + y**2 + z**2)
                theta = np.arccos(z/r)
                phi = np.arctan2(y,x)

                func = waveFunction(n,l,m,r,theta,phi)
                probability = func*np.conj(func)
                print probability
                ran = [random.random() for i in range(20)]
                if min(ran) < probability:
                    ax.scatter3D(x,y,z,s=10)

    plt.show()

if __name__=='__main__':
    plotGraph(2,1,-1,hydrogen)