In [ ]:
from pylab import *
from scipy.linalg import solve_banded
from scipy.special import hermite
from scipy.misc import factorial
In [37]:
%matplotlib inline
def V(x):
return .5 * np.sum(x ** 2, axis=0)
def Psi_h(x, n, t=0):
H=hermite(n)
return 1/sqrt(2**n*factorial(n))*1/pi**0.25*exp(-x**2/2)*H(x)*exp(-1j*(n+0.5)*t)
'''
N = 256
x0, x1 = -5., 5.
x = linspace(x0, x1, N)
dx = (x1 - x0) / (N-1)
'''
dt = 1. / 64 * 1
n=1 # mean energy [( (n+1) +1/2 ) + (n+1/2) + ( (n+1) +1/2 )]/2
E_mean= 3 / 2
p_max=sqrt(2*E_mean)
dx=1/512*pi/p_max
# computational grid depends on n
x0, x1=-sqrt(n)*5, sqrt(n)*5
N=int(round((x1-x0)/dx))
x=linspace(x0, x1, N)
dx=(x1-x0)/(N-1)
Psi = Psi_h(x, 0, t=0).reshape((len(Psi_h(x, 0, t=0)), 1)).dot(Psi_h(x, 0, t=0).reshape(len(Psi_h(x, 0, t=0)), 1).T)
x = np.mgrid[0:N,0:N] * dx - 5.
dp = 2 * pi / (N * dx)
p = linspace(0, (N - 1) * dp, N)
p[p > .5 * N * dp] - N * dp
U_1 = exp(-.5 * 1j * V(x) * dt)
U_2 = exp(-1j * .5 * p ** 2 * dt)
H=zeros((3, N))
H[0, 1: ]=-0.5/dx**2
H[1, : ]=+1.0/dx**2
H[2, :-1]=-0.5/dx**2
Dp=+0.5j*dt*H
Dp[1, :]+=1
Dm=-0.5j*dt*H
Dm[1, :]+=1
for k in range(512 // 1):
Psi *= U_1
#Psi = fft(Psi)
#Psi *= U_2
#Psi = ifft(Psi)
b=Dm[1, :]*Psi
b[1:]+=Dm[0, 1:]*Psi[:-1, :]
b[:-1]+=Dm[2, :-1]*Psi[1:, :]
Psi=solve_banded((1, 1), Dp, b)
b=Dm[1, :]*Psi.T
b[1:]+=Dm[0, 1:]*Psi.T[:-1, :]
b[:-1]+=Dm[2, :-1]*Psi.T[1:, :]
Psi=solve_banded((1, 1), Dp, b).T
Psi *= U_1
plt.imshow(x, np.absolute(Psi))
In [ ]:
Dm[0, 1:].shape, Psi[:-1, :].shape
In [8]:
Out[8]:
In [18]:
n=4 # mean energy [( (n+1) +1/2 ) + (n+1/2) + ( (n+1) +1/2 )]/2
E_mean=1/2
p_max=sqrt(2*E_mean)
dx=1/512*pi/p_max
In [11]:
plot((Psi.real))
Out[11]:
In [15]:
dx
Out[15]:
In [16]:
dx=1/512*pi/p_max
In [19]:
dx
Out[19]:
In [15]:
tmp = np.arange(10) +1
In [10]:
gu = np.zeros((10,10))
In [19]:
gu[:, -1] += 1
In [24]:
tmp * np.ones(gu.shape) * gu, tmp
Out[24]:
In [27]:
(tmp * np.ones(gu.shape)).T
Out[27]:
In [5]:
w = np.mgrid[0:N,0:N] * dx - 5.
In [39]:
w[:, -1, 0] - 5.
Out[39]:
In [38]:
x[-1]
Out[38]:
In [4]:
w
Out[4]:
In [12]:
z = .5 * np.sum(w ** 2, axis=0)
z.shape
Out[12]:
In [17]:
Psi = Psi_h(x, 0, t=0).reshape(len(Psi_h(x, 0, t=0)), 1).dot(Psi_h(x, 0, t=0).reshape(len(Psi_h(x, 0, t=0)), 1).T)
In [20]:
plt.imshow(np.absolute(Psi))
Out[20]:
In [19]:
Psi
Out[19]:
In [38]:
np.mgrid[0:10,0:10]
Out[38]:
In [41]:
np.arange(10)[:, newaxis].shape
Out[41]:
In [ ]: