entangle.py

"""Two particles trapped inside a 1D box.

This script is based on this article by Daniel Schroeder:

https://physics.weber.edu/schroeder/quantum/EntanglementIsntJustForSpin.pdf

as well as this accompanying interactive web page:

https://physics.weber.edu/schroeder/software/CollidingPackets.html


"""
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from scipy.signal import convolve2d


HBAR, M_E = 1.0, 1.0
NX, NY = 256, 256
LX, LY = 256.0, 256.0
DELTA_T = 0.25
DELTA_X, DELTA_Y = LX/NX, LY/NY
X_1D = np.linspace(0.0, LX - DELTA_X, NX)
Y_1D = np.linspace(0.0, LY - DELTA_Y, NY)
X, Y = np.meshgrid(X_1D, Y_1D)
V = np.exp(-0.5*(X - Y)**2/1.5**2) + 4.0*((X/LX - 0.5)**2 + (Y/LY - 0.5)**2)
# V = 1.0/np.abs(X - Y) + 4.0*((X/LX - 0.5)**2 + (Y/LY - 0.5)**2)
LAPLACIAN2D_9PT_4TH_OR = np.array([
    [0.0, 0.0, -1.0/(12.0*DELTA_Y**2), 0.0, 0.0],
    [0.0, 0.0, 4.0/(3.0*DELTA_Y**2), 0.0, 0.0],
    [-1.0/(12.0*DELTA_X**2), 4.0/(3.0*DELTA_X**2),
     -5.0/(2.0*DELTA_X**2) - 5.0/(2.0*DELTA_Y**2),
     4.0/(3.0*DELTA_X**2), -1.0/(12.0*DELTA_X**2)],
    [0.0, 0.0, 4.0/(3.0*DELTA_Y**2), 0.0, 0.0],
    [0.0, 0.0, -1.0/(12.0*DELTA_Y**2), 0.0, 0.0],])


def hamiltonian(psi: np.ndarray) -> np.ndarray:
    return -HBAR**2/(2.0*M_E)*\
        convolve2d(psi, LAPLACIAN2D_9PT_4TH_OR, mode='same', boundary='wrap')\
        + V*psi


def time_step(psi0: np.ndarray, psi1: np.ndarray, dt: float) -> np.ndarray:
    return psi0 - 1.0j*dt/HBAR*hamiltonian(psi1)


def normalize(psi: np.ndarray) -> np.ndarray:
    return psi/np.sqrt(np.sum(np.abs(psi)**2))


def compute_purity(psi_t: np.ndarray) -> float:
    psi = normalize(psi_t)
    trace_x_psi_density = np.einsum('ij,kj->ik', psi, np.conj(psi))
    # return np.sum(np.abs(trace_x_psi_density)**2)
    return np.real(np.einsum('ik,ki->', *2*[trace_x_psi_density]))


def simulation():
    psi0 = (np.exp(-0.25*(X/LX - 0.25)**2/0.05**2)
            *np.exp(-0.25*(Y/LY - 0.75)**2/0.05**2)
            )*np.exp(2.0j*np.pi*(15.0*X/LX - 2.0*Y/LY))
    psi1 = time_step(psi0, psi0, DELTA_T/2.0)
    psi = [psi0, psi1, psi1]
    fig = plt.figure()
    ax1, ax2 = fig.subplots(1, 2)
    line_x, = ax1.plot(X_1D,
                       np.sum(np.abs(normalize(psi[0]))**2/DELTA_X, axis=0),
                       label='Particle 1')
    line_y, = ax1.plot(Y_1D,
                       np.sum(np.abs(normalize(psi[0]))**2/DELTA_Y, axis=1),
                       label='Particle 2')
    ax1.set_xlim(min(X_1D[0], Y_1D[0]), max(X_1D[-1], Y_1D[-1]))
    ax1.set_xlabel('x (a.u.)')
    ax1.set_ylabel('Probability density')
    ax1.legend(title='Probability density of')
    im = ax2.imshow(np.flip(np.abs(psi0), axis=0), 
                    extent=(0, LX, 0, LY), interpolation='bilinear')
    ax2.set_title('$|\psi(x_1, x_2)|^2$')
    ax2.set_xlabel('$x_1$ (a.u.)')
    ax2.set_ylabel('$x_2$ (a.u.)')
    ax2.set_xlim(X_1D[0], X_1D[-1])
    ax2.set_ylim(Y_1D[0], Y_1D[-1])


    def animation_func(*_args):
        for _ in range(6):
            psi[2] = time_step(psi[0], psi[1], DELTA_T)
            psi[2], psi[0], psi[1] = psi
        psi_t = psi[0]
        # psi_t = psi_t - psi_t.T
        print(compute_purity(psi_t))
        im.set_data(np.flip(np.abs(psi_t), axis=0))
        line_x.set_ydata(np.sum(np.abs(normalize(psi_t))**2/DELTA_X, axis=0))
        line_y.set_ydata(np.sum(np.abs(normalize(psi_t))**2/DELTA_Y, axis=1))
        return [im, line_x, line_y]


    _ = animation.FuncAnimation(
        fig, animation_func, blit=True, interval=1000.0/60.0)
    # plt.title('Two Particles in a 1D Box')
    plt.show()


