#!/usr/bin/env python3 # -*- coding: utf-8 -*- """ 1D coupled oscillator wave-packet collision / scattering with argparse options and optional GIF save. Site mass: m_site[i] Bond spring constants: k2_bond[i] : spring between site i and i+1 k4_bond[i] : nonlinear spring between site i and i+1 Potential: V = sum_i [ 1/2 k2_bond[i] (u[i+1]-u[i])^2 + 1/4 k4_bond[i] (u[i+1]-u[i])^4 ] Force: Let du[i] = u[i+1] - u[i] Let g[i] = k2_bond[i] du[i] + k4_bond[i] du[i]^3 Then for an interior site j: F[j] = g[j] - g[j-1] """ import argparse import re import sys from pathlib import Path import numpy as np import matplotlib.pyplot as plt from matplotlib.animation import FuncAnimation # ============================================================ # argparse # ============================================================ def build_argparser(): p = argparse.ArgumentParser( description=( "1D coupled oscillator wave-packet scattering simulation " "with linear + quartic springs." ), formatter_class=argparse.ArgumentDefaultsHelpFormatter, ) # -------------------------- # basic simulation parameters # -------------------------- p.add_argument("--n", type=int, default=200, help="number of masses") p.add_argument("--a", type=float, default=1.0, help="lattice spacing") p.add_argument("--m0", type=float, default=1.0, help="background mass") p.add_argument("--k2", type=float, default=5.0, help="background linear spring constant") p.add_argument("--k4", type=float, default=500.0, help="background nonlinear spring constant") p.add_argument( "--method", type=str, default="verlet", choices=["euler", "verlet"], help="time-integration method", ) p.add_argument("--dt", type=float, default=0.05, help="time step") p.add_argument("--nstep", type=int, default=1000, help="number of MD/animation steps") p.add_argument("--interval", type=int, default=1, help="animation interval [ms]") p.add_argument( "--save_interval", type=int, default=10, help="interval for FFT and energy snapshots", ) p.add_argument( "--boundary", type=str, default="fixed", choices=["fixed"], help="boundary condition", ) # -------------------------- # wave packet parameters # -------------------------- p.add_argument("--A", type=float, default=0.20, help="packet amplitude") p.add_argument("--sigma", type=float, default=2.0, help="packet width") p.add_argument("--k0", type=float, default=0.1, help="central wave number") p.add_argument( "--initial_mode", type=str, default="left_only", choices=["left_only", "two_packets"], help="initial packet configuration", ) p.add_argument( "--x0_left", type=float, default=None, help="left packet center. If omitted, n*a*0.25 is used.", ) p.add_argument( "--x0_right", type=float, default=None, help="right packet center. If omitted, n*a*0.75 is used.", ) # -------------------------- # scatterer parameters # -------------------------- p.add_argument( "--scatterer_mode", type=str, default="heavy_mass", choices=[ "none", "heavy_mass", "light_mass", "soft_bond", "hard_bond", "nonlinear_bond", "barrier", ], help="scatterer type", ) p.add_argument( "--scatterer_center", type=int, default=None, help="scatterer center index. If omitted, n//2 is used.", ) # The following values reproduce the current hard-coded settings by default. p.add_argument("--mass_width", type=int, default=3, help="mass-defect width") p.add_argument("--heavy_mass_factor", type=float, default=10.0, help="heavy mass factor") p.add_argument("--light_mass_factor", type=float, default=0.2, help="light mass factor") p.add_argument("--bond_width", type=int, default=4, help="bond-defect width") p.add_argument("--soft_bond_k2_factor", type=float, default=0.2, help="soft-bond k2 factor") p.add_argument("--soft_bond_k4_factor", type=float, default=1.0, help="soft-bond k4 factor") p.add_argument("--hard_bond_k2_factor", type=float, default=5.0, help="hard-bond k2 factor") p.add_argument("--hard_bond_k4_factor", type=float, default=1.0, help="hard-bond k4 factor") p.add_argument("--nonlinear_bond_k2_factor", type=float, default=1.0, help="nonlinear-bond k2 factor") p.add_argument("--nonlinear_bond_k4_factor", type=float, default=10.0, help="nonlinear-bond k4 factor") p.add_argument("--barrier_half_width", type=int, default=3, help="barrier half width") p.add_argument("--barrier_mass", type=float, default=8.0, help="barrier mass value") p.add_argument("--barrier_k2", type=float, default=20.0, help="barrier k2 value") p.add_argument( "--barrier_k4", type=float, default=None, help="barrier k4 value. If omitted, background k4 is used.", ) # -------------------------- # output / display # -------------------------- p.add_argument( "--save", type=int, default=0, choices=[0, 1], help="save GIF with ani.save after the animation window is closed", ) p.add_argument( "--gif_fps", type=int, default=30, help="frames per second for GIF save", ) p.add_argument( "--gif_dpi", type=int, default=100, help="DPI for GIF save", ) return p def fmt_param(x): """Convert a parameter value to a filename-friendly short string.""" s = f"{x:g}" s = s.replace("-", "m").replace("+", "") s = s.replace(".", "p") return s def safe_stem(s): """Sanitize a script stem for Windows-safe output filenames.""" return re.sub(r"[^0-9A-Za-z_\-]+", "_", s).strip("_") or "md_packets" def make_gif_filename(args): stem = safe_stem(Path(sys.argv[0]).stem) return ( f"{stem}_{args.method}_{args.initial_mode}_{args.scatterer_mode}_" f"{fmt_param(args.k2)}_{fmt_param(args.k4)}.gif" ) def main(argv=None): args = build_argparser().parse_args(argv) # ============================================================ # parameters # ============================================================ n = args.n if n < 3: raise ValueError("--n must be >= 3") a = args.a m0 = args.m0 k2_0 = args.k2 k4_0 = args.k4 method = args.method dt = args.dt nstep = args.nstep interval = args.interval save_interval = args.save_interval boundary = args.boundary A = args.A sigma = args.sigma k0 = args.k0 initial_mode = args.initial_mode x0_left = args.x0_left if args.x0_left is not None else n * a * 0.25 x0_right = args.x0_right if args.x0_right is not None else n * a * 0.75 scatterer_mode = args.scatterer_mode scatterer_center = args.scatterer_center if args.scatterer_center is not None else n // 2 barrier_k4 = args.barrier_k4 if args.barrier_k4 is not None else k4_0 # ============================================================ # equilibrium positions # ============================================================ x_eq = np.arange(n) * a # ============================================================ # site and bond arrays # ============================================================ m_site = np.full(n, m0, dtype=float) # k2_bond[i], k4_bond[i] are for the bond between site i and i+1 k2_bond = np.full(n - 1, k2_0, dtype=float) k4_bond = np.full(n - 1, k4_0, dtype=float) scatterer_spans = [] # ============================================================ # scatterer helper functions # ============================================================ def add_mass_defect(center, width=1, m_factor=5.0): """ Add a mass defect. Parameters ---------- center : int Center site index. width : int Number of sites in the defect region. m_factor : float Multiplication factor for the mass. Example: m_factor=10 makes the defect sites 10 times heavier. """ i1 = max(0, int(center - width // 2)) i2 = min(n, int(i1 + width)) if i2 <= i1: return m_site[i1:i2] *= m_factor x1 = x_eq[i1] x2 = x_eq[i2 - 1] scatterer_spans.append((x1, x2, f"mass x{m_factor:g}")) def set_mass_region(i1, i2, m_value): """ Set m_site[i1:i2] = m_value. i2 is exclusive. """ i1 = max(0, int(i1)) i2 = min(n, int(i2)) if i2 <= i1: return m_site[i1:i2] = m_value x1 = x_eq[i1] x2 = x_eq[i2 - 1] scatterer_spans.append((x1, x2, f"m={m_value:g}")) def add_bond_defect(center, width=1, k2_factor=1.0, k4_factor=1.0): """ Add a bond defect. This modifies bonds, not sites. k2_bond[i] and k4_bond[i] are the spring constants between site i and site i+1. Parameters ---------- center : int Center bond index. width : int Number of bonds in the defect region. k2_factor : float Multiplication factor for k2. k4_factor : float Multiplication factor for k4. """ i1 = max(0, int(center - width // 2)) i2 = min(n - 1, int(i1 + width)) if i2 <= i1: return k2_bond[i1:i2] *= k2_factor k4_bond[i1:i2] *= k4_factor x1 = x_eq[i1] x2 = x_eq[i2] scatterer_spans.append( (x1, x2, f"k2 x{k2_factor:g}, k4 x{k4_factor:g}") ) def set_bond_region(i1, i2, k2_value=None, k4_value=None): """ Set spring constants in a bond region. Bonds i1, i1+1, ..., i2-1 are modified. Bond i connects site i and site i+1. Parameters ---------- i1, i2 : int Bond range. i2 is exclusive. k2_value : float or None If not None, set k2_bond[i1:i2] to this value. k4_value : float or None If not None, set k4_bond[i1:i2] to this value. """ i1 = max(0, int(i1)) i2 = min(n - 1, int(i2)) if i2 <= i1: return label_parts = [] if k2_value is not None: k2_bond[i1:i2] = k2_value label_parts.append(f"k2={k2_value:g}") if k4_value is not None: k4_bond[i1:i2] = k4_value label_parts.append(f"k4={k4_value:g}") if label_parts: x1 = x_eq[i1] x2 = x_eq[i2] scatterer_spans.append((x1, x2, ", ".join(label_parts))) # ============================================================ # scatterer settings # ============================================================ def setup_scatterer(): if scatterer_mode == "none": pass elif scatterer_mode == "heavy_mass": # A local heavy-mass defect near the center. # This gives reflection and transmission. add_mass_defect( center=scatterer_center, width=args.mass_width, m_factor=args.heavy_mass_factor, ) elif scatterer_mode == "light_mass": # A local light-mass defect. add_mass_defect( center=scatterer_center, width=args.mass_width, m_factor=args.light_mass_factor, ) elif scatterer_mode == "soft_bond": # A soft spring region. add_bond_defect( center=scatterer_center, width=args.bond_width, k2_factor=args.soft_bond_k2_factor, k4_factor=args.soft_bond_k4_factor, ) elif scatterer_mode == "hard_bond": # A hard spring region. add_bond_defect( center=scatterer_center, width=args.bond_width, k2_factor=args.hard_bond_k2_factor, k4_factor=args.hard_bond_k4_factor, ) elif scatterer_mode == "nonlinear_bond": # Only nonlinear spring constant is enhanced. add_bond_defect( center=scatterer_center, width=args.bond_width, k2_factor=args.nonlinear_bond_k2_factor, k4_factor=args.nonlinear_bond_k4_factor, ) elif scatterer_mode == "barrier": # A wider barrier-like region. hw = args.barrier_half_width set_mass_region(scatterer_center - hw, scatterer_center + hw + 1, m_value=args.barrier_mass) set_bond_region(scatterer_center - hw, scatterer_center + hw, k2_value=args.barrier_k2, k4_value=barrier_k4) else: raise ValueError(f"Unknown scatterer_mode: {scatterer_mode}") setup_scatterer() # ============================================================ # reference dispersion relation for the homogeneous incident region # ============================================================ def omega_k(k): """ Linear-chain dispersion relation for the background region. omega(k) = 2 sqrt(k2_0/m0) |sin(k a / 2)| This is used only to construct the initial incident packet. """ return 2.0 * np.sqrt(k2_0 / m0) * np.abs(np.sin(0.5 * k * a)) def vg_k(k): """ Group velocity d omega / d k in the background region. For 0 < |k| < pi/a: vg = sqrt(k2_0/m0) a cos(k a / 2) sign(k) """ if k == 0: return 0.0 return ( np.sqrt(k2_0 / m0) * a * np.cos(0.5 * np.abs(k) * a) * np.sign(k) ) # ============================================================ # traveling wave packet # ============================================================ def packet_u(x, x0, k): """ Initial displacement of a Gaussian wave packet. u(x,0) = A exp[-(x-x0)^2/(2 sigma^2)] cos[k(x-x0)] """ dx = x - x0 env = np.exp(-dx**2 / (2.0 * sigma**2)) return A * env * np.cos(k * dx) def packet_v(x, x0, k): """ Initial velocity of a traveling Gaussian wave packet. We assume approximately u(x,t) = A exp[-(x-x0-vg t)^2/(2 sigma^2)] cos[k(x-x0) - omega t] Therefore, at t=0, v(x,0) = du/dt = A env [ vg * dx/sigma^2 * cos(k dx) + omega * sin(k dx) ] For k > 0, the packet moves to +x. For k < 0, the packet moves to -x. """ dx = x - x0 env = np.exp(-dx**2 / (2.0 * sigma**2)) omega0 = omega_k(k) vg0 = vg_k(k) return A * env * ( vg0 * dx / sigma**2 * np.cos(k * dx) + omega0 * np.sin(k * dx) ) # ============================================================ # force calculation # ============================================================ def force(u): """ Force for nearest-neighbor linear + quartic nonlinear springs. k2_bond[i], k4_bond[i] are for the bond between site i and i+1. Fixed-end boundary: u[0] and u[-1] are fixed to zero. Forces at endpoints are set to zero. """ f = np.zeros_like(u) du = u[1:] - u[:-1] # bond force-like quantity # g[i] = dV_i / d(u[i+1]-u[i]) g = k2_bond * du + k4_bond * du**3 # interior sites f[1:-1] = g[1:] - g[:-1] # endpoints are fixed, so forces there are not used f[0] = 0.0 f[-1] = 0.0 return f def apply_boundary(u, v): """ Fixed boundary condition. """ if boundary == "fixed": u[0] = 0.0 u[-1] = 0.0 v[0] = 0.0 v[-1] = 0.0 else: raise ValueError(f"Unknown boundary condition: {boundary}") return u, v # ============================================================ # energy calculation # ============================================================ def total_energy(u, v): """ Total energy = kinetic + potential. """ kinetic = 0.5 * np.sum(m_site * v**2) du = u[1:] - u[:-1] potential2 = 0.5 * np.sum(k2_bond * du**2) potential4 = 0.25 * np.sum(k4_bond * du**4) return kinetic + potential2 + potential4 # ============================================================ # initial condition # ============================================================ def make_initial_condition(): if initial_mode == "left_only": # One packet from left to right. # This is useful for seeing scattering by a defect. u0 = packet_u(x_eq, x0_left, +k0) v0 = packet_v(x_eq, x0_left, +k0) elif initial_mode == "two_packets": # Left packet moves right: k = +k0 # Right packet moves left: k = -k0 u0 = ( packet_u(x_eq, x0_left, +k0) + packet_u(x_eq, x0_right, -k0) ) v0 = ( packet_v(x_eq, x0_left, +k0) + packet_v(x_eq, x0_right, -k0) ) else: raise ValueError(f"Unknown initial_mode: {initial_mode}") u0, v0 = apply_boundary(u0, v0) return u0, v0 u_initial, v_initial = make_initial_condition() u = u_initial.copy() v = v_initial.copy() E0 = total_energy(u, v) # ============================================================ # console output # ============================================================ print("=== Simulation settings ===") print(f"n = {n}") print(f"a = {a}") print(f"m0 = {m0}") print(f"k2_0 = {k2_0}") print(f"k4_0 = {k4_0}") print(f"dt = {dt}") print(f"nstep = {nstep}") print(f"method = {method}") print(f"boundary = {boundary}") print(f"scatterer_mode = {scatterer_mode}") print(f"initial_mode = {initial_mode}") print(f"save = {args.save}") print() print("=== Array summary ===") print(f"m_site min/max = {m_site.min():.6g} / {m_site.max():.6g}") print(f"k2_bond min/max = {k2_bond.min():.6g} / {k2_bond.max():.6g}") print(f"k4_bond min/max = {k4_bond.min():.6g} / {k4_bond.max():.6g}") print() print("=== Wave packet settings ===") print(f"A = {A}") print(f"sigma = {sigma}") print(f"k0 = {k0}") print(f"x0_left = {x0_left}") print(f"x0_right = {x0_right}") print(f"omega0 = {omega_k(k0):.8f}") print(f"vg0 = {vg_k(k0):.8f}") print() print(f"Initial energy = {E0:.12e}") print() # ============================================================ # storage for FFT and energy # ============================================================ snapshots = [] snapshot_times = [] energy_list = [] energy_times = [] animation_phase = "display" def clear_observables(): snapshots.clear() snapshot_times.clear() energy_list.clear() energy_times.clear() def reset_state(): nonlocal u, v u = u_initial.copy() v = v_initial.copy() line.set_ydata(u) title_text.set_text("1D wave packet scattering") clear_observables() # ============================================================ # one MD step # ============================================================ def step_euler(u, v): f = force(u) v_new = v + dt * f / m_site u_new = u + dt * v_new u_new, v_new = apply_boundary(u_new, v_new) return u_new, v_new def step_verlet(u, v): f = force(u) v_half = v + 0.5 * dt * f / m_site u_new = u + dt * v_half # Apply boundary before calculating new force if boundary == "fixed": u_new[0] = 0.0 u_new[-1] = 0.0 f_new = force(u_new) v_new = v_half + 0.5 * dt * f_new / m_site u_new, v_new = apply_boundary(u_new, v_new) return u_new, v_new def md_step(u, v): if method == "euler": return step_euler(u, v) elif method == "verlet": return step_verlet(u, v) else: raise ValueError(f"Unknown method: {method}") # ============================================================ # animation update # ============================================================ def update(frame): nonlocal u, v u, v = md_step(u, v) step_no = frame + 1 t = step_no * dt if save_interval > 0 and step_no % save_interval == 0: snapshots.append(u.copy()) snapshot_times.append(t) E = total_energy(u, v) energy_list.append(E) energy_times.append(t) line.set_ydata(u) if step_no % 20 == 0 or step_no == 1: E = total_energy(u, v) drift = (E - E0) / E0 title_text.set_text( f"1D wave packet scattering step={step_no}, t={t:.2f}, " f"energy drift={drift:.2e}" ) if step_no % 50 == 0 or step_no == nstep: E = total_energy(u, v) drift = (E - E0) / E0 print( f"[{animation_phase}] step={step_no:6d}/{nstep}, " f"t={t:.4g}, energy drift={drift:.3e}", flush=True, ) return line, title_text # ============================================================ # plot setup # ============================================================ fig, ax = plt.subplots(figsize=(10, 4)) line, = ax.plot(x_eq, u, lw=2) title_text = ax.set_title("1D wave packet scattering") ax.set_xlim(x_eq[0], x_eq[-1]) ax.set_ylim(-0.5, 0.5) ax.set_xlabel("x") ax.set_ylabel("u") ax.grid(True) # mark scatterer regions for x1, x2, label in scatterer_spans: ax.axvspan(x1, x2, alpha=0.20) ax.text( 0.5 * (x1 + x2), 0.43, label, ha="center", va="center", fontsize=9, rotation=90, ) ani = FuncAnimation( fig, update, frames=nstep, interval=interval, blit=False, ) plt.show() # ============================================================ # optional GIF save # ============================================================ if args.save: # FuncAnimation with a stateful update function replays frames during ani.save. # Reset first so the saved GIF starts from the same initial condition. animation_phase = "save" reset_state() gif_name = make_gif_filename(args) print(f"Saving GIF: {gif_name}") try: ani.save(gif_name, writer="pillow", fps=args.gif_fps, dpi=args.gif_dpi) except Exception: print("ERROR: failed to save GIF.") print("Please check that pillow is installed:") print(" pip install pillow") raise print(f"Saved GIF: {gif_name}") # ============================================================ # FFT visualization # ============================================================ if len(snapshots) > 1: snapshots_arr = np.array(snapshots) snapshot_times_arr = np.array(snapshot_times) # Remove spatial average before FFT snapshots0 = snapshots_arr - snapshots_arr.mean(axis=1, keepdims=True) Uk = np.fft.rfft(snapshots0, axis=1) P = np.abs(Uk)**2 k_axis = 2.0 * np.pi * np.fft.rfftfreq(n, d=a) plt.figure(figsize=(10, 5)) plt.imshow( P.T, aspect="auto", origin="lower", extent=[ snapshot_times_arr[0], snapshot_times_arr[-1], k_axis[0], k_axis[-1], ], ) plt.xlabel("time") plt.ylabel("k") plt.ylim([0, min(np.pi / a, k0 * 5)]) plt.title(r"Mode power $|u_k(t)|^2$") plt.colorbar(label="Power") plt.tight_layout() plt.show() # ============================================================ # energy drift visualization # ============================================================ if len(energy_list) > 1: energy_arr = np.array(energy_list) energy_times_arr = np.array(energy_times) plt.figure(figsize=(10, 4)) plt.plot(energy_times_arr, (energy_arr - E0) / E0) plt.xlabel("time") plt.ylabel("relative energy drift") plt.title("Energy conservation check") plt.grid(True) plt.tight_layout() plt.show() # ============================================================ # final report # ============================================================ E1 = total_energy(u, v) print() print("=== Final energy check ===") print(f"Initial energy = {E0:.12e}") print(f"Final energy = {E1:.12e}") print(f"Relative drift = {(E1 - E0) / E0:.12e}") if __name__ == "__main__": main()