量子的トンネル効果のシュミレーション

# This program is made by AI(claude) at 2026. 10. 28.
#  Step_Scatter_fft_momentum.py に入射波の運動量k0値を入力できるようにしたもの．

"""
階段ポテンシャルによる波束散乱 (PDF §9.4, 式 9.42) のシミュレーション
  上段: 位置の確率密度 |psi(x,t)|^2
  下段: 運動量の確率密度 |psi_hat(k,t)|^2

平均運動量 k0 (E0 = k0^2/2) と階段の高さ V0 を変えて、
  E0 > V0 (古典的には全透過) と E0 < V0 (古典的には全反射)
での反射の様子の違いを比較できる。

(使い方)
  $ conda activate myenv
  $ python Step_Scatter_k0_input.py                # 既定値 (k0=6, V0=12.5)
  $ python Step_Scatter_k0_input.py --k0 3 --v0 12.5
    実行中は k0 スライダー / k0 入力欄 / V0 スライダー で変更可能 (変更すると t=0 から再開)
"""

import argparse
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from matplotlib.widgets import Slider, TextBox, Button

# =========================================================
# コマンドライン引数
# =========================================================
parser = argparse.ArgumentParser()
parser.add_argument('--k0', type=float, default=6.0, help='平均運動量 k0 (E0=k0^2/2)')
parser.add_argument('--v0', type=float, default=12.5, help='階段の高さ V0')
parser.add_argument('--x0', type=float, default=-4.0, help='初期位置')
args, _ = parser.parse_known_args()

x0 = args.x0

# =========================================================
# 数値パラメータ
# =========================================================
emin = 0.01          # エネルギー積分の下限 (0付近の不安定性を避ける)
n_energy = 600       # エネルギー積分の分割数
k_spread = 5.0       # 積分範囲: k0 + k_spread まで (ガウス幅 ~1 に対して十分大)

x_left, x_right, dx = -12.0, 12.0, 0.05
x_grid = np.arange(x_left, x_right + dx, dx)
X = x_grid[:, np.newaxis]

k_plot_min, k_plot_max = -10.0, 10.0
n_fft = 2**14

dt, t_max = 0.01, 4.0   # アニメーションの時間刻みと周期 dt=0.02, t_max=6.0

# =========================================================
# 物理関数
# =========================================================
def k(e):
    return np.sqrt(2 * e)

def k_trans(e, v0):
    # E < V0 のとき純虚数 (指数減衰)
    return np.sqrt(2 * (e - v0 + 0j))

def reflection_coeff(e, v0):
    k1, k2 = k(e), k_trans(e, v0)
    return (k1 - k2) / (k1 + k2)

def transmission_coeff(e, v0):
    k1, k2 = k(e), k_trans(e, v0)
    return 2 * np.sqrt(k1 * k2) / (k1 + k2)

def g_energy(e, k0, x0):
    kv = k(e)
    return (1.0 / np.sqrt(kv)) * (1.0 / np.pi)**0.25 \
        * np.exp(-(kv - k0)**2 / 2.0) * np.exp(-1j * x0 * kv)

# =========================================================
# k0, V0 が決まったときの基底関数・係数をまとめて作る
# =========================================================
def build(k0, v0):
    emax = (k0 + k_spread)**2 / 2
    e_grid = np.linspace(emin, emax, n_energy)
    de = e_grid[1] - e_grid[0]
    E = e_grid[np.newaxis, :]

    K1 = k(E)
    K2 = k_trans(E, v0)
    G = g_energy(E, k0, x0)

    in_wave = (1.0 / np.sqrt(2 * np.pi * K1)) * np.exp(1j * K1 * X)
    re_wave = reflection_coeff(E, v0) * (1.0 / np.sqrt(2 * np.pi * K1)) * np.exp(-1j * K1 * X)
    tr_wave = transmission_coeff(E, v0) * (1.0 / np.sqrt(2 * np.pi * K2)) * np.exp(1j * K2 * X)
    phi_basis = np.where(X < 0, (in_wave + re_wave) * G, tr_wave * G)

    # 反射確率 P_R = ∫|R(E)|^2 |g(E)|^2 dE,  透過確率 P_T = 1 - P_R  (E0 に限った積分なので、全体の規格化 ∫|g|^2 dE で割って比をとる)
    g2 = np.abs(G[0])**2
    norm = np.sum(g2) * de
    R2 = np.abs(reflection_coeff(e_grid, v0))**2
    p_r = np.sum(R2 * g2) * de / norm
    return dict(e_grid=e_grid, de=de, phi=phi_basis, p_r=p_r, p_t=1 - p_r, k0=k0, v0=v0)

# =========================================================
# 運動量表示 (ゼロパディング FFT)
#   psi_hat(k) = 1/sqrt(2π) ∫ psi(x) e^{-ikx} dx
# =========================================================
k_fft = np.fft.fftshift(2 * np.pi * np.fft.fftfreq(n_fft, d=dx))
k_mask = (k_fft >= k_plot_min) & (k_fft <= k_plot_max)
k_axis = k_fft[k_mask]

def momentum_density(psi):
    psi_hat = np.fft.fftshift(np.fft.fft(psi, n=n_fft)) * dx / np.sqrt(2 * np.pi)
    return np.abs(psi_hat[k_mask])**2

# =========================================================
# 描画の準備
# =========================================================
state = dict(sim=build(args.k0, args.v0), t=0.0)

