splitstep1d.py

"""
Single particle quantum mechanics in 1D using the split-operator method.

References:
https://www.algorithm-archive.org/contents/
split-operator_method/split-operator_method.html

https://en.wikipedia.org/wiki/Split-step_method

"""
import numpy as np
import scipy.constants as const
import matplotlib.pyplot as plt
import matplotlib.animation as animation


# Constants
N = 1024 # Number of points to use
L = 1e-8 # Extent of simulation (in meters)
X = L*np.linspace(-0.5, 0.5 - 1.0/N, N)
DX = X[1] - X[0] # Spatial step size
DT = 1e-17 # Time step size


class TimeEvolutionOperator:

    def __init__(self, potential):
        self.V: np.ndarray = potential
        # self.V[0] = 1e-12
        # self.V[-1] = 1e-12
        self._exp_potential: np.ndarray = np.zeros([N])
        self._exp_kinetic: np.ndarray = np.zeros([N])
        self._norm = False
        self.set_timestep(DT)

    def normalize_at_each_step(self, conditional):
        self._norm = conditional

    def set_timestep(self, timestep):
        self._exp_potential = np.exp(-0.25j*timestep*self.V/const.hbar)
        p = np.fft.fftfreq(N)*N*2.0*np.pi*const.hbar/L
        self._exp_kinetic = np.exp(-0.5j*timestep*p**2/
                                   (2*const.m_e*const.hbar))

    def __mul__(self, psi):
        psi_p = np.fft.fft(psi*self._exp_potential)
        psi_p *= self._exp_kinetic
        psi = np.fft.ifft(psi_p)*self._exp_potential
        if self._norm:
            psi *= 1.0/np.sqrt(np.dot(psi, np.conj(psi)))
        return psi


# The wavefunction
k = 0.0 # Modify this value to change the initial momentum of the wavefunction
sigma = 0.056568
psi = np.exp(-((X/L+0.25)/sigma)**2/2.0)*np.exp(2.0j*np.pi*k*X/L)
norm_factor = 1.0/np.sqrt(np.dot(psi, np.conj(psi)))
psi *= norm_factor

# The potential
V = 6*1e-18*(X/L)**2 # Simple Harmonic Oscillator
# V = np.zeros([N]) # Zero potential everywhere (except at boundaries)
# V = 0.9*1e-18*np.array([1.0 if i >= 0.48*N and i < 0.52*N else 0.0 
#                         for i in range(N)]) # Potential Barrier
# V = 0.9*1e-18*np.array([1.0 if i <= 0.2*N or i > 0.8*N else 0.0 
#                         for i in range(N)]) # Potential Barrier

fig = plt.figure()
ax = fig.add_subplot(1, 1, 1)
ax.set_xlim(X[0], X[-1])
ax.set_ylim(-6.0, 6.0)
ax.set_xlabel('Position (m)')
ax.set_ylabel('Potential Energy (eV)')
data = {'psi': psi, 't': 0.0}
real_plot, = ax.plot(X, np.real(psi), label=r'Re($\psi(x)$)')
imag_plot, = ax.plot(X, np.imag(psi), label=r'Im($\psi(x)$)')
abs_plot, = ax.plot(X, np.abs(psi), color='black', linewidth=1.0, 
                    label=r'|$\psi(x)$|')
potential_plot, = ax.plot(X, V/const.electron_volt, color='gray', 
                          label='Potential V(x)', linewidth=2.0)
U = TimeEvolutionOperator(V)
# U.normalize_at_each_step(True)
# U.set_timestep(-DT*1.0j)

def animation_func(*args):
    psi = data['psi']
    psi = U*psi
    psi_scaled = psi*40.0
    real_plot.set_ydata(np.real(psi_scaled))
    imag_plot.set_ydata(np.imag(psi_scaled))
    abs_plot.set_ydata(np.abs(psi_scaled))
    data['psi'] = psi
    data['t'] += DT
    return real_plot, imag_plot, abs_plot

ani = animation.FuncAnimation(fig, animation_func, blit=True, interval=1.0)
plt.grid(linestyle='--', linewidth=0.5)
plt.legend()
plt.show()
plt.close()
print(f"Total time elapsed: {data['t']}s")

splitstep2d.py

"""
Single particle quantum mechanics in 2D using the split-operator method.

References:
https://www.algorithm-archive.org/contents/
split-operator_method/split-operator_method.html

https://en.wikipedia.org/wiki/Split-step_method

"""
import numpy as np
import scipy.constants as const
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from numba import jit, c16
from time import perf_counter
import os


