"""
水素原子の波動関数プロットスクリプト
概要:
このスクリプトは、水素原子の電子の波動関数、その絶対値、位相、確率密度分布などを様々な形式で可視化します。
コマンドライン引数でプロットモードと関数タイプを指定し、静的プロット、アニメーション、
等値面、3D散布図など多様な表現で波動関数を分析することができます。
詳細説明:
本スクリプトは、以下の主要な機能を提供します。
- rplot: 球面調和関数の実部、虚部、絶対値、絶対値の二乗を3D表面としてプロット。
- ranim: rplot の3D表面プロットをアニメーション化。
- dot, dotc: 波動関数の確率密度に比例した3D散布図をプロット（位相による色分け）。
- dotanim: dot プロットをアニメーション化。
- iso3d, iso3danim: 波動関数の等値面を3Dでプロットし、アニメーション化も可能。
- Rnl: 動径波動関数と動径確率密度分布を2Dでプロット。
- Ylm, Phi: 球面調和関数の角度依存部分（ThetaとPhi）を2Dでプロット。
- normalize: 波動関数の正規化を数値的に検証。
- torus, isosurface3d, contour3d, 3dMC_animation: その他の汎用的な3Dプロットやアニメーション。
利用可能な関数タイプ (functype):
- 1s, 2s, 2px, 3dxy などの軌道指定。
- Rnl, Ylm, Theta, Phi, normalize などの波動関数コンポーネント。
- c (複素数), r (実部), i (虚部), a (絶対値), a2 (絶対値の2乗) の表示タイプ。
関連リンク:
- 元記事: https://zenn.dev/shittoku_xxx/articles/13afd6fdfac44e
"""
import sys
import argparse
import numpy as np
from numpy import abs, angle, exp, sqrt, cos, sin, arccos, arctan2, pi, min, max
from scipy.integrate import quad
from scipy.constants import physical_constants
from scipy.special import sph_harm, lpmv, assoc_laguerre
import matplotlib
import matplotlib.pyplot as plt
from matplotlib import cm
from matplotlib.animation import FuncAnimation
from matplotlib.colors import LightSource, Normalize, LinearSegmentedColormap
import tkWavefunction_H as wf
import tklib.tkgraphic.tkPlot3d as plt3d
a0 = physical_constants['Bohr radius'][0] * 1e10 # オングストローム単位に変換
pi = np.pi
mode = 'rplot'
functype = '2px'
# for Monte Carlo sampling
nsamples = 30000
figsize = (12, 10)
fontsize = 16
nbins = 100
colors_wf = [(0, 0, 1), (0, 1, 0), (1, 0, 0)] # 緑、青、赤
colors_phase = [(1, 0, 0), (0, 1, 0), (0, 0, 1), (1, 1, 0), (1, 0, 0)] # 赤, 緑, 青, 黄, 赤
cmap_name_wf = 'custom cmap for wavefunction'
cmap_name_phase = 'custom cmap for phase'
[ドキュメント]
class CustomArgumentParser(argparse.ArgumentParser):
"""
概要:
argparse.ArgumentParser を拡張し、カスタムのヘルプメッセージを提供するクラス。
詳細説明:
標準のヘルプ出力に加えて、一般的な使用方法と具体的な使用例を追記します。
"""
[ドキュメント]
def print_help(self, *args, **kwargs):
"""
概要:
カスタムのヘルプメッセージを出力します。
詳細説明:
親クラスの print_help メソッドの実行前後に、
本スクリプトの一般的な使用法と具体的な使用例をコンソールに表示します。
引数:
:param args: argparse.ArgumentParser.print_help に渡される可変長引数。
:type args: tuple
:param kwargs: argparse.ArgumentParser.print_help に渡されるキーワード引数。
:type kwargs: dict
戻り値:
:returns: None
"""
print()
print("General usage:")
super().print_help(*args, **kwargs)
print()
print("Examples:")
print(r"usage: python .\plot_wf.py [rplot|ranim|dot|dotc|dotanim|iso3d|iso3danim|rsurface] [100r|2pxa etc]")
print(r"usage: python .\plot_wf.py [Rnl|Ylm|Theta|Phi|normalize]")
print(r"usage: python .\plot_wf.py [3dMC_animation|contour3d|torus]")
[ドキュメント]
def usage():
"""
概要:
コマンドライン引数のヘルプメッセージを表示します。
詳細説明:
CustomArgumentParser クラスの print_help メソッドを呼び出し、
プログラムの使用法と例をコンソールに出力します。
戻り値:
:returns: None
"""
parser.print_help()
[ドキュメント]
def torus(theta: np.ndarray, phi: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
"""
概要:
トーラスの3D座標を計算します。
詳細説明:
トーラスの大半径Rと小半径rを用いて、与えられた角度からX, Y, Z座標を計算します。
引数:
:param theta: トーラスの縦方向の角度 (0から2π) を表すNumPy配列。
:type theta: numpy.ndarray
:param phi: トーラスの横方向の角度 (0から2π) を表すNumPy配列。
:type phi: numpy.ndarray
戻り値:
:returns: X, Y, Z座標のタプル。各要素はNumPy配列。
:rtype: tuple of numpy.ndarray
"""
# トーラスのパラメータ
R = 3 # 大半径
r = 1 # 小半径
X = (R + r * np.cos(theta)) * np.cos(phi)
Y = (R + r * np.cos(theta)) * np.sin(phi)
Z = r * np.sin(theta)
return X, Y, Z
[ドキュメント]
def sphere(X: np.ndarray, Y: np.ndarray, Z: np.ndarray) -> np.ndarray:
"""
概要:
3次元空間における点 (X, Y, Z) の半径の二乗を計算します。
詳細説明:
各座標成分の二乗和 X^2 + Y^2 + Z^2 を計算し、
単位球や一般的な球の等値面を定義する際に利用できます。
引数:
:param X: X座標のNumPy配列。
:type X: numpy.ndarray
:param Y: Y座標のNumPy配列。
:type Y: numpy.ndarray
:param Z: Z座標のNumPy配列。
:type Z: numpy.ndarray
戻り値:
:returns: 各点における X^2 + Y^2 + Z^2 の値。NumPy配列。
:rtype: numpy.ndarray
"""
return (X**2 + Y**2 + Z**2) # 単位球: x^2 + y^2 + z^2 = r^2
[ドキュメント]
def exp_r2(X: np.ndarray, Y: np.ndarray, Z: np.ndarray) -> np.ndarray:
"""
概要:
exp(-r^2) の値を計算します。
詳細説明:
3次元空間における点 (X, Y, Z) から原点までの距離rの二乗を用いて、
ガウス型の減衰関数 exp(-(X^2 + Y^2 + Z^2)) を計算します。
引数:
:param X: X座標のNumPy配列。
:type X: numpy.ndarray
:param Y: Y座標のNumPy配列。
:type Y: numpy.ndarray
:param Z: Z座標のNumPy配列。
:type Z: numpy.ndarray
戻り値:
:returns: 各点における exp(-(X^2 + Y^2 + Z^2)) の値。NumPy配列。
:rtype: numpy.ndarray
"""
r2 = X**2 + Y**2 + Z**2
return np.exp(-r2)
[ドキュメント]
def plot_torus() -> matplotlib.pyplot:
"""
概要:
トーラスの3D表面プロットを作成します。
詳細説明:
torus 関数を用いて座標を生成し、カスタムカラーマップで色付けされた
トーラスの表面をMatplotlibの3Dプロットで描画します。
戻り値:
:returns: プロットが描画された matplotlib.pyplot モジュール。
:rtype: matplotlib.pyplot
"""
# パラメトリック変数
theta = np.linspace(0, 2 * np.pi, 100)
phi = np.linspace(0, 2 * np.pi, 100)
theta, phi = np.meshgrid(theta, phi)
X, Y, Z = torus(theta, phi)
# カスタムカラーマップの定義
colors = [(1, 0.5, 0.5), (0, 1, 0), (1, 0.5, 0.5)] # 薄い赤から青、符水赤へのグラデーション
color_map = plt3d.make_color_map(phi, colors, cmap_name = 'custom_cmap', nbins = nbins)
# プロット
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
plt3d.plot_surface3d(ax, X, Y, Z, facecolors = color_map, edgecolor = 'none', alpha = 0.7, shade = True)
return plt
[ドキュメント]
def plot_isosurface3d() -> matplotlib.pyplot:
"""
概要:
定義された関数に基づき、複数の等値面を3Dでプロットします。
詳細説明:
3Dグリッド上でスカラー場 f (グローバル変数) を計算し、
指定されたレベルで等値面を抽出し、Matplotlibの3Dプロットで描画します。
戻り値:
:returns: プロットが描画された matplotlib.pyplot モジュール。
:rtype: matplotlib.pyplot
"""
# 3次元グリッドを作成
x = np.linspace(-1.2, 1.2, 40)
y = np.linspace(-1.2, 1.2, 40)
z = np.linspace(-1.2, 1.2, 40)
X, Y, Z = np.meshgrid(x, y, z, indexing='ij')
# 3Dスカラー場: 球の等値面 (中心を原点にする)
F = f(X, Y, Z)
# 等値面のレベル
levels = [0.1, 0.2, 0.3]
colors = ['cyan', 'magenta', 'yellow'] # 等値面の色
# 3D プロット
fig = plt.figure(figsize=(8, 8))
ax = fig.add_subplot(111, projection='3d')
dx, dy, dz = x[1] - x[0], y[1] - y[0], z[1] - z[0]
plt3d.plot_isosurface3d(ax, X, Y, Z, F, levels = levels, origin = [x.min(), y.min(), z.min()], spacing = [dx, dy, dz],
colors = colors, edgecolor = 'k', alpha = 0.3, linewidth = 0.1)
return plt
[ドキュメント]
def plot_contour3d() -> matplotlib.pyplot:
"""
概要:
exp(-r^2) 関数のスライスされた等高線を3Dでプロットします。
詳細説明:
3Dグリッド上で exp_r2 関数を計算し、特定のZ座標でのスライスにおける
等高線 (F[:, :, 50]) をMatplotlibの3Dプロットに描画します。
戻り値:
:returns: プロットが描画された matplotlib.pyplot モジュール。
:rtype: matplotlib.pyplot
"""
x = np.linspace(-5, 5, 100)
y = np.linspace(-5, 5, 100)
z = np.linspace(-5, 5, 100)
X, Y, Z = np.meshgrid(x, y, z)
F = exp_r2(X, Y, Z)
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
plt3d.contour3d(ax, X[:, :, 50], Y[:, :, 50], F[:, :, 50], nbins = nbins, cmap = cm.viridis)
return plt
[ドキュメント]
def plot_3dMC_animation() -> None:
"""
概要:
モンテカルロサンプリングを用いた3Dアニメーションをプロットします。
詳細説明:
確率分布 abs(x) * 4π r^3 exp(-r**2) に基づいてランダムな点を3D空間に順次追加し、
その分布の変化をアニメーションとして表示します。
点のX座標の符号に基づいて赤または青に色分けされます。
アニメーションはユーザーがEnterキーを押すまで継続します。
戻り値:
:returns: None
"""
# サンプル点の総数と一度に追加する点の数
num_points = 100000
nadd = 100
# アニメーションの設定
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
sc = ax.scatter([], [], [], marker='o', s=1, alpha=0.5)
# 初期化関数
def init():
sc._offsets3d = ([], [], [])
return sc,
# 更新関数
def update(frame):
new_points = nadd # 追加する新しい点の数
u = np.random.rand(new_points)
r = np.cbrt(-np.log(1 - u)) # 立方根で変換
phi = np.random.rand(new_points) * 2.0 * np.pi
cos_theta = np.random.rand(new_points) * 2 - 1 # cos(θ) を一様分布させる
theta = np.arccos(cos_theta)
x = r * np.sin(theta) * np.cos(phi)
y = r * np.sin(theta) * np.sin(phi)
z = r * np.cos(theta)
# x, y, z の範囲を -10 から 10 にスケーリング
x = (x / np.max(np.abs(x))) * 10
y = (y / np.max(np.abs(y))) * 10
z = (z / np.max(np.abs(z))) * 10
# 分布関数 abs(x) * 4π r^3 exp(-r**2)
weights = np.abs(x) * 4 * np.pi * r**3 * np.exp(-r**2)
weights /= np.sum(weights) # 正規化
indices = np.random.choice(np.arange(new_points), size=new_points, p=weights)
x_sampled = x[indices]
y_sampled = y[indices]
z_sampled = z[indices]
offsets = np.array(sc._offsets3d)
offsets = np.vstack([offsets.T, np.column_stack((x_sampled, y_sampled, z_sampled))]).T
sc._offsets3d = (offsets[0], offsets[1], offsets[2])
# 色の設定
color_map = np.where(offsets[0] > 0, 'red', 'blue') # x > 0 を赤、x < 0 を青
sc.set_color(color_map)
ax.set_xlabel('X')
ax.set_ylabel('Y')
ax.set_zlabel('Z')
# 軸の範囲を設定
ax.set_xlim([-10, 10])
ax.set_ylim([-10, 10])
ax.set_zlim([-10, 10])
# タイトルに点の総数を表示
ax.set_title(f'Total points: {len(offsets[0])}')
return sc,
# アニメーションの作成
ani = FuncAnimation(fig, update, frames=num_points // nadd, init_func=init, blit=False, repeat=False)
if fplot: plt.pause(0.1)
#ani.save('3dMC_animation.gif', writer='imagemagick')
input("Press ENTER to terminate>>")
return None
[ドキュメント]
def plot_Phi() -> matplotlib.pyplot:
"""
概要:
球面調和関数の角度依存部分 exp(j*m*phi) の実部と虚部をプロットします。
詳細説明:
さまざまな磁気量子数 m (0から4) に対して、角度 phi に応じた
exp(j*m*phi) の実部と虚部の変化を2つのサブプロットに分けてグラフ化します。
戻り値:
:returns: プロットが描画された matplotlib.pyplot モジュール。
:rtype: matplotlib.pyplot
"""
N = 101
phi0, phi1 = 0.0, 2.0 * pi
phi = np.linspace(phi0, phi1, N)
flist = []
for m in range(0, 5):
f = wf.Ylm_phi(m, phi)
flist.append([m, f.real, f.imag])
fig, axes = plt.subplots(2, 1, figsize = (10, 10))
phi_deg = 180.0 / pi * phi
for inf in flist:
axes[0].plot(phi_deg, inf[1], label = f"m={inf[0]},real")
axes[1].plot(phi_deg, inf[2], label = f"m={inf[0]},imag")
axes[0].set_xlabel('phi')
axes[0].set_ylabel('exp(j*m*phi),real')
axes[0].legend()
axes[1].set_xlabel('phi')
axes[1].set_ylabel('exp(j*m*phi),imag')
axes[1].legend()
return plt
[ドキュメント]
def plot_Ylm() -> matplotlib.pyplot:
"""
概要:
球面調和関数の角度依存部分 Ylm_theta(l, m, theta) をプロットします。
詳細説明:
さまざまな方位量子数 l (0から2) と磁気量子数 m に対して、
角度 theta に応じた Ylm_theta の変化をグラフに描画します。
戻り値:
:returns: プロットが描画された matplotlib.pyplot モジュール。
:rtype: matplotlib.pyplot
"""
N = 101
theta0, theta1 = 0.0, pi
theta = np.linspace(theta0, theta1, N)
flist = []
for l in range(0, 3):
for m in range(-l, l + 1):
f = wf.Ylm_theta(l, m, theta)
flist.append([l, m, f])
fig, ax = plt.subplots(1, 1, figsize = (10, 10))
theta_deg = 180.0 / pi * theta
for inf in flist:
ax.plot(theta_deg, inf[2], label = f"l,m={inf[0]},{inf[1]}")
ax.set_title('Spherical harmonic (theta dependence)')
ax.set_xlabel('theta')
ax.set_ylabel('Ylm')
ax.legend()
return plt
[ドキュメント]
def plot_Rnl() -> matplotlib.pyplot:
"""
概要:
水素原子の動径波動関数 Rnl(r) と動径確率密度分布 Pnl(r) をプロットします。
詳細説明:
主量子数 n と方位量子数 l の異なる組み合わせ (1s, 2s, 2p, ..., 4f) について、
動径波動関数 Rnl(r) と r^2 * abs(Rnl(r))^2 (動径確率密度分布) を計算し、
それぞれを9つのサブプロットを持つ2つの図に描画します。
fplot および fsave 引数に応じて表示または保存されます。
戻り値:
:returns: プロットが描画された matplotlib.pyplot モジュール。
:rtype: matplotlib.pyplot
"""
Nr = 101
rmin = 0.0
rmax = 40.0
a00 = 1.0
r = np.linspace(rmin, rmax, Nr)
f1list = []
f2list = []
P1list = []
P2list = []
for inf in [['1s', 1, 0], ['2s', 2, 0], ['2p', 2, 1], ['3s', 3, 0], ['3p', 3, 1], ['3d', 3, 2],
['4s', 4, 0], ['4p', 4, 1], ['4d', 4, 2], ['4f', 4, 3]]:
f1 = wf.Rnl(r, inf[1], inf[2])
f2 = wf.Rnl_if(r, inf[1], inf[2])
f1list.append([inf[0], inf[1], inf[2], f1])
f2list.append([inf[0], inf[1], inf[2], f2])
if f1 is not None:
P1list.append([inf[0], inf[1], inf[2], r**2 * f1**2])
else:
P1list.append([inf[0], inf[1], inf[2], None])
if f2 is not None:
P2list.append([inf[0], inf[1], inf[2], r**2 * f2**2])
else:
P2list.append([inf[0], inf[1], inf[2], None])
#plot
# plt.rcParams["font.size"] = 18
fig1, ax1 = plt.subplots(3, 3, figsize = (10, 10))
ax1 = ax1.flatten()
fig1.subplots_adjust(hspace=0.4, wspace=0.3)
fig1.suptitle('$R_{nl}(r)$: Radial wave function')
for ax, inf1, inf2 in zip(ax1, f1list, f2list):
if inf1[3] is not None:
ax.plot(r, inf1[3], label='Module')
if inf2[3] is not None:
ax.plot(r, inf2[3], label='explicit', linestyle = 'None', marker='o', markersize = 2.0)
ax.set_title(inf1[0])
ax.grid()
ax.legend()
ax.set_xlabel('r/$a_0$')
ax.set_ylabel('$R_{nl}(r)$')
if fsave: plt.savefig('plot_wf_Rnl.png')
#動径確率密度分布の描画
fig2, ax2 = plt.subplots(3, 3, figsize = (10, 10))
ax2 = ax2.flatten()
fig2.subplots_adjust(hspace=0.4, wspace=0.3)
fig2.suptitle('$P_{nl}(r)$: Radial probability density distribution')
for ax, inf1, inf2 in zip(ax2, P1list, P2list):
if inf1[3] is not None:
ax.plot(r, inf1[3], label='Module')
if inf2[3] is not None:
ax.plot(r, inf2[3], label='explicit', linestyle = 'None', marker='o', markersize = 2.0)
ax.set_title(inf1[0])
ax.grid()
ax.legend()
ax.set_xlabel('r/$a_0$')
ax.set_ylabel('$P_{nl}(r)$')
if fsave: plt.savefig('plot_wf_Pnl.png')
return plt
[ドキュメント]
def plot_rplot() -> matplotlib.pyplot:
"""
概要:
指定された水素原子の波動関数 (球面調和関数) の3D表面プロットを作成します。
詳細説明:
グローバル変数 nlm で指定された量子数 n, l, m と orb_type に基づいて
波動関数を計算し、その実部、虚部、絶対値、または絶対値の2乗を3D空間にプロットします。
プロットの色は波動関数の値または位相に基づいてカスタムカラーマップで色分けされます。
XY, YZ, ZX平面に波動関数の2D等高線も重ねて描画されます。
波動関数のエネルギー En も表示されます。
戻り値:
:returns: プロットが描画された matplotlib.pyplot モジュール。
:rtype: matplotlib.pyplot
"""
n, l, m = nlm
orb_name = wf.get_orb_name(n, l, m)
n, l, m = wf.get_qnumbers(n, l, m)
print()
print(f"Plot R(theta, phi) for {orb_name} n, l, m=({n}, {l}, {m})")
print(f" orbital type: {orb_type}")
En = wf.En(n)
print(f" En={En:.6g} eV")
theta = np.linspace(0, pi, 101)
phi = np.linspace(0, 2.0 * pi, 101)
Theta2d, Phi2d = np.meshgrid(theta, phi)
# r = wf.Ylm_real(Theta2d, Phi2d, l, m)
f = wf.Ylm(Theta2d, Phi2d, l, m)
r, phase = wf.get_by_type(f, orb_type)
nfigures = 2
if orb_type == 'a':
title = r"|$\psi$|"
elif orb_type == 'a2':
title = r"|$\psi$|$^2$"
elif orb_type == 'r':
title = r"$\psi_r$"
elif orb_type == 'i':
title = r"$\psi_i$"
ra = abs(r)
z = ra * cos(Theta2d)
rxy = ra * sin(Theta2d)
x = rxy * cos(Phi2d)
y = rxy * sin(Phi2d)
alpha = 0.7
# カスタムカラーマップの作成
color_list = colors_wf
cmap_name = cmap_name_wf
custom_cmap = LinearSegmentedColormap.from_list(cmap_name, color_list, N = nbins)
# rの値にcmapを対応させる
norm = Normalize(vmin = min(r), vmax = max(r))
face_colors = custom_cmap(norm(r))
# ax = fig.add_subplot(111, projection='3d')
fig, axes = plt.subplots(1, nfigures, figsize = (8 * nfigures, 10), subplot_kw = {'projection': '3d'})
if nfigures == 1:
ax = axes
else:
ax = axes[0]
ax2 = axes[1]
if phase is None:
options = {"facecolors": face_colors, "edgecolor": None}
else:
phase_min = np.min(phase)
phase_max = np.max(phase)
if phase_min == phase_max:
phase_min -= 0.1
phase_max += 0.1
normalized_phase = (phase - phase_min) / (phase_max - phase_min)
# カスタムカラーマップを定義
colors = colors_phase
cmap_name = cmap_name_phase
custom_cmap = LinearSegmentedColormap.from_list(cmap_name, colors, N = nbins)
facecolors = custom_cmap(normalized_phase)
options = {"facecolors": facecolors, "edgecolor": None}
plt3d.plot_surface3d(ax, x, y, z, **options, alpha = alpha, shade = False,
xlabel = 'x', ylabel = 'y', zlabel = 'z')
def contour_func(x, y, z):
f = wf.Ylm_xyz(x, y, z, l, m)
r, phase = wf.get_by_type(f, orb_type)
return r
maxx = plt3d.get_max_xyz(x, y, z)
plt3d.plot_contours_xyz_by_func(ax, contour_func, minx = -maxx, maxx = maxx, nmesh = 100,
posx = 0.1, posy = 0.1, posz = 0.1,
offsetx = -maxx, offsety = maxx, offsetz = -maxx,
cmap = custom_cmap, levels = 20, alpha = 0.3)
plt3d.set_cubic_scale(ax, maxx)
ax.set_title(title)
if phase is None:
cbar = plt3d.show_color_bar(fig, ax, scale = z, cmap = custom_cmap, shrink = 0.5, aspect = 10.0)
if orb_type == 'i':
cbar.set_label(r'$\psi_i$')
else:
cbar.set_label(r'$\psi_r$')
else:
vmin = phase_min * 180.0 / pi
vmax = phase_max * 180.0 / pi
scmap = plt.cm.ScalarMappable(cmap = custom_cmap, norm = plt.Normalize(vmin = vmin, vmax = vmax))
scmap.set_array(normalized_phase)
cbar = plt.colorbar(scmap, ax=ax, shrink=0.5, aspect=10.0)
cbar.set_label('phase')
if phase is not None and nfigures > 1:
ra = abs(phase)
z = ra * cos(Theta2d)
rxy = ra * sin(Theta2d)
x = rxy * cos(Phi2d)
y = rxy * sin(Phi2d)
plt3d.plot_surface3d(ax2, x * 180.0/pi, y * 180.0/pi, z * 180.0/pi, **options, alpha = alpha, shade = False,
xlabel = 'x', ylabel = 'y', zlabel = 'z')
# plt3d.set_cubic_scale(ax2, maxx)
ax2.set_title("phase")
return plt
[ドキュメント]
def plot_ranim() -> None:
"""
概要:
水素原子の波動関数の3D表面プロットをアニメーションとして表示します。
詳細説明:
plot_rplot と同様に、グローバル変数 nlm と orb_type に基づいて
波動関数を計算し、時間依存性 (phase_m パラメータ) を加えてアニメーション化します。
波動関数の位相変化を表現し、GIFファイルとして保存することも可能です。
アニメーションはユーザーがEnterキーを押すまで継続します。
戻り値:
:returns: None
"""
n, l, m = nlm
orb_name = wf.get_orb_name(n, l, m)
n, l, m = wf.get_qnumbers(n, l, m)
print()
print(f"Animation of R(theta, phi)=Ylm(theta, phi) for {orb_name} n, l, m=({n}, {l}, {m})")
print(f" orbital type: {orb_type}")
En = wf.En(n)
print(f" En={En:.6g} eV")
theta = np.linspace(0, pi, 101)
phi = np.linspace(0, 2.0 * pi, 101)
Theta2d, Phi2d = np.meshgrid(theta, phi)
# r = wf.Ylm_real(Theta2d, Phi2d, l, m)
f = wf.Ylm(Theta2d, Phi2d, l, m)
r, phase = wf.get_by_type(f, orb_type)
nfigures = 1
if orb_type == 'a':
title = rf"|$\psi$| ({n}{l}{m})"
elif orb_type == 'a2':
title = rf"|$\psi$|$^2$ ({n}{l}{m})"
elif orb_type == 'r':
title = rf"$\psi_r$ ({n}{l}{m})"
elif orb_type == 'i':
title = rf"$\psi_i$ ({n}{l}{m})"
ra = abs(r)
z = ra * cos(Theta2d)
rxy = ra * sin(Theta2d)
x = rxy * cos(Phi2d)
y = rxy * sin(Phi2d)
alpha = 0.7
fig, ax = plt.subplots(1, nfigures, figsize = (8, 8), subplot_kw = {'projection': '3d'})
phase_min = -pi
phase_max = pi
normalized_phase = (phase - phase_min) / (phase_max - phase_min)
# カスタムカラーマップを定義
colors = colors_phase
cmap_name = cmap_name_phase
custom_cmap = LinearSegmentedColormap.from_list(cmap_name, colors, N = nbins)
facecolors = custom_cmap(normalized_phase)
options = {"facecolors": facecolors, "edgecolor": None}
plt3d.plot_surface3d(ax, x, y, z, **options, alpha = alpha, shade = False,
xlabel = 'x', ylabel = 'y', zlabel = 'z')
maxx = plt3d.get_max_xyz(x, y, z)
plt3d.set_cubic_scale(ax, maxx)
ax.set_title(title)
vmin = -180.0
vmax = 180.0
scmap = plt.cm.ScalarMappable(cmap = custom_cmap, norm = plt.Normalize(vmin = vmin, vmax = vmax))
scmap.set_array(normalized_phase)
cbar = plt.colorbar(scmap, ax = ax, shrink = 0.5, aspect = 10.0)
cbar.set_label('phase')
nmaxiter = 100
wait = 5 # ms
ani = None
window_closed = False
def update_animation(frame):
nonlocal window_closed, ani
if window_closed :
if ani.event_source: ani.event_source.stop()
return ax
if frame % 10 == 0: print(f"frame: {frame}/{nmaxiter}")
ax.cla() # 軸の設定を保持するために全体をクリア
f = wf.Ylm(Theta2d, Phi2d, l, m, phase_m = frame * 0.1)
r, phase = wf.get_by_type(f, orb_type)
ra = abs(r)
z = ra * cos(Theta2d)
rxy = ra * sin(Theta2d)
x = rxy * cos(Phi2d)
y = rxy * sin(Phi2d)
phase_min = -pi
phase_max = pi
normalized_phase = (phase - phase_min) / (phase_max - phase_min)
custom_cmap = LinearSegmentedColormap.from_list(cmap_name, colors, N = nbins)
facecolors = custom_cmap(normalized_phase)
options = {"facecolors": facecolors, "edgecolor": None}
plt3d.plot_surface3d(ax, x, y, z, **options, alpha = alpha, shade = False,
xlabel = 'x', ylabel = 'y', zlabel = 'z')
plt3d.set_cubic_scale(ax, maxx)
ax.set_title(f"{title} frame: {frame}/{nmaxiter}")
if fplot: plt.pause(0.1)
# if frame >= nmaxiter - 1:
# ani.event_source.stop()
# return ax
return ax
def on_close(event):
nonlocal window_closed
window_closed = True
fig.canvas.mpl_connect('close_event', on_close)
ani = FuncAnimation(fig, update_animation, frames = nmaxiter, interval = wait)
#ani = FuncAnimation(fig, update, frames=nmaxiter, init_func=init, blit=False, repeat=False)
outfile = f'plot_wf_{mode}_{functype}.gif'
if fsave:
print(f"saving to {outfile}")
ani.save(outfile, writer='imagemagick')
if fplot:
input("Press Enter to close the window...")
# plt.show()
return None
[ドキュメント]
def plot_rsurface() -> matplotlib.pyplot:
"""
概要:
水素原子の動径波動関数の実部を3D表面としてプロットします。
詳細説明:
球面座標系における動径波動関数の実部を計算し、それを3D空間の表面として描画します。
X軸、Y軸、Z軸はそれぞれ直交座標に対応し、色は波動関数の値に基づいて色分けされます。
XY, YZ, ZX平面に波動関数の2D等高線も重ねて描画されます。
戻り値:
:returns: プロットが描画された matplotlib.pyplot モジュール。
:rtype: matplotlib.pyplot
"""
n, l, m = nlm
orb_name = wf.get_orb_name(n, l, m)
n, l, m = wf.get_qnumbers(n, l, m)
print()
print(f"Plot R(theta, phi) for {orb_name} n, l, m=({n}, {l}, {m})")
theta = np.linspace(0, pi, 101)
phi = np.linspace(0, 2.0 * pi, 101)
Theta2d, Phi2d = np.meshgrid(theta, phi)
r = wf.Ylm_real(Theta2d, Phi2d, l, m)
ra = abs(r)
z = ra * cos(Theta2d)
rxy = ra * sin(Theta2d)
x = rxy * cos(Phi2d)
y = rxy * sin(Phi2d)
alpha = 0.7
# カスタムカラーマップの作成
color_list = colors_wf
cmap_name = cmap_name_wf
custom_cmap = LinearSegmentedColormap.from_list(cmap_name, color_list, N = nbins)
# rの値にcmapを対応させる
norm = Normalize(vmin = min(r), vmax = max(r))
face_colors = custom_cmap(norm(r))
fig = plt.figure()
ax = fig.add_subplot(111, projection = '3d')
# options = {"colors": colors}
# options = {"color": color}
# options = {"cmap": cmap, "edgecolor": None}
options = {"facecolors": face_colors, "edgecolor": None}
# options = {"facecolors": colors, "edgecolor": None}
plt3d.plot_surface3d(ax, x, y, r, **options, alpha = alpha, shade = False,
xlabel = 'x', ylabel = 'y', zlabel = 'z')
maxx = plt3d.get_max_xyz(x, y, z)
plt3d.plot_contours_xyz_by_func(ax, lambda x, y, z: wf.Ylm_xyz_real(x, y, z, l, m), minx = -maxx, maxx = maxx, nmesh = 100,
posx = 0.1, posy = 0.1, posz = 0.1,
offsetx = -maxx, offsety = maxx, offsetz = -maxx,
cmap = custom_cmap, levels = 20, alpha = 0.3)
plt3d.set_cubic_scale(ax, maxx)
plt3d.show_color_bar(fig, ax, scale = z, cmap = custom_cmap, shrink = 0.5, aspect = 10.0)
return plt
[ドキュメント]
def estimate_max_f(func, minx: float = -3, maxx: float = -3, nsamples_estimate: int = 1000) -> float:
"""
概要:
3次元空間で定義された関数の絶対値の最大値を推定します。
詳細説明:
指定された範囲内でランダムにサンプリングを行い、
関数の最大値を概算します。これはリジェクションサンプリングの受け入れ率Mを
決定する際に使用されます。
引数:
:param func: 最大値を推定する対象の関数 (3つの引数x, y, zを取る)。
:type func: callable
:param minx: X, Y, Z座標の最小値。
:type minx: float
:param maxx: X, Y, Z座標の最大値。
:type maxx: float
:param nsamples_estimate: 最大値推定のためのサンプル数。
:type nsamples_estimate: int
戻り値:
:returns: 推定された関数の最大値。
:rtype: float
"""
x_samples, y_samples, z_samples = np.random.uniform(minx, maxx, (3, nsamples_estimate))
f_vals = func(x_samples, y_samples, z_samples)
maxf = np.max(f_vals)
return maxf
[ドキュメント]
def rejection_sampling(func, func2, M: float | None = None, minx: float = -3, maxx: float = 3, nsamples: int = 10000, nsamples_estimate: int = 1000) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, float]:
"""
概要:
リジェクションサンプリングを用いて、指定された関数に従う3D点をサンプリングします。
詳細説明:
3次元の関数 func (絶対値) と func2 (値と位相) に基づいて、
ランダムな点を生成し、そのうち条件を満たす点のみを採択します。
これにより、波動関数の確率密度に比例する点の分布を得ます。
M が指定されない場合、estimate_max_f を用いて自動で推定されます。
引数:
:param func: サンプリングの確率密度として使用する関数（絶対値）。
:type func: callable
:param func2: 波動関数の値と位相を返す関数。
:type func2: callable
:param M: サンプリングの受け入れ率の上限値。Noneの場合、estimate_max_f で推定される。
:type M: float or None
:param minx: サンプリング範囲の最小X, Y, Z座標。
:type minx: float
:param maxx: サンプリング範囲の最大X, Y, Z座標。
:type maxx: float
:param nsamples: 試行するサンプル点の総数。
:type nsamples: int
:param nsamples_estimate: M がNoneの場合に estimate_max_f で使われるサンプル数。
:type nsamples_estimate: int
戻り値:
:returns: サンプリングされたX, Y, Z座標、関数値、位相のNumPy配列、
および最終的に使用されたM値のタプル。
:rtype: tuple of (numpy.ndarray, numpy.ndarray, numpy.ndarray, numpy.ndarray, numpy.ndarray, float)
"""
if M is None: M = estimate_max_f(lambda x, y, z: abs(func(x, y, z)), minx, maxx, nsamples_estimate)
x_list = []
y_list = []
z_list = []
f_list = []
phase_list = []
for _ in range(nsamples):
x, y, z = np.random.uniform(minx, maxx, 3)
f_val, p_val = func2(x, y, z)
if np.random.uniform(0, M) < abs(f_val):
x_list.append(x)
y_list.append(y)
z_list.append(z)
f_list.append(f_val)
phase_list.append(p_val)
return np.array(x_list), np.array(y_list), np.array(z_list), np.array(f_list), np.array(phase_list), M
[ドキュメント]
def plot_dot(plot_contours: bool = True) -> matplotlib.pyplot:
"""
概要:
リジェクションサンプリングで生成された点を用いて、水素原子の波動関数を3D散布図でプロットします。
詳細説明:
グローバル変数 nlm と orb_type で指定された波動関数の確率密度に比例して
分布する点を3D空間に描画し、点の色は波動関数の位相によって色分けされます。
plot_contours がTrueの場合、波動関数の値に基づく2D等高線もプロットされます。
各軸のスケールは立方体になるように調整されます。
引数:
:param plot_contours: 波動関数の等高線プロットを表示するかどうかのブール値。
:type plot_contours: bool
戻り値:
:returns: プロットが描画された matplotlib.pyplot モジュール。
:rtype: matplotlib.pyplot
"""
nsamples_estimate = 1000
n, l, m = nlm
maxx_list = [None, 3.0, 5.0, 10.0, 15.0, 25.0]
kmaxx_list = [0.5, 0.7, 1.0, 1.0, 1.0, 1.0]
maxx = maxx_list[n] * kmaxx_list[l]
minx = -maxx
offsetx = -maxx
offsety = maxx
offsetz = -maxx
posx = 0.1
posy = 0.1
posz = 0.1
if orb_type == 'r':
title = r'$\psi_r$'
elif orb_type == 'i':
title = r'$\psi_i$'
elif orb_type == 'a':
title = r'|$\psi$|'
elif orb_type == 'a2':
title = r'|$\psi_i$|$^2$'
print()
print(f"Plot wave function with dots: n,l,m=({n},{l},{m})")
print(f"nsamples={nsamples}")
# func = lambda x, y, z: wf.psi_xyz_real(x, y, z, n, l, m)
def func(x, y, z):
f = wf.psi_xyz(x, y, z, n, l, m)
f, phase = wf.get_by_type(f, orb_type)
return f
def func2(x, y, z):
f = wf.psi_xyz(x, y, z, n, l, m)
f, phase = wf.get_by_type(f, orb_type)
return f, phase
x_sampled, y_sampled, z_sampled, f_sampled, phase_sampled, M = rejection_sampling(func, func2, None, minx, maxx, nsamples, nsamples_estimate)
nsampled = f_sampled.size
# 正規化の重み
weights = np.abs(f_sampled)
weights /= np.sum(weights)
print(f"nsampled={nsampled}")
print("Estimated max func:", M)
alpha_s = 0.2
alpha_c = 0.1
# 3Dプロットで色を変える
fig = plt.figure(figsize = (10, 10))
ax = fig.add_subplot(111, projection = '3d')
# phaseに基づいて色を設定
colors = colors_phase
custom_cmap = LinearSegmentedColormap.from_list(cmap_name_phase, colors, N = nbins)
norm = Normalize(vmin = -pi, vmax = pi)
scatter = plt3d.plot_scatter3d(ax, x_sampled, y_sampled, z_sampled, minx, maxx, minx, maxx, minx, maxx,
size = 1.0, c = phase_sampled, cmap = custom_cmap, norm = norm, alpha = alpha_s)
pad = 0.1
if True:
scmap = plt.cm.ScalarMappable(cmap = custom_cmap, norm = norm)
scmap.set_array(phase_sampled)
cbar = fig.colorbar(scmap, ax = ax, shrink = 0.5, aspect = 10.0, pad = pad)
cbar.set_label('phase')
if plot_contours:
# maxx = plt3d.get_max_xyz(x1, y1, z1)
# カスタムカラーマップの作成
color_list = colors_wf
custom_cmap = LinearSegmentedColormap.from_list(cmap_name_wf, color_list, N = nbins)
fr1d = f_sampled.flatten()
minF = min(fr1d)
maxF = max(fr1d)
# rの値にcmapを対応させる
norm = Normalize(vmin = minF, vmax = maxF)
plt3d.plot_contours_xyz_by_func(ax, func, minx = -maxx, maxx = maxx, nmesh = 100,
posx = posx, posy = posy, posz = posz,
offsetx = -maxx, offsety = maxx, offsetz = -maxx,
cmap = custom_cmap, levels = 20, alpha = alpha_c)
scmap = plt.cm.ScalarMappable(cmap = custom_cmap, norm = norm)
scmap.set_array(f_sampled)
cbar = fig.colorbar(scmap, ax = ax, shrink = 0.5, aspect = 10.0, pad = pad)
cbar.set_label(title)
plt3d.set_cubic_scale(ax, maxx)
plt.title(f"dot plot of {title} for nlm={n}{l}{m}")
return plt
[ドキュメント]
def plot_dotanim(plot_contours: bool = True) -> matplotlib.pyplot:
"""
概要:
リジェクションサンプリングで生成された点を用いて、水素原子の波動関数の3D散布図アニメーションをプロットします。
詳細説明:
plot_dot と同様に、グローバル変数 nlm と orb_type に基づいて
波動関数の確率密度に比例した3D散布図を、時間依存性 (phase パラメータ) を加えてアニメーション化します。
点の色は波動関数の位相によって色分けされ、GIFファイルとして保存することも可能です。
アニメーションはユーザーがEnterキーを押すまで継続します。
plot_contours 引数は現在この関数では使用されていません。
引数:
:param plot_contours: 等高線プロットを表示するかどうかのブール値 (アニメーションでは未使用)。
:type plot_contours: bool
戻り値:
:returns: プロットが描画された matplotlib.pyplot モジュール。
:rtype: matplotlib.pyplot
"""
nsamples_estimate = 1000
n, l, m = nlm
maxx_list = [None, 3.0, 5.0, 10.0, 15.0, 25.0]
kmaxx_list = [0.5, 0.7, 1.0, 1.0, 1.0, 1.0]
maxx = maxx_list[n] * kmaxx_list[l]
minx = -maxx
if orb_type == 'r':
title = r'$\psi_r$'
elif orb_type == 'i':
title = r'$\psi_i$'
elif orb_type == 'a':
title = r'|$\psi$|'
elif orb_type == 'a2':
title = r'|$\psi_i$|$^2$'
print()
print(f"Animation of wave function with dots: n,l,m=({n},{l},{m})")
print(f"nsamples={nsamples}")
def func(x, y, z):
f = wf.psi_xyz(x, y, z, n, l, m)
f, phase = wf.get_by_type(f, orb_type)
return f
def func2(x, y, z):
f = wf.psi_xyz(x, y, z, n, l, m)
f, phase = wf.get_by_type(f, orb_type)
return f, phase
x_sampled, y_sampled, z_sampled, f_sampled, phase_sampled, M = rejection_sampling(func, func2, None, minx, maxx, nsamples, nsamples_estimate)
nsampled = f_sampled.size
print(f"nsampled={nsampled}")
print("Estimated max func:", M)
# plot
fig = plt.figure(figsize = (10, 10))
ax = fig.add_subplot(111, projection = '3d')
# phaseに基づいて色を設定
colors = colors_phase
custom_cmap = LinearSegmentedColormap.from_list(cmap_name_phase, colors, N = nbins)
norm = Normalize(vmin = -pi, vmax = pi)
scatter = plt3d.plot_scatter3d(ax, x_sampled, y_sampled, z_sampled, minx, maxx, minx, maxx, minx, maxx,
size = 1.0, c = phase_sampled, cmap = custom_cmap, norm = norm)
scmap = plt.cm.ScalarMappable(cmap = custom_cmap, norm = norm)
scmap.set_array(phase_sampled)
cbar = fig.colorbar(scmap, ax = ax, shrink = 0.5, aspect = 10.0)
cbar.set_label('phase')
plt3d.set_cubic_scale(ax, maxx)
plt.title(f"dot plot of {title} for nlm={n}{l}{m}")
nmaxiter = 62
wait = 5 # ms
ani = None
window_closed = False
def update_animation(frame):
nonlocal window_closed, ani
if window_closed :
if ani.event_source: ani.event_source.stop()
return ax
if frame == 0 and not hasattr(update_animation, 'initialized'):
update_animation.initialized = True
return
if frame % 10 == 0: print(f"frame: {frame}/{nmaxiter}")
ax.cla() # 軸の設定を保持するために全体をクリア
def func2(x, y, z):
f = wf.psi_xyz(x, y, z, n, l, m, phase = frame * 0.1)
f, phase = wf.get_by_type(f, orb_type)
return f, phase
x_sampled, y_sampled, z_sampled, f_sampled, phase_sampled, _M \
= rejection_sampling(None, func2, M, minx, maxx, nsamples, nsamples_estimate)
plt3d.plot_scatter3d(ax, x_sampled, y_sampled, z_sampled, minx, maxx, minx, maxx, minx, maxx,
size = 1.0, c = phase_sampled, cmap = custom_cmap, norm = norm)
plt3d.set_cubic_scale(ax, maxx)
ax.set_title(f"{title} frame: {frame}/{nmaxiter}")
if fplot: plt.pause(0.1)
return ax
def on_close(event):
nonlocal window_closed
window_closed = True
fig.canvas.mpl_connect('close_event', on_close)
ani = FuncAnimation(fig, update_animation, frames = nmaxiter, interval = wait)
outfile = f'plot_wf_{mode}_{functype}.gif'
if fsave:
print(f"saving to {outfile}")
ani.save(outfile, writer='imagemagick')
if fplot:
# plt.show()
plt.pause(0.1)
input("Press Enter to close the window...")
return plt
[ドキュメント]
def plot_iso3d() -> matplotlib.pyplot:
"""
概要:
水素原子の波動関数の等値面を3Dでプロットします。
詳細説明:
グローバル変数 nlm と orb_type で指定された量子数に基づいて
3次元グリッド上で水素原子の波動関数を計算し、指定されたレベルで等値面を抽出してプロットします。
波動関数の値に基づく2D等高線も重ねて描画されます。
戻り値:
:returns: プロットが描画された matplotlib.pyplot モジュール。
:rtype: matplotlib.pyplot
"""
n, l, m = nlm
orb_name = wf.get_orb_name(n, l, m)
n, l, m = wf.get_qnumbers(n, l, m)
maxx_list = [None, 3.0, 5.0, 10.0, 15.0, 25.0]
kmaxx_list = [0.5, 0.7, 1.0, 1.0, 1.0, 1.0]
maxx = maxx_list[n] * kmaxx_list[l]
x = np.linspace(-maxx, maxx, 50)
y = np.linspace(-maxx, maxx, 50)
z = np.linspace(-maxx, maxx, 50)
X, Y, Z = np.meshgrid(x, y, z)
def func(x, y, z):
f = wf.psi_xyz(x, y, z, n, l, m)
f, phase = wf.get_by_type(f, orb_type)
return f, phase
def func1(x, y, z):
return func(x, y, z)[0]
F, phase = func(X, Y, Z)
# 等値面のレベル
fr1d = F.flatten()
minF = min(fr1d)
maxF = max(fr1d)
print(f"F range: {minF} - {maxF}")
levels = [-0.6 * maxF, -0.3 * maxF, 0.3 * maxF, 0.6 * maxF]
colors = ['blue', 'cyan', 'pink', 'red'] # 等値面の色
# 3D プロット
fig = plt.figure(figsize = (8, 8))
ax = fig.add_subplot(111, projection='3d')
dx, dy, dz = x[1] - x[0], y[1] - y[0], z[1] - z[0]
plt3d.plot_isosurface3d(ax, X, Y, Z, F, levels = levels, origin = [x.min(), y.min(), z.min()], spacing = [dx, dy, dz],
colors = colors, edgecolor = 'k', alpha = 0.3, linewidth = 0.1)
# 2D等高線
# カスタムカラーマップの作成
color_list = colors_wf
cmap_name = cmap_name_wf
custom_cmap = LinearSegmentedColormap.from_list(cmap_name, color_list, N = nbins)
# rの値にcmapを対応させる
# norm = Normalize(vmin = -maxF, vmax = maxF)
norm = Normalize(vmin = minF, vmax = maxF)
face_colors = custom_cmap(norm(F))
maxx = plt3d.get_max_xyz(X, Y, Z)
plt3d.plot_contours_xyz_by_func(ax, func1, minx = -maxx, maxx = maxx, nmesh = 100,
posx = 0.1, posy = 0.1, posz = 0.1,
offsetx = -maxx, offsety = maxx, offsetz = -maxx,
cmap = custom_cmap, levels = 20, alpha = 0.3)
plt3d.set_cubic_scale(ax, maxx)
scmap = plt.cm.ScalarMappable(cmap = custom_cmap, norm = norm)
scmap.set_array(F)
cbar = plt.colorbar(scmap, ax=ax, shrink=0.5, aspect=10.0)
plt.title(f"isosurfaces for nlm={n}{l}{m}")
return plt
[ドキュメント]
def plot_iso3danim() -> matplotlib.pyplot:
"""
概要:
水素原子の波動関数の等値面アニメーションを3Dでプロットします。
詳細説明:
plot_iso3d と同様に、グローバル変数 nlm と orb_type に基づいて
波動関数の等値面を計算し、時間依存性 (phase パラメータ) を加えてアニメーション化します。
波動関数の位相変化を表現し、GIFファイルとして保存することも可能です。
アニメーションはユーザーがEnterキーを押すまで継続します。
戻り値:
:returns: プロットが描画された matplotlib.pyplot モジュール。
:rtype: matplotlib.pyplot
"""
n, l, m = nlm
orb_name = wf.get_orb_name(n, l, m)
n, l, m = wf.get_qnumbers(n, l, m)
if orb_type == 'r':
title = r'$\psi_r$'
elif orb_type == 'i':
title = r'$\psi_i$'
elif orb_type == 'a':
title = r'|$\psi$|'
elif orb_type == 'a2':
title = r'|$\psi_i$|$^2$'
maxx_list = [None, 3.0, 5.0, 10.0, 15.0, 25.0]
kmaxx_list = [0.5, 0.7, 1.0, 1.0, 1.0, 1.0]
maxx = maxx_list[n] * kmaxx_list[l]
x = np.linspace(-maxx, maxx, 50)
y = np.linspace(-maxx, maxx, 50)
z = np.linspace(-maxx, maxx, 50)
X, Y, Z = np.meshgrid(x, y, z)
def func(x, y, z):
f = wf.psi_xyz(x, y, z, n, l, m)
f, phase = wf.get_by_type(f, orb_type)
return f, phase
def func1(x, y, z):
return func(x, y, z)[0]
F, phase = func(X, Y, Z)
# 等値面のレベル
fr1d = F.flatten()
minF = min(fr1d)
maxF = max(fr1d)
print(f"F range: {minF} - {maxF}")
levels = [-0.6 * maxF, -0.3 * maxF, 0.3 * maxF, 0.6 * maxF]
colors = ['blue', 'cyan', 'pink', 'red'] # 等値面の色
# 3D プロット
fig = plt.figure(figsize = (8, 8))
ax = fig.add_subplot(111, projection='3d')
dx, dy, dz = x[1] - x[0], y[1] - y[0], z[1] - z[0]
plt3d.plot_isosurface3d(ax, X, Y, Z, F, levels = levels, origin = [x.min(), y.min(), z.min()], spacing = [dx, dy, dz],
colors = colors, edgecolor = 'k', alpha = 0.3, linewidth = 0.1)
plt3d.set_cubic_scale(ax, maxx)
plt.title(f"isosurfaces for nlm={n}{l}{m}")
nmaxiter = 62
wait = 5 # ms
ani = None
window_closed = False
def update_animation(frame):
nonlocal window_closed, ani
if window_closed :
if ani.event_source: ani.event_source.stop()
return ax
if frame == 0 and not hasattr(update_animation, 'initialized'):
update_animation.initialized = True
return
if frame % 10 == 0: print(f"frame: {frame}/{nmaxiter}")
ax.cla() # 軸の設定を保持するために全体をクリア
def func(x, y, z):
f = wf.psi_xyz(x, y, z, n, l, m, phase = frame * 0.1)
f, phase = wf.get_by_type(f, orb_type)
return f, phase
F, phase = func(X, Y, Z)
plt3d.plot_isosurface3d(ax, X, Y, Z, F, levels = levels, origin = [x.min(), y.min(), z.min()], spacing = [dx, dy, dz],
colors = colors, edgecolor = 'k', alpha = 0.3, linewidth = 0.1)
plt3d.set_cubic_scale(ax, maxx)
ax.set_title(f"{title} frame: {frame}/{nmaxiter}")
if fplot: plt.pause(0.1)
return ax
def on_close(event):
nonlocal window_closed
window_closed = True
fig.canvas.mpl_connect('close_event', on_close)
ani = FuncAnimation(fig, update_animation, frames = nmaxiter, interval = wait)
# plt.show()
outfile = f'plot_wf_{mode}_{functype}.gif'
if fsave:
print(f"saving to {outfile}")
ani.save(outfile, writer='imagemagick')
if fplot:
# plt.show()
# plt.pause(0.1)
input("Press Enter to close the window...")
return plt
[ドキュメント]
def make_3d_density_map(output_file: str = "EDens.vasp", Z: float = 1.0, n: int = 1, l: int = 0, m: int = 0, lx: float = 10.0, ly: float = 10.0, lz: float = 10.0, step: float = 0.2) -> int:
"""
概要:
水素原子の電子密度を3DグリッドデータとしてVASP形式で出力します。
詳細説明:
指定された量子数 n, l, m と原子番号 Z の波動関数に基づいて、
X, Y, Z方向のグリッド上に電子密度 (波動関数の絶対値の2乗) を計算し、
VASP (LOCPOT) ファイル形式で出力します。
引数:
:param output_file: 出力ファイル名。Orbital名が埋め込まれる（例: "EDens_{Orb}.vasp"）。
:type output_file: str
:param Z: 原子番号 (陽子数)。
:type Z: float
:param n: 主量子数。
:type n: int
:param l: 方位量子数。
:type l: int
:param m: 磁気量子数。
:type m: int
:param lx: X方向のボックスサイズ。
:type lx: float
:param ly: Y方向のボックスサイズ。
:type ly: float
:param lz: Z方向のボックスサイズ。
:type lz: float
:param step: グリッドのステップサイズ。
:type step: float
戻り値:
:returns: 成功時は 0、ファイル書き込みエラー時は -1、
無効な量子数の場合は -2 を返します。
:rtype: int
"""
print("\nCalc 3D Electron Density map\n")
OrbName = wf.GetOrbName(n, l, m)
print(f" Z={Z}")
print(f" n,l,m={n},{l},{m}")
print(f" Orbital: {OrbName}")
output_file = output_file.replace("{Orb}", OrbName)
print(f"OutFile: [{output_file}]")
try:
out = open(output_file, "w")
except IOError:
print(f"Error: Can not write to [{output_file}].")
return -1
nx = int(lx / step + 0.1)
ny = int(ly / step + 0.1)
nz = int(lz / step + 0.1)
print(f"Range: {lx} {ly} {lz}")
print(f"nMesh: {nx} {ny} {nz}")
out.write(f"Electron density for Z={Z}, n={n}, l={l}, m={m}\n")
out.write(f" {lx:12.6f}\n")
out.write(f" {lx/lx:12.6f} {0.0:12.6f} {0.0:12.6f}\n")
out.write(f" {0.0:12.6f} {ly/lx:12.6f} {0.0:12.6f}\n")
out.write(f" {0.0:12.6f} {0.0:12.6f} {lz/lx:12.6f}\n")
out.write(" 1\n")
out.write("Direct\n")
out.write(" 0.500000 0.500000 0.500000\n")
out.write("\n")
out.write(f" {nx:3d} {ny:3d} {nz:3d}\n")
count = 0
for ix in range(nx):
for iy in range(ny):
for iz in range(nz):
x = ix * step
y = iy * step
z = iz * step
# 中心は(0.5*lx, 0.5*ly, 0.5*lz)
f = wf.cal_wave_function(Z, n, l, m, x, y, z, 0.5*lx, 0.5*ly, 0.5*lz)
if f is None:
print(f"Invalid n,l,m({n},{l},{m}).")
out.close()
return -2
out.write(f" {f:11.5g}")
count += 1
if count % 10 == 0:
out.write("\n")
out.close()
[ドキュメント]
def integ_MC3d(f, x_min: float, x_max: float, y_min: float, y_max: float, z_min: float, z_max: float, n_samples: int) -> float:
"""
概要:
モンテカルロ法（リジェクトサンプリングではない）を用いて、3次元空間の関数を積分します。
詳細説明:
指定された関数の3次元積分を、指定された範囲内でランダムにサンプリングした点での
関数値の平均に体積を掛けることで近似します。
引数:
:param f: 積分する対象の関数 (3つの引数x, y, zを取る)。
:type f: callable
:param x_min: X軸の最小値。
:type x_min: float
:param x_max: X軸の最大値。
:type x_max: float
:param y_min: Y軸の最小値。
:type y_min: float
:param y_max: Y軸の最大値。
:type y_max: float
:param z_min: Z軸の最小値。
:type z_min: float
:param z_max: Z軸の最大値。
:type z_max: float
:param n_samples: サンプリング点の数。
:type n_samples: int
戻り値:
:returns: 近似された積分の値。
:rtype: float
"""
x_samples = np.random.uniform(x_min, x_max, n_samples)
y_samples = np.random.uniform(y_min, y_max, n_samples)
z_samples = np.random.uniform(z_min, z_max, n_samples)
S = sum(f(x_samples, y_samples, z_samples))
volume = (x_max - x_min) * (y_max - y_min) * (z_max - z_min)
integral = S * volume / n_samples
return integral
[ドキュメント]
def normalize() -> None:
"""
概要:
水素原子の波動関数が正規化されているかを確認します。
詳細説明:
動径波動関数 Rnl(r) と全波動関数 psi(x,y,z) の両方について、
abs(psi)^2 (確率密度) が全空間で1に積分されることを検証します。
動径部分の積分には scipy.integrate.quad を、3次元部分の積分には
モンテカルロ積分を使用します。異なる量子数 n, l の組み合わせについて検証結果を出力します。
戻り値:
:returns: None
"""
print("\nCalculate |psi(x,y,z)|^2 to check normalization\n")
print("Rnl:")
for n in range(1, 5):
for l in range(0, n):
def func(r):
# return r**2 * abs(wf.psi_r(r, 0.0, 0.0, n, l, 0))**2
return 4 * pi * r**2 * abs(wf.psi_r(r, 0.0, 0.0, n, l, 0))**2
# return 4 * pi * r**2 * wf.Rnl(r, n, 0)**2
result, error = quad(func, 0, np.inf)
print(f" n,l={n},{l}: S={result} +- {error}")
nsampling = 1000000
print()
print("psi:")
for n in range(1, 5):
lmax = n - 1
for l in range(0, lmax + 1):
print(f"Orbital: {n}{l}0: En={wf.En(n):.6g} eV", end = '')
def psi2(x, y, z):
f = wf.psi_xyz(x, y, z, n, l, 0)
return f.real**2 + f.imag**2
if n == 1:
maxx = 5.0
else:
maxx = 20.0
integ = integ_MC3d(psi2, -maxx, maxx, -maxx, maxx, -maxx, maxx, nsampling)
print(f" |psi|^2={integ}")
[ドキュメント]
def main(mode: str) -> None:
"""
概要:
プログラムの主要な実行ロジックを制御します。
詳細説明:
コマンドライン引数から受け取った mode に基づいて、
対応するプロット関数または処理関数 (例: plot_torus, plot_Rnl, normalize) を呼び出します。
結果のプロット表示 (fplot) やファイル保存 (fsave) も管理します。
無効な mode が指定された場合はエラーメッセージを出力し終了します。
引数:
:param mode: 実行モードを指定する文字列。
:type mode: str
戻り値:
:returns: None
"""
if mode == 'torus':
plt = plot_torus()
elif mode == 'Rnl':
plt = plot_Rnl()
elif mode == 'Ylm' or mode == 'Theta':
plt = plot_Ylm()
elif mode == 'Phi':
plt = plot_Phi()
elif mode == 'normalize':
normalize()
exit()
elif mode == 'rplot':
plt = plot_rplot()
elif mode == 'ranim':
plt = plot_ranim()
elif mode == 'dot':
plt = plot_dot(plot_contours = False)
elif mode == 'dotc':
plt = plot_dot(plot_contours = True)
elif mode == 'dotanim':
plt = plot_dotanim(plot_contours = False)
elif mode == 'iso3d':
plt = plot_iso3d()
elif mode == 'iso3danim':
plt = plot_iso3danim()
elif mode == 'rsurface':
plt = plot_rsurface()
elif mode == 'isosurface3d':
plt = plot_isosurface3d()
elif mode == 'contour3d':
plt = plot_contour3d()
elif mode == '3dMC_animation':
plt = plot_3dMC_animation()
else:
print(f"\nError: Invalid mode [{mode}]")
exit()
if plt:
if fsave:
output_file = f"plot_wf_{mode}_{functype}.png"
print()
print(f"Save plot to {output_file}")
plt.savefig(output_file)
if fplot: plt.show()
if __name__ == "__main__":
# argparse のセットアップ
parser = CustomArgumentParser(description='Plot wave functions of H\n\n')
parser.add_argument('mode', nargs='?', default=mode, help='Mode of operation')
parser.add_argument('functype', nargs='?', default=functype, help='Function type')
parser.add_argument('nsamples', nargs='?', type=int, default=nsamples, help='Number of samples for Monte Carlo sampling')
parser.add_argument('--fplot', type=int, default=1, help='Flag to plot graphs (1: enable, 0: disable)')
parser.add_argument('--fsave', type=int, default=1, help='Flag to save figure file (1: enable, 0: disable)')
# 引数の解析
args = parser.parse_args()
# 引数の設定
mode = args.mode
functype = args.functype
nsamples = args.nsamples
fplot = args.fplot
fsave = args.fsave
# 使用例の出力
print()
print(f"mode: {mode}")
print(f"function type: {functype}")
ftype, nlmt = wf.analyze_function(functype)
nlm, orb_type = nlmt[:3], nlmt[3]
print(f" type: {ftype} nlm:", nlm, f" orb_type: {orb_type}")
print(f"nsamples = {nsamples}")
print(f"fplot = {fplot}")
if mode == 'torus':
f = torus
elif ftype == 'c' or ftype == 'r':
f = wf.psi_r
else:
if ftype == 'f':
if nlm == 'sphere':
f = sphere
elif nlm == 'exp_r2':
f = exp_r2
elif mode == 'contour3d' or mode == 'iso3d':
f = sphere
else:
f = sphere
main(mode)
usage()