In [1]:
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import animation, rc
import sys
from IPython.display import HTML, Image
%matplotlib inline

m = 1
hbar = 1
sigma = 1
p0 = 30
param  = 1

def wavefunction(p,p0,sigma,t):
    return np.exp(- (p - p0) ** 2 / (4 * sigma ** 2)) * np.exp(-1.j * p ** 2 * t / (2 * m * hbar))

def inverse_fourier(t):
    p = np.fft.ifftshift(np.linspace(p0 - 3 * sigma, p0 +3 * sigma, 10000))
    y = np.fft.ifft(wavefunction(p,p0,sigma,t))
    #x = np.fft.ifftshift(np.linspace(-sigma, sigma, 10000))
    y = (1 / (2 * np.pi * sigma ** 4) ** (.25)) * np.exp(-1.j * param * sigma / hbar) * y
    return np.absolute(y) ** 2




def init():
    line.set_data([], [])
    return (line,)

def animate(i):
    x = np.fft.ifftshift(np.linspace(-sigma, sigma, 10000))
    y = inverse_fourier(i)
    line.set_data(x, y)
    return (line,)

rc('animation', html='html5')

p = np.fft.ifftshift(np.linspace(p0 - 3 * sigma, p0 +3 * sigma, 10000))

fig, ax = plt.subplots()

ax.set_xlim([-sigma,sigma]) 
ax.set_ylim([-0.001, 0.02]) 

line, = ax.plot([], [], lw=2)

anim = animation.FuncAnimation(fig, animate, np.arange(0, 60), init_func=init,
                               interval=200, blit=True)



In [17]:
anim.save('animation3.gif', writer='imagemagick', fps=30)

In [1]:
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import animation, rc
import sys
from IPython.display import HTML, Image


%matplotlib inline

m = 1
hbar = 1
sigma = 3
p0 = 100
c = 1

def wavefunction(p,p0,sigma,t):
    waves = []
    E = np.sqrt(p ** 2 * c**2  + m*m * c**4)
    waves.append(np.exp(- (p - p0) ** 2 / (4 * sigma ** 2)) * np.exp(-1.j * E * t))
    E = -E
    waves.append(np.exp(- (p - p0) ** 2 / (4 * sigma ** 2)) * np.exp(-1.j * E * t))
    return waves

def inverse_fourier(t):
    p = np.fft.ifftshift(np.linspace(p0 - 3 * sigma, p0 +3 * sigma, 10000))
    ys = []
    for wave in wavefunction(p,p0,sigma,t):
        y = np.fft.ifft(wave)
        x = np.fft.ifftshift(np.linspace(-sigma, sigma, 10000))
        ys.append((1 / (2 * np.pi * sigma ** 4) ** (.25)) * np.exp(-1.j * x * sigma / hbar) * y)
    res = 0
    for wave in ys:
        res += np.absolute(wave) ** 2
    return res


def init():
    line.set_data([], [])
    return (line,)

def animate(i):
    x = np.fft.ifftshift(np.linspace(-sigma, sigma, 10000))
    y = inverse_fourier(i)
    line.set_data(x, y)
    return (line,)

rc('animation', html='html5')

p = np.fft.ifftshift(np.linspace(p0 - 3 * sigma, p0 +3 * sigma, 10000))

fig, ax = plt.subplots()

ax.set_xlim([-sigma,sigma]) 
ax.set_ylim([-0.001, 0.02]) 

line, = ax.plot([], [], lw=2)

anim = animation.FuncAnimation(fig, animate, np.arange(0, 100), init_func=init,
                               interval=200, blit=True)



In [2]:
anim.save('animation4.gif', writer='imagemagick', fps=30)

In [ ]: