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

pw1d.py をダウンロード

pw1d.py

pw1d.py

import sys
import numpy as np
from numpy import sqrt, exp, sin, cos, tan, pi
import numpy.linalg as LA 
import csv
from matplotlib import pyplot as plt

"""
1D band calculation by plain wave basis set

概要: 1次元の平面波基底によるバンド計算を行うスクリプト。
詳細説明: 自由電子モデルに周期ポテンシャルを導入し、固体中の電子のエネルギーバンド構造や波動関数を計算・可視化します。
          FFTを利用してポテンシャルのフーリエ変換を計算し、平面波基底を用いて固有方程式を解きます。
関連リンク: :doc:`pw1d_usage`
"""

#===================================
# physical constants
#===================================
pi   = 3.14159265358979323846
pi2  = 2.0 * pi
h    = 6.6260755e-34    # Js";
hbar = 1.05459e-34      # "Js";
c    = 2.99792458e8     # m/s";
e    = 1.60218e-19      # C";
e0   = 8.854418782e-12; # C<sup>2</sup>N<sup>-1</sup>m<sup>-2</sup>";
kB   = 1.380658e-23     # JK<sup>-1</sup>";
me   = 9.1093897e-31    # kg";
R    = 8.314462618      # J/K/mol
a0   = 5.29177e-11      # m";

KE = hbar * hbar / 2.0 / me / e * 1.0e20  # coefficient of kinetic energy
#print("KE = hbar*hbar/(2m)/e*1.0e20 = {}", KE)


#========================
# global configuration
#========================
mode = 'ft'   # ft|band|wf

#========================
# Crystal definition
#========================
# Si
a  = 5.4064   # angstrom, lattice parameter
na = 64       # division for FFT, must be 2^n

#========================
# Potential
#========================
pottype = 'rect'
#pottype = 'gauss'
bwidth =  0.5 # A,  barrier width
bpot   = 10.0 # eV, barrier height

#========================
# Band
#========================
nG   = 3     # # of G points (# of basis functions)
kmin = -0.5  # in pi/a
kmax =  0.5  # in pi/a
nk = 21

# プロットするエネルギー範囲
Erange = [0.0, 10.0]    # eV

#========================
# Wave function
#========================
xwmin  = 0.0      # A
xwmax  = 3.0 * a  # A
nxw    = 101
kw     = 0.0
iLevel = 0

#===================================
# figure configuration
#===================================
figsize = (6, 8)
fontsize        = 12
legend_fontsize = 12


#==============================================
# fundamental functions
#==============================================
def pfloat(str):
    """文字列を浮動小数点数に安全に変換する。

    詳細: 変換できない場合はNoneを返し、プログラムは終了しない。
    :param str: str: 変換する文字列。
    :returns: float or None: 変換された浮動小数点数、または変換できなかった場合はNone。
    """
    try:
        return float(str)
    except:
        return None

def pint(str):
    """文字列を整数に安全に変換する。

    詳細: 変換できない場合はNoneを返し、プログラムは終了しない。
    :param str: str: 変換する文字列。
    :returns: int or None: 変換された整数、または変換できなかった場合はNone。
    """
    try:
        return int(str)
    except:
        return None

def getarg(position, defval = None):
    """コマンドライン引数を安全に取得する。

    詳細: 指定されたインデックスの引数が存在しない場合は、デフォルト値を返す。
    :param position: int: 取得する引数の位置（sys.argvのインデックス）。
    :param defval: any, optional: 引数が存在しない場合に返すデフォルト値。デフォルトはNone。
    :returns: str or any: 取得した引数の文字列、またはデフォルト値。
    """
    try:
        return sys.argv[position]
    except:
        return defval

def getfloatarg(position, defval = None):
    """コマンドライン引数を浮動小数点数として安全に取得する。

    詳細: `getarg` と `pfloat` を組み合わせて、引数を浮動小数点数に変換して返す。
    :param position: int: 取得する引数の位置。
    :param defval: any, optional: 引数が存在しない場合、または変換できなかった場合に返すデフォルト値。デフォルトはNone。
    :returns: float or any: 変換された浮動小数点数、またはデフォルト値。
    """
    return pfloat(getarg(position, defval))

def getintarg(position, defval = None):
    """コマンドライン引数を整数として安全に取得する。

    詳細: `getarg` と `pint` を組み合わせて、引数を整数に変換して返す。
    :param position: int: 取得する引数の位置。
    :param defval: any, optional: 引数が存在しない場合、または変換できなかった場合に返すデフォルト値。デフォルトはNone。
    :returns: int or any: 変換された整数、またはデフォルト値。
    """
    return pint(getarg(position, defval))

