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

refineE_schrodinger1d.py をダウンロード

refineE_schrodinger1d.py
refineE_schrodinger1d.py
  1"""
  2概要: シュレーディンガー方程式の固有エネルギーをBrent法で精密化するスクリプト。
  3
  4詳細説明:
  5本スクリプトは、schrodinger1dモジュールを使用してシュレーディンガー方程式を解き、
  6その結果に基づいて境界条件が満たされない残差関数を定義します。
  7与えられた初期エネルギー値から出発し、残差関数の符号反転区間を自動的に探索します。
  8その後、SciPyライブラリのBrent法 (scipy.optimize.brentq) を利用して、
  9残差がゼロとなるエネルギー（固有エネルギー）を高精度で求めます。
 10精密化された固有エネルギーと計算の詳細は、refineE.csvファイルに記録されます。
 11
 12関連リンク:
 13refineE_schrodinger1d_usage
 14"""
 15import os
 16import argparse
 17import numpy as np
 18from scipy.optimize import brentq
 19from schrodinger1d import solve_schrodinger
 20
 21# ----------------------------------------
 22# residual(E): ψ(L) を返す関数
 23# ----------------------------------------
 24def residual(E, args, print_level = 0):
 25    """
 26    概要: シュレーディンガー方程式を解き、波動関数の右端での値 (psi(L)) を返す。
 27
 28    詳細説明:
 29    schrodinger1d.solve_schrodinger 関数を呼び出し、与えられたエネルギー E
 30    に対するシュレーディンガー方程式の波動関数を計算する。
 31    この関数は、Brent法などの根探索アルゴリズムで使用される残差関数として機能し、
 32    波動関数の右端 (x=L) での値がゼロになるエネルギーを探索するために用いられる。
 33
 34    引数:
 35        :param E: 試行するエネルギー値。
 36        :type E: float
 37        :param args: main 関数で定義されたコマンドライン引数を格納するオブジェクト。
 38                     シュレーディンガー方程式のパラメータ (L, nx, psi_max, bc, eps) を含む。
 39        :type args: argparse.Namespace
 40        :param print_level: 0以外の場合、計算されたエネルギーと psi(L) の値を出力する。
 41        :type print_level: int
 42    戻り値:
 43        :returns: 計算された波動関数の右端 (x=L) での値。
 44        :rtype: float
 45    """
 46    x, psi, info = solve_schrodinger(
 47        E=E,
 48        L=args.L,
 49        nx=args.nx,
 50        psi_max=args.psi_max,
 51        report_step=0,
 52        bc=args.bc,
 53        eps=args.eps
 54    )
 55    
 56    if print_level:
 57        print(f"E={E:12.8g}  psi(L)={psi[-1]:14.6e}")
 58
 59    return psi[-1]  # ψ(L)
 60
 61
 62# ----------------------------------------
 63# 初期値 E0 から符号反転区間を自動探索
 64# ----------------------------------------
 65def find_bracket(E0, args, max_expand=20, direction='both'):
 66    """
 67    概要: 与えられた初期エネルギー E0 を中心に、残差関数の符号が反転するエネルギー区間を探索する。
 68
 69    詳細説明:
 70    E0 を出発点とし、初期の delta 値を用いて E0 - delta と E0 + delta
 71    の範囲で residual 関数の符号をチェックする。
 72    もし符号反転が見つからない場合、delta を2倍にして探索範囲を広げ、
 73    max_expand 回数までこのプロセスを繰り返す。
 74    符号反転区間が見つかった場合、その下限 E_low と上限 E_high を返す。
 75
 76    引数:
 77        :param E0: 探索を開始する初期エネルギー値。
 78        :type E0: float
 79        :param args: residual 関数に渡すシュレーディンガー方程式のパラメータ。
 80        :type args: argparse.Namespace
 81        :param max_expand: 探索範囲を拡大する最大回数。この回数を超えても符号反転区間が見つからない場合はエラーとなる。
 82        :type max_expand: int
 83        :param direction: 探索する方向 ('upper', 'lower', 'both')。
 84                          'upper': E0より大きい側のみ探索。
 85                          'lower': E0より小さい側のみ探索。
 86                          'both': E0を挟んで両側を探索。
 87        :type direction: str
 88    戻り値:
 89        :returns: 符号反転区間の下限 E_low と上限 E_high のタプル。
 90        :rtype: tuple[float, float]
 91    例外:
 92        :raises ValueError: direction が 'upper', 'lower', 'both' のいずれでもない場合。
 93        :raises RuntimeError: max_expand 回数内に符号反転区間が見つからなかった場合。
 94    """
 95    delta = 0.05 * abs(E0) + 0.01
 96
 97    # 初期値の符号
 98    f0 = residual(E0, args)
 99
