I am simulating spring-mass system. This is the code.
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
k = 1
m = 1
def f(x, u, t):
return -k/m*x
x_graph = []
t_graph = []
u_graph = []
y_graph = []
def func(x_0, u_0, t_0, h):
for i in range(1, 1000):
m1 = h*u_0
k1 = h*f(x_0, u_0, t_0)
m2 = h*(u_0 + 0.5*k1)
k2 = h*f(x_0+0.5*m1, u_0+0.5*k1, t_0+0.5*h)
m3 = h*(u_0 + 0.5*k2)
k3 = h*f(x_0+0.5*m2, u_0+0.5*k2, t_0+0.5*h)
m4 = h*(u_0 + k3)
k4 = h*f(x_0+m3, u_0+k3, t_0+h)
x_0 += (m1 + 2*m2 + 2*m3 + m4)/6
u_0 += (k1 + 2*k2 + 2*k3 + k4)/6
t_0 += h
x_graph.append(x_0)
t_graph.append(t_0)
u_graph.append(u_0)
y_graph.append(0)
return x_0
print(func(0, 5, np.pi, 0.01))
fig, ax = plt.subplots()
ax.set(xlim=(-5.1, 5.1), ylim=(-0.05,1))
# ax.grid()
def animate(i):
l1.set_data(x_graph[:i],y_graph[:i])
return l1,
l1, = ax.plot([],[], 'o-',markevery=[-1])
ani = animation.FuncAnimation(fig, animate,frames=len(t_graph),interval =5, blit=True)
# plt.show()
ani.save('pandemic.gif', writer='ffmpeg')
I want to have a spring at point (-5,0) which changes its position as the block moves. My current output is this
Edit1:- I have added code which includes data for x_graph, y_graph and t_graph.
Correct answer has already been provided. I am adding a different solution. This solution will suit if the spring is attached on the left to a rigid wall and the ball is moving. The spring stretches and compresses based on the distance between the wall and the ball. This looks like below.
I changed your code very slightly to do this. Please beware that some of the values are hard coded so one has to change them to suit the requirement. It is very easy to change them depending on the x_lim values.
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
k = 1
m = 1
def f(x, u, t):
return -k/m*x
x_graph = []
t_graph = []
u_graph = []
y_graph = []
def func(x_0, u_0, t_0, h):
for i in range(1, 1000):
m1 = h*u_0
k1 = h*f(x_0, u_0, t_0)
m2 = h*(u_0 + 0.5*k1)
k2 = h*f(x_0+0.5*m1, u_0+0.5*k1, t_0+0.5*h)
m3 = h*(u_0 + 0.5*k2)
k3 = h*f(x_0+0.5*m2, u_0+0.5*k2, t_0+0.5*h)
m4 = h*(u_0 + k3)
k4 = h*f(x_0+m3, u_0+k3, t_0+h)
x_0 += (m1 + 2*m2 + 2*m3 + m4)/6
u_0 += (k1 + 2*k2 + 2*k3 + k4)/6
t_0 += h
x_graph.append(x_0)
t_graph.append(t_0)
u_graph.append(u_0)
y_graph.append(0)
return x_0
print(func(0, 5, np.pi, 0.01))
fig, ax = plt.subplots()
ax.set(xlim=(-5.1, 5.1), ylim=(-1,1))
# ax.grid()
def animate(i):
l1.set_data(x_graph[:i],y_graph[:i])
sin_x = []
sin_y = []
for j in np.linspace(-5,x_graph[i],250):
sin_x.append(j)
sin_y.append(0.15*np.sin((j+5)*(2*np.pi)*10/(x_graph[i]+5)))
l2.set_data(sin_x, sin_y)
return l1,l2,
l1, = ax.plot([],[], 'o',markevery=[-1])
l2, = ax.plot([],[], '-')
ani = animation.FuncAnimation(fig, animate, frames=len(t_graph),interval =0.5, blit=True)
# plt.show()
ani.save('pandemic.gif', writer='imagemagic')
I think you are trying to simulate a mass osculating around a point.
This small edit should work
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
k = 1
m = 1
def f(x, u, t):
return -k/m*x
def func(x_0, u_0, t_0, h):
x_graph = []
t_graph = []
u_graph = []
y_graph = []
for i in range(1, 1000):
m1 = h*u_0
k1 = h*f(x_0, u_0, t_0)
m2 = h*(u_0 + 0.5*k1)
k2 = h*f(x_0+0.5*m1, u_0+0.5*k1, t_0+0.5*h)
m3 = h*(u_0 + 0.5*k2)
k3 = h*f(x_0+0.5*m2, u_0+0.5*k2, t_0+0.5*h)
m4 = h*(u_0 + k3)
k4 = h*f(x_0+m3, u_0+k3, t_0+h)
x_0 += (m1 + 2*m2 + 2*m3 + m4)/6
u_0 += (k1 + 2*k2 + 2*k3 + k4)/6
t_0 += h
x_graph.append(x_0)
t_graph.append(t_0)
u_graph.append(u_0)
y_graph.append(0)
return t_graph, x_graph, y_graph, u_graph
t, x, y, u = func(0, 5, np.pi, 0.01)
fig, ax = plt.subplots()
ax.set(xlim=(-5.1, 5.1), ylim=(-0.05,1))
# ax.grid()
def animate(i):
i, j = [0, x[i]], [0, y[i]]
l1.set_data(i, j)
return l1,
l1, = ax.plot([],[], 'o-',markevery=[-1])
l2, = ax.plot([],[], 'o-',markevery=[-1])
l3, = ax.plot([],[], 'o-',markevery=[-1])
ani = animation.FuncAnimation(fig, animate,frames=len(t),interval =5, blit=True)
fig.show()
ani.save('pandemic.mp4', writer='ffmpeg')
Related
I've got a set of points (coordinates X and Y) generated by a mathematical expression and I want to draw the resulting figure in a specific position of the screen (I'd like to determine the position in which to center the drawn figure).
I tried using the following code to test if the formula resulted in the correct figure. But now I need to draw the same contour on a pre-existing image at a specific position.
B = 185
L = 250
W = (L-B)/6
D = (L/2)-L/4
x = np.linspace(-L/2, L/2, 500)
y1 = []
y2 = []
for X in x:
termo1 = sqrt((L**2 - 4*X**2) / (L**2 + 8*W*X + 4*W**2))
termo2 = ((sqrt(5.5*L**2 + 11*L*W + 4*W**2) * (sqrt(3)*B*B - 2*D*sqrt(L**2 + 2*W*L + 4*W**2))
) / (sqrt(3)*B*L*(sqrt(5.5*L**2 + 11*L*W + 4*W**2) - 2*sqrt(L**2 + 2*W*L + 4*W**2))))
termo3 = 1 - sqrt((L*(L**2 + 8*W*X + 4*W**2)) / (2*(L - 2*W)*X**2 +
(L**2 + 8*L*W - 4*W**2)*X + 2*L*W**2 + L**2*W + L**2*W + L**3))
calculo = B/2 * termo1 * (1-termo2 * termo3)
y1.append(calculo)
calculo = -B/2 * termo1 * (1-termo2 * termo3)
y2.append(calculo)
fig = plt.figure()
ax = fig.add_subplot(1, 1, 1)
ax.spines['left'].set_position('center')
ax.spines['bottom'].set_position('zero')
ax.spines['right'].set_color('none')
ax.spines['top'].set_color('none')
ax.xaxis.set_ticks_position('bottom')
ax.yaxis.set_ticks_position('left')
plt.plot(x, y1, 'r')
plt.plot(x, y2, 'r')
plt.show()
You can do this by creating an onclick event; it will take the mouse cords when click and use them as an offset...I think that's what you are asking for? Though with the current plot x/y limits, it won't show up depending on where you click so I added those in the configuration of the plot.
import numpy as np
from math import sqrt
import matplotlib.pyplot as plt
import os
import matplotlib.image as mpimg
def onclick(event):
global ix, iy
ix, iy = event.xdata, event.ydata
plt.plot(x + ix, y1+ iy, 'r')
plt.plot(x + ix, y2+ iy, 'r')
plt.show()
fig.canvas.mpl_disconnect(cid)
return
B = 185
L = 250
W = (L-B)/6
D = (L/2)-L/4
x = np.linspace(-L/2, L/2, 500)
y1 = []
y2 = []
for X in x:
termo1 = sqrt((L**2 - 4*X**2) / (L**2 + 8*W*X + 4*W**2))
termo2 = ((sqrt(5.5*L**2 + 11*L*W + 4*W**2) * (sqrt(3)*B*B - 2*D*sqrt(L**2 + 2*W*L + 4*W**2))
) / (sqrt(3)*B*L*(sqrt(5.5*L**2 + 11*L*W + 4*W**2) - 2*sqrt(L**2 + 2*W*L + 4*W**2))))
termo3 = 1 - sqrt((L*(L**2 + 8*W*X + 4*W**2)) / (2*(L - 2*W)*X**2 +
(L**2 + 8*L*W - 4*W**2)*X + 2*L*W**2 + L**2*W + L**2*W + L**3))
calculo = B/2 * termo1 * (1-termo2 * termo3)
y1.append(calculo)
calculo = -B/2 * termo1 * (1-termo2 * termo3)
y2.append(calculo)
fig = plt.figure()
ax = fig.add_subplot(1, 1, 1)
ax.set_xlim([-500,500])
ax.set_ylim([-500,500])
ax.spines['left'].set_position('center')
ax.spines['bottom'].set_position('zero')
ax.spines['right'].set_color('none')
ax.spines['top'].set_color('none')
ax.xaxis.set_ticks_position('bottom')
ax.yaxis.set_ticks_position('left')
coords = []
cid = fig.canvas.mpl_connect('button_press_event', onclick)
plt.show()
My program evaluates error in solving a linear differential equation. It uses only numpy arrays. When I try to use numba's jit decorator for the functions I define, I just get errors. Can you please help me use it properly?
My code:
import numpy as np
from numba import jit
def rk4(t_prev, x_prev, derivs, dt):
k1 = dt * derivs(t_prev, x_prev)
k2 = dt * derivs(t_prev + 1/2*dt, x_prev + 1/2*k1)
k3 = dt * derivs(t_prev + 1/2*dt, x_prev + 1/2*k2)
k4 = dt * derivs(t_prev + dt, x_prev + k3)
x_next = x_prev + 1/6*k1 + 1/3*k2 + 1/3*k3 + 1/6*k4
return x_next
global k, x_0, v_0, t_0, t_f
k = 1
x_0 = 0
v_0 = np.sqrt(k)
t_0 = 0
t_f = 10
dtList = np.logspace(0, -5, 1000)
def derivs(t, X):
deriv = np.zeros([2])
deriv[0] = X[1]
deriv[1] = -k * X[0]
return deriv
def err(dt):
tList = np.arange(t_0, t_f + dt, dt)
N = tList.shape[0]
XList = np.zeros([N,2])
XList[0][0], XList[0][1] = x_0, v_0
for i in range(N-1):
XList[i+1] = rk4(tList[i], XList[i], derivs, dt)
error = np.abs(XList[-1][0] - np.sin(10))
return error
print(err(.001))
The following works for me:
import numpy as np
from numba import jit
#jit(nopython=True)
def rk4(t_prev, x_prev, derivs, dt):
k1 = dt * derivs(t_prev, x_prev)
k2 = dt * derivs(t_prev + 1/2*dt, x_prev + 1/2*k1)
k3 = dt * derivs(t_prev + 1/2*dt, x_prev + 1/2*k2)
k4 = dt * derivs(t_prev + dt, x_prev + k3)
x_next = x_prev + 1/6*k1 + 1/3*k2 + 1/3*k3 + 1/6*k4
return x_next
global k, x_0, v_0, t_0, t_f
k = 1
x_0 = 0
v_0 = np.sqrt(k)
t_0 = 0
t_f = 10
dtList = np.logspace(0, -5, 1000)
#jit(nopython=True)
def derivs(t, X):
deriv = np.zeros(2)
deriv[0] = X[1]
deriv[1] = -k * X[0]
return deriv
#jit(nopython=True)
def err(dt):
tList = np.arange(t_0, t_f + dt, dt)
N = tList.shape[0]
XList = np.zeros((N,2))
XList[0][0], XList[0][1] = x_0, v_0
for i in range(N-1):
XList[i+1] = rk4(tList[i], XList[i], derivs, dt)
error = np.abs(XList[-1][0] - np.sin(10))
return error
print(err(.001))
Note, the only two changes I made to your code was to replace the calls to np.zeros that passed in lists to either a tuple in the 2d case, or just the bare integer in the 1d case. See the following issue for an explanation of why this is:
https://github.com/numba/numba/issues/3993
So the purpose of my code is to use inputted data points to give a gaussian plot distribution. I figured out how to make it work with a double gaussian but I'm having a lot of trouble adding a third. Im not quite sure what I'm doing wrong. 1 of the errors I keep getting is an Index Error saying that the list index is out of range. I would appreciate any help with this.
Heres my code:
from pylab import *
import numpy as np
from numpy import loadtxt
from scipy.optimize import leastsq
from scipy.optimize import least_squares
from scipy.stats import iqr
import math
import matplotlib.pyplot as plt
import sys
matplotlib.rcParams['mathtext.default'] = 'regular'
fitfunc_triple = lambda p, x: np.abs(p[0]) * exp(-0.5 * ((x - p[1]) / p[2]) ** 2) + np.abs(p[3]) * exp(
-0.5 * ((x - p[4]) / p[5]) ** 2) + np.abs(p[6]) * exp(-0.5 * ((x - p[7])/ p[8] **2 ))
fitfunc_double = lambda p, x: np.abs(p[0]) * exp(-0.5 * ((x - p[1]) / p[2]) ** 2) + np.abs(p[3]) * exp(
-0.5 * ((x - p[4]) / p[5]) ** 2)
fitfunc_single = lambda p, x: np.abs(p[0]) * exp(-0.5 * ((x - p[1]) / p[2]) ** 2)
errfunc = lambda p, x, y: (y - fitfunc(p, x))
dataR = np.loadtxt("/Users/Safi/Library/Preferences/PyCharmCE2018.1/scratches/rspecial1385.a2261.dat5", skiprows=0)
RA = dataR[:, 0]
DEC = dataR[:, 1]
VELR = dataR[:, 2]
REDSH = dataR[:, 3]
RADD = dataR[:, 4]
sl = 3E5
zbar = np.mean(REDSH)
vc = zbar * sl
VEL = vc + sl * ((REDSH - zbar) / (1 + zbar))
wdith = 200
iters = 10
sig2 = 500
binN = int(math.ceil((np.max(VEL) - np.min(VEL)) / wdith))
sys.stdout = open(str(wdith) + "_1sigma_" + str(iters) + "_sig2_" + str(sig2) + ".txt", "w")
plt.figure(1)
y, x, _ = plt.hist(VEL, binN, alpha=0.5, label='data')
x = (x[1:] + x[:-1]) / 2 # for len(x)==len(y)
data = np.vstack((x, y)).T
xdata = data[:, 0]
ydata = data[:, 1]
yerr = ydata ** 0.5
init = [10, 69500, 1200, 5, 68000, sig2]
bds = ([0, 66000, 800, 0, 66000, sig2], [50, 70000, 1750, 20, 70000, sig2 + 0.01])
def index_outlier(data):
inter_quart = iqr(data) * 1.5
bd2 = np.percentile(data, 75) + inter_quart
bd1 = np.percentile(data, 25) - inter_quart
index = []
for i in [i for i, x in enumerate(data) if x < bd1 or x > bd2]:
index.append(i)
return (index)
#### Bootstrapping Estimation Function ####
def fit_bootstrap(fitfunc, datax, datay, init, bds, sigma, iterations=iters):
errfunc = lambda p, x, y: (y - fitfunc(p, x))
# Fit first time
pfit = least_squares(errfunc, init, bounds=bds, args=(datax, datay), max_nfev=10000)
model = fitfunc(pfit.x, datax)
residuals = pfit.fun
# Random data sets are generated and fitted
ps = []
for i in range(iterations):
randomdataY = []
for k in range(len(sigma)):
randomDelta = np.random.normal(0., sigma[k], 1)
randomdataY.append(datay[k] + randomDelta)
out = np.concatenate(randomdataY)
randomfit = least_squares(errfunc, init, bounds=bds, args=(datax, out))
ps.append(randomfit.x)
# Removing outliers
# Finding outliers and indexing them
master_list = []
indexed = []
for k in range(len(ps[0])): # 0-6
it = []
for i in range(len(ps)): # 0-1000
it.append(ps[i][k])
master_list.append(it)
# indexed.append(index_outlier(master_list[k]))
# # List of outlier indicies
# flat_list=[item for sublist in indexed for item in sublist]
# no_dups= list(set(flat_list))
# # Removing bad fits
# for k in range(len(master_list)):
# for i in sorted(no_dups,reverse=True):
# del master_list[k][i]
pfit_bootstrap = []
perr_bootstrap = []
for i in master_list:
pfit_bootstrap.append(np.median(i))
perr_pos = np.round(np.percentile(i, 84) - np.median(i), 4)
perr_neg = np.round(np.median(i) - np.percentile(i, 16), 4)
perr_bootstrap.append(str('[+') + str(perr_pos) + str(',-') + str(perr_neg) + str(']'))
return (pfit_bootstrap, perr_bootstrap, residuals, pfit.nfev, master_list)
pfit, perr, residuals, nfev, master_list = fit_bootstrap(fitfunc_double, xdata, ydata, init, bds, yerr)
pfit1, perr1, residuals1, nfev1, master_list1 = fit_bootstrap(fitfunc_single, xdata, ydata, init, bds, yerr)
more_data = np.linspace(np.min(xdata), np.max(xdata), 1000)
real_func = fitfunc_double(pfit, more_data)
real_func1 = fitfunc_single(pfit1, more_data)
######## Saving Coefficients #########
A1 = pfit[0]
m1 = pfit[1]
s1 = pfit[2]
A2 = pfit[3]
m2 = pfit[4]
s2 = pfit[5]
A3 = pfit[6]
m3 = pfit[7]
s3 = pfit[8]
pecp = VEL - vc
m1p = m1 - vc
m2p = m2 - vc
m3p = m3 - vc
xdatap = xdata - vc
plt.figure(6)
plt.hist(pecp, binN, alpha=.5, label='data', color='skyblue')
xhmax = np.amax(pecp + 1500)
xhmin = np.amin(pecp - 1500)
xh = np.linspace(xhmin, xhmax, 50)
# yh1=(mlab.normpdf(xh, c[1], c[2]))
yh1 = np.abs(A1) * exp(-0.5 * (((xh - m1p) / (s1)) ** 2))
yh2 = np.abs(A2) * exp(-0.5 * (((xh - m2p) / (s2)) ** 2))
yh3 = np.abs(A3) * exp(-0.5 * (((xh - m3p) / (s3)) ** 2))
plt.plot(xh, yh1, color='b', linewidth=2)
plt.plot(xh, yh2, color='r', linewidth=2)
plt.plot(xh, yh3, color='g', linewidth=2)
plt.plot(xh, yh1 + yh2 + yh3, color='purple', linewidth=3)
# plt.errorbar(xdatap,y,xerr=wdith/2,ls='none', yerr=yerr,color='k',linewidth=2)
# plt.plot(xdatap, ydata,'.',color='k')
plt.ylim(0, np.max(ydata) + 2)
plt.xlabel('Peculiar Velocity (km/s)')
plt.ylabel('N$_{gal}$')
plt.text(-4800, 15, '$\mu_{2}$-$\mu_{1}$ = ' + str(int(m2 - m1)) + ' km/s')
plt.savefig(str(wdith) + "_1sigma_" + str(iters) + "_sig2_" + str(sig2) + "_hist.ps")
divi = -1800
memlow = np.array([[0 for x in range(2)] for y in range(1)])
memhigh = np.array([[0 for x in range(2)] for y in range(1)])
j = 0
k = 0
plt.show()
I am studding automatic and have a matlab/SIMULINK model of nonlinear vibration. I am trying to resolve the same problem using python instead.
Here we have a schema. The black frame is a part that I have already. I also have a Model tłumnika MR (eng. Magnetorheological Damper). Just I don't know how to connect it together?
There are welcome any suggestion how to resolve this including modifying model.
Code for part 1 (black frame) [just pseudo code ]
def rms(v): # ROOT MEAN SQRT
return np.sqrt(np.mean(np.square(v)))
def transmissibility(_in, _out):
# Transmissibility (T) = output/input
return rms(_out) / rms(_in)
def q(ss, t, Ampituda, freqs):
y2 = []
for f in freqs:
u = Ampituda * np.sin(2 * np.pi * f * t) # wektor wejscia
t1, y1, x1 = signal.lsim(ss, u, t) # wyliczenie modelu w dziedzinie czasu
y2.append(transmissibility(u, y1))
return y2
vec_c = np.arange(1000, 6000 , 1000)
for c in vec_c:
A = [[0, 1], [-(k/m), -(c/m)]]
B = [[-c/m], [(k/m) - (c/m)**2]]
C = [1, 0]
D = 0
ss = signal.StateSpace(A,B,C,D) # State Space model
y2 = q(ss, t, Ampituda, freqs)
plt.plot(freqs, y2, label=r'$c = {} \frac{}{} $'.format(c, "{Ns}", "{m}"), linewidth=2.0)
plt.legend()
Code for MR (Model tłumnika MR) [copy paste example]
from scipy import signal
import numpy as np
def I(to_intagrate, dt, y0=0):
i = [y0]
for v in to_intagrate:
i.append(i[-1] + v * dt)
del i[0]
return i
def MR(v, i, dt):
""" #v -- prędkość (z' -w') wektor przyśpieszeń
#i -- natężenie prądu w amperach
return Siła generowan przez tłumnik w N
Slajd 18
http://home.agh.edu.pl/~jastrzeb/images/SSD/SSD_2DOF_v1.pdf
"""
b1 = 3415.7
b2 = 93.324
b3 = 74.487
F0 = b1 * (i**2) + b2 * i + b3
b4 = 2534.1
b5 = 19.55
b6 = 643.1
C1 = b4 * (i**2) + b5*i + b6
beta = 50
p1 = 4
p2 = 0.2
x = I(v, dt)
Ft = []
for x1, v1 in zip(x, v):
part1 = F0 * np.tanh( beta * (v1 + (p1 * x1)))
part2 = C1 * (v1 + (p2 * x1))
Ft.append(part1 + part2)
return Ft, x
import matplotlib.pyplot as plt
plt.rcParams['figure.figsize'] = (20, 16)
plt.rcParams['font.size'] = 18
for f in [2, 5]: # wybrane częstotliwości
# f = 5
i = 0.2
# x_sin wektor wartości x dla wymuszenia sinusoidalnego 201 pkt
# na każdy okres sinusa. Rozpoczęcie pkt pracy w -0.2
x_sin = np.linspace(-np.pi/2, (np.pi * f) - np.pi/2, num=201 * f)
u = np.sin(x_sin) * 0.2 # przeskalowanie przyśpieszenia
dt = 1/(f*201)
ft, x = MR(u, i, dt) # sila
plt.plot(u, ft, label='Freq = {}Hz, I={}A'.format(f, i))
plt.legend()
plt.grid()
plt.show()
Sorry that this post is long but I am trying to simulate two dimensional Schrodinger equation in python using split-step method.
One dimensional problem of this equation has been explained in this post:
https://jakevdp.github.io/blog/2012/09/05/quantum-python/
I tried adding another dimension and modified the operators following the above post but I am confused as how to plot the psi function now that I added another dimension, here is my modified class equation by adding and extra dimension, I also modified the helper functions now that I added extra dimension:
import numpy as np
from scipy.fftpack import fft, ifft
class Schrodinger(object):
def __init__(self, x,y,psi_x0,psi_y0,V_x,V_y,k0=None, hbar=1, m=1, t0=0.0):
self.x,self.y,psi_x0,psi_y0,self.V_x,V_y=map(np.asarray(x,y,psi_x0,psi_y0,V_x,V_y))
N=self.x.size
assert self.x.shape==(N,)
assert self.y.shape==(N,)
assert psi_x0.shape==(N,)
assert self.V_x.shape==(N,)
assert self.V_y.shape==(N,)
self.hbar = hbar
self.m = m
self.t = t0
self.dt_ = None
self.N = len(x)
self.dx=self.x[1]-self.x[0]
self.dy=self.y[1]-self.y[0]
self.dk = 2 * np.pi / (self.N * self.dx)
if k0 == None:
self.k0 = -0.5 * self.N * self.dk
else:
self.k0 = k0
self.k = self.k0 + self.dk * np.arange(self.N)
self.psi_x = psi_x0
self.psi_y=psi_y0
self.compute_k_from_x()
self.compute_k_from_y()
self.x_evolve_half = None
self.x_evolve = None
self.y_evolve_half = None
self.y_evolve = None
self.k_evolve = None
self.psi_x_line = None
self.psi_y_line = None
self.psi_k_line = None
self.V_x_line = None
self.V_y_line = None
def _set_psi_x(self,psi_x):
self.psi_mode_x=(psi_x*np.exp(-1j*self.k[0]*self.x)*self.dx/np.sqrt(2*np.pi))
def _get_psi_x(self):
return (self.psi_mode_x*np.exp(1j*self.k[0])*np.sqrt(2*np.pi)/self.dx)
def _set_psi_y(self,psi_y):
self.psi_mode_y=(psi_y*np.exp(-1j*self.k[0]*self.y)*self.dy/np.sqrt(2*np.pi))
def _get_psi_y(self):
return (self.psi_mode_y*np.exp(1j*self.k[0])*np.sqrt(2*np.pi)/self.dy)
def _set_psi_k_x(self,psi_k_x):
self.psi_mode_k_x=psi_k_x*np.exp(1j*self.x[0]*self.dk*np.arange(self.N))
def _get_psi_k_x(self):
return self.psi_mode_k_x*np.exp(-1j*self.x[0]*self.dk*np.arange(self.N))
def _set_psi_k_y(self,psi_k_y):
self.psi_mode_k_y=psi_k_y*np.exp(1j*self.y[0]*self.dk*np.arange(self.N))
def _get_psi_k_y(self):
return self.psi_mode_k_y*np.exp(-1j*self.y[0]*self.dk*np.arange(self.N))
def _get_dt(self):
return self.dt_
def _set_dt(self,dt):
if dt!=self.dt_:
self.dt_=dt
self.x_evolve_half=np.exp(-0.5*1j*self.V_x/self.hbar*dt)
self.y_evolve_half=np.exp(-0.5*1j*self.V_y/self.hbar*dt)
self.x_evolve=self.x_evolve_half*self.x_evolve_half
self.y_evolve=self.y_evolve_half*self.y_evolve_half
self.k_evolve=np.exp(-0.5*1j*self.hbar/self.m*(self.k*self.k)*dt)
psi_x=property(_get_psi_x,_set_psi_x)
psi_y=property(_get_psi_y,_set_psi_y)
psi_k_x=property(_get_psi_k_x,_set_psi_k_x)
psi_k_y=property(_get_psi_k_y,_set_psi_k_y)
dz=property(_get_dt,_set_dt)
def compute_k_from_x(self):
self.psi_mode_k_x=fft(self.psi_mode_x)
def compute_k_from_y(self):
self.psi_mode_k_y=fft(self.psi_mode_y)
def compute_x_from_k(self):
self.psi_mode_x=ifft(self.psi_mode_k_x)
def compute_y_from_k(self):
self.psi_mode_y=ifft(self.psi_mode_k_y)
def time_step(self,dt,Nsteps=1):
self.dt=dt
if Nsteps>0:
self.psi_mode_x*=self.x_evolve_half
self.psi_mode_y*=self.y_evolve_half
for i in range(Nsteps-1):
self.compute_k_from_x()
self.compute_k_from_y()
self.psi_mode_k_x*=self.k_evolve
self.psi_mode_k_y*=self.k_evolve
self.compute_x_from_k()
self.compute_y_from_k()
self.psi_mode_x*=self.x_evolve
self.psi_mode_y*=self.y_evolve
self.compute_k_from_x()
self.compute_k_from_y()
self.psi_mode_k_x*=self.k_evolve
self.psi_mode_k_y*=self.k_evolve
self.compute_x_from_k()
self.psi_mode_x*=self.x_evolve_half
self.compute_y_from_k()
self.psi_mode_y*=self.y_evolve_half
self.compute_k_from_x()
self.compute_k_from_y()
self.t+=dt*Nsteps
def gauss_x(x, a, x0, k0):
return ((a*np.sqrt(np.pi))**(-0.5)* np.exp(-0.5*((x-x0)* 1./a)**2 +1j*x*k0))
def gauss_y(y, a, y0, k0):
return ((a*np.sqrt(np.pi))**(-0.5)* np.exp(-0.5*((y-y0)* 1./a)**2 +1j*y*k0))
def gauss_k_x(k,a,x0,k0):
return ((a/np.sqrt(np.pi))**0.5* np.exp(-0.5*(a*(k- k0))** 2- 1j*(k- k0)*x0))
def gauss_k_y(k,a,y0,k0):
return ((a/np.sqrt(np.pi))**0.5* np.exp(-0.5*(a*(k- k0))** 2- 1j*(k- k0)*y0))
def theta(x):
x = np.asarray(x)
y = np.zeros(x.shape)
y[x > 0] = 1.0
return y
def square_barrier(x, width, height):
return height * (theta(x) - theta(x - width))
dt = 0.01
N_steps = 50
t_max = 120
frames = int(t_max / float(N_steps * dt))
hbar = 1.0
m = 1.9
N = 2 ** 11
dx = 0.1
dy=0.1
x = dx * (np.arange(N) - 0.5 * N)
y = dy * (np.arange(N) - 0.5 * N)
V0 = 1.5
L = hbar / np.sqrt(2 * m * V0)
a = 3 * L
x0 = -60 * L
y0 = -60 * L
V_x = square_barrier(x, a, V0)
V_y = square_barrier(y, a, V0)
V_x[x < -98] = 1E6
V_x[x > 98] = 1E6
V_y[y < -98] = 1E6
V_y[y > 98] = 1E6
p0 = np.sqrt(2 * m * 0.2 * V0)
dp2 = p0 * p0 * 1./80
d = hbar / np.sqrt(2 * dp2)
k0 = p0 / hbar
v0 = p0 / m
psi_x0 = gauss_x(x, d, x0, k0)
psi_y0 = gauss_y(y, d, y0, k0)
S =Schrodinger(x=x,y=y,psi_x0=psi_x0,psi_y0=psi_y0,V_x=V,V_y=V_y,hbar=hbar,m=m,k0=-28)
the code for one-dimensional plot and the animation is as below:
from matplotlib import pyplot as pl
from matplotlib import animation
fig = pl.figure()
xlim = (-100, 100)
klim = (-5, 5)
ymin = 0
ymax = V0
ax1 = fig.add_subplot(211, xlim=xlim,
ylim=(ymin - 0.2 * (ymax - ymin),
ymax + 0.2 * (ymax - ymin)))
psi_x_line, = ax1.plot([], [], c='r', label=r'$|\psi(x)|$')
V_x_line, = ax1.plot([], [], c='k', label=r'$V(x)$')
center_line = ax1.axvline(0, c='k', ls=':',
label = r"$x_0 + v_0t$")
title = ax1.set_title("")
ax1.legend(prop=dict(size=12))
ax1.set_xlabel('$x$')
ax1.set_ylabel(r'$|\psi(x)|$')
ymin = abs(S.psi_k).min()
ymax = abs(S.psi_k).max()
ax2 = fig.add_subplot(212, xlim=klim,
ylim=(ymin - 0.2 * (ymax - ymin),
ymax + 0.2 * (ymax - ymin)))
psi_k_line, = ax2.plot([], [], c='r', label=r'$|\psi(k)|$')
p0_line1 = ax2.axvline(-p0 / hbar, c='k', ls=':', label=r'$\pm p_0$')
p0_line2 = ax2.axvline(p0 / hbar, c='k', ls=':')
mV_line = ax2.axvline(np.sqrt(2 * V0) / hbar, c='k', ls='--',
label=r'$\sqrt{2mV_0}$')
ax2.legend(prop=dict(size=12))
ax2.set_xlabel('$k$')
ax2.set_ylabel(r'$|\psi(k)|$')
V_x_line.set_data(S.x, S.V_x)
# Animate plot
def init():
psi_x_line.set_data([], [])
V_x_line.set_data([], [])
center_line.set_data([], [])
psi_k_line.set_data([], [])
title.set_text("")
return (psi_x_line, V_x_line, center_line, psi_k_line, title)
def animate(i):
S.time_step(dt, N_steps)
psi_x_line.set_data(S.x, 4 * abs(S.psi_x))
V_x_line.set_data(S.x, S.V_x)
center_line.set_data(2 * [x0 + S.t * p0 / m], [0, 1])
psi_k_line.set_data(S.k, abs(S.psi_k))
title.set_text("t = %.2f" % S.t)
return (psi_x_line, V_x_line, center_line, psi_k_line, title)
anim = animation.FuncAnimation(fig, animate, init_func=init,
frames=frames, interval=30, blit=True)
pl.show()