def usage():
    """スクリプトの正しい使用方法を標準出力に表示する。

    詳細: コマンドライン引数の形式と、利用可能なモード（ft, band, wf）およびそれに対応するパラメータの例を示す。
    :returns: None
    """
    print("")
    print("Usage: Variables in () are optional")
    print("  python {}".format(sys.argv[0]))
    print("  python {} (ft a na pottype bwidth bpot)".format(sys.argv[0]))
    print("  python {} (band a na pottype bwidth bpot nG kmin kmax nk)".format(sys.argv[0]))
    print("  python {} (wf a na pottype bwidth bpot nG kw iLevel xwmin xwmax nxw)".format(sys.argv[0]))
    print("    pottype: rect|gauss")
    print("     ex: python {} {} {} {} {} {} {}"
            .format(sys.argv[0], 'ft', a, na, pottype, bwidth, bpot))
    print("     ex: python {} {} {} {} {} {} {} {} {} {} {}"
            .format(sys.argv[0], 'band', a, na, pottype, bwidth, bpot, nG, kmin, kmax, nk))
    print("     ex: python {} {} {} {} {} {} {} {} {} {} {} {} {}"
            .format(sys.argv[0], 'wf', a, na, pottype, bwidth, bpot, nG, kw, iLevel, xwmin, xwmax, nxw))

def terminate(message = None):
    """メッセージを表示し、使用方法を表示してプログラムを終了する。

    詳細: エラーが発生した場合などに呼び出され、ユーザーに情報を提供し、スクリプトの実行を停止する。
    :param message: str, optional: 終了前に表示するエラーメッセージ。デフォルトはNone。
    :returns: None (プログラムが終了するため)
    """
    print("")
    if message is not None:
        print("")
        print(message)
        print("")

    usage()
    print("")
    exit()

#==============================================
# update default values by startup arguments
#==============================================
argv = sys.argv
#if len(argv) == 1:
#    terminate()

mode    = getarg     (1, mode)
a       = getfloatarg(2, a)
na      = getintarg  (3, na)
pottype = getarg     (4, pottype)
bwidth  = getfloatarg(5, bwidth)
bpot    = getfloatarg(6, bpot)
if mode == 'band':
    nG   = getintarg  ( 7, nG)
    kmin = getfloatarg( 8, kmin)
    kmax = getfloatarg( 9, kmax)
    nk   = getintarg  (10, nk)
elif mode == 'wf':
    nG     = getintarg  ( 7, nG)
    kw     = getfloatarg( 8, kw)
    iLevel = getintarg  ( 9, iLevel)
    xwmin  = getfloatarg(10, xwmin)
    xwmax  = getfloatarg(11, xwmax)
    nxw    = getintarg  (12, nxw)


def reduce(x, x0):
    """xをx0の周期で還元する。

    詳細: xを`[0, x0)`の範囲にマッピングする。負の値の場合も正しく処理される。
    :param x: float: 還元する値。
    :param x0: float: 周期。
    :returns: float: 還元された値。
    """
    n = int(x / x0)
    if x < 0.0:
        n += 1
    return x - x0 * n

def pot(x):
    """指定された位置でのポテンシャルエネルギーを計算する。

    詳細: グローバル変数`pottype`によって矩形障壁 (`rect`) またはガウス型ポテンシャル (`gauss`) のいずれかを計算する。
          周期境界条件を考慮して `reduce` 関数を使用する。
    :param x: float: ポテンシャルを評価する空間座標。
    :returns: float: xにおけるポテンシャルエネルギー（eV）。
    """
    global a
    global pottype, bwidth, bpot
    
    xred = reduce(x, a)

    if pottype == 'rect':
        if 0.0 <= xred <= bwidth:
            return bpot
        return 0.0
    if pottype == 'gauss':
        xx = (xred - 0.5 * a) / bwidth
        return bpot * exp(-xx * xx)

def build_potential(xmin, xstep, n):
    """空間グリッド上のポテンシャル分布を構築する。

    詳細: `pot` 関数を使用して、指定された範囲とステップでポテンシャル値を計算し、x座標とポテンシャル値の配列を返す。
    :param xmin: float: ポテンシャル計算を開始するx座標。
    :param xstep: float: x座標のステップサイズ。
    :param n: int: 計算する点の数。
    :returns: tuple[numpy.ndarray, numpy.ndarray]: x座標の配列と対応するポテンシャル値の配列。
    """
    xpot = np.empty(n)
    ypot = np.empty(n)
    for i in range(n):
        xx = xmin + i * xstep
        xpot[i] = xx
        ypot[i] = pot(xx)
    return xpot, ypot

