adiabatic2rapid.py

import numpy as np
from scipy.fft import dst, idst
from scipy.integrate import trapezoid
import matplotlib.pyplot as plt
import matplotlib.animation as animation


# Constants
N: int = 512
X = np.arange(0, N, dtype=np.float64)
DX: float = X[1] - X[0]
DT: float = 0.5  # time step

# Energies of the stationary states for the infinite square well.
# Note that the discrete sin transformation used for getting the wave function
# in terms of the stationary states of the infinite square well does not
# include the two boundary points where the wave function vanishes: 
# hence the extra 2*DX term in the denominator.
# Formula for infinite square well energies from
# https://en.wikipedia.org/wiki/Particle_in_a_box
E = np.pi**2*np.arange(1, N+1)**2/(2.0*(float(N)+2.0*DX)**2)


def time_evolve_infinite_square_well(psi: np.ndarray, t: float) -> np.ndarray:
    """
    Get the wave function in terms of infinite square well eigenfunctions by
    doing a discrete sin transformation on it, then time evolve it using the
    infinite square well energies. Inverse discrete sin transform to get
    it back in its position representation. 
    
    This will eventually be used in the split operator method;
    see eq 2.22 and 2.23 of [1].

    [1] Antoine X., Bao W., Besse C.
        Computational methods for the dynamics of the 
        nonlinear Schrodinger/Gross-Pitaevskii equations.
        https://arxiv.org/abs/1305.1093

    """
    # Use the definition of the type I DST to get the wave function in terms
    # of the stationary states of the infinite square well. See
    # https://docs.scipy.org/doc/scipy/reference/generated/scipy.fft.dst.html
    return idst(np.exp(-1.0j*E*t)*dst(psi, type=1), type=1)


def phi(x: np.ndarray, t: float) -> np.ndarray:
    """
    The time varying potential phi(x, t).
    """
    # This function currently implements a harmonic oscillator 
    # whose minimum oscillates in time. Note that since the DST is used
    # in the integration method, the wave function must always
    # vanish at the boundaries, no matter what potential is used.
    # Formula for the harmonic oscillator Hamiltonian from
    # https://en.wikipedia.org/wiki/Quantum_harmonic_oscillator
    kt: float = 0.00001*t
    # Corresponds to omega = sqrt(2)/N
    return ((x/float(N)) - (0.5 + 0.05*np.sin(kt + 150.0*kt**2)))**2


def split_step(psi: np.ndarray,  # wave function 
               phi: 'function',  # potential
               t: float, dt: float) -> np.ndarray:
    """
    Advance the wave function in time using the Split step method.
    For a time dependent potential this follows II.1 of [1], as well
    as eq. 2.22 - 2.23 of [2].  
    
    [1] Bauke H., Keitel C. Accelerating the Fourier split operator 
        method via graphics processing units.
        https://arxiv.org/abs/1012.3911

    [2] Antoine X., Bao W., Besse C.
        Computational methods for the dynamics of the 
        nonlinear Schrodinger/Gross-Pitaevskii equations.
        https://arxiv.org/abs/1305.1093

    [3] Schloss J. The Split Operator Method - Arcane Algorithm Archive.
        https://www.algorithm-archive.org/contents/split-operator_method/
        split-operator_method.html

    """
    # Numerically integrate the time dependent potential from t to t + dt.
    int_phi_dt = trapezoid([phi(X, t + i*dt/3.0) 
                            for i in range(4)], dx=dt/3.0, axis=0)
    ux = np.exp(-0.5j*int_phi_dt)
    return ux*time_evolve_infinite_square_well(ux*psi, dt)


# Currently initialized as the ground state of the harmonic oscillator.
# Formula for the harmonic oscillator eigenstates from
# https://en.wikipedia.org/wiki/Quantum_harmonic_oscillator
psi = np.exp(-0.5*float(N)*np.sqrt(2.0)*(X/float(N) - 0.5)**2)*(0.0 + 1j)


# Now do some fancy plots and animations #####################################

fig = plt.figure(figsize=(5, 7)
                 # figsize=(6, 8)
                 )
ax, ax2 = fig.subplots(2, 1, height_ratios=[1, 0.2])
ax.set_title('Position space')
ax2.set_title('Momentum space')
VIEW_SCALE = 4.0
line1, = ax.plot(X, VIEW_SCALE*phi(X, 0.0), 
                 label='V(x, t)', color='gray', linewidth=2.0)
line2, = ax.plot(X, np.real(psi), label=r'$Re(\psi(x))$')
c_re = line2.get_color()
line3, = ax.plot(X, np.imag(psi), label=r'$Im(\psi(x))$')
c_im = line3.get_color()
line4, = ax.plot(X, np.abs(psi), label=r'$|\psi(x)|$', color='black')
freq = np.fft.fftshift(np.fft.fftfreq(N))
psi_p = np.fft.fftshift(np.fft.fft(psi))
line5, = ax2.plot(freq, np.real(psi_p), color=c_re)
line6, = ax2.plot(freq, np.imag(psi_p), color=c_im)
line7, = ax2.plot(freq, np.abs(psi_p), color='gray')
ax.set_xlim(X[0], X[-1])
ax.set_ylim(-1.2, 1.2)
ax.set_ylabel(r'$\psi(x)$')
ax.legend(loc='lower left')
ax2.set_facecolor('black')
ax2.set_ylabel(r'$\psi(p)$')
ax2.set_xlim(np.amin(freq)/4.0, np.amax(freq)/4.0)
ax2.set_xlabel('Frequency')
ax2.grid(linestyle='--', linewidth=0.5)
ax.set_xlabel(r'Position $x$ (a.u.)')
fig.tight_layout()

animation_data = {'psi': psi, 
                  'lines': [line1, line2, line3, 
                            line4, line5, line6, line7],
                  't': 0.0}

def animation_func(*_):
    lines = animation_data['lines']
    for _ in range(100):
        t = animation_data['t']
        psi = animation_data['psi']
        animation_data['psi'] = split_step(psi, phi, t, DT)
        animation_data['t'] += DT
    lines[0].set_ydata(VIEW_SCALE*phi(X, t))
    lines[1].set_ydata(np.real(psi))
    lines[2].set_ydata(np.imag(psi))
    lines[3].set_ydata(np.abs(psi))
    psi_p = np.fft.fftshift(np.fft.fft(psi))
    lines[4].set_ydata(np.real(psi_p))
    lines[5].set_ydata(np.imag(psi_p))
    lines[6].set_ydata(np.abs(psi_p))
    return lines

a = animation.FuncAnimation(fig, animation_func,
                            # frames=1800,
                            blit=True, interval=1000//60)
# a.save('animation2.gif', writer='ffmpeg', fps=60, bitrate=1800)
plt.show()
plt.close()