Draw 2D wave functions
Download script from .\wavefunction2D.py
Related files:
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
"""
Draw 2D wave functions
"""
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
#=============================
figsize = (12, 8)
fontsize = 10
legend_fontsize = 10
cmap_r = 'bwr' # 'plasma', 'coolwarm', 'rainbow'
#===================
# 起動時引数
#===================
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):
print("")
print(message)
exit()
def reduce2unitcell(x):
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):
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):
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):
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):
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):
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):
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):
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):
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):
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():
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()