100    for _ in range(max_expand):
101
102        if direction == 'upper':
103            E_low  = E0
104            E_high = E0 + delta
105
106        elif direction == 'lower':
107            E_low  = E0 - delta
108            E_high = E0
109
110        elif direction == 'both':
111            E_low  = E0 - delta
112            E_high = E0 + delta
113
114        else:
115            raise ValueError("direction must be 'upper', 'lower', or 'both'")
116
117        f_low  = residual(E_low, args)
118        f_high = residual(E_high, args)
119
120        if f_low * f_high < 0:
121            print(f"[INFO] Found bracket: [{E_low}, {E_high}]")
122            return E_low, E_high
123
124        delta *= 2
125
126    raise RuntimeError("Could not find sign change interval.")
127
128# ----------------------------------------
129# Brent 法の maxiter を自動計算
130# ----------------------------------------
131def compute_maxiter(E_low, E_high, E_tol):
132    """
133    概要: Brent法の最大反復回数を自動的に計算する。
134
135    詳細説明:
136    Brent法は、探索区間の幅を指数関数的に減少させる。
137    この関数は、初期区間 [E_low, E_high] の幅と目標精度 E_tol に基づいて、
138    必要な反復回数を推定する。これにより、不要な計算を避けつつ、
139    十分な精度に達するための反復回数を設定できる。
140    経験的に5回の余裕を持たせている。
141
142    引数:
143        :param E_low: 符号反転区間の下限。
144        :type E_low: float
145        :param E_high: 符号反転区間の上限。
146        :type E_high: float
147        :param E_tol: ターゲットとするエネルギーの許容誤差。
148        :type E_tol: float
149    戻り値:
150        :returns: Brent法に推奨される最大反復回数。
151        :rtype: int
152    """
153    width = abs(E_high - E_low)
154    return int(np.ceil(np.log2(width / E_tol))) + 5
155
156
157# ----------------------------------------
158# main
159# ----------------------------------------
160def main():
161    """
162    概要: コマンドライン引数を受け取り、Brent法を用いてシュレーディンガー方程式の固有エネルギーを精密化する。
163
164    詳細説明:
165    このメイン関数は以下の手順で動作する:
166    1.  argparse を用いてコマンドライン引数を解析し、初期エネルギー E0、許容誤差 E_tol、
167        およびシュレーディンガー方程式の物理パラメータなどを取得する。
168    2.  find_bracket 関数を呼び出し、与えられた初期エネルギー args.E を含む符号反転区間
169        [E_low, E_high] を探索する。
170    3.  もし maxiter がコマンドラインで指定されていない場合、compute_maxiter 関数を用いて、
171        目標精度に基づいて自動的に最大反復回数を設定する。
172    4.  scipy.optimize.brentq 関数を使用して、residual 関数がゼロとなるエネルギー値
173        （すなわち固有エネルギー）を E_low と E_high の間で精密に探索する。
174        探索中には residual 関数の呼び出しごとに途中経過が出力される。
175    5.  精密化された固有エネルギーと、その探索結果（反復回数、関数呼び出し回数など）を出力する。
176    6.  最終的な結果は refineE.csv ファイルに追記される。ファイルが存在しない場合は、
177        最初にヘッダー行が追加される。
178    """
179    parser = argparse.ArgumentParser(
180        description="Refine eigenvalue using Brent method"
181    )
182
183    parser.add_argument("--E", type=float, required=True,
184                        help="Initial guess for eigenvalue")
185    parser.add_argument("--E_tol", type=float, default=1e-6,
186                        help="Target precision for eigenvalue")
187    parser.add_argument("--maxiter", type=int, default=None,
188                        help="Max iterations for Brent method (optional)")
189
190    # schrodinger1d.py と同じ引数を流用
191    parser.add_argument("--L", type=float, default=10.0)
192    parser.add_argument("--nx", type=int, default=2000)
193    parser.add_argument("--psi_max", type=float, default=1e10)
194    parser.add_argument("--bc", choices=["asymptotic", "zero"], default="asymptotic")
195    parser.add_argument("--eps", type=float, default=1e-6)
196
197    args = parser.parse_args()
198
199    # 1. 初期値 E0 から符号反転区間を探す
200    E_low, E_high = find_bracket(args.E, args, direction = 'upper')
201
202    # 2. maxiter を自動設定
203    if args.maxiter is None:
204        args.maxiter = compute_maxiter(E_low, E_high, args.E_tol)
205        print(f"[INFO] maxiter set to {args.maxiter}")
206
207    # 3. Brent 法で固有値を精密化
208    E_root, result = brentq(
209        lambda E: residual(E, args, print_level=1),
210        E_low, E_high,
211        xtol=args.E_tol,
212        maxiter=args.maxiter,
213        full_output=True
214    )
215
216    print("====================================")
217    print(" Refined eigenvalue:")
218    print(f"   E = {E_root:.12f}")
219    print("====================================")
220    
221    # f_low, f_high を計算
222    f_low  = residual(E_low, args)
223    f_high = residual(E_high, args)
224
225    log_csv = "refineE.csv"
226    print()
227    if not os.path.exists(log_csv):
228        print(f"{log_csv} does not exist. Add header labels.")
229        with open(log_csv, "a") as fp:
230            fp.write("E0, E_low, E_high, E_root, f_low, f_high, iterations, funcalls\n")
231
232    print(f"Append the result to {log_csv}")
233    with open(log_csv, "a") as fp:
234        fp.write(
235            f"{args.E},"
236            f"{E_low},{E_high},"
237            f"{E_root},"
238            f"{f_low},{f_high},"
239            f"{result.iterations},{result.function_calls}\n"
240        )
241
242
243if __name__ == "__main__":
244    main()