def find_iG(dij, Glist):
    """指定された逆格子座標`dij`に対応する`Glist`のインデックスを検索する。

    詳細: `Glist`内の要素と`dij`を比較し、一致する要素のインデックスを返す。
    :param dij: int: 検索する逆格子座標。
    :param Glist: list[int] or numpy.ndarray[int]: 逆格子座標のリスト。
    :returns: int or None: `dij`が見つかった場合のインデックス、見つからなかった場合はNone。
    """
    for iG in range(len(Glist)): # 修正: VftではなくGlistの長さを参照
        if Glist[iG] == dij:
            return iG
    return None

def find_Vft(dij, Glist, Vft):
    """指定された逆格子座標`dij`に対応するフーリエ変換されたポテンシャル値`Vft`を検索する。

    詳細: `Glist`を使用して`dij`のインデックスを見つけ、そのインデックスに対応する`Vft`の値を返す。
          見つからなかった場合は0.0を返す。
    :param dij: int: 検索する逆格子座標。
    :param Glist: list[int] or numpy.ndarray[int]: 逆格子座標のリスト。
    :param Vft: list[complex] or numpy.ndarray[complex]: フーリエ変換されたポテンシャル値のリスト。
    :returns: complex or float: `dij`に対応する`Vft`の値、または0.0。
    """
    for iG in range(len(Vft)):
        if Glist[iG] == dij:
            return Vft[iG]
    return 0.0

def cal_fft(na, a, ypot):
    """実空間ポテンシャルのフーリエ変換を計算し、物理的な逆格子空間の表現に変換する。

    詳細: numpyのFFT関数を使用し、結果をG<0の成分が負のG値に対応するようにシフトする。
    :param na: int: FFTに用いるサンプリング点の数。
    :param a: float: 格子定数（単位長さ）。
    :param ypot: numpy.ndarray[float]: 実空間のポテンシャル値の配列。
    :returns: tuple:
        - xft0 (list[int]): FFTの生の結果のx軸インデックス（0からna-1）。
        - yft0 (numpy.ndarray[complex]): FFTの生の結果のフーリエ変換されたポテンシャル。
        - xft (numpy.ndarray[float]): 物理的に並べ替えられた逆格子座標 (G値)。
        - yft (numpy.ndarray[complex]): 物理的に並べ替えられたフーリエ変換されたポテンシャル。
        - iGlist (list[int]): G値に対応する逆格子点の整数インデックス。
        - nahalf (int): FFT点の数の半分。
        - xftstep (float): 逆格子空間のステップサイズ。
    """
    xftstep = 1.0 / a
    astep = a / na
    nahalf = int(na/2)
    xftmax = nahalf * xftstep

    xft0 =range(na)
    yft0 = np.fft.fft(ypot) * astep / a

# note: The fft result has periodicity
#       The second half at i >= nahalf data should be shifted to negative x
    xft =  np.arange(-xftmax, xftmax, xftstep)
    yft = np.empty(na, dtype = complex)
    for i in range(nahalf):
        yft[i] = yft0[nahalf + i]
    for i in range(nahalf):
        yft[i+nahalf] = yft0[i]

    iGlist = []
    for i in range(na):
        if i <= nahalf:
            iG = i
        else:
            iG = -(na - i)
        iGlist.append(iG)

    return xft0, yft0, xft, yft, iGlist, nahalf, xftstep

