"""
1次元シュレーディンガー方程式ソルバー
このスクリプトは、1次元の定常シュレーディンガー方程式を射撃法を用いて数値的に解きます。
原子単位系 (atomic units) を使用し、Verlet積分を用いて波動関数を計算します。
境界条件として「漸近条件 (asymptotic)」または「ゼロ条件 (zero)」を選択でき、
結果はExcelファイルに保存され、matplotlibで可視化されます。
:doc:`schrodinger1d_usage`
"""
import argparse
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
# =========================
# Potential (atomic units)
# =========================
[ドキュメント]
def V(x):
"""
ポテンシャルエネルギーを計算します。
詳細説明:
現在の実装では、調和振動子ポテンシャル V(x) = 0.5 * x^2 (Ha) を使用しています。
他のポテンシャルを使用する場合は、この関数を変更してください。
:param x: 位置 (原子単位)
:type x: float or numpy.ndarray
:returns: ポテンシャルエネルギー V(x) (Ha)
:rtype: float or numpy.ndarray
"""
return 0.5 * x**2
# return 0.5 * x**2 + 1e-2 * x**4
# =========================
# Argument initialization
# =========================
[ドキュメント]
def initialize():
"""
コマンドライン引数を初期化し、パースします。
詳細説明:
`argparse` モジュールを使用して、エネルギーの推定値、計算領域、メッシュ点数、
波動関数の発散閾値、レポート間隔、出力ファイル名、境界条件、初期微分値といった
パラメータを定義します。
:returns:
- parser (argparse.ArgumentParser): 引数パーサーオブジェクト
- args (argparse.Namespace): パースされた引数の名前空間
:rtype: tuple[argparse.ArgumentParser, argparse.Namespace]
"""
parser = argparse.ArgumentParser(
description="Solve 1D Schrodinger equation by shooting method (atomic units)"
)
parser.add_argument("--E", type=float, required=True,
help="Energy eigenvalue guess (Ha)")
parser.add_argument("--L", type=float, default=10.0,
help="Half domain size [-L, L]")
parser.add_argument("--nx", type=int, default=2000,
help="Number of mesh points")
parser.add_argument("--psi_max", type=float, default=1e10,
help="Divergence threshold for psi")
parser.add_argument("--report_step", type=int, default=200,
help="Interval for progress report")
parser.add_argument("--outfile", type=str, default="schrodinger.xlsx",
help="Output Excel filename")
# --- boundary condition ---
parser.add_argument("--bc", choices=["asymptotic", "zero"],
default="asymptotic",
help="Boundary condition at x=-L")
# used only for bc=zero
parser.add_argument("--eps", type=float, default=1e-6,
help="Initial derivative dpsi/dx at x=-L (bc=zero only)")
args = parser.parse_args()
return parser, args
# =========================
# Schrodinger solver
# =========================
[ドキュメント]
def solve_schrodinger(E, L, nx, psi_max, report_step, bc, eps):
"""
1次元シュレーディンガー方程式を射撃法で解きます。
詳細説明:
Verlet積分を用いて波動関数を数値的に伝播させます。
指定された境界条件 (`bc`) に基づいて初期条件を設定し、
計算中に波動関数が `psi_max` を超えた場合、発散と判断して計算を中断します。
計算の進行状況は `report_step` 間隔で出力されます。
:param E: エネルギー固有値の推定値 (Ha)
:type E: float
:param L: 計算領域の半分のサイズ。xは[-L, L]の範囲になります。
:type L: float
:param nx: メッシュ点の数
:type nx: int
:param psi_max: 波動関数の発散を検出するための閾値
:type psi_max: float
:param report_step: 計算の進行状況を報告するステップ間隔。0以下の場合は報告しません。
:type report_step: int
:param bc: 境界条件のタイプ ("asymptotic" または "zero")
:type bc: str
:param eps: `bc="zero"` の場合に使用される、x=-Lにおける波動関数の初期微分値
:type eps: float
:returns:
- x_full (numpy.ndarray): 計算領域のx座標の配列
- psi_full (numpy.ndarray): 計算された波動関数の配列。発散した場合は、発散点までの配列。
- info (dict): 計算結果に関する情報を含む辞書。
- "success" (bool): 積分が正常に完了したか。
- "diverged" (bool): 波動関数が発散したか。
- "diverged_x" (float or None): 波動関数が発散し始めたx座標。発散しなかった場合はNone。
- "message" (str): 計算結果に関する簡潔なメッセージ。
:rtype: tuple[numpy.ndarray, numpy.ndarray, dict]
"""
x_full = np.linspace(-L, L, nx)
dx = x_full[1] - x_full[0]
psi_full = np.zeros(nx)
# =========================
# Initial conditions
# =========================
if bc == "zero":
# ψ(-L)=0, ψ'(-L)=eps
psi_full[0] = 0.0
psi_full[1] = eps * dx
elif bc == "asymptotic":
# ψ'(-L) = κ ψ(-L)
psi0 = 1e-6
kappa = np.sqrt(2.0 * (V(-L) - E))
psi_full[0] = psi0
psi_full[1] = psi0 * (1.0 + kappa * dx)
else:
raise ValueError("Unknown boundary condition")
diverged_x = None
diverge_index = None
# =========================
# Verlet integration
# =========================
for i in range(1, nx - 1):
k = 2.0 * (V(x_full[i]) - E)
psi_full[i + 1] = (
2.0 * psi_full[i]
- psi_full[i - 1]
+ dx**2 * k * psi_full[i]
)
if abs(psi_full[i + 1]) > psi_max:
diverge_index = i + 1
diverged_x = x_full[i + 1]
break
if report_step > 0 and i % report_step == 0:
print(f"[INFO] step={i}, x={x_full[i]:.3f}, psi={psi_full[i]:.3e}")
info = {}
if diverge_index is None:
info["success"] = True
info["diverged"] = False
info["diverged_x"] = None
info["message"] = "integration completed successfully"
return x_full, psi_full, info
else:
info["success"] = False
info["diverged"] = True
info["diverged_x"] = diverged_x
info["message"] = f"diverged at x={diverged_x:.3f}"
return x_full[:diverge_index], psi_full[:diverge_index], info
# =========================
# Save & plot
# =========================
[ドキュメント]
def save_and_plot(x, psi, outfile, diverged=False, diverged_x=None):
"""
計算結果の波動関数データをExcelファイルに保存し、matplotlibでプロットします。
詳細説明:
波動関数 `psi(x)` と確率密度 `|psi(x)|^2` をExcelファイル (`outfile`) に保存します。
matplotlibを用いて `psi(x)` と `|psi(x)|^2` をグラフ表示します。
波動関数が発散した場合は、発散が始まったx座標に垂直な破線が追加され、y軸が対数スケールになります。
:param x: x座標の配列
:type x: numpy.ndarray
:param psi: 計算された波動関数の配列
:type psi: numpy.ndarray
:param outfile: データ保存先のExcelファイル名
:type outfile: str
:param diverged: 波動関数が計算中に発散したかどうかを示すフラグ。デフォルトはFalse。
:type diverged: bool
:param diverged_x: 波動関数が発散し始めたx座標。`diverged` がTrueの場合にのみ使用されます。
デフォルトはNone。
:type diverged_x: float or None
:returns: なし
:rtype: None
"""
df = pd.DataFrame({
"x": x,
"psi": psi,
"psi2": psi**2
})
df.to_excel(outfile, index=False)
print(f"[INFO] Data saved to {outfile}")
plt.figure(figsize=(8, 5))
if not diverged:
plt.plot(x, psi, label=r"$\psi(x)$")
plt.plot(x, psi**2, label=r"$|\psi(x)|^2$")
if diverged and diverged_x is not None:
plt.axvline(diverged_x, color="r", linestyle="--",
label="divergence start")
plt.xlabel("x (a.u.)")
plt.ylabel("Wavefunction")
if diverged:
plt.yscale("log")
plt.legend()
plt.grid()
plt.tight_layout()
plt.show()
# =========================
# Main
# =========================
[ドキュメント]
def main():
"""
プログラムのメイン処理を実行します。
詳細説明:
1. コマンドライン引数をパースして初期設定を行います。
2. パースされた引数を使用して `solve_schrodinger` 関数を呼び出し、
1次元シュレーディンガー方程式を解きます。
3. 計算結果のメッセージと終了点の波動関数値を出力します。
4. `save_and_plot` 関数を呼び出し、結果をExcelに保存し、グラフを表示します。
:returns: なし
:rtype: None
"""
_, args = initialize()
print("=== 1D Schrodinger Solver (atomic units) ===")
print(f"E = {args.E} Ha, L = {args.L}, nx = {args.nx}")
print(f"Boundary condition: {args.bc}")
x, psi, info = solve_schrodinger(
E=args.E,
L=args.L,
nx=args.nx,
psi_max=args.psi_max,
report_step=args.report_step,
bc=args.bc,
eps=args.eps
)
print("[RESULT]", info["message"])
print(f"psi(end) = {psi[-1]:.3e}")
save_and_plot(
x,
psi,
args.outfile,
diverged=info["diverged"],
diverged_x=info["diverged_x"]
)
if __name__ == "__main__":
main()