import sys
import os
from numpy import sin, cos, tan, arcsin, arccos, arctan, arctan2, exp, log, sqrt
import numpy as np
from numpy import linalg as la
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
from matplotlib.colors import Normalize
"""
2次元波動関数描画ツール
このスクリプトは、様々な物理モデル（自由電子、無限井戸型量子ドット、水素原子）における2次元波動関数を計算し、視覚化します。
波動関数の実部、虚部、および確率密度を1Dプロットと2Dカラーマップ/3Dサーフェスプロットで表示します。
起動時引数により、モード、量子数、ブロッホ波ベクトル、メッシュ数を設定できます。
詳細説明:
- 'pw'モード: 自由電子（平面波）の波動関数を計算します。
- 'qbox'モード: 無限井戸型量子ドット内の粒子の波動関数を計算します。
- 'H'モード: 水素原子様の波動関数を計算します。
- 波動関数はブロッホ波の位相因子 'phi' と組み合わせて計算されます。
- 表示範囲は単位セル内外を指定でき、単位セル境界も明示されます。
関連リンク:
:doc:'wavefunction2D_usage'
"""
pi = 3.1415926535
pi2 = pi + pi
pi2j = 1.0j * pi2
#===================
# 初期値
#===================
"""
スクリプトの実行モードと物理パラメータの初期設定。
これらの値はコマンドライン引数によって上書きされる可能性があります。
"""
# mode: 'pw', 'qbox', 'H', 'harmonic'
#mode = 'pw'
mode = 'qbox'
# lattice parameter
ax, ay, az = 1.0, 1.0, 1.0
nlxrange, nlyrange, nlzrange = (-1.5, 1.5), (-1.5, 1.5), (-1.5, 1.5)
# number of mesh
nmeshx, nmeshy, nmeshz = 120, 120, 120
# quantum number: kx, ky, kz in pi/a for mode = 'pw
n1, n2, n3 = 0.0, 0.0, 0.0
# k vector in pi/a for mode = 'pw
kx, ky, kz = 0.0, 0.0, 0.0
# Elementary charge
e = 1.602176634e-19 # C
# Boha radius
aB = 0.529177 # A
# Nuclear charge
Z = 1.0
# Expansion ratio for psi(x) of qdt model
#kqbox = [0.3, 0.3, 0.8, 0.8, 0.8]
kqbox = [1.0, 1.0, 1.0, 1.0, 1.0]
# Expansion ratio for R(r) of H model
kR = [0.5, 6.0, 15.0]
#=============================
# Graph configuration
#=============================
"""
グラフ描画に関する設定パラメータ。
matplotlibによるプロットの外観を制御します。
"""
figsize = (12, 8)
fontsize = 10
legend_fontsize = 10
cmap_r = 'bwr' # 'plasma', 'coolwarm', 'rainbow'
#===================
# 起動時引数
#===================
if __name__ == "__main__":
argv = sys.argv
narg = len(argv)
if narg >= 2:
mode = argv[1]
if narg >= 3:
n1 = float(argv[2])
if narg >= 4:
n2 = float(argv[3])
if narg >= 5:
n3 = float(argv[4])
if narg >= 6:
kx = float(argv[5])
if narg >= 7:
ky = float(argv[6])
if narg >= 8:
kz = float(argv[7])
if narg >= 9:
nmeshx = float(argv[8])
if narg >= 10:
nmeshy = float(argv[9])
if narg >= 11:
nmeshz = float(argv[10])
[ドキュメント]
def usage():
"""
スクリプトの正しい使用法と引数の説明を表示します。
詳細説明:
スクリプトの実行時に指定可能なモード（自由電子、無限井戸型量子ドット、水素原子）と、
それぞれのモードで必要な量子数や波ベクトル、メッシュ数の引数の形式を示します。
"""
print("")
print("kx, ky, kz: Bloch's wave vector")
print("nmeshx, nmeshy, nmeshz: Number of x, y, z mesh to calculate wavefunction values")
print("usage: python {} pw kx0 ky0 kz0 kx ky kz nmeshx nmeshy nmeshz".format(argv[0]))
print(" For free electron (plain wave)")
print("usage: python {} qbox nx ny nz kx ky kz nmeshx nmeshy nmeshz".format(argv[0]))
print(" For quantum dot separated by infite potential barrier")
print("usage: python {} H n l m kx ky kz nmeshx nmeshy nmeshz".format(argv[0]))
print(" For hydrogen like model. Radius of wavefunctions are adjusted for clear visualization")
print("")
[ドキュメント]
def terminate(message, usage = usage):
"""
エラーメッセージを表示し、スクリプトを終了します。
詳細説明:
指定されたメッセージを出力した後、'usage' 関数を呼び出して使用法を表示し、
システムを終了します。
:param message: str: 表示するエラーメッセージ。
:param usage: callable, optional: 使用法を表示するための関数。デフォルトはスクリプト内の 'usage' 関数。
"""
print("")
print(message)
exit()
[ドキュメント]
def reduce2unitcell(x):
"""
座標を単位セル内に収まるように変換します。
詳細説明:
与えられた座標 'x' を、周期境界条件を考慮して [-0.5, 0.5) の範囲にマッピングします。
これは、物理的な系が周期性を持ち、単位セル内の振る舞いに関心がある場合に有用です。
:param x: float: 変換する座標値。
:returns: float: 単位セル内に収まるように変換された座標値。
"""
x += 0.5
if x < 0.0:
x = x - int(x) + 1.0
if x >= 1.0:
x = x - int(x)
return x - 0.5
[ドキュメント]
def phi(x, y, z, ax, ay, az, kx, ky, kz):
"""
ブロッホ波の位相因子を計算します。
詳細説明:
3次元空間におけるブロッホ波の複素位相因子 'exp(i * 2 * pi * (kx*x/ax + ky*y/ay + kz*z/az))' を計算します。
これは、周期的なポテンシャル下での波動関数に現れる位相項を表します。
:param x: float: x座標。
:param y: float: y座標。
:param z: float: z座標。
:param ax: float: x方向の格子定数。
:param ay: float: y方向の格子定数。
:param az: float: z方向の格子定数。
:param kx: float: x方向の波のベクトル成分 (π/a単位)。
:param ky: float: y方向の波のベクトル成分 (π/a単位)。
:param kz: float: z方向の波のベクトル成分 (π/a単位)。
:returns: complex: 計算されたブロッホ波の位相因子。
"""
return exp(pi2j * (kx * x / ax + ky * y / ay + kz * z / az))
[ドキュメント]
def psi_pw(x, y, z, ax, ay, az, n1, n2, n3, kx, ky, kz, xonly = False, yonly = False, zonly = False):
"""
自由電子（平面波）の波動関数を計算します。
詳細説明:
ブロッホの定理に基づき、平面波 'exp(i * 2 * pi * (n1*x/ax + ...))' に
ブロッホ波の位相因子 'phi' を乗じた形で波動関数を計算します。
'xonly', 'yonly', 'zonly' フラグにより、特定の方向の成分のみを返すことができます。
:param x: float: x座標。
:param y: float: y座標。
:param z: float: z座標。
:param ax: float: x方向の格子定数。
:param ay: float: y方向の格子定数。
:param az: float: z方向の格子定数。
:param n1: float: x方向の量子数 (π/a単位)。
:param n2: float: y方向の量子数 (π/a単位)。
:param n3: float: z方向の量子数 (π/a単位)。
:param kx: float: x方向のブロッホ波ベクトル成分 (π/a単位)。
:param ky: float: y方向のブロッホ波ベクトル成分 (π/a単位)。
:param kz: float: z方向のブロッホ波ベクトル成分 (π/a単位)。
:param xonly: bool, optional: x方向の成分のみを計算する場合にTrue。デフォルトはFalse。
:param yonly: bool, optional: y方向の成分のみを計算する場合にTrue。デフォルトはFalse。
:param zonly: bool, optional: z方向の成分のみを計算する場合にTrue。デフォルトはFalse。
:returns: tuple[float, float, float]: 波動関数の実部、虚部、および確率密度 (|Ψ|^2)。
"""
fx = exp(pi2j * (n1 * x / ax))
fy = exp(pi2j * (n2 * y / ay))
fz = exp(pi2j * (n3 * z / az))
# phi = exp(pi2j * (kx * x / ax + ky * y / ay + kz * z / az))
_phi = phi(x, y, z, ax, ay, az, kx, ky, kz)
if xonly:
f = fx * _phi
elif yonly:
f = fy * _phi
elif zonly:
f = fz * _phi
else:
f = fx * fy * fz * _phi
return f.real, f.imag, f.real * f.real + f.imag * f.imag
[ドキュメント]
def psi_qbox1(x, y, z, ax, ay, az, n1, n2, n3, kx, ky, kz, xonly = False, yonly = False, zonly = False):
"""
無限井戸型量子ドットモデルの波動関数の基本成分を計算します。
詳細説明:
周期的な境界条件下での無限井戸型量子ドットの波動関数を計算します。
'reduce2unitcell' を使用して座標を単位セル内にマッピングし、
奇数/偶数の量子数に応じてcos/sin関数を適用します。
ブロッホ波の位相因子 'phi' と乗算されます。
:param x: float: x座標。
:param y: float: y座標。
:param z: float: z座標。
:param ax: float: x方向の格子定数。
:param ay: float: y方向の格子定数。
:param az: float: z方向の格子定数。
:param n1: float: x方向の量子数 (整数)。
:param n2: float: y方向の量子数 (整数)。
:param n3: float: z方向の量子数 (整数)。
:param kx: float: x方向のブロッホ波ベクトル成分 (π/a単位)。
:param ky: float: y方向のブロッホ波ベクトル成分 (π/a単位)。
:param kz: float: z方向のブロッホ波ベクトル成分 (π/a単位)。
:param xonly: bool, optional: x方向の成分のみを計算する場合にTrue。デフォルトはFalse。
:param yonly: bool, optional: y方向の成分のみを計算する場合にTrue。デフォルトはFalse。
:param zonly: bool, optional: z方向の成分のみを計算する場合にTrue。デフォルトはFalse。
:returns: complex: 計算された波動関数の複素数値。
"""
n1 = int(n1 + 1.0e-3)
n2 = int(n2 + 1.0e-3)
n3 = int(n3 + 1.0e-3)
x0 = reduce2unitcell(x) * kqbox[n1-1]
y0 = reduce2unitcell(y) * kqbox[n2-1]
z0 = reduce2unitcell(z) * kqbox[n3-1]
if n1 % 2 == 1:
fx = cos(pi * x0 / ax * n1)
else:
fx = sin(pi * x0 / ax * n1)
if n2 % 2 == 1:
fy = cos(pi * y0 / ay * n2)
else:
fy = sin(pi * y0 / ay * n2)
if n3 % 2 == 1:
# fz = cos(pi * z0 / az * n3)
fz = 1.0
else:
# fz = sin(pi * z0 / az * n3)
fz = 1.0
_phi = phi(x, y, z, ax, ay, az, kx, ky, kz)
if xonly:
f = fx * _phi
elif yonly:
f = fy * _phi
elif zonly:
f = fz * _phi
else:
f = fx * fy * fz * _phi
return f
[ドキュメント]
def psi_qbox(x, y, z, ax, ay, az, n1, n2, n3, kx, ky, kz, xonly = False, yonly = False, zonly = False):
"""
無限井戸型量子ドットモデルの波動関数を周期境界条件を考慮して計算します。
詳細説明:
'psi_qbox1' を用いて、中心セルとその隣接する6つのセル
（x, y, z方向の+/-1）からの寄与を合計することで、周期性を模倣した波動関数を構築します。
これにより、全空間にわたって連続的な波動関数を表現します。
:param x: float: x座標。
:param y: float: y座標。
:param z: float: z座標。
:param ax: float: x方向の格子定数。
:param ay: float: y方向の格子定数。
:param az: float: z方向の格子定数。
:param n1: float: x方向の量子数 (整数)。
:param n2: float: y方向の量子数 (整数)。
:param n3: float: z方向の量子数 (整数)。
:param kx: float: x方向のブロッホ波ベクトル成分 (π/a単位)。
:param ky: float: y方向のブロッホ波ベクトル成分 (π/a単位)。
:param kz: float: z方向のブロッホ波ベクトル成分 (π/a単位)。
:param xonly: bool, optional: x方向の成分のみを計算する場合にTrue。デフォルトはFalse。
:param yonly: bool, optional: y方向の成分のみを計算する場合にTrue。デフォルトはFalse。
:param zonly: bool, optional: z方向の成分のみを計算する場合にTrue。デフォルトはFalse。
:returns: tuple[float, float, float]: 波動関数の実部、虚部、および確率密度 (|Ψ|^2)。
"""
f = 0.0
for v in [[0, 0, 0], [-1, 0, 0], [1, 0, 0], [0, -1, 0], [0, 1, 0], [0, 0, -1], [0, 0, 1]]:
f += psi_qbox1(x, y, z, ax, ay, az, n1, n2, n3, kx + v[0], ky + v[1], kz + v[2], xonly, yonly, zonly)
return f.real, f.imag, f.real * f.real + f.imag * f.imag
[ドキュメント]
def psi_R_H(r, n, l):
"""
水素様原子の動径波動関数 R_nl(r) を計算します。
詳細説明:
主量子数 'n' と方位量子数 'l' に基づいて、水素原子の動径波動関数を計算します。
'kR' 配列は、視覚化のために動径方向のスケールを調整するために使用されます。
サポートされている量子数は (n, l) = (1,0), (2,0), (2,1), (3,0), (3,1), (3,2) です。
:param r: float: 動径座標。
:param n: int: 主量子数。
:param l: int: 方位量子数。
:returns: float: 計算された動径波動関数 R_nl(r) の値。
"""
n = int(n + 1.0e-3)
kZaB = kR[n-1] * Z / aB
if n == 1:
return 2.0 * pow(kZaB, 1.5) * exp(-kZaB * r)
elif n == 2 and l == 0:
return 2.0 * pow(kZaB / 2.0, 1.5) * (2.0 - kZaB * r) * exp(-0.5 * kZaB * r)
elif n == 2 and l == 1:
return 1.0 / sqrt(3.0) * pow(kZaB / 2.0, 1.5) * kZaB * r * exp(-0.5 * kZaB * r)
elif n == 3 and l == 0:
return 2.0 / 3.0 * pow(kZaB / 3.0, 1.5) * (3.0 - 2.0 * kZaB * r + 2.0 / 9.0 * kZaB * kZaB * r * r) * exp(-kZaB / 3.0 * r)
elif n == 3 and l == 1:
return 1.0 / 81.0 / sqrt(3.0) * pow(2.0 * kZaB, 1.5) * (6.0 - Z / aB * r) * kZaB * r * exp(-kZaB / 3.0 * r)
elif n == 3 and l == 2:
return 1.0 / 81.0 / sqrt(15.0) * pow(2.0 * kZaB, 1.5) * kZaB * kZaB * r * r * exp(-kZaB / 3.0 * r)
else:
terminate(f"Invalid quantum numbers [n={n}, l={l}]", usage = usage)
[ドキュメント]
def psi_Y_H(Theta, Phi, l, m):
"""
水素様原子の球面調和関数 Y_lm(Theta, Phi) を計算します。
詳細説明:
方位量子数 'l' と磁気量子数 'm' に基づいて、球面調和関数を計算します。
ここでは、一般的な原子軌道（s, pz, px, py, dz2, dxz, dyz, dx2-y2, dxy）に対応する
実数形式の球面調和関数を返します。
サポートされている量子数は (l, m) = (0,0), (1,0), (1,1), (1,-1), (2,0), (2,1), (2,-1), (2,2), (2,-2) です。
:param Theta: float: 極角 (ラジアン)。
:param Phi: float: 方位角 (ラジアン)。
:param l: int: 方位量子数。
:param m: int: 磁気量子数。
:returns: float: 計算された球面調和関数 Y_lm(Theta, Phi) の値。
"""
l = int(l + 1.0e-3)
if m < 0:
m = int(m - 1.0e-3)
else:
m = int(m + 1.0e-3)
if l == 0 and m == 0:
return sqrt(1.0 / 4.0 / pi)
elif l == 1 and m == 0:
return sqrt(3.0 / 4.0 / pi) * cos(Theta)
elif l == 1 and m == 1:
return +sqrt(3.0 / 8.0 / pi) * sin(Theta) * sqrt(2.0) * cos(Phi) #exp(1.0j * Phi)
elif l == 1 and m == -1:
return +sqrt(3.0 / 8.0 / pi) * sin(Theta) * sqrt(2.0) * sin(Phi) #exp(-1.0j * Phi)
elif l == 2 and m == 0:
return +sqrt(5.0 / 16.0 / pi) * (cos(Theta)**2 - 1.0)
elif l == 2 and m == 1:
return +sqrt(15.0 / 8.0 / pi) * cos(Theta) * sin(Theta) * sqrt(2.0) * cos(Phi) #exp(1.0j * Phi)
elif l == 2 and m == -1:
return +sqrt(15.0 / 8.0 / pi) * cos(Theta) * sin(Theta) * sqrt(2.0) * sin(Phi) #exp(-1.0j * Phi)
elif l == 2 and m == 2:
return sqrt(15.0 / 32.0 / pi) * sin(Theta)**2 * sqrt(2.0) * cos(2.0 * Phi) #exp(2.0j * Phi)
elif l == 2 and m == -2:
return sqrt(15.0 / 32.0 / pi) * sin(Theta)**2 * sqrt(2.0) * sin(2.0 * Phi) #exp(-2.0j * Phi)
else:
terminate(f"psi_Y_H(): Invalid quantum number [m={m}]", usage = usage)
[ドキュメント]
def wfname_H(n, l, m):
"""
水素様原子の波動関数に対応する軌道名（例: "1s", "2pz"）を生成します。
詳細説明:
主量子数 'n'、方位量子数 'l'、磁気量子数 'm' に基づいて、
一般的な原子軌道の表記（s, pz, px, py, dz2, dxz, dyz, dx2-y2, dxy）を生成します。
:param n: int: 主量子数。
:param l: int: 方位量子数。
:param m: int: 磁気量子数。
:returns: str: 生成された軌道名。
"""
print("n=", n, l, m)
n = int(n+1.0e-3)
l = int(l+1.0e-3)
if m < 0:
m = int(m-1.0e-3)
else:
m = int(m+1.0e-3)
print("n=", n, l, m)
s = f"{n}"
if l == 0:
s += 's'
elif l == 1 and m == 0:
s += 'pz'
elif l == 1 and m == 1:
s += 'px'
elif l == 1 and m == -1:
s += 'py'
elif l == 2 and m == 0:
s += 'dz2'
elif l == 2 and m == 1:
s += 'dxz'
elif l == 2 and m == -1:
s += 'dyz'
elif l == 2 and m == 2:
s += 'dx2-y2'
elif l == 2 and m == -2:
s += 'dxy'
else:
terminate(f"wfname_H(): Invalid quantum number [m={m}]", usage = usage)
return s
[ドキュメント]
def psi_H1(x, y, z, ax, ay, az, n, l, m, kx, ky, kz, xonly = False, yonly = False, zonly = False):
"""
水素様原子モデルの波動関数の基本成分を計算します。
詳細説明:
動径波動関数 'psi_R_H' と球面調和関数 'psi_Y_H' を組み合わせて波動関数を計算し、
さらにブロッホ波の位相因子 'phi' を乗じます。
座標は 'reduce2unitcell' で単位セルにマッピングされ、極座標に変換されます。
:param x: float: x座標。
:param y: float: y座標。
:param z: float: z座標。
:param ax: float: x方向の格子定数。
:param ay: float: y方向の格子定数。
:param az: float: z方向の格子定数。
:param n: int: 主量子数。
:param l: int: 方位量子数。
:param m: int: 磁気量子数。
:param kx: float: x方向のブロッホ波ベクトル成分 (π/a単位)。
:param ky: float: y方向のブロッホ波ベクトル成分 (π/a単位)。
:param kz: float: z方向のブロッホ波ベクトル成分 (π/a単位)。
:param xonly: bool, optional: x方向の成分のみを計算する場合にTrue。デフォルトはFalse。
:param yonly: bool, optional: y方向の成分のみを計算する場合にTrue。デフォルトはFalse。
:param zonly: bool, optional: z方向の成分のみを計算する場合にTrue。デフォルトはFalse。
:returns: complex: 計算された波動関数の複素数値。
"""
x0 = reduce2unitcell(x)
y0 = reduce2unitcell(y)
z0 = reduce2unitcell(z)
rxy = sqrt(x0 * x0 + y0 * y0)
r = sqrt(x0 * x0 + y0 * y0 + z0 * z0)
Phi = arctan2(y0, x0)
if rxy == 0.0:
Theta = 0.0
else:
Theta = arcsin(rxy / r)
# print("r=", x, y, z, x0, y0, z0, r, rxy, Theta, Phi)
Rnl = psi_R_H(r, n, l)
Yml = psi_Y_H(Theta, Phi, l, m)
# print("r=", r, Theta, Phi, Rnl, Yml)
# print("phi=", Theta, Phi, Rnl, Yml)
_phi = phi(x, y, z, ax, ay, az, kx, ky, kz)
f = Rnl * Yml * _phi
return f
[ドキュメント]
def psi_H(x, y, z, ax, ay, az, n, l, m, kx, ky, kz, xonly = False, yonly = False, zonly = False):
"""
水素様原子モデルの波動関数を周期境界条件を考慮して計算します。
詳細説明:
'psi_H1' を用いて、中心セルとその隣接する6つのセル
（x, y, z方向の+/-1）からの寄与を合計することで、周期性を模倣した波動関数を構築します。
これにより、全空間にわたって連続的な波動関数を表現します。
:param x: float: x座標。
:param y: float: y座標。
:param z: float: z座標。
:param ax: float: x方向の格子定数。
:param ay: float: y方向の格子定数。
:param az: float: z方向の格子定数。
:param n: int: 主量子数。
:param l: int: 方位量子数。
:param m: int: 磁気量子数。
:param kx: float: x方向のブロッホ波ベクトル成分 (π/a単位)。
:param ky: float: y方向のブロッホ波ベクトル成分 (π/a単位)。
:param kz: float: z方向のブロッホ波ベクトル成分 (π/a単位)。
:param xonly: bool, optional: x方向の成分のみを計算する場合にTrue。デフォルトはFalse。
:param yonly: bool, optional: y方向の成分のみを計算する場合にTrue。デフォルトはFalse。
:param zonly: bool, optional: z方向の成分のみを計算する場合にTrue。デフォルトはFalse。
:returns: tuple[float, float, float]: 波動関数の実部、虚部、および確率密度 (|Ψ|^2)。
"""
f = 0.0
for v in [[0, 0, 0], [-1, 0, 0], [1, 0, 0], [0, -1, 0], [0, 1, 0], [0, 0, -1], [0, 0, 1]]:
f += psi_H1(x, y, z, ax, ay, az, n, l, m, kx + v[0], ky + v[1], kz + v[2], xonly, yonly, zonly)
return f.real, f.imag, f.real * f.real + f.imag * f.imag
[ドキュメント]
def main():
"""
スクリプトのメイン処理を実行します。
詳細説明:
1. コマンドライン引数に基づき、計算モードと物理パラメータ（格子定数、プロット範囲、メッシュ数、量子数、波ベクトル）を設定します。
2. 選択されたモード（''pw'', ''qbox'', ''H'') に応じて波動関数計算関数を決定します。
3. x方向、y方向、およびxy平面上（z=0固定）の波動関数（実部、虚部、確率密度）を計算し、データを出力します。
4. 計算結果をmatplotlibを用いて1Dプロットおよび2Dカラーマップ、3Dサーフェスプロットとして表示します。
5. ユーザーがEnterキーを押すまでグラフウィンドウを保持し、その後スクリプトを終了します。
:param: なし
:returns: なし
"""
global n1, n2, n3
print("")
print(f"mode: {mode}")
print(f"lattice parameters: {ax}, {ay}, {az} A")
print(f"plot range in unit cell: {nlxrange}, {nlyrange}, {nlzrange}")
print(f"number of mesh: {nmeshx}, {nmeshy}, {nmeshz}")
if mode == 'pw':
psi = psi_pw
elif mode == 'qbox':
psi = psi_qbox
elif mode == 'H':
psi = psi_H
else:
terminate(f"Error: Invalid mode [{mode}]", usage = usage)
xstart = nlxrange[0]
xend = nlxrange[1]
xstep = (xend - xstart) / nmeshx
print("")
print(f"plot psi(x) for x range ({xstart:6.2f} - {xend:6.2f}) at {xstep:8.4f} step")
x = []
fxr = []
fxi = []
fx2 = []
phixr = []
phixi = []
print("x\ty\tz\tPsi_r\tPsi_i\t|Psi|2\tphi_r\tphi_i")
for i in range(nmeshx + 1):
_x = xstart + i * xstep
_yr, _yi, _y2 = psi(_x, 0.0, 0.0, ax, ay, az, n1, n2, n3, kx, ky, kz, xonly = True)
_phi = phi(_x, 0.0, 0.0, ax, ay, az, kx, ky, kz)
x.append(_x)
fxr.append(_yr)
fxi.append(_yi)
fx2.append(_y2)
phixr.append(_phi.real)
phixi.append(_phi.imag)
print(f"{_x:8.3f}\t{_yr:8.3f}\t{_yi:8.3f}\t{_y2:8.3f}\t{_phi.real:8.3f}\t{_phi.imag:8.3f}")
ystart = nlyrange[0]
yend = nlyrange[1]
ystep = (yend - ystart) / nmeshy
print("")
print(f"plot psi(y) for y range ({ystart:6.2f} - {yend:6.2f}) at {ystep:8.4f} step")
y = []
fyr = []
fyi = []
fy2 = []
phiyr = []
phiyi = []
print("x\ty\tz\tPsi_r\tPsi_i\t|Psi|2\tphi_r\tphi_i")
for i in range(nmeshy + 1):
_y = ystart + i * ystep
_yr, _yi, _y2 = psi(0.0, _y, 0.0, ax, ay, az, n1, n2, n3, kx, ky, kz, yonly = True)
_phi = phi(0.0, _y, 0.0, ax, ay, az, kx, ky, kz)
y.append(_y)
fyr.append(_yr)
fyi.append(_yi)
fy2.append(_y2)
phiyr.append(_phi.real)
phiyi.append(_phi.imag)
print(f"{_y:8.3f}\t{_yr:8.3f}\t{_yi:8.3f}\t{_y2:8.3f}\t{_phi.real:8.3f}\t{_phi.imag:8.3f}")
print("")
print(f"plot psi(x, y) for x range ({xstart:6.2f} - {xend:6.2f}) at {xstep:8.4f} step and y range ({ystart:6.2f} - {yend:6.2f}) at {ystep:8.4f} step")
x3d = np.empty([nmeshx+1, nmeshy+1], dtype = float)
y3d = np.empty([nmeshx+1, nmeshy+1], dtype = float)
f3dr = np.empty([nmeshx+1, nmeshy+1], dtype = float)
f3di = np.empty([nmeshx+1, nmeshy+1], dtype = float)
f3d2 = np.empty([nmeshx+1, nmeshy+1], dtype = float)
print("x\ty\tz\tPsi_r\tPsi_i\t|Psi|2")
for iy in range(nmeshy + 1):
_y = ystart + iy * ystep
for ix in range(nmeshx + 1):
_x = xstart + ix * xstep
_yr, _yi, _y2 = psi(_x, _y, 0.0, ax, ay, az, n1, n2, n3, kx, ky, kz)
x3d[ix][iy] = _x
y3d[ix][iy] = _y
f3dr[ix][iy] = _yr
f3di[ix][iy] = _yi
f3d2[ix][iy] = _y2
if -0.5 <= _x <= 0.5 and -0.5 <= _y <= 0.5:
print(f"{_x:8.3f}\t{_y:8.3f}\t{_yr:8.3f}\t{_yi:8.3f}\t{_y2:8.3f}")
print("")
print("Quantum number:")
if mode == 'pw':
print(f"kx0, ky0, kz0 = {n1}, {n2}, {n3} * pi/a")
elif mode == 'qbox':
print(f"nx, ny, nz = {n1}, {n2}, {n3}")
elif mode == 'H':
print(f"n, l, m = {n1}, {n2}, {n3}")
else:
terminate(f"Error: Invalid mode [{mode}]", usage = usage)
print(f"kx, ky, kz = {kx}, {ky}, {kz} * pi/a")
#=============================
# グラフの表示
#=============================
maxfr = 0.0
for list in f3dr:
for v in list:
if abs(v) > maxfr:
maxfr = abs(v)
if maxfr < 1.0e-3:
maxfr = 1.0e-3
print("maxfr=", maxfr)
maxfi = 0.0
for list in f3di:
for v in list:
if abs(v) > maxfi:
maxfi = abs(v)
if maxfi < 1.0e-3:
maxfi = 1.0e-3
print("maxfi=", maxfi)
if maxfi > 1.0e-3:
maxfi = 1.0e-3
print("")
fig = plt.figure(figsize = figsize)
if mode == 'H':
wfn = wfname_H(n1, n2, n3)
else:
wfn = ''
if mode != 'pw':
n1 = int(n1 + 1.0e-3)
n2 = int(n2 + 1.0e-3)
if n3 < 0:
n3 = int(n3 - 1.0e-3)
else:
n3 = int(n3 + 1.0e-3)
title = f"mode:{mode} n=({n1},{n2},{n3}){wfn} k=({kx},{ky},{kz})"
plt.suptitle(title)
axcr = fig.add_subplot(2, 3, 1)
axci = fig.add_subplot(2, 3, 2)
ax3dr = fig.add_subplot(2, 3, 3, projection='3d')
ax3di = fig.add_subplot(2, 3, 6, projection='3d')
axfx = fig.add_subplot(4, 3, 7)
axphix = fig.add_subplot(4, 3, 10)
axfy = fig.add_subplot(4, 3, 8)
axphiy = fig.add_subplot(4, 3, 11)
axfx.plot(x, fxr, label = '$\psi$$_x$$_r$', linestyle = 'dashed', linewidth = 0.5, color = 'red')
axfx.plot(x, fxi, label = '$\psi$$_x$$_i$', linestyle = 'dashed', linewidth = 0.5, color = 'blue')
# axfx.plot(x, fx2, label = '|$\psi$$|$^2$', linestyle = '-', linewidth = 0.5, color = 'green')
xlim = nlxrange #axfx.get_xlim()
ylim = axfx.get_ylim()
axfx.plot(xlim, [0.0, 0.0], linestyle = 'dashed', linewidth = 0.3, color = 'red')
for i in range(int(nlxrange[0]), int(nlxrange[1]) + 1):
axfx.plot([i + 0.5, i + 0.5], ylim, linestyle = '-', linewidth = 0.5, color = 'black')
axfx.set_xlim(xlim)
axfx.set_ylim(ylim)
axfx.set_xlabel("$x$ (A)", fontsize = fontsize)
axfx.set_ylabel("$\psi$$_x$($x$)", fontsize = fontsize)
axfx.legend(fontsize = legend_fontsize, loc = 'best')
axfx.tick_params(labelsize = fontsize)
axphix.plot(x, phixr, label = '$\phi_r$', linestyle = 'dashed', linewidth = 0.5, color = 'red')
axphix.plot(x, phixi, label = '$\phi_i$', linestyle = 'dashed', linewidth = 0.5, color = 'blue')
xlim = nlxrange #axfx.get_xlim()
ylim = axphix.get_ylim()
axphix.plot(xlim, [0.0, 0.0], linestyle = 'dashed', linewidth = 0.3, color = 'red')
for i in range(int(nlxrange[0]), int(nlxrange[1]) + 1):
axphix.plot([i + 0.5, i + 0.5], ylim, linestyle = '-', linewidth = 0.5, color = 'black')
axphix.set_xlim(xlim)
axphix.set_ylim(ylim)
axphix.set_xlabel("$x$ (A)", fontsize = fontsize)
axphix.set_ylabel("$\phi$($x$)", fontsize = fontsize)
axphix.legend(fontsize = legend_fontsize, loc = 'best')
axphix.tick_params(labelsize = fontsize)
axfy.plot(y, fyr, label = '$\psi$$_y$$_r$', linestyle = 'dashed', linewidth = 0.5, color = 'red')
axfy.plot(y, fyi, label = '$\psi$$_y$$_i$', linestyle = 'dashed', linewidth = 0.5, color = 'blue')
# axfy.plot(y, fy2, label = '|$\psi$$_y$|$^2$', linestyle = '-', linewidth = 0.5, color = 'black')
xlim = nlyrange #axfx.get_ylim()
ylim = axfy.get_ylim()
axfy.plot(xlim, [0.0, 0.0], linestyle = 'dashed', linewidth = 0.3, color = 'red')
for i in range(int(nlxrange[0]), int(nlxrange[1]) + 1):
axfy.plot([i + 0.5, i + 0.5], ylim, linestyle = '-', linewidth = 0.5, color = 'black')
axfy.set_xlim(xlim)
axfy.set_ylim(ylim)
axfy.set_xlabel("$y$ (A)", fontsize = fontsize)
axfy.set_ylabel("$\psi$$_y$($y$)", fontsize = fontsize)
axfy.legend(fontsize = legend_fontsize, loc = 'best')
axfy.tick_params(labelsize = fontsize)
axphiy.plot(y, phiyr, label = '$\phi_r$', linestyle = 'dashed', linewidth = 0.5, color = 'red')
axphiy.plot(y, phiyi, label = '$\phi_i$', linestyle = 'dashed', linewidth = 0.5, color = 'blue')
xlim = nlyrange #axfx.get_ylim()
ylim = axphiy.get_ylim()
axphiy.plot(xlim, [0.0, 0.0], linestyle = 'dashed', linewidth = 0.3, color = 'red')
for i in range(int(nlxrange[0]), int(nlxrange[1]) + 1):
axphiy.plot([i + 0.5, i + 0.5], ylim, linestyle = '-', linewidth = 0.5, color = 'black')
axphiy.set_xlim(xlim)
axphiy.set_ylim(ylim)
axphiy.set_xlabel("$y$ (A)", fontsize = fontsize)
axphiy.set_ylabel("$\phi$($y$)", fontsize = fontsize)
axphiy.legend(fontsize = legend_fontsize, loc = 'best')
axphiy.tick_params(labelsize = fontsize)
surf = ax3dr.plot_surface(x3d, y3d, f3dr, label = '$\Psi_r$', cmap = cmap_r)
ax3dr.contour(x3d, y3d, f3dr, cmap = cmap_r, offset = -1,
vmin = -maxfr, vmax = maxfr, levels = 100)
# norm = Normalize(vmin = -maxfr, vmax = maxfr), levels = 100)
# fig.colorbar(surf, ax = ax3dr)
# ax3dr.set_aspect('equal')
ax3dr.set_xlabel("$x$ (A)", fontsize = fontsize)
ax3dr.set_ylabel("$y$ (A)", fontsize = fontsize)
ax3dr.set_zlabel("$\Psi$$_r$", fontsize = fontsize)
# ax3dr.legend(fontsize = legend_fontsize, loc = 'best')
ax3dr.tick_params(labelsize = fontsize)
surf = ax3di.plot_surface(x3d, y3d, f3di, label = '$\Psi_i$', cmap = cmap_r)
ax3di.contour(x3d, y3d, f3di, cmap = cmap_r, offset = -1,
norm = Normalize(vmin = -maxfi, vmax = maxfi), levels = 100)
# fig.colorbar(surf, ax = ax3di)
# ax3di.set_aspect('equal')
ax3di.set_xlabel("$x$ (A)", fontsize = fontsize)
ax3di.set_ylabel("$y$ (A)", fontsize = fontsize)
ax3di.set_zlabel("$\Psi$$_i$", fontsize = fontsize)
# ax3di.legend(fontsize = legend_fontsize, loc = 'best')
ax3di.tick_params(labelsize = fontsize)
axcr.set_title("$\Psi_r$")
# contour = axcr.pcolormesh(x3d, y3d, f3dr, label = '$\Psi_r$', cmap = cmap_r, shading='auto')
if maxfr > 1.0e-3:
# contour = axcr.contour(x3d, y3d, f3dr, cmap = cmap_r,
# levels = np.arange(-maxfr, maxfr, 0.01 * maxfr))
contour = axcr.contourf(x3d, y3d, f3dr,
cmap = cmap_r, norm = Normalize(vmin = -maxfr, vmax = maxfr), levels = 100)
fig.colorbar(contour, ax = axcr, shrink = 0.5)
axcr.set_aspect('equal')
xlim = nlxrange #axfy.get_xlim()
ylim = nlyrange #axfy.get_ylim()
for i in range(int(nlxrange[0]), int(nlxrange[1]) + 1):
# if i == 0:
# continue
axcr.plot([i + 0.5, i + 0.5], ylim, linestyle = '-', linewidth = 1.0, color = 'black')
for i in range(int(nlyrange[0]), int(nlyrange[1]) + 1):
# if i == 0:
# continue
axcr.plot(xlim, [i + 0.5, i + 0.5], linestyle = '-', linewidth = 1.0, color = 'black')
axcr.plot(xlim, [0.0, 0.0], linestyle = '-', linewidth = 0.5, color = 'red')
axcr.plot([0.0, 0.0], ylim, linestyle = '-', linewidth = 0.5, color = 'red')
axcr.set_xlim(xlim)
axcr.set_ylim(ylim)
axcr.set_xlabel("$x$ (A)", fontsize = fontsize)
axcr.set_ylabel("$y$ (A)", fontsize = fontsize)
# axcr.legend(fontsize = legend_fontsize, loc = 'best')
axcr.tick_params(labelsize = fontsize)
axci.set_title("$\Psi_i$")
# contour = axci.pcolormesh(x3d, y3d, f3di, label = '$\Psi_i$', cmap = cmap_r, shading='auto')
if maxfi > 1.0e-3:
# contour = axci.contour(x3d, y3d, f3di, cmap = cmap_r,
# levels = np.arange(-maxfi, maxfi, 0.01 * maxfi))
contour = axci.contourf(x3d, y3d, f3di,
cmap = cmap_r, norm = Normalize(vmin = -maxfi, vmax = maxfi), levels = 100)
fig.colorbar(contour, ax = axci, shrink = 0.5)
axci.set_aspect('equal')
xlim = nlxrange #axfy.get_xlim()
ylim = nlyrange #axfy.get_ylim()
for i in range(int(nlxrange[0]), int(nlxrange[1]) + 1):
# if i == 0:
# continue
axci.plot([i + 0.5, i + 0.5], ylim, linestyle = '-', linewidth = 1.0, color = 'black')
for i in range(int(nlyrange[0]), int(nlyrange[1]) + 1):
# if i == 0:
# continue
axci.plot(xlim, [i + 0.5, i + 0.5], linestyle = '-', linewidth = 1.0, color = 'black')
axci.plot(xlim, [0.0, 0.0], linestyle = '-', linewidth = 0.5, color = 'red')
axci.plot([0.0, 0.0], ylim, linestyle = '-', linewidth = 0.5, color = 'red')
axci.set_xlim(xlim)
axci.set_ylim(ylim)
axci.set_xlabel("$x$ (A)", fontsize = fontsize)
axci.set_ylabel("$y$ (A)", fontsize = fontsize)
# axci.legend(fontsize = legend_fontsize, loc = 'best')
axci.tick_params(labelsize = fontsize)
plt.tight_layout()
plt.pause(0.1)
print("Press ENTER to exit>>", end = '')
input()
terminate("", usage = usage)
if __name__ == '__main__':
main()