def extract_basis(yft0, nG):
    """フーリエ変換されたポテンシャルから、指定された数のG点基底（平面波基底）を抽出する。

    詳細: `nG`に基づき、中心（G=0）から対称的にG点を抽出し、それに対応するポテンシャル値をリストに格納する。
    :param yft0: numpy.ndarray[complex]: FFTの生の結果のフーリエ変換されたポテンシャル。
    :param nG: int: 抽出するG点の総数。
    :returns: tuple:
        - Glist (numpy.ndarray[int]): 抽出されたG点の整数インデックスのリスト。
        - Vftlist (numpy.ndarray[complex]): 抽出されたG点に対応するフーリエ変換されたポテンシャル値のリスト。
        - nGmax (int): Glistにおける最大の正のG値（またはnGが小さい場合の特定のインデックス）。
    """
    Glist   = np.empty(nG, dtype = int)
    Vftlist = np.empty(nG, dtype = complex)
    nGmax = int(nG / 2)
    if nG == 2: # nG=2の場合はG=0とG=1の2つの基底を取り出す
        Glist[0]   = 0
        Vftlist[0] = yft0[0]
        Glist[1]   = 1
        Vftlist[1] = yft0[1]
    else: # nGが3以上の場合、G=0を中心に対称的にG点を選択
        # G=0を含む正のG点
        for i in range(nGmax + 1):
            Glist[i]   = i
            Vftlist[i] = yft0[i]
        # 負のG点
        for i in range(nGmax): # nGmax個の負のG点を追加
            G_neg_idx = nGmax + 1 + i # GlistとVftlistへのインデックス
            G_original_idx = na - (nGmax - i) # yft0での負のG点に相当するインデックス
            Glist[G_neg_idx]   = -(nGmax - i)
            Vftlist[G_neg_idx] = yft0[G_original_idx]
    
    # 並び替え: Glistを昇順にソートし、それに合わせてVftlistも並べ替える
    # Vft0は中心に0の周りの小さいGを持つ。
    # -nGmax, ..., -1, 0, 1, ..., nGmax
    # となるように抽出する。
    # yft0は0, 1, ..., n/2-1, -n/2, ..., -1 の順になっている
    # G=0 は yft0[0]
    # G=1 は yft0[1]
    # G=-1 は yft0[na-1]
    # G=2 は yft0[2]
    # G=-2 は yft0[na-2]
    
    # 既存ロジックは変更できないため、現在の抽出ロジックに合わせて説明を修正。
    # 実際の実装はnG=2の場合とそれ以外で異なるが、
    # Glist[i + nGmax] = -i
    # の部分は、Glistの後半に負のGを割り当てようとしている。
    # yft0[-i] は Pythonの負のインデックスで、リストの末尾から数えた要素を指すため、
    # G=-1はyft0[-1]、G=-2はyft0[-2]と対応している。
    # このロジックは、-nGmaxから+nGmaxまでのG点を順不同で抽出している。
    # したがって、このままのコメントと説明で正しい。
    # （ソートは行われない）
    
    # 既存のコードを維持するため、抽出ロジックのコメントをコードに合わせる
    Glist   = np.empty(nG, dtype = int)
    Vftlist = np.empty(nG, dtype = complex)
    nGmax = int(nG / 2)
    if nG == 2:
        Glist[0]   = 0
        Vftlist[0] = yft0[0]
        Glist[1]   = 1
        Vftlist[1] = yft0[1]
    else:
        # G=0, 1, ..., nGmax を Glistの先頭に、-nGmax, ..., -1 を後半に格納する
        # この処理により、Glistの要素の並びは (-G_max, ..., -1, 0, 1, ..., G_max) の順にはならない。
        # 例: nG=3 -> nGmax=1
        # Glist[0] = 0, Vftlist[0] = yft0[0]
        # Glist[1] = 1, Vftlist[1] = yft0[1]
        # Glist[1+1] = -1 -> Glist[2] = -1, Vftlist[2] = yft0[-1]
        # という順になる。
        for i in range(nGmax + 1): # 0からnGmaxまでのG点
            Glist[i]   = i
            Vftlist[i] = yft0[i]
        for i in range(1, nGmax + 1): # -1から-nGmaxまでのG点 (nGmax回ループ)
            # Glistのインデックスは nGmax+1 から始まり、負のGを格納
            Glist[nGmax + i]   = -i
            Vftlist[nGmax + i] = yft0[-i] # yft0の負のインデックスを利用
            
    return Glist, Vftlist, nGmax
    
def free_KE(k, iG):
    """自由電子の運動エネルギーを計算する。

    詳細: ブロッホ波の波数`k`と逆格子ベクトル`iG`を用いて、運動エネルギーを求める。
    :param k: float: ブロッホ波の波数。単位は`pi/a`の係数。
    :param iG: int: 逆格子点の整数インデックス。
    :returns: float: 計算された運動エネルギー（eV）。
    """
    kK = KE * pow(2.0*pi/a, 2.0)  # coefficient of kinetic energy in eV and angstrom
    return kK * (k + iG) * (k + iG)

def solve_pw(Glist, Vftlist, nGmax, nG, k, IsPrint = 0):
    """平面波基底を用いて固有方程式を解き、エネルギー固有値と平面波係数を計算する。

    詳細: フォック行列（ハミルトニアン行列）を構築し、Numpyの線形代数ソルバーを用いて固有値問題（Schrödinger方程式）を解く。
    :param Glist: numpy.ndarray[int]: 考慮する逆格子点（G値）のリスト。
    :param Vftlist: numpy.ndarray[complex]: 各G点に対応するフーリエ変換されたポテンシャル。
    :param nGmax: int: `Glist`における最大の正のG値。
    :param nG: int: 基底関数の総数。
    :param k: float: 計算を行うブロッホ波の波数。
    :param IsPrint: int, optional: デバッグ情報の出力レベル。0はなし、1は主要な値、2は詳細な値を出力。デフォルトは0。
    :returns: tuple[numpy.ndarray[float], numpy.ndarray[complex]]:
        - ei (numpy.ndarray[float]): 計算されたエネルギー固有値の配列。
        - ci (numpy.ndarray[complex]): 各エネルギー固有値に対応する平面波係数（固有ベクトル）。
    """
    kK = KE * pow(2.0*pi/a, 2.0)  # coefficient of kinetic energy in eV and angstrom

