import numpy as np

import matplotlib.pyplot as plt

from scipy.integrate import simps

L, x0, k, a, N = 40.0, 4.0, 1.0, 2.0, 2000

xs, h = np.linspace(0, L, N, retstep = True)

us = np.exp(-(xs-x0)**2/(2*a))*np.exp(1.j*k*xs)

Vs = xs[1:-1]*0

Vs[800:1000]=5

Vs1 = np.zeros(N)

Vs1[800:1000]=0.5

dt = 0.01

mu = (0. + 1.j)*dt/(2*h**2)

t, tf = 0, 5.0

hbar, m = 1, 1

E = (hbar**2/2.0/m)*(k**2 + 0.5/a)

print('wavelenth energy:', E)

A1 = np.eye(N-2)*(1.0 + 2.0*mu + 1.j*dt/2*Vs) - mu*np.eye(N-2,k=1) - mu*np.eye(N-2,k=-1)

A2 = np.eye(N-2)*(1.0 - 2.0*mu + 1.j*dt/2*Vs) + mu*np.eye(N-2,k=1) + mu*np.eye(N-2,k=-1)

B=np.dot(np.linalg.inv(A1),A2)

while(t<tf):

plt.cla()

us[1:-1] = np.dot(B,us[1:-1])

t+=dt

normus = us/np.sqrt(simps(us*np.conj(us),xs))

plt.ylim(-0.6, 0.6)

plt.plot(xs,Vs1)

plt.plot(xs, np.real(normus))

plt.plot(xs, abs(normus)**2)

plt.pause(0.001)

plt.show()
import numpy as np

import matplotlib.pyplot as plt

from scipy.integrate import simps

L, x0, k, a, N = 40.0, 4.0, 5.0, 2.0, 2000

xs, h = np.linspace(0, L, N, retstep = True)

us = np.exp(-(xs-x0)**2/a)*np.exp(1.j*k*xs)

Vs = xs[1:-1]*0.0

dt = 0.01

mu = (0. + 1.j)*dt/(2*h**2)

t, tf = 0, 5.0

A1 = np.eye(N-2)*(1.0 + 2.0*mu + 1.j*dt/2*Vs) - mu*np.eye(N-2,k=1) - mu*np.eye(N-2,k=-1)

A2 = np.eye(N-2)*(1.0 - 2.0*mu + 1.j*dt/2*Vs) + mu*np.eye(N-2,k=1) + mu*np.eye(N-2,k=-1)

B=np.dot(np.linalg.inv(A1),A2)

while(t<tf):

plt.cla()

us[1:-1] = np.dot(B,us[1:-1])

t+=dt

normus = us/np.sqrt(simps(us*np.conj(us),xs))

plt.ylim(-0.6, 0.6)

plt.plot(xs, np.real(normus))

plt.plot(xs, abs(normus)**2)

plt.pause(0.001)

plt.show()