animation3d.py

from splitstep import SplitStepMethod
import numpy as np

# Constants (Metric Units)
N = 64  # Number of points to use
L = 1e-8  # Extent of simulation (in meters)
X, Y, Z = np.meshgrid(L*np.linspace(-0.5, 0.5 - 1.0/N, N),
                        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  # timestep in seconds


# The wavefunction
k = 0.0
sigma = 0.056568
# sigma = 3.0
e = (1.0 + 0.0j) #*np.exp(2.0j*np.pi*k*X/L)
psi1 = e*np.exp(-((X/L+0.25)/sigma)**2/2.0
                -((Y/L-0.25)/sigma)**2/2.0
                -((Z/L+0.25)/sigma)**2/2.0)
psi2 = e*np.exp(-((X/L-0.25)/sigma)**2/2.0
                -((Y/L+0.25)/sigma)**2/2.0
                -((Z/L-0.25)/sigma)**2/2.0)
psi12 = psi1  + psi2
psi = psi12/np.sqrt(np.sum(psi12*np.conj(psi12)))

V = 6*1e-18*((X/L)**2 + (Y/L)**2 + (Z/L)**2) # Simple Harmonic Oscillator
U = SplitStepMethod(V, (L, L, L), DT)

data = {'psi': psi}

# import matplotlib.pyplot as plt
# import matplotlib.animation as animation
# fig = plt.figure()
# axes = fig.subplots(1, 3)
# im_x = axes[0].imshow(np.sum(np.abs(psi), axis=0))
# im_y = axes[1].imshow(np.sum(np.abs(psi), axis=1))
# im_z = axes[2].imshow(np.sum(np.abs(psi), axis=2))

# def animation_func(*_):
#     data['psi'] = U(data['psi'])
#     im_x.set_data(np.sum(np.abs(data['psi']), axis=0))
#     im_y.set_data(np.sum(np.abs(data['psi']), axis=1))
#     im_z.set_data(np.sum(np.abs(data['psi']), axis=2))
#     return im_x, im_y, im_z

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


from mayavi import mlab
plot_data = mlab.contour3d(np.abs(psi))

@mlab.animate
def animation():
    while (1):
        for _ in range(3):
            data['psi'] = U(data['psi'])
        plot_data.mlab_source.scalars = np.abs(data['psi'])
        yield

animation()
mlab.show()

splitstep.py

"""
Single particle quantum mechanics simulation
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

"""
from typing import Union, Any, Tuple
import numpy as np
import scipy.constants as const


class SplitStepMethod:
    """
    Class for the split step method.
    """

    def __init__(self, potential: np.ndarray,
                 dimensions: Tuple[float, ...],
                 timestep: Union[float, np.complex128] = 1e-17):
        if len(potential.shape) != len(dimensions):
            raise Exception('Potential shape does not match dimensions')
        self.V = potential
        self._dim = dimensions
        self._exp_potential = None
        self._exp_kinetic = None
        self._norm = False
        self._dt = 0
        self.set_timestep(timestep)

    def set_timestep(self, timestep: Union[float, np.complex128]) -> None:
        """
        Set the timestep. It can be real or complex.
        """
        self._dt = timestep
        self._exp_potential = np.exp(-0.25j*(self._dt/const.hbar)*self.V)
        p = np.meshgrid(*[2.0*np.pi*const.hbar*np.fft.fftfreq(d)*d/self._dim[i]
                          for i, d in enumerate(self.V.shape)])
        self._exp_kinetic = np.exp(-0.5j*(self._dt/(2*const.m_e*const.hbar))
                                   * sum([p_i**2 for p_i in p]))

    def __call__(self, psi: np.ndarray) -> np.ndarray:
        """
        Step the wavefunction in time.
        """
        psi_p = np.fft.fftn(psi*self._exp_potential)
        psi_p = psi_p*self._exp_kinetic
        psi = np.fft.ifftn(psi_p)*self._exp_potential
        if self._norm:
            psi = psi/np.sqrt(np.sum(psi*np.conj(psi)))
        return psi

    def normalize_at_each_step(self, norm: bool) -> None:
        """
        Whether to normalize the wavefunction at each time step or not.
        """
        self._norm = norm


if __name__ == '__main__':

    # Do an animation.
    import matplotlib.pyplot as plt
    import matplotlib.animation as animation

    # Constants (Metric Units)
    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  # timestep in seconds

    # The wavefunction
    SIGMA = 0.056568
    wavefunc = np.exp(-((X/L+0.25)/SIGMA)**2/2.0
                      - ((Y/L-0.25)/SIGMA)**2/2.0)*(1.0 + 0.0j)
    wavefunc = wavefunc/np.sqrt(np.sum(wavefunc*np.conj(wavefunc)))

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

    fig = plt.figure()
    ax = fig.add_subplot(1, 1, 1)
    # print(np.amax(np.angle(psi)))
    max_val = np.amax(np.abs(wavefunc))
    im = ax.imshow(np.angle(X + 1.0j*Y),
                   alpha=np.abs(wavefunc)/max_val,
                   extent=(X[0, 1], X[0, -1], Y[0, 0], Y[-1, 0]),
                   interpolation='none',
                   cmap='hsv')
    potential_im_data = np.transpose(np.array([V, V, V, np.amax(V)*np.ones([N, N])])
                                     /np.amax(V), (2, 1, 0))
    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': wavefunc}

    def animation_func(*_: Tuple[Any]):
        """
        Animation function
        """
        data['psi'] = U(data['psi'])
        im.set_data(np.angle(data['psi']))
        im.set_alpha(np.abs(data['psi'])/max_val)
        return (im2, im)

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