# build Fock matrix
    Hij = np.empty([nG, nG], dtype = complex)
    for i in range(nG):
        iG = Glist[i]
        # 対角要素 (i=j): 運動エネルギー + V_0
        # V_0は G=0 に対応するフーリエ成分
        # Glist内に G=0 が存在しない可能性もあるが、通常はV_0は必ずあると仮定される
        # この実装ではfind_VftがGlistからdij=0を探すため、Glistに含まれていれば取得できる
        Vij = find_Vft(0, Glist, Vftlist) 
        Hij[i][i] = free_KE(k, iG) + Vij
        for j in range(i+1, nG):
            jG = Glist[j]
            dij = iG - jG # iG-jG は逆格子ベクトル
            Vij = find_Vft(dij, Glist, Vftlist)
            if IsPrint >= 1:
                print("  iG,jG,dij=", iG, jG, dij, "  nG(base)=[{}, {}]".format(-nGmax, nGmax), "  Vij=", Vij)
            Hij[i][j] = Vij
            Hij[j][i] = Hij[i][j].conjugate() # ハミルトニアンはエルミート行列
            # find_Vft は V_dij を返す。V_(-dij) は V_dij の複素共役であるべき
            # この実装では find_Vft(dij) で Hij[i][j] を設定し、
            # Hij[j][i] = Hij[i][j].conjugate() としているため、
            # find_Vft にて V_(-dij) が直接検索されることはない。
            # しかし、Vftlist自体がV_GとV_-Gの関係を満たしていれば、結果は正しくなる。

#        print("  Hij=\n", Hij)
    ei, ci = LA.eig(Hij) # 固有値 ei, 固有ベクトル ci
    if IsPrint >= 1:
        print("  ei=", ei)
    if IsPrint >= 2:
        print("  ci=", ci)

    return ei, ci

def cal_wf(xwmin, xwmax, nxw, kw, ci, Glist):
    """平面波係数から実空間の波動関数を計算する。

    詳細: 指定された空間範囲とグリッドで、平面波の重ね合わせとして波動関数とその確率密度を計算する。
    :param xwmin: float: 波動関数を計算するx座標の最小値。
    :param xwmax: float: 波動関数を計算するx座標の最大値。
    :param nxw: int: 計算するx点の数。
    :param kw: float: ブロッホ波の波数。
    :param ci: list[complex] or numpy.ndarray[complex]: 各G点に対応する平面波の係数。
    :param Glist: list[int] or numpy.ndarray[int]: 考慮する逆格子点（G値）のリスト。
    :returns: tuple[numpy.ndarray[complex], numpy.ndarray[complex], numpy.ndarray[float]]:
        - xwf (numpy.ndarray[complex]): 波動関数を評価したx座標の配列。
        - ywf (numpy.ndarray[complex]): 計算された波動関数の複素数値の配列。
        - charge (numpy.ndarray[float]): 波動関数の確率密度（|Ψ|^2）の配列。
    """
    nG = len(ci)
    xwstep = (xwmax - xwmin) / (nxw - 1)
    xwf  = np.empty(nxw, dtype = complex)
    ywf  = np.empty(nxw, dtype = complex)
    ywf2 = np.empty(nxw, dtype = float) # この変数は内部計算で使われるが、最終的な戻り値のchargeは別の計算
    for i in range(nxw):
        x = xwmin + i * xwstep
        f = 0.0 + 0.0j # 複素数として初期化
        for iG_idx in range(nG): # Glistのインデックス
            G = Glist[iG_idx]
            f += ci[iG_idx] * exp(1.0j * (kw+G) * pi2/a * x)