fig, (ax_x, ax_k) = plt.subplots(2, 1, figsize=(10, 8.5))
fig.subplots_adjust(left=0.1, right=0.95, top=0.92, bottom=0.2, hspace=0.35)

# --- 上段 ---
line_x, = ax_x.plot([], [], lw=2, color='blue', label=r'$|\psi(x,t)|^2$')
pot_line, = ax_x.plot([x_left, 0, 0, x_right], [0, 0, 0, 0], color='gray', lw=1.5,
                      label='Potential step (scaled)')
e0_line = ax_x.axhline(0, color='red', ls='--', lw=1, label='E0 (scaled)')
ax_x.set_xlim(x_left, x_right)
ax_x.set_ylim(-0.05, 1.3)
ax_x.set_xlabel('x')
ax_x.set_ylabel('Probability Density')
ax_x.legend(loc='upper left', fontsize=8)
info_text = ax_x.text(0.99, 0.97, '', transform=ax_x.transAxes, ha='right', va='top',
                      fontsize=9, family='monospace',
                      bbox=dict(boxstyle='round', fc='white', ec='gray', alpha=0.9))

# --- 下段 ---
line_k, = ax_k.plot([], [], lw=2, color='green', label=r'$|\hat\psi(k,t)|^2$')
vl_k0p = ax_k.axvline(0, color='red', ls=':', lw=1, label=r'$\pm k_0$')
vl_k0m = ax_k.axvline(0, color='red', ls=':', lw=1)
vl_v0 = ax_k.axvline(0, color='gray', ls='--', lw=1, label=r'$\sqrt{2V_0}$')
vl_kp = ax_k.axvline(0, color='purple', ls=':', lw=1, label=r"$k_0'=\sqrt{k_0^2-2V_0}$")
ax_k.set_xlim(k_plot_min, k_plot_max)
ax_k.set_ylim(-0.02, 1.0)
ax_k.set_xlabel('k')
ax_k.set_ylabel('Momentum Density')
ax_k.legend(loc='upper left', fontsize=8)

def refresh_static():
    """k0, V0 に依存する静的な要素(ポテンシャル・目印の線・説明文)を更新"""
    s = state['sim']
    k0, v0 = s['k0'], s['v0']
    E0 = k0**2 / 2
    scale = 0.8 / max(v0, E0)
    pot_line.set_ydata([0, 0, v0 * scale, v0 * scale])
    e0_line.set_ydata([E0 * scale, E0 * scale])

    vl_k0p.set_xdata([k0, k0])
    vl_k0m.set_xdata([-k0, -k0])
    vl_v0.set_xdata([np.sqrt(2 * v0)] * 2)
    if k0**2 - 2 * v0 > 0:
        kp = np.sqrt(k0**2 - 2 * v0)
        vl_kp.set_xdata([kp, kp])
        vl_kp.set_visible(True)
    else:
        vl_kp.set_visible(False)

    if E0 > v0:
        regime = 'E0 > V0 : classically full transmission'
    elif E0 < v0:
        regime = 'E0 < V0 : classically full reflection'
    else:
        regime = 'E0 = V0'
    info_text.set_text(f"k0={k0:.2f}  E0={E0:.2f}  V0={v0:.2f}\n{regime}\n"
                       f"Reflection prob. R = {s['p_r']:.3f}\n"
                       f"Transmission prob. T = {s['p_t']:.3f}")

refresh_static()

# =========================================================
# ウィジェット (k0 スライダー / k0 入力欄 / V0 スライダー / Restart)
# =========================================================
ax_sk = fig.add_axes([0.15, 0.10, 0.55, 0.03])
ax_sv = fig.add_axes([0.15, 0.05, 0.55, 0.03])
ax_tb = fig.add_axes([0.82, 0.095, 0.10, 0.04])
ax_bt = fig.add_axes([0.82, 0.04, 0.10, 0.04])

sl_k0 = Slider(ax_sk, 'k0', 1.5, 8.0, valinit=args.k0)
sl_v0 = Slider(ax_sv, 'V0', 0.5, 40.0, valinit=args.v0)
tb_k0 = TextBox(ax_tb, 'k0 input ', initial=f'{args.k0:g}')
btn = Button(ax_bt, 'Restart')

def on_param_change(_=None):
    state['sim'] = build(sl_k0.val, sl_v0.val)
    state['t'] = 0.0
    refresh_static()

def on_text_submit(text):
    try:
        val = float(text)
    except ValueError:
        return
    val = min(max(val, sl_k0.valmin), sl_k0.valmax)
    sl_k0.set_val(val)   # -> on_param_change が呼ばれる

def on_restart(_):
    state['t'] = 0.0

sl_k0.on_changed(on_param_change)
sl_v0.on_changed(on_param_change)
tb_k0.on_submit(on_text_submit)
btn.on_clicked(on_restart)

# =========================================================
# アニメーション
# =========================================================
def update(_frame):
    s = state['sim']
    t = state['t']
    # psi(x,t) = ∫ psi(E,x) g(E) e^{-iEt} dE
    psi = s['phi'] @ (np.exp(-1j * s['e_grid'] * t) * s['de'])

    line_x.set_data(x_grid, np.abs(psi)**2)
    line_k.set_data(k_axis, momentum_density(psi))
    fig.suptitle(f"Time t = {t:.2f}")

    state['t'] = t + dt
    if state['t'] > t_max:
        state['t'] = 0.0
    return line_x, line_k

ani = FuncAnimation(fig, update, frames=None, blit=False, interval=40,
                    cache_frame_data=False)

plt.show()