# Constants
N = 256 # Number of points to use
L = 1e-8 # Extent of simulation (in meters)
X, Y = np.meshgrid(L*np.linspace(-0.5, 0.5 - 1.0/N, N),
                   L*np.linspace(-0.5, 0.5 - 1.0/N, N))
DX = X[1] - X[0] # Spatial step size
DT = 5e-17


@jit('c16[:,:](c16[:,:], c16[:,:])', 
     parallel=True, nogil=True, nopython=True)
def mul2d(a, b):
    return a*b


def norm(v):
    return v/np.sqrt(np.sum(v*np.conj(v)))


class TimeEvolutionOperator:

    def __init__(self, potential):
        self.V = potential
        self._exp_potential = np.zeros([N])
        self._exp_kinetic = np.zeros([N])
        self._norm = False
        self._dt = 0
        self.set_timestep(DT)

    def normalize_at_each_step(self, conditional):
        self._norm = conditional

    def set_timestep(self, timestep):
        self._dt = timestep
        self._exp_potential = np.exp(-0.25j*self._dt*self.V/const.hbar)
        fx, fy = np.meshgrid(np.fft.fftfreq(N),np.fft.fftfreq(N))
        p = np.sqrt(fx**2 + fy**2)*N*2.0*np.pi*const.hbar/L
        self._exp_kinetic = np.exp(-0.5j*self._dt*p**2/
                                   (2*const.m_e*const.hbar))

    def __mul__(self, psi):
        psi_p = np.fft.fft2(mul2d(psi, self._exp_potential))
        psi_p = mul2d(psi_p, self._exp_kinetic)
        psi = mul2d(np.fft.ifft2(psi_p), self._exp_potential)
        if self._norm:
            psi = norm(psi)
        return psi


def complex_to_colour(values: np.ndarray, dyn_alpha=True) -> np.ndarray:
    """
    Get a colouring array given a complex array of data.

    How to map an angle on the colour wheel to a colour value:
    https://en.wikipedia.org/wiki/Hue
    https://en.wikipedia.org/wiki/File:HSV-RGB-comparison.svg

    """
    arg = np.angle(values)
    br = np.abs(values)
    coeffs = [0.99502488/2.0, 0.60809906, -0.00411985, -0.135546618]
    # This coefficients are found by Fourier transforming the hue angle to colour
    # relation and only considering the first four frequencies.
    r = sum([c*np.cos(n*arg) for n, c in enumerate(coeffs)])
    b = sum([c*np.cos(n*(2.0*np.pi/3.0 + arg)) for n, c in enumerate(coeffs)])
    g = sum([c*np.cos(n*(arg - 2.0*np.pi/3.0)) for n, c in enumerate(coeffs)])
    return np.transpose(np.array([2.0*br*r, 2.0*br*g, 2.0*br*b, 
                        br/np.amax(br) if dyn_alpha else r/r]),
                        (1, 2, 0))


# The wavefunction
k = 0.0 # Modify this value to change the initial momentum of the wavefunction
sigma = 0.056568
psi = np.exp(-((X/L+0.25)/sigma)**2/2.0
             -((Y/L-0.25)/sigma)**2/2.0)*(1.0 + 0.0j) #*np.exp(2.0j*np.pi*k*X/L)
psi = norm(psi)

# The potential
V = 6*1e-18*((X/L)**2 + (Y/L)**2) # Simple Harmonic Oscillator
U = TimeEvolutionOperator(V)

fig = plt.figure()
ax = fig.add_subplot(1, 1, 1)
br = 10.0*(N/128.0)
potential_im_data = np.transpose(np.array([V, V, V, V])/np.amax(V),
                                 (2, 1, 0))
im = ax.imshow(br*(complex_to_colour(psi)),
               extent=(X[0, 1], X[0, -1], Y[0, 0], Y[-1, 0]),
               interpolation='bilinear')
im2 = ax.imshow(potential_im_data,
                extent=(X[0, 1], X[0, -1], Y[0, 0], Y[-1, 0]),
                interpolation='bilinear')
ax.set_xlabel('x (m)')
ax.set_ylabel('y (m)')
ax.set_title('Wavefunction')
data = {'psi': psi}

def animation_func(*args):
    data['psi'] = U*data['psi']
    im.set_data(br*(complex_to_colour(data['psi'], dyn_alpha=False)))
    return im, im2

ani = animation.FuncAnimation(fig, animation_func, blit=True, interval=1.0)
plt.show()