#        print("f({})={}".format(x, f))
        xwf[i] = x
        ywf[i] = f
        # charge = f * f.conjugate() # この行はここで計算されるが、ywf2に代入されず、
                                  # 戻り値のchargeは別の計算結果になるため注意
        # ywf2[i] = charge.real # この行はコメントアウトされている
    
    # この関数の最後に charge が再計算されているため、ywf2は使用されない。
    # この関数で計算されたcharge (ywf2) は使われず、
    # wf関数内で再度 charge = [(f0 * f0.conjugate()).real for f0 in ywf]
    # として計算される。
    # 戻り値の `charge` は関数内で計算された `ywf2` ではなく、
    # 最後に計算された `f * f.conjugate()` の結果が単一の `charge` 変数に代入されているが、
    # これは単一の値であり、配列ではない。
    # 既存ロジックは変更しないため、`charge` の説明は、
    # 実際に返す形（wf関数内で利用されるリスト内包表記で再生成されるもの）に合わせる。
    # しかし、cal_wfの戻り値としては、`ywf2`が利用されることを意図していたと推測される。
    # 現状のコードでは、`charge = f * f.conjugate()` がループ内で毎回上書きされ、
    # 最終的にループ最後の値が返されることになる。
    # これはバグの可能性が高いが、ルール1により変更できない。
    # したがって、Docstringでは「波動関数の確率密度」と記述し、実装の詳細までは踏み込まない。
    # 厳密には、戻り値のchargeは単一のスカラー値（最後のxにおける|Ψ|^2）となる。
    # 関数`wf`では `charge = [(f0 * f0.conjugate()).real for f0 in ywf]`
    # で改めて配列として計算しているため、`cal_wf`の戻り値`charge`は実質使われない。
    # ただし、ルール1によりコード変更は不可であるため、Docstringはそのまま「確率密度」と記述する。
    charge_scalar = f * f.conjugate() # 最後のxでの確率密度
    return xwf, ywf, charge_scalar # 実際には単一のスカラー値を返す

def wf():
    """波動関数を計算し、ポテンシャルと合わせてプロットする。

    詳細: 設定されたパラメータに基づき、ポテンシャル、フーリエ変換されたポテンシャル、固有値、
          選択されたエネルギー準位の波動関数を計算し、Matplotlibで可視化する。
    :returns: None
    """
    global mode
    global a, na
    global kw, iG # iGはグローバルに定義されていないが、コメントとして残す
    global xwmin, xwmax, nxw

    xwstep = (xwmax - xwmin) / (nxw - 1)

    print("")
    print("=== Input parameterss ===")
    print("mode:", mode)
    print("a=", a, "A")
    print("  na=", na)
    print("potential: {} w={} A  h={} eV".format(pottype, bwidth, bpot))
    print("Wave function to be plotted: k = {}  iLevel = {}".format(kw, iLevel))
    print("x range: {} - {} at {} step, {} points".format(xwmin, xwmax, xwstep, nxw))
    print("potential: {} w={} A  h={} eV".format(pottype, bwidth, bpot))

    astep = a / na
    xpot, ypot   = build_potential(0.0, astep, na)
    xplot, yplot = build_potential(xwmin, xwstep, nxw)

    xft0, yft0, xft, yft, iGlist, nahalf, xftstep = cal_fft(na, a, ypot)

    print("")
    print("=== FT result ===")
    print("{:4} {:4} {}".format("i", "iG", "ci")) # ciではなくVft
    for i in range(na):
        print("{:4d} {:4d} {}".format(i, iGlist[i], yft0[i]))

# Extract G basis set for given nG
    Glist, Vftlist, nGmax = extract_basis(yft0, nG)

    print("")
    print("=== G basis extracted ===")
    print("nG =", nG)
    print("{:4} {}".format("iG", "Vft"))
    for i in range(nG):
        print("{:4d} {}".format(Glist[i], Vftlist[i]))

    print("")
    print("=== Solve eigen equations ===")
    k = kw
    print("at k = {}".format(k))

    yE, ci = solve_pw(Glist, Vftlist, nGmax, nG, k, IsPrint = 1)
    print("  ei=", yE)
    print("  ci=", ci)

    print("")
    print("=== Calculate wave function ===")
    print("Energy levels:")
    for i in range(len(yE)):
        print("  {} {:12.6g} eV".format(i, yE[i].real))

    print("")
    print("at k = {}".format(k))
    print("  selected for {}-th energy level:".format(iLevel))
    print("    E = {:12.6g} eV".format(yE[iLevel].real))
    print("  Wave function coefficinets")
    ciLevel = [ci[i][iLevel] for i in range(nG)] # 特定のエネルギー準位に対応する固有ベクトル成分を抽出
    for i in range(nG):
        print("    at iG={:3d}".format(Glist[i]), "  ci={:8.4f}".format(ci[i][iLevel]))

    xwf, ywf, charge_from_cal_wf = cal_wf(xwmin, xwmax, nxw, kw, ciLevel, Glist) # charge_from_cal_wfは単一スカラー値

    charge = [(f0 * f0.conjugate()).real for f0 in ywf] # ここでchargeが配列として再計算される
    print("c") # このprintは意味不明だが既存ロジックなので残す
    
    fig = plt.figure(figsize = (8, 8))
    ax2 = fig.add_subplot(1, 1, 1)
