wavefunction2D.py ダウンロード/コピー

wavefunction2D.py をダウンロード

wavefunction2D.py
wavefunction2D.py
  1import sys
  2import os
  3import traceback
  4from numpy import sin, cos, tan, arcsin, arccos, arctan, arctan2, exp, log, sqrt
  5import numpy as np
  6from numpy import linalg as la
  7from mpl_toolkits.mplot3d import Axes3D
  8import matplotlib.pyplot as plt
  9from matplotlib.colors import Normalize
 10
 11"""
 122次元波動関数を描画するスクリプトです。
 13
 14詳細説明:
 15様々な物理モデル(平面波、量子箱、水素原子モデルなど)の波動関数を計算し、
 16その実部、虚部、および確率密度を1次元プロットと2次元カラーマップ/3Dサーフェスプロットで可視化します。
 17コマンドライン引数により、モードと量子数を指定できます。
 18
 19:doc:`wavefunction2D_usage`
 20"""
 21
 22
 23pi   = 3.1415926535
 24pi2  = pi + pi
 25pi2j = 1.0j * pi2
 26
 27
 28#===================
 29# 初期値
 30#===================
 31# mode: 'pw', 'qbox', 'H', 'harmonic'
 32#mode = 'pw'
 33#mode = 'qbox'
 34mode = ''
 35
 36# lattice parameter
 37ax, ay, az = 1.0, 1.0, 1.0
 38nlxrange, nlyrange, nlzrange = (-1.5, 1.5), (-1.5, 1.5), (-1.5, 1.5)
 39# number of mesh
 40nmeshx, nmeshy, nmeshz = 120, 120, 120
 41
 42# quantum number: kx, ky, kz in pi/a for mode = 'pw
 43n1, n2, n3 = 0.0, 0.0, 0.0
 44# k vector in pi/a for mode = 'pw
 45kx, ky, kz = 0.0, 0.0, 0.0
 46
 47# Elementary charge
 48e = 1.602176634e-19  # C
 49# Boha radius
 50aB = 0.529177         # A
 51# Nuclear charge
 52Z  = 1.0
 53
 54# Expansion ratio for psi(x) of qdt model
 55#kqbox = [0.3, 0.3, 0.8, 0.8, 0.8]
 56kqbox = [1.0, 1.0, 1.0, 1.0, 1.0]
 57
 58# Expansion ratio for R(r) of H model
 59kR = [0.5, 6.0, 15.0]
 60
 61#=============================
 62# Graph configuration
 63#=============================
 64figsize         = (12, 8)
 65fontsize        = 10
 66legend_fontsize = 10
 67cmap_r = 'bwr'  # 'plasma', 'coolwarm', 'rainbow'
 68
 69
 70def usage():
 71    """
 72    スクリプトの利用方法と引数を表示します。
 73
 74    詳細説明:
 75    標準出力にコマンドライン引数の指定方法や利用可能なモードについての情報を出力します。
 76
 77    :returns: None - 戻り値はありません。
 78    """
 79    print("")
 80    print("kx, ky, kz: Bloch's wave vector")
 81    print("nmeshx, nmeshy, nmeshz: Number of x, y, z mesh to calculate wavefunction values")
 82    print("usage: python {} pw kx0 ky0 kz0 kx ky kz nmeshx nmeshy nmeshz".format(sys.argv[0]))
 83    print("   For free electron (plain wave)")
 84    print("usage: python {} qbox nx ny nz kx ky kz nmeshx nmeshy nmeshz".format(sys.argv[0]))
 85    print("   For quantum dot separated by infite potential barrier")
 86    print("usage: python {} H n l m kx ky kz nmeshx nmeshy nmeshz".format(sys.argv[0]))
 87    print("   For hydrogen like model. Radius of wavefunctions are adjusted for clear visualization")
 88    print("")
 89
 90
 91def terminate(message: str, usage = usage):
 92    """
 93    エラーメッセージを表示し、必要に応じて利用方法を表示してプログラムを終了します。
 94
 95    詳細説明:
 96    指定されたエラーメッセージを表示後、プログラムを強制終了(exit)します。
 97
 98    :param message: str - 表示するエラーまたは終了メッセージ。
 99    :param usage: function - 呼び出す利用方法表示関数(デフォルトは `usage`)。
100    :returns: None - プログラムが終了するため戻り値はありません。
101    """
102    print("")
103    print(message)
104    exit()
105
106
107def reduce2unitcell(x: float) -> float:
108    """
109    与えられた座標値を単位セル内の[-0.5, 0.5)の範囲に変換します。
110
111    詳細説明:
112    座標 x を +0.5 して [0, 1) の範囲に移動させ、整数部分を引いて [0, 1) に収め、
113    最後に -0.5 して [-0.5, 0.5) の範囲に戻します。
114
115    :param x: float - 変換する座標値。
116    :returns: float - 単位セル内の変換された座標。
117    """
118    x += 0.5
119    if x < 0.0:
120        x = x - int(x) + 1.0
121    if x >= 1.0:
122        x = x - int(x)
123    return x - 0.5
124    
125
126def phi(x: float, y: float, z: float, ax: float, ay: float, az: float, kx: float, ky: float, kz: float) -> complex:
127    """
128    ブロッホ因子(平面波部分)を計算します。
129
130    詳細説明:
131    `exp(i * 2 * pi * (kx*x/ax + ky*y/ay + kz*z/az))` の形式でブロッホ因子を返します。
132
133    :param x: float - x座標。
134    :param y: float - y座標。
135    :param z: float - z座標。
136    :param ax: float - x方向の格子定数。
137    :param ay: float - y方向の格子定数。
138    :param az: float - z方向の格子定数。
139    :param kx: float - x方向の波数ベクトル成分 (単位: pi/a)。
140    :param ky: float - y方向の波数ベクトル成分 (単位: pi/a)。
141    :param kz: float - z方向の波数ベクトル成分 (単位: pi/a)。
142    :returns: complex - ブロッホ因子の複素数値。
143    """
144    return exp(pi2j * (kx * x / ax + ky * y / ay + kz * z / az))
145
146
147def psi_pw(x: float, y: float, z: float, ax: float, ay: float, az: float, n1: float, n2: float, n3: float, kx: float, ky: float, kz: float, xonly: bool = False, yonly: bool = False, zonly: bool = False) -> tuple[float, float, float]:
148    """
149    自由電子(平面波)モデルの波動関数を計算します。
150
151    詳細説明:
152    n1, n2, n3 で指定される波数ベクトルを持つ平面波と、kx, ky, kz で指定されるブロッホ因子を
153    組み合わせた波動関数を計算します。`xonly`, `yonly`, `zonly` が True の場合、
154    対応する方向の成分のみを考慮します。
155
156    :param x: float - x座標。
157    :param y: float - y座標。
158    :param z: float - z座標。
159    :param ax: float - x方向の格子定数。
160    :param ay: float - y方向の格子定数。
161    :param az: float - z方向の格子定数。
162    :param n1: float - x方向の量子数 (波数ベクトル成分)。
163    :param n2: float - y方向の量子数 (波数ベクトル成分)。
164    :param n3: float - z方向の量子数 (波数ベクトル成分)。
165    :param kx: float - x方向のブロッホ波数ベクトル成分 (単位: pi/a)。
166    :param ky: float - y方向のブロッホ波数ベクトル成分 (単位: pi/a)。
167    :param kz: float - z方向のブロッホ波数ベクトル成分 (単位: pi/a)。
168    :param xonly: bool - Trueの場合、x方向成分のみを考慮。
169    :param yonly: bool - Trueの場合、y方向成分のみを考慮。
170    :param zonly: bool - Trueの場合、z方向成分のみを考慮。
171    :returns: tuple[float, float, float] - 波動関数の実部、虚部、および確率密度(絶対値の2乗)。
172    """
173    fx = exp(pi2j * (n1 * x / ax))
174    fy = exp(pi2j * (n2 * y / ay))
175    fz = exp(pi2j * (n3 * z / az))
176
177#    phi = exp(pi2j * (kx * x / ax + ky * y / ay + kz * z / az))
178    _phi = phi(x, y, z, ax, ay, az, kx, ky, kz)
179
180    if xonly:
181        f = fx * _phi
182    elif yonly:
183        f = fy * _phi
184    elif zonly:
185        f = fz * _phi
186    else:
187        f = fx * fy * fz * _phi
188    
189    return f.real, f.imag, f.real * f.real + f.imag * f.imag
190
191
192def psi_qbox1(x: float, y: float, z: float, ax: float, ay: float, az: float, n1: float, n2: float, n3: float, kx: float, ky: float, kz: float, xonly: bool = False, yonly: bool = False, zonly: bool = False) -> complex:
193    """
194    量子箱モデルの波動関数の基本成分を計算します。
195
196    詳細説明:
197    無限ポテンシャル障壁を持つ量子箱内の粒子の波動関数を計算します。
198    `reduce2unitcell` を使用して座標を単位セル内に収め、量子数 `n1, n2, n3` に応じて
199    cosまたはsin関数を適用します。ブロッホ因子も考慮されます。
200    `kqbox` は波動関数の広がりを調整する係数です。
201
202    :param x: float - x座標。
203    :param y: float - y座標。
204    :param z: float - z座標。
205    :param ax: float - x方向の格子定数。
206    :param ay: float - y方向の格子定数。
207    :param az: float - z方向の格子定数。
208    :param n1: float - x方向の量子数。
209    :param n2: float - y方向の量子数。
210    :param n3: float - z方向の量子数。
211    :param kx: float - x方向のブロッホ波数ベクトル成分 (単位: pi/a)。
212    :param ky: float - y方向のブロッホ波数ベクトル成分 (単位: pi/a)。
213    :param kz: float - z方向のブロッホ波数ベクトル成分 (単位: pi/a)。
214    :param xonly: bool - Trueの場合、x方向成分のみを考慮。
215    :param yonly: bool - Trueの場合、y方向成分のみを考慮。
216    :param zonly: bool - Trueの場合、z方向成分のみを考慮。
217    :returns: complex - 量子箱モデル波動関数の複素数値。
218    """
219    n1 = int(n1 + 1.0e-3)
220    n2 = int(n2 + 1.0e-3)
221    n3 = int(n3 + 1.0e-3)
222    x0 = reduce2unitcell(x) * kqbox[n1-1]
223    y0 = reduce2unitcell(y) * kqbox[n2-1]
224    z0 = reduce2unitcell(z) * kqbox[n3-1]
225
226    if n1 % 2 == 1:
227        fx = cos(pi * x0 / ax * n1)
228    else:
229        fx = sin(pi * x0 / ax * n1)
230    if n2 % 2 == 1:
231        fy = cos(pi * y0 / ay * n2)
232    else:
233        fy = sin(pi * y0 / ay * n2)
234    if n3 % 2 == 1:
235#        fz = cos(pi * z0 / az * n3)
236        fz = 1.0
237    else:
238#        fz = sin(pi * z0 / az * n3)
239        fz = 1.0
240
241    _phi = phi(x, y, z, ax, ay, az, kx, ky, kz)
242
243    if xonly:
244        f = fx * _phi
245    elif yonly:
246        f = fy * _phi
247    elif zonly:
248        f = fz * _phi
249    else:
250        f = fx * fy * fz * _phi
251    
252    return f
253
254
255def psi_qbox(x: float, y: float, z: float, ax: float, ay: float, az: float, n1: float, n2: float, n3: float, kx: float, ky: float, kz: float, xonly: bool = False, yonly: bool = False, zonly: bool = False) -> tuple[float, float, float]:
256    """
257    量子箱モデルの周期的な波動関数を計算します。
258
259    詳細説明:
260    中心セルとその隣接する6つのセルからの `psi_qbox1` 成分を合計することで、
261    周期境界条件を模倣した量子箱波動関数を生成します。
262
263    :param x: float - x座標。
264    :param y: float - y座標。
265    :param z: float - z座標。
266    :param ax: float - x方向の格子定数。
267    :param ay: float - y方向の格子定数。
268    :param az: float - z方向の格子定数。
269    :param n1: float - x方向の量子数。
270    :param n2: float - y方向の量子数。
271    :param n3: float - z方向の量子数。
272    :param kx: float - x方向のブロッホ波数ベクトル成分 (単位: pi/a)。
273    :param ky: float - y方向のブロッホ波数ベクトル成分 (単位: pi/a)。
274    :param kz: float - z方向のブロッホ波数ベクトル成分 (単位: pi/a)。
275    :param xonly: bool - Trueの場合、x方向成分のみを考慮。
276    :param yonly: bool - Trueの場合、y方向成分のみを考慮。
277    :param zonly: bool - Trueの場合、z方向成分のみを考慮。
278    :returns: tuple[float, float, float] - 波動関数の実部、虚部、および確率密度(絶対値の2乗)。
279    """
280    f = 0.0
281    for v in [[0, 0, 0], [-1, 0, 0], [1, 0, 0], [0, -1, 0], [0, 1, 0], [0, 0, -1], [0, 0, 1]]:
282        f += psi_qbox1(x, y, z, ax, ay, az, n1, n2, n3, kx + v[0], ky + v[1], kz + v[2], xonly, yonly, zonly)
283    
284    return f.real, f.imag, f.real * f.real + f.imag * f.imag
285
286
287def psi_R_H(r: float, n: float, l: float) -> float:
288    """
289    水素原子モデルの動径波動関数部分を計算します。
290
291    詳細説明:
292    主量子数 `n` と方位量子数 `l` に基づいて、水素様原子の動径波動関数 `R_nl(r)` を計算します。
293    `kR` は波動関数の広がりを調整する係数です。
294
295    :param r: float - 中心からの動径距離。
296    :param n: float - 主量子数。
297    :param l: float - 方位量子数。
298    :returns: float - 動径波動関数の実数値。
299    """
300    n = int(n + 1.0e-3)
301    kZaB = kR[n-1] * Z / aB
302
303    if n == 1:
304        return 2.0 * pow(kZaB, 1.5) * exp(-kZaB * r)
305    elif n == 2 and l == 0:
306        return 2.0 * pow(kZaB / 2.0, 1.5) * (2.0 - kZaB * r) * exp(-0.5 * kZaB * r)
307    elif n == 2 and l == 1:
308        return 1.0 / sqrt(3.0) * pow(kZaB / 2.0, 1.5) * kZaB * r * exp(-0.5 * kZaB * r)
309    elif n == 3 and l == 0:
310        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)
311    elif n == 3 and l == 1:
312        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)
313    elif n == 3 and l == 2:
314        return 1.0 / 81.0 / sqrt(15.0) * pow(2.0 * kZaB, 1.5) * kZaB * kZaB * r * r * exp(-kZaB / 3.0 * r)
315    else:
316        terminate(f"Invalid quantum numbers [n={n}, l={l}]", usage = usage)
317
318
319def psi_Y_H(Theta: float, Phi: float, l: float, m: float) -> float:
320    """
321    水素原子モデルの角度波動関数(球面調和関数)部分を計算します。
322
323    詳細説明:
324    方位量子数 `l` と磁気量子数 `m` に基づいて、球面調和関数 `Y_lm(Theta, Phi)` の
325    実数部分を計算します。
326
327    :param Theta: float - 天頂角 (ラジアン)。
328    :param Phi: float - 方位角 (ラジアン)。
329    :param l: float - 方位量子数。
330    :param m: float - 磁気量子数。
331    :returns: float - 角度波動関数の実数値。
332    """
333    l = int(l + 1.0e-3)
334    if m < 0:
335        m = int(m - 1.0e-3)
336    else:
337        m = int(m + 1.0e-3)
338
339    if l == 0 and m == 0:
340        return sqrt(1.0 / 4.0 / pi)
341    elif l == 1 and m == 0:
342        return sqrt(3.0 / 4.0 / pi) * cos(Theta)
343    elif l == 1 and m == 1:
344        return +sqrt(3.0 / 8.0 / pi) * sin(Theta) * sqrt(2.0) * cos(Phi) #exp(1.0j * Phi)
345    elif l == 1 and m == -1:
346        return +sqrt(3.0 / 8.0 / pi) * sin(Theta) * sqrt(2.0) * sin(Phi) #exp(-1.0j * Phi)
347    elif l == 2 and m == 0:
348        return +sqrt(5.0 / 16.0 / pi) * (cos(Theta)**2 - 1.0)
349    elif l == 2 and m == 1:
350        return +sqrt(15.0 / 8.0 / pi) * cos(Theta) * sin(Theta) * sqrt(2.0) * cos(Phi) #exp(1.0j * Phi)
351    elif l == 2 and m == -1:
352        return +sqrt(15.0 / 8.0 / pi) * cos(Theta) * sin(Theta) * sqrt(2.0) * sin(Phi) #exp(-1.0j * Phi)
353    elif l == 2 and m == 2:
354        return sqrt(15.0 / 32.0 / pi) * sin(Theta)**2 * sqrt(2.0) * cos(2.0 * Phi) #exp(2.0j * Phi)
355    elif l == 2 and m == -2:
356        return sqrt(15.0 / 32.0 / pi) * sin(Theta)**2 * sqrt(2.0) * sin(2.0 * Phi) #exp(-2.0j * Phi)
357    else:
358        terminate(f"psi_Y_H(): Invalid quantum number [m={m}]", usage = usage)    
359
360
361def wfname_H(n: float, l: float, m: float) -> str:
362    """
363    水素原子モデルの波動関数名(軌道名)を生成します。
364
365    詳細説明:
366    与えられた量子数 `n, l, m` に対応する軌道名(例: 1s, 2pz, 3dx2-y2など)を
367    文字列として返します。
368
369    :param n: float - 主量子数。
370    :param l: float - 方位量子数。
371    :param m: float - 磁気量子数。
372    :returns: str - 波動関数の軌道名。
373    """
374    print("n=", n, l, m)
375    n = int(n+1.0e-3)
376    l = int(l+1.0e-3)
377    if m < 0:
378        m = int(m-1.0e-3)
379    else:
380        m = int(m+1.0e-3)
381    print("n=", n, l, m)
382
383    s = f"{n}"
384    if l == 0:
385        s += 's'
386    elif l == 1 and m == 0:
387        s += 'pz'
388    elif l == 1 and m == 1:
389        s += 'px'
390    elif l == 1 and m == -1:
391        s += 'py'
392    elif l == 2 and m == 0:
393        s += 'dz2'
394    elif l == 2 and m == 1:
395        s += 'dxz'
396    elif l == 2 and m == -1:
397        s += 'dyz'
398    elif l == 2 and m == 2:
399        s += 'dx2-y2'
400    elif l == 2 and m == -2:
401        s += 'dxy'
402    else:
403        terminate(f"wfname_H(): Invalid quantum number [m={m}]", usage = usage)    
404    return s
405    
406
407def psi_H1(x: float, y: float, z: float, ax: float, ay: float, az: float, n: float, l: float, m: float, kx: float, ky: float, kz: float, xonly: bool = False, yonly: bool = False, zonly: bool = False) -> complex:
408    """
409    水素原子モデルの波動関数の基本成分を計算します。
410
411    詳細説明:
412    座標 `(x, y, z)` を球座標に変換し、`psi_R_H` と `psi_Y_H` を使用して動径部分と
413    角度部分を計算します。これにブロッホ因子を乗じて波動関数を生成します。
414    `reduce2unitcell` で座標を調整します。
415
416    :param x: float - x座標。
417    :param y: float - y座標。
418    :param z: float - z座標。
419    :param ax: float - x方向の格子定数。
420    :param ay: float - y方向の格子定数。
421    :param az: float - z方向の格子定数。
422    :param n: float - 主量子数。
423    :param l: float - 方位量子数。
424    :param m: float - 磁気量子数。
425    :param kx: float - x方向のブロッホ波数ベクトル成分 (単位: pi/a)。
426    :param ky: float - y方向のブロッホ波数ベクトル成分 (単位: pi/a)。
427    :param kz: float - z方向のブロッホ波数ベクトル成分 (単位: pi/a)。
428    :param xonly: bool - Trueの場合、x方向成分のみを考慮。
429    :param yonly: bool - Trueの場合、y方向成分のみを考慮。
430    :param zonly: bool - Trueの場合、z方向成分のみを考慮。
431    :returns: complex - 水素原子モデル波動関数の複素数値。
432    """
433    x0 = reduce2unitcell(x)
434    y0 = reduce2unitcell(y)
435    z0 = reduce2unitcell(z)
436    rxy = sqrt(x0 * x0 + y0 * y0)
437    r   = sqrt(x0 * x0 + y0 * y0 + z0 * z0)
438    Phi = arctan2(y0, x0)
439    if rxy == 0.0:
440        Theta = 0.0
441    else:
442        Theta = arcsin(rxy / r)
443#    print("r=", x, y, z, x0, y0, z0, r, rxy, Theta, Phi)
444
445    Rnl = psi_R_H(r, n, l)
446    Yml = psi_Y_H(Theta, Phi, l, m)
447#    print("r=", r, Theta, Phi, Rnl, Yml)
448#    print("phi=", Theta, Phi, Rnl, Yml)
449
450    _phi = phi(x, y, z, ax, ay, az, kx, ky, kz)
451
452    f = Rnl * Yml * _phi
453    
454    return f
455
456
457def psi_H(x: float, y: float, z: float, ax: float, ay: float, az: float, n: float, l: float, m: float, kx: float, ky: float, kz: float, xonly: bool = False, yonly: bool = False, zonly: bool = False) -> tuple[float, float, float]:
458    """
459    水素原子モデルの周期的な波動関数を計算します。
460
461    詳細説明:
462    中心セルとその隣接する6つのセルからの `psi_H1` 成分を合計することで、
463    周期境界条件を模倣した水素原子モデル波動関数を生成します。
464
465    :param x: float - x座標。
466    :param y: float - y座標。
467    :param z: float - z座標。
468    :param ax: float - x方向の格子定数。
469    :param ay: float - y方向の格子定数。
470    :param az: float - z方向の格子定数。
471    :param n: float - 主量子数。
472    :param l: float - 方位量子数。
473    :param m: float - 磁気量子数。
474    :param kx: float - x方向のブロッホ波数ベクトル成分 (単位: pi/a)。
475    :param ky: float - y方向のブロッホ波数ベクトル成分 (単位: pi/a)。
476    :param kz: float - z方向のブロッホ波数ベクトル成分 (単位: pi/a)。
477    :param xonly: bool - Trueの場合、x方向成分のみを考慮。
478    :param yonly: bool - Trueの場合、y方向成分のみを考慮。
479    :param zonly: bool - Trueの場合、z方向成分のみを考慮。
480    :returns: tuple[float, float, float] - 波動関数の実部、虚部、および確率密度(絶対値の2乗)。
481    """
482    f = 0.0
483    for v in [[0, 0, 0], [-1, 0, 0], [1, 0, 0], [0, -1, 0], [0, 1, 0], [0, 0, -1], [0, 0, 1]]:
484        f += psi_H1(x, y, z, ax, ay, az, n, l, m, kx + v[0], ky + v[1], kz + v[2], xonly, yonly, zonly)
485    
486    return f.real, f.imag, f.real * f.real + f.imag * f.imag
487
488
489def main():
490    """
491    スクリプトの主要な実行ロジック。波動関数の計算とグラフ描画を行います。
492
493    詳細説明:
494    コマンドライン引数に基づいてモードと量子数を設定し、選択されたモデル
495    (平面波、量子箱、水素原子モデル)の波動関数を計算します。
496    計算結果は、1次元プロット(x方向、y方向)と、2次元平面(xy平面)における
497    波動関数の実部、虚部、確率密度のカラーマップ/3Dサーフェスプロットで可視化されます。
498
499    :returns: None - 戻り値はありません。
500    """
501    global mode, n1, n2, n3, kx, ky, kz, nmeshx, nmeshy, nmeshz
502    
503    #===================
504    # 起動時引数
505    #===================
506    argv = sys.argv
507    narg = len(argv)
508    if narg >= 2:
509        mode = argv[1]
510    if narg >= 3:
511        n1 = float(argv[2])
512    if narg >= 4:
513        n2 = float(argv[3])
514    if narg >= 5:
515        n3 = float(argv[4])
516    if narg >= 6:
517        kx = float(argv[5])
518    if narg >= 7:
519        ky = float(argv[6])
520    if narg >= 8:
521        kz = float(argv[7])
522    if narg >= 9:
523        nmeshx = float(argv[8])
524    if narg >= 10:
525        nmeshy = float(argv[9])
526    if narg >= 11:
527        nmeshz = float(argv[10])
528
529    if mode == "":
530        exit()
531
532    print("")
533    print(f"mode: {mode}")
534    print(f"lattice parameters: {ax}, {ay}, {az} A")
535    print(f"plot range in unit cell: {nlxrange}, {nlyrange}, {nlzrange}")
536    print(f"number of mesh: {nmeshx}, {nmeshy}, {nmeshz}")
537
538    if mode == 'pw':
539        psi = psi_pw
540    elif mode == 'qbox':
541        psi = psi_qbox
542    elif mode == 'H':
543        psi = psi_H
544    else:
545        terminate(f"Error: Invalid mode [{mode}]", usage = usage)
546
547    xstart = nlxrange[0]
548    xend   = nlxrange[1]
549    xstep  = (xend - xstart) / nmeshx
550    print("")
551    print(f"plot psi(x) for x range ({xstart:6.2f} - {xend:6.2f}) at {xstep:8.4f} step")
552    x     = []
553    fxr   = []
554    fxi   = []
555    fx2   = []
556    phixr = []
557    phixi = []
558    print("x\ty\tz\tPsi_r\tPsi_i\t|Psi|2\tphi_r\tphi_i")
559    for i in range(int(nmeshx) + 1):
560        _x = xstart + i * xstep
561        _yr, _yi, _y2 = psi(_x, 0.0, 0.0, ax, ay, az, n1, n2, n3, kx, ky, kz, xonly = True)
562        _phi = phi(_x, 0.0, 0.0, ax, ay, az, kx, ky, kz)
563        x.append(_x)
564        fxr.append(_yr)
565        fxi.append(_yi)
566        fx2.append(_y2)
567        phixr.append(_phi.real)
568        phixi.append(_phi.imag)
569        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}")
570
571    ystart = nlyrange[0]
572    yend   = nlyrange[1]
573    ystep  = (yend - ystart) / nmeshy
574    print("")
575    print(f"plot psi(y) for y range ({ystart:6.2f} - {yend:6.2f}) at {ystep:8.4f} step")
576    y  = []
577    fyr = []
578    fyi = []
579    fy2 = []
580    phiyr = []
581    phiyi = []
582    print("x\ty\tz\tPsi_r\tPsi_i\t|Psi|2\tphi_r\tphi_i")
583    for i in range(int(nmeshy) + 1):
584        _y = ystart + i * ystep
585        _yr, _yi, _y2 = psi(0.0, _y, 0.0, ax, ay, az, n1, n2, n3, kx, ky, kz, yonly = True)
586        _phi = phi(0.0, _y, 0.0, ax, ay, az, kx, ky, kz)
587        y.append(_y)
588        fyr.append(_yr)
589        fyi.append(_yi)
590        fy2.append(_y2)
591        phiyr.append(_phi.real)
592        phiyi.append(_phi.imag)
593        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.8f}")
594
595    print("")
596    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")
597    x3d  = np.empty([int(nmeshx)+1, int(nmeshy)+1], dtype = float)
598    y3d  = np.empty([int(nmeshx)+1, int(nmeshy)+1], dtype = float)
599    f3dr = np.empty([int(nmeshx)+1, int(nmeshy)+1], dtype = float)
600    f3di = np.empty([int(nmeshx)+1, int(nmeshy)+1], dtype = float)
601    f3d2 = np.empty([int(nmeshx)+1, int(nmeshy)+1], dtype = float)
602    print("x\ty\tz\tPsi_r\tPsi_i\t|Psi|2")
603    for iy in range(int(nmeshy) + 1):
604        _y = ystart + iy * ystep
605        for ix in range(int(nmeshx) + 1):
606            _x = xstart + ix * xstep
607            _yr, _yi, _y2 = psi(_x, _y, 0.0, ax, ay, az, n1, n2, n3, kx, ky, kz)
608            x3d[ix][iy]  = _x
609            y3d[ix][iy]  = _y
610            f3dr[ix][iy] = _yr
611            f3di[ix][iy] = _yi
612            f3d2[ix][iy] = _y2
613            if -0.5 <= _x <= 0.5 and -0.5 <= _y <= 0.5:
614                print(f"{_x:8.3f}\t{_y:8.3f}\t{_yr:8.3f}\t{_yi:8.3f}\t{_y2:8.3f}")
615
616    print("")
617    print("Quantum number:")
618    if mode == 'pw':
619        print(f"kx0, ky0, kz0 = {n1}, {n2}, {n3} * pi/a")
620    elif mode == 'qbox':
621        print(f"nx, ny, nz = {n1}, {n2}, {n3}")
622    elif mode == 'H':
623        print(f"n, l, m = {n1}, {n2}, {n3}")
624    else:
625        terminate(f"Error: Invalid mode [{mode}]", usage = usage)
626        
627    print(f"kx, ky, kz = {kx}, {ky}, {kz} * pi/a")
628
629#=============================
630# グラフの表示
631#=============================
632    maxfr = 0.0
633    for list in f3dr:
634        for v in list:
635            if abs(v) > maxfr:
636                maxfr = abs(v)
637    if maxfr < 1.0e-3:
638        maxfr = 1.0e-3
639    print("maxfr=", maxfr)
640
641    maxfi = 0.0
642    for list in f3di:
643        for v in list:
644            if abs(v) > maxfi:
645                maxfi = abs(v)
646    if maxfi < 1.0e-3:
647        maxfi = 1.0e-3
648    print("maxfi=", maxfi)
649    if maxfi > 1.0e-3:
650        maxfi = 1.0e-3
651
652    print("")
653    fig = plt.figure(figsize = figsize)
654
655    if mode == 'H':
656        wfn = wfname_H(n1, n2, n3)
657    else:
658        wfn = ''
659    if mode != 'pw':
660        n1 = int(n1 + 1.0e-3)
661        n2 = int(n2 + 1.0e-3)
662        if n3 < 0:
663            n3 = int(n3 - 1.0e-3)
664        else:
665            n3 = int(n3 + 1.0e-3)
666    title = f"mode:{mode} n=({n1},{n2},{n3}){wfn} k=({kx},{ky},{kz})"
667    plt.suptitle(title)
668
669    axcr   = fig.add_subplot(2, 3, 1)
670    axci   = fig.add_subplot(2, 3, 2)
671    ax3dr  = fig.add_subplot(2, 3, 3, projection='3d')
672    ax3di  = fig.add_subplot(2, 3, 6, projection='3d')
673    axfx   = fig.add_subplot(4, 3,  7)
674    axphix = fig.add_subplot(4, 3, 10)
675    axfy   = fig.add_subplot(4, 3,  8)
676    axphiy = fig.add_subplot(4, 3, 11)
677
678    axfx.plot(x, fxr, label = r'$\psi$$_x$$_r$', linestyle = 'dashed', linewidth = 0.5, color = 'red')
679    axfx.plot(x, fxi, label = r'$\psi$$_x$$_i$', linestyle = 'dashed', linewidth = 0.5, color = 'blue')
680#    axfx.plot(x, fx2, label = r'|$\psi$$|$^2$', linestyle = '-',     linewidth = 0.5, color = 'green')
681    xlim = nlxrange #axfx.get_xlim()
682    ylim = axfx.get_ylim()
683    axfx.plot(xlim, [0.0, 0.0], linestyle = 'dashed', linewidth = 0.3, color = 'red')
684    for i in range(int(nlxrange[0]), int(nlxrange[1]) + 1):
685        axfx.plot([i + 0.5, i + 0.5], ylim, linestyle = '-', linewidth = 0.5, color = 'black')
686    axfx.set_xlim(xlim)
687    axfx.set_ylim(ylim)
688    axfx.set_xlabel(r"$x$ (A)", fontsize = fontsize)
689    axfx.set_ylabel(r"$\psi$$_x$($x$)", fontsize = fontsize)
690    axfx.legend(fontsize = legend_fontsize, loc = 'best')
691    axfx.tick_params(labelsize = fontsize)
692
693    axphix.plot(x, phixr, label = r'$\phi_r$', linestyle = 'dashed', linewidth = 0.5, color = 'red')
694    axphix.plot(x, phixi, label = r'$\phi_i$', linestyle = 'dashed', linewidth = 0.5, color = 'blue')
695    xlim = nlxrange #axfx.get_xlim()
696    ylim = axphix.get_ylim()
697    axphix.plot(xlim, [0.0, 0.0], linestyle = 'dashed', linewidth = 0.3, color = 'red')
698    for i in range(int(nlxrange[0]), int(nlxrange[1]) + 1):
699        axphix.plot([i + 0.5, i + 0.5], ylim, linestyle = '-', linewidth = 0.5, color = 'black')
700    axphix.set_xlim(xlim)
701    axphix.set_ylim(ylim)
702    axphix.set_xlabel(r"$x$ (A)", fontsize = fontsize)
703    axphix.set_ylabel(r"$\phi$($x$)", fontsize = fontsize)
704    axphix.legend(fontsize = legend_fontsize, loc = 'best')
705    axphix.tick_params(labelsize = fontsize)
706
707    axfy.plot(y, fyr, label = r'$\psi$$_y$$_r$', linestyle = 'dashed', linewidth = 0.5, color = 'red')
708    axfy.plot(y, fyi, label = r'$\psi$$_y$$_i$', linestyle = 'dashed', linewidth = 0.5, color = 'blue')
709#    axfy.plot(y, fy2, label = r'|$\psi$$_y$|$^2$', linestyle = '-',      linewidth = 0.5, color = 'black')
710    xlim = nlyrange #axfx.get_ylim()
711    ylim = axfy.get_ylim()
712    axfy.plot(xlim, [0.0, 0.0], linestyle = 'dashed', linewidth = 0.3, color = 'red')
713    for i in range(int(nlxrange[0]), int(nlxrange[1]) + 1):
714        axfy.plot([i + 0.5, i + 0.5], ylim, linestyle = '-', linewidth = 0.5, color = 'black')
715    axfy.set_xlim(xlim)
716    axfy.set_ylim(ylim)
717    axfy.set_xlabel(r"$y$ (A)", fontsize = fontsize)
718    axfy.set_ylabel(r"$\psi$$_y$($y$)", fontsize = fontsize)
719    axfy.legend(fontsize = legend_fontsize, loc = 'best')
720    axfy.tick_params(labelsize = fontsize)
721
722    axphiy.plot(y, phiyr, label = r'$\phi_r$', linestyle = 'dashed', linewidth = 0.5, color = 'red')
723    axphiy.plot(y, phiyi, label = r'$\phi_i$', linestyle = 'dashed', linewidth = 0.5, color = 'blue')
724    xlim = nlyrange #axfx.get_ylim()
725    ylim = axphiy.get_ylim()
726    axphiy.plot(xlim, [0.0, 0.0], linestyle = 'dashed', linewidth = 0.3, color = 'red')
727    for i in range(int(nlxrange[0]), int(nlxrange[1]) + 1):
728        axphiy.plot([i + 0.5, i + 0.5], ylim, linestyle = '-', linewidth = 0.5, color = 'black')
729    axphiy.set_xlim(xlim)
730    axphiy.set_ylim(ylim)
731    axphiy.set_xlabel(r"$y$ (A)", fontsize = fontsize)
732    axphiy.set_ylabel(r"$\phi$($y$)", fontsize = fontsize)
733    axphiy.legend(fontsize = legend_fontsize, loc = 'best')
734    axphiy.tick_params(labelsize = fontsize)
735
736    surf = ax3dr.plot_surface(x3d, y3d, f3dr, label = r'$\Psi_r$', cmap = cmap_r)
737    ax3dr.contour(x3d, y3d, f3dr, cmap = cmap_r, offset = -1,
738                vmin = -maxfr, vmax = maxfr, levels = 100) 
739#                norm = Normalize(vmin = -maxfr, vmax = maxfr), levels = 100) 
740#    fig.colorbar(surf, ax = ax3dr)
741#    ax3dr.set_aspect('equal')
742    ax3dr.set_xlabel(r"$x$ (A)", fontsize = fontsize)
743    ax3dr.set_ylabel(r"$y$ (A)", fontsize = fontsize)
744    ax3dr.set_zlabel(r"$\Psi$$_r$", fontsize = fontsize)
745#    ax3dr.legend(fontsize = legend_fontsize, loc = 'best')
746    ax3dr.tick_params(labelsize = fontsize)
747
748    surf = ax3di.plot_surface(x3d, y3d, f3di, label = r'$\Psi_i$', cmap = cmap_r)
749    ax3di.contour(x3d, y3d, f3di, cmap = cmap_r, offset = -1,
750                norm = Normalize(vmin = -maxfi, vmax = maxfi), levels = 100) 
751#    fig.colorbar(surf, ax = ax3di)
752#    ax3di.set_aspect('equal')
753    ax3di.set_xlabel(r"$x$ (A)", fontsize = fontsize)
754    ax3di.set_ylabel(r"$y$ (A)", fontsize = fontsize)
755    ax3di.set_zlabel(r"$\Psi$$_i$", fontsize = fontsize)
756#    ax3di.legend(fontsize = legend_fontsize, loc = 'best')
757    ax3di.tick_params(labelsize = fontsize)
758
759    axcr.set_title(r"$\Psi_r$")
760#    contour = axcr.pcolormesh(x3d, y3d, f3dr, label = '$\Psi_r$', cmap = cmap_r, shading='auto')
761    if maxfr > 1.0e-3:
762#        contour = axcr.contour(x3d, y3d, f3dr, cmap = cmap_r, 
763#                levels = np.arange(-maxfr, maxfr, 0.01 * maxfr))
764        contour = axcr.contourf(x3d, y3d, f3dr, 
765                cmap = cmap_r, norm = Normalize(vmin = -maxfr, vmax = maxfr), levels = 100) 
766        fig.colorbar(contour, ax = axcr, shrink = 0.5)
767    axcr.set_aspect('equal')
768    xlim = nlxrange #axfy.get_xlim()
769    ylim = nlyrange #axfy.get_ylim()
770    for i in range(int(nlxrange[0]), int(nlxrange[1]) + 1):
771#        if i == 0:
772#            continue
773        axcr.plot([i + 0.5, i + 0.5], ylim, linestyle = '-', linewidth = 1.0, color = 'black')
774    for i in range(int(nlyrange[0]), int(nlyrange[1]) + 1):
775#        if i == 0:
776#            continue
777        axcr.plot(xlim, [i + 0.5, i + 0.5], linestyle = '-', linewidth = 1.0, color = 'black')
778    axcr.plot(xlim, [0.0, 0.0], linestyle = '-', linewidth = 0.5, color = 'red')
779    axcr.plot([0.0, 0.0], ylim, linestyle = '-', linewidth = 0.5, color = 'red')
780    axcr.set_xlim(xlim)
781    axcr.set_ylim(ylim)
782    axcr.set_xlabel(r"$x$ (A)", fontsize = fontsize)
783    axcr.set_ylabel(r"$y$ (A)", fontsize = fontsize)
784#    axcr.legend(fontsize = legend_fontsize, loc = 'best')
785    axcr.tick_params(labelsize = fontsize)
786
787    axci.set_title(r"$\Psi_i$")
788#    contour = axci.pcolormesh(x3d, y3d, f3di, label = r'$\Psi_i$', cmap = cmap_r, shading='auto')
789    if maxfi > 1.0e-3:
790#        contour = axci.contour(x3d, y3d, f3di, cmap = cmap_r, 
791#                levels = np.arange(-maxfi, maxfi, 0.01 * maxfi))
792        contour = axci.contourf(x3d, y3d, f3di, 
793                cmap = cmap_r, norm = Normalize(vmin = -maxfi, vmax = maxfi), levels = 100) 
794        fig.colorbar(contour, ax = axci, shrink = 0.5)
795    axci.set_aspect('equal')
796    xlim = nlxrange #axfy.get_xlim()
797    ylim = nlyrange #axfy.get_ylim()
798    for i in range(int(nlxrange[0]), int(nlxrange[1]) + 1):
799#        if i == 0:
800#            continue
801        axci.plot([i + 0.5, i + 0.5], ylim, linestyle = '-', linewidth = 1.0, color = 'black')
802    for i in range(int(nlyrange[0]), int(nlyrange[1]) + 1):
803#        if i == 0:
804#            continue
805        axci.plot(xlim, [i + 0.5, i + 0.5], linestyle = '-', linewidth = 1.0, color = 'black')
806    axci.plot(xlim, [0.0, 0.0], linestyle = '-', linewidth = 0.5, color = 'red')
807    axci.plot([0.0, 0.0], ylim, linestyle = '-', linewidth = 0.5, color = 'red')
808    axci.set_xlim(xlim)
809    axci.set_ylim(ylim)
810    axci.set_xlabel("$x$ (A)", fontsize = fontsize)
811    axci.set_ylabel("$y$ (A)", fontsize = fontsize)
812#    axci.legend(fontsize = legend_fontsize, loc = 'best')
813    axci.tick_params(labelsize = fontsize)
814
815    plt.tight_layout()
816    plt.pause(0.1)
817
818    print("Press ENTER to exit>>", end = '')
819    input()
820
821    terminate("", usage = usage)
822
823
824if __name__ == '__main__':
825    main()