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

schrodinger1d.py をダウンロード

schrodinger1d.py
schrodinger1d.py
  1"""
  2概要: 1次元シュレーディンガー方程式ソルバー
  3
  4詳細説明:
  5    このスクリプトは、1次元の定常シュレーディンガー方程式を射撃法を用いて数値的に解きます。
  6    原子単位系 (atomic units) を使用し、Verlet積分を用いて波動関数を計算します。
  7    境界条件として「漸近条件 (asymptotic)」または「ゼロ条件 (zero)」を選択でき、
  8    結果はExcelファイルに保存され、matplotlibで可視化されます。
  9
 10関連リンク:
 11    schrodinger1d_usage
 12"""
 13import argparse
 14import numpy as np
 15import pandas as pd
 16import matplotlib.pyplot as plt
 17
 18# =========================
 19# Potential (atomic units)
 20# =========================
 21def V(x):
 22    """
 23    概要: ポテンシャルエネルギーを計算します。
 24
 25    詳細説明:
 26        現在の実装では、調和振動子ポテンシャル V(x) = 0.5 * x**2 (Ha) を使用しています。
 27        他のポテンシャルを使用する場合は、この関数を変更してください。
 28
 29    引数:
 30        :param x: 位置 (原子単位)
 31        :type x: float or numpy.ndarray
 32    戻り値:
 33        :returns: ポテンシャルエネルギー V(x) (Ha)
 34        :rtype: float or numpy.ndarray
 35    """
 36    return 0.5 * x**2
 37#    return 0.5 * x**2 + 1e-2 * x**4
 38
 39
 40# =========================
 41# Argument initialization
 42# =========================
 43def initialize():
 44    """
 45    概要: コマンドライン引数を初期化し、パースします。
 46
 47    詳細説明:
 48        argparse モジュールを使用して、エネルギーの推定値、計算領域、メッシュ点数、
 49        波動関数の発散閾値、レポート間隔、出力ファイル名、境界条件、初期微分値といった
 50        パラメータを定義します。
 51
 52    戻り値:
 53        :returns: argparse.ArgumentParserオブジェクトとargparse.Namespaceオブジェクトを含むタプルです。
 54        :rtype: tuple[argparse.ArgumentParser, argparse.Namespace]
 55    """
 56    parser = argparse.ArgumentParser(
 57        description="Solve 1D Schrodinger equation by shooting method (atomic units)"
 58    )
 59
 60    parser.add_argument("--E", type=float, required=True,
 61                        help="Energy eigenvalue guess (Ha)")
 62    parser.add_argument("--L", type=float, default=10.0,
 63                        help="Half domain size [-L, L]")
 64    parser.add_argument("--nx", type=int, default=2000,
 65                        help="Number of mesh points")
 66    parser.add_argument("--psi_max", type=float, default=1e10,
 67                        help="Divergence threshold for psi")
 68    parser.add_argument("--report_step", type=int, default=200,
 69                        help="Interval for progress report")
 70    parser.add_argument("--outfile", type=str, default="schrodinger.xlsx",
 71                        help="Output Excel filename")
 72
 73    # --- boundary condition ---
 74    parser.add_argument("--bc", choices=["asymptotic", "zero"],
 75                        default="asymptotic",
 76                        help="Boundary condition at x=-L")
 77
 78    # used only for bc=zero
 79    parser.add_argument("--eps", type=float, default=1e-6,
 80                        help="Initial derivative dpsi/dx at x=-L (bc=zero only)")
 81
 82    args = parser.parse_args()
 83    return parser, args
 84
 85
 86# =========================
 87# Schrodinger solver
 88# =========================
 89def solve_schrodinger(E, L, nx, psi_max, report_step, bc, eps):
 90    """
 91    概要: 1次元シュレーディンガー方程式を射撃法で解きます。
 92
 93    詳細説明:
 94        Verlet積分を用いて波動関数を数値的に伝播させます。
 95        指定された境界条件 (bc) に基づいて初期条件を設定し、
 96        計算中に波動関数が psi_max を超えた場合、発散と判断して計算を中断します。
 97        計算の進行状況は report_step 間隔で出力されます。
 98
 99    引数:
100        :param E: エネルギー固有値の推定値 (Ha)
101        :type E: float
102        :param L: 計算領域の半分のサイズ。xは[-L, L]の範囲になります。
103        :type L: float
104        :param nx: メッシュ点の数
105        :type nx: int
106        :param psi_max: 波動関数の発散を検出するための閾値
107        :type psi_max: float
108        :param report_step: 計算の進行状況を報告するステップ間隔。0以下の場合は報告しません。
109        :type report_step: int
110        :param bc: 境界条件のタイプ ("asymptotic" または "zero")
111        :type bc: str
112        :param eps: bc="zero" の場合に使用される、x=-Lにおける波動関数の初期微分値
113        :type eps: float
114    戻り値:
115        :returns:
116            x座標の配列 x_full、波動関数の配列 psi_full、および計算情報を含む辞書 info のタプルを返します。
117            info辞書には、"success" (bool)、"diverged" (bool)、"diverged_x" (float or None)、
118            "message" (str) のキーが含まれています。
119        :rtype: tuple[numpy.ndarray, numpy.ndarray, dict]
120    """
121    x_full = np.linspace(-L, L, nx)
122    dx = x_full[1] - x_full[0]
123
124    psi_full = np.zeros(nx)
125
126    # =========================
127    # Initial conditions
128    # =========================
129    if bc == "zero":
130        # ψ(-L)=0, ψ'(-L)=eps
131        psi_full[0] = 0.0
132        psi_full[1] = eps * dx
133
134    elif bc == "asymptotic":
135        # ψ'(-L) = κ ψ(-L)
136        psi0 = 1e-6
137        kappa = np.sqrt(2.0 * (V(-L) - E))
138
139        psi_full[0] = psi0
140        psi_full[1] = psi0 * (1.0 + kappa * dx)
141
142    else:
143        raise ValueError("Unknown boundary condition")
144
145    diverged_x = None
146    diverge_index = None
147
148    # =========================
149    # Verlet integration
150    # =========================
151    for i in range(1, nx - 1):
152        k = 2.0 * (V(x_full[i]) - E)
153        psi_full[i + 1] = (
154            2.0 * psi_full[i]
155            - psi_full[i - 1]
156            + dx**2 * k * psi_full[i]
157        )
158
159        if abs(psi_full[i + 1]) > psi_max:
160            diverge_index = i + 1
161            diverged_x = x_full[i + 1]
162            break
163
164        if report_step > 0 and i % report_step == 0:
165            print(f"[INFO] step={i}, x={x_full[i]:.3f}, psi={psi_full[i]:.3e}")
166
167    info = {}
168
169    if diverge_index is None:
170        info["success"] = True
171        info["diverged"] = False
172        info["diverged_x"] = None
173        info["message"] = "integration completed successfully"
174        return x_full, psi_full, info
175    else:
176        info["success"] = False
177        info["diverged"] = True
178        info["diverged_x"] = diverged_x
179        info["message"] = f"diverged at x={diverged_x:.3f}"
180        return x_full[:diverge_index], psi_full[:diverge_index], info
181
182
183# =========================
184# Save & plot
185# =========================
186def save_and_plot(x, psi, outfile, diverged=False, diverged_x=None):
187    """
188    概要: 計算結果の波動関数データをExcelファイルに保存し、matplotlibでプロットします。
189
190    詳細説明:
191        波動関数 psi と確率密度 abs(psi)**2 をExcelファイル (outfile) に保存します。
192        matplotlibを用いて psi と abs(psi)**2 をグラフ表示します。
193        波動関数が発散した場合は、発散が始まったx座標に垂直な破線が追加され、y軸が対数スケールになります。
194
195    引数:
196        :param x: x座標の配列
197        :type x: numpy.ndarray
198        :param psi: 計算された波動関数の配列
199        :type psi: numpy.ndarray
200        :param outfile: データ保存先のExcelファイル名
201        :type outfile: str
202        :param diverged: 波動関数が計算中に発散したかどうかを示すフラグ。デフォルトはFalse。
203        :type diverged: bool
204        :param diverged_x: 波動関数が発散し始めたx座標。diverged がTrueの場合にのみ使用されます。
205                           デフォルトはNone。
206        :type diverged_x: float or None
207    戻り値:
208        :returns: なし
209        :rtype: None
210    """
211    df = pd.DataFrame({
212        "x": x,
213        "psi": psi,
214        "psi2": psi**2
215    })
216    df.to_excel(outfile, index=False)
217    print(f"[INFO] Data saved to {outfile}")
218
219    plt.figure(figsize=(8, 5))
220    if not diverged:
221        plt.plot(x, psi, label=r"$\psi(x)$")
222    plt.plot(x, psi**2, label=r"$|\psi(x)|^2$")
223
224    if diverged and diverged_x is not None:
225        plt.axvline(diverged_x, color="r", linestyle="--",
226                    label="divergence start")
227
228    plt.xlabel("x (a.u.)")
229    plt.ylabel("Wavefunction")
230    if diverged:
231        plt.yscale("log")
232    plt.legend()
233    plt.grid()
234    plt.tight_layout()
235    plt.show()
236
237
238# =========================
239# Main
240# =========================
241def main():
242    """
243    概要: プログラムのメイン処理を実行します。
244
245    詳細説明:
246        本関数は、まずコマンドライン引数をパースして初期設定を行います。
247        次に、パースされた引数を使用して solve_schrodinger 関数を呼び出し、
248        1次元シュレーディンガー方程式を解きます。
249        その後、計算結果のメッセージと終了点の波動関数値を出力し、
250        最後に save_and_plot 関数を呼び出して結果をExcelに保存し、グラフを表示します。
251
252    戻り値:
253        :returns: なし
254        :rtype: None
255    """
256    _, args = initialize()
257
258    print("=== 1D Schrodinger Solver (atomic units) ===")
259    print(f"E = {args.E} Ha,  L = {args.L},  nx = {args.nx}")
260    print(f"Boundary condition: {args.bc}")
261
262    x, psi, info = solve_schrodinger(
263        E=args.E,
264        L=args.L,
265        nx=args.nx,
266        psi_max=args.psi_max,
267        report_step=args.report_step,
268        bc=args.bc,
269        eps=args.eps
270    )
271
272    print("[RESULT]", info["message"])
273    print(f"psi(end) = {psi[-1]:.3e}")
274
275    save_and_plot(
276        x,
277        psi,
278        args.outfile,
279        diverged=info["diverged"],
280        diverged_x=info["diverged_x"]
281    )
282
283
284if __name__ == "__main__":
285    main()