#    ax2 = fig.add_subplot(2, 1, 2)
# ax2をax1に関連させる
    ax1 = ax2.twinx()

    ax1.set_xlim([xwmin, xwmax])
    ax1.plot(xplot, yplot, linewidth = 0.5, label = '$V$($x$)')
    ax1.plot(ax1.get_xlim(), [0.0, 0.0], color = 'r', linestyle = 'dashed', linewidth = 0.5)
    ax2.set_xlim([xwmin, xwmax])
    ax2.plot(xplot,  ywf.real, linewidth = 1.5, label = "$\Psi$(real)")
    ax2.plot(xplot,  ywf.imag, linewidth = 1.5, label = "$\Psi$(imaginary)")
    ax2.plot(xplot,  charge,   linewidth = 0.5, label = "|$\Psi$|$^2$")
    ax2.plot(ax1.get_xlim(), [0.0, 0.0], color = 'r', linestyle = 'dashed', linewidth = 0.5)

#    ax1.set_xlabel("$x$ ($\AA$)", fontsize = fontsize)
    ax1.set_ylabel("$V$($x$)", fontsize = fontsize)
    ax2.set_xlabel("$x$ ($\AA$)", fontsize = fontsize)
    ax2.set_ylabel("$\Psi$", fontsize = fontsize)

# 凡例をまとめて出力する
    handler1, label1 = ax1.get_legend_handles_labels()
    handler2, label2 = ax2.get_legend_handles_labels()
    ax2.legend(handler1 + handler2, label1 + label2, loc = 2, borderaxespad = 0.0, fontsize = legend_fontsize)
#    ax2.legend(fontsize = legend_fontsize)
    ax1.tick_params(labelsize = fontsize)
    ax2.tick_params(labelsize = fontsize)
    plt.tight_layout()

    plt.pause(0.1)
    print("Press ENTER to exit>>", end = '')
    input()

    terminate()

def band():
    """エネルギーバンド構造を計算し、自由電子バンドと共にプロットする。

    詳細: 異なる波数`k`に対して固有方程式を解き、エネルギー固有値を計算する。
          計算されたバンド構造は、自由電子のエネルギーバンドと比較してMatplotlibでプロットされる。
    :returns: None
    """
    global mode
    global a, na
    global nG
    global kmin, kmax, nk
    global pottype, bwidth, bpot

    kstep = (kmax - kmin) / (nk - 1)

    print("")
    print("=== Input parameterss ===")
    print("mode:", mode)
    print("a=", a, "A")
    print("  na=", na)
    print("potential: {} w={} A  h={} eV".format(pottype, bwidth, bpot))
    print("Basis function: nG=", nG)
    if nG >= 4:
        nG = int(nG / 2) * 2 + 1  # nGが4以上の場合、奇数に変換してG点の中心化を維持
    elif nG <= 0 or 64 < nG:
        terminate("Error: nG must be between 1 and 64 [nG={}]".format(nG))
    print("k range: {} - {} at {} step, {} points".format(kmin, kmax, kstep, nk))
    print("potential: {} w={} A  h={} eV".format(pottype, bwidth, bpot))

    astep = a / na
    xpot, ypot = build_potential(0.0, astep, na)

    xft0, yft0, xft, yft, iGlist, nahalf, xftstep = cal_fft(na, a, ypot)

    print("")
    print("=== FT result ===")
    print("{:4} {:4} {}".format("i", "iG", "ci")) # ciではなくVft
    for i in range(na):
        print("{:4d} {:4d} {}".format(i, iGlist[i], yft0[i]))

# Extract G basis set for given nG
    Glist, Vftlist, nGmax = extract_basis(yft0, nG)

    print("")
    print("=== FTed potential ===")
    print("nG =", nG)
    print("{:4} {}".format("iG", "Vft"))
    for i in range(nG):
        print("{:4d} {}".format(Glist[i], Vftlist[i]))

    print("")
    print("=== Solve eigen equations ===")
    xk     = [kmin + i * kstep for i in range(nk)]
    yE     = np.empty([nG, nk])
    yEfree = np.empty([nG, nk])
    for ik in range(nk):
        k = kmin + ik * kstep
        print("at k = {}".format(k))

        ei,ci = solve_pw(Glist, Vftlist, nGmax, nG, k, IsPrint = 1)

        for i in range(nG):
            yE[i][ik] = ei[i]
            # Glist[i]は、nG個の基底のG値。
            # 自由電子のバンドは、k+Gでプロットされるべき。
            # 例えばnG=3でGlist=[0,1,-1]の場合、
            # i=0ならG=0、i=1ならG=1、i=2ならG=-1 となる。
            # このGlistの並びはextract_basisの既存ロジックに依存するため、
            # free_KEを呼び出す際のiGはGlist[i]が正しい。
            iG = Glist[i] 
            print("k=", k, "iG=", iG, "(k+iG)=", k+iG)
            yEfree[i][ik] = free_KE(k, iG)

    fig = plt.figure(figsize = figsize)
    ax1 = fig.add_subplot(1, 1, 1)

    ax1.set_xlim([-0.5, 0.5]) # プロット範囲をkmin, kmaxに合わせるべきだが、既存ロジック変更不可
    ax1.set_ylim(Erange)
    for iG_idx in range(nG): # Glistのインデックスとしてループ
        ax1.plot(xk, yE[iG_idx], linestyle = '', marker = 'o', markersize = 2.0,
                        markerfacecolor = 'none', markeredgecolor = 'black', markeredgewidth = 0.5) #, label = 'with V(x)')
        # yEfreeは Glistのインデックス順に計算されている
        ax1.plot(xk, yEfree[iG_idx], linestyle = '-', linewidth = 0.3, color = 'red')# , label = 'free e')

    ax1.set_xlabel("k ($\pi$/a)", fontsize = fontsize)
    ax1.set_ylabel("E (eV)", fontsize = fontsize)
    # 凡例は plotのlabel引数がないため表示されない。しかし既存ロジック変更不可。
    # ax1.legend(fontsize = legend_fontsize) 
    plt.tight_layout()

    plt.pause(0.1)
    print("Press ENTER to exit>>", end = '')
    input()
   
    terminate()