simulation()

entangle3.py

"""Three particles trapped inside a 1D box.
"""
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from scipy.ndimage import convolve


HBAR, M_E = 1.0, 1.0
NX, NY, NZ = 64, 64, 64
LX, LY, LZ = 64.0, 64.0, 64.0
DELTA_T = 0.1
DELTA_X, DELTA_Y, DELTA_Z = LX/NX, LY/NY, LZ/NZ
X_1D = np.linspace(0.0, LX - DELTA_X, NX)
Y_1D = np.linspace(0.0, LY - DELTA_Y, NY)
Z_1D = np.linspace(0.0, LZ - DELTA_Z, NZ)
X, Y, Z = np.meshgrid(X_1D, Y_1D, Z_1D)
V = (
    np.exp(-0.5*(X - Y)**2/1.5**2)
    + np.exp(-0.5*(X - Z)**2/1.5**2)
    + np.exp(-0.5*(Y - Z)**2/1.5**2)
    + 4.0*((X/LX - 0.5)**2 + (Y/LY - 0.5)**2 + (Z/LZ - 0.5)**2)
)


def get_discrete_1d_laplacian_stencil(dx: float) -> np.ndarray:
    return np.array([-1.0/12.0, 4.0/3.0, -5.0/2.0, 4.0/3.0, -1.0/12.0])/dx**2


def get_discrete_3d_laplacian_stencil() -> np.ndarray:
    stencil = np.zeros([5, 5, 5])
    stencil[2, 2, :] += get_discrete_1d_laplacian_stencil(DELTA_X)
    stencil[2, :, 2] += get_discrete_1d_laplacian_stencil(DELTA_Y)
    stencil[:, 2, 2] += get_discrete_1d_laplacian_stencil(DELTA_Z)
    return stencil


LAPLACIAN = get_discrete_3d_laplacian_stencil()


def hamiltonian(psi: np.ndarray) -> np.ndarray:
    return (
        -HBAR**2/(2.0*M_E)*convolve(psi, LAPLACIAN, mode='wrap') + V*psi)


def time_step(psi0: np.ndarray, psi1: np.ndarray, dt: float) -> np.ndarray:
    return psi0 - 1.0j*dt/HBAR*hamiltonian(psi1)


def normalize(psi: np.ndarray) -> np.ndarray:
    return psi/np.sqrt(np.sum(np.abs(psi)**2))


def compute_purities(psi: np.ndarray) -> tuple[float]:
    density_x = np.einsum('ijk,ijl->kl', psi, np.conj(psi))
    density_y = np.einsum('ijk,ilk->jl', psi, np.conj(psi))
    density_z = np.einsum('ijk,ljk->il', psi, np.conj(psi))
    return [np.sum(np.abs(density_x)**2), 
            np.sum(np.abs(density_y)**2), 
            np.sum(np.abs(density_z)**2)]


def projected_prob_density(psi: np.ndarray, d: int) -> np.ndarray:
    return np.sum(np.abs(psi)**2, axis=tuple(i for i in (0, 1, 2) if i != d))


def simulation():
    psi0 = (np.exp(-0.25*(X/LX - 0.25)**2/0.05**2)
            *np.exp(-0.25*(Y/LY - 0.75)**2/0.05**2)
            *np.exp(-0.25*(Z/LY - 0.5)**2/0.05**2)
            )*np.exp(2.0j*np.pi*(4.0*Z/LZ))
    psi0 = normalize(psi0)
    psi1 = time_step(psi0, psi0, DELTA_T/2.0)
    psi = [psi0, psi1, psi1]
    fig = plt.figure()
    ax = fig.add_subplot(1, 1, 1)
    ax.set_title("Three 1D Particles in a Box")
    line_x, = ax.plot(X_1D, projected_prob_density(psi[0], 2),
                      label='Particle 1')
    line_y, = ax.plot(Y_1D, projected_prob_density(psi[0], 1),
                      label='Particle 2')
    line_z, = ax.plot(Z_1D, projected_prob_density(psi[0], 0),
                      label='Particle 3')
    ax.set_xlim(min(X_1D[0], Y_1D[0], Z_1D[0]),
                max(X_1D[-1], Y_1D[-1], Z_1D[-1]))
    ax.set_xlabel('x')
    ax.set_ylabel('Probability density')

    def animation_func(*_args):
        for _ in range(6):
            psi[2] = time_step(psi[0], psi[1], DELTA_T)
            psi[2], psi[0], psi[1] = psi
        # print(compute_purities(psi[0]))
        line_x.set_ydata(projected_prob_density(psi[0], 2))
        line_y.set_ydata(projected_prob_density(psi[0], 1))
        line_z.set_ydata(projected_prob_density(psi[0], 0))
        return [line_x, line_y, line_z]


    _ = animation.FuncAnimation(
        fig, animation_func, blit=True, interval=1000.0/60.0)
    plt.legend()
    plt.show()

simulation()