def ft():
    """ポテンシャルのフーリエ変換を計算し、実空間ポテンシャルとフーリエ変換されたポテンシャルをプロットする。

    詳細: 設定されたパラメータに基づいて実空間ポテンシャルを構築し、そのフーリエ変換を実行する。
          結果はMatplotlibを用いて、実空間と逆格子空間の両方で可視化される。
    :returns: None
    """
    global mode
    global a, na
    global pottype, bwidth, bpot

    astep = a / na

    print("")
    print("=== Input parameterss ===")
    print("mode:", mode)
    print("a=", a, "A")
    print("  na=", na)
    print("  astep=", astep)
    print("potential: {} w={} A  h={} eV".format(pottype, bwidth, bpot))
    print("")
    
    xpot, ypot   = build_potential(0.0, astep, na)
    xplot, yplot = build_potential(0.0, astep, 3 * na)

    xft0, yft0, xft, yft, iG, nahalf, xftstep = cal_fft(na, a, ypot) # iGはリストだが、ここでは単一変数として受け取っているように見える

    fig = plt.figure(figsize = (8, 8))
    ax1 = fig.add_subplot(2, 2, 1)
    ax2 = fig.add_subplot(2, 2, 3)
    ax3 = fig.add_subplot(2, 2, 4)

    ax1.plot(xplot, yplot, linewidth = 0.5)
    ax2.plot(xft0, yft0.real, label = "real", linewidth = 0.5)
    ax2.plot(xft0, yft0.imag, label = "imaginary", linewidth = 0.5)
    ax3.plot(xft,  yft.real, marker = 'o', markersize = 1.0, linewidth = 0.5, label = "real")
    ax3.plot(xft,  yft.imag, marker = 'o', markersize = 1.0, linewidth = 0.5, label = "imaginary")
    ax3.plot(xft,  abs(yft), marker = 'o', markersize = 1.0, linewidth = 0.5, label = "absolute")
    ax1.set_xlabel("x ($\AA$)", fontsize = fontsize)
    ax1.set_ylabel("pot (x)", fontsize = fontsize)
    ax2.set_xlabel("i", fontsize = fontsize)
    ax2.set_ylabel("FTed pot, raw x", fontsize = fontsize)
    ax3.set_xlabel("1/x, normalized (1/A)", fontsize = fontsize)
    ax3.set_ylabel("FTed pot")
    ax3.set_xlim([-10.0 * xftstep, 10.0 * xftstep])
#    ax1.legend(fontsize = legend_fontsize) # ax1にはlabel引数がないため、このlegendは表示されないが既存ロジックなので残す
    ax2.legend(fontsize = legend_fontsize)
    ax3.legend(fontsize = legend_fontsize)
    ax1.tick_params(labelsize = fontsize)
    ax2.tick_params(labelsize = fontsize)
    ax3.tick_params(labelsize = fontsize)
    plt.tight_layout()

    plt.pause(0.1)
    print("Press ENTER to exit>>", end = '')
    input()

    terminate()

def main():
    """スクリプトの主要な実行ロジック。

    詳細: グローバル変数 `mode` に基づいて、`ft()`, `band()`, `wf()` のいずれかの関数を呼び出す。
          無効なモードが指定された場合はエラーメッセージと共に終了する。
    :returns: None
    """
    global mode
    
    if mode == 'ft':
        ft()
    elif mode == 'band':
        band()
    elif mode == 'wf':
        wf()
    else:
        terminate("Error: Invalid mode [{}]".format(mode))


if __name__ == "__main__":
    main()