So I'm trying to draw a circle instead of a quarter circle does anyone knows how do I change the commands ?
I want the center of the circle to be in the coordinates (0.5,0.5) , but I don't know how to change it. I know that the I need to change the radius to 0.5 but I can't figure out where I need to change the values for the center. Thanks!
import matplotlib.pyplot as plt
from matplotlib.pyplot import *
matplotlib.use("TkAgg")
fig = figure(figsize=(8, 8), dpi=120)
nTrials = int(100000)
radius = 1
nInside = 0
nDrops = 0
XrandCoords = np.random.default_rng().uniform(0, 1, (nTrials,))
YrandCoords = np.random.default_rng().uniform(0, 1, (nTrials,))
fig1 = plt.figure(1)
plt.get_current_fig_manager().window.wm_geometry("+00+00") # move the window
plt.xlim(0, 1)
plt.ylim(0, 1)
plt.legend()
isFirst1 = True
isFirst2 = True
piValueI = []
nDrops_arr = []
insideX = []
outsideX = []
insideY = []
outsideY = []
for i in range(nTrials):
x = XrandCoords[i]
y = YrandCoords[i]
nDrops = nDrops + 1
if (x ** 2 + y ** 2) <= radius ** 2:
nInside = nInside + 1
insideX.append(x)
insideY.append(y)
else:
outsideX.append(x)
outsideY.append(y)
if i % 100 == 0:
plt.figure(1)
if isFirst1:
plt.scatter(insideX, insideY, c='blue', s=50, label='În interior')
isFirst1 = False
plt.legend(loc=(0.75, 0.9))
else:
plt.scatter(insideX, insideY, c='blue', s=50)
plt.figure(1)
if isFirst2:
plt.scatter(outsideX, outsideY, c='black', s=50, label='În exterior')
isFirst2 = False
plt.legend(loc=(0.75, 0.9))
else:
plt.scatter(outsideX, outsideY, c='black', s=50)
area = 4 * nInside / nDrops
plt.figure(1)
plt.title('Nr de puncte luate = ' + str(nDrops) + '; Nr punctelor din interiorul cercului = ' + str(
nInside) + r'; π ≈ $4\frac{Nr_\mathrm{interior}}{Nr_\mathrm{total}}=$ ' + str(np.round(area, 6)))
piValueI.append(area)
nDrops_arr.append(nDrops)
plt.pause(0.1)
area = 4 * nInside / nTrials
print("Valoare estimată a lui Pi: ", area)
plt.show()
Please try :
radius = 0.5
XrandCoords = np.random.uniform(0.5, 1, (nTrials,))
YrandCoords = np.random.uniform(0.5, 1, (nTrials,))
radius = 0.5
if ((x-0.5) ** 2 + (y-0.5) ** 2) <= radius ** 2:
....
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()
I have the code below. If you run that program a scatter plot is created. The plot is very messy. I want every point to represent one car. I only want to see the point of the current position of each car in the animation. I am not very good at matplotlib, but how would I do that?
from matplotlib import pyplot as plt
import random
import math
from matplotlib import animation
import numpy as np
roadLength = 100
numFrames = 200
numCars = 60
velocities = []
positions = []
theta = []
r = []
color = []
probability = 0.5
vmax = 1
flowr = []
fig = plt.figure()
ax = fig.add_subplot(111, projection='polar')
ax.axis('off')
for i in range(numCars):
positions.append(i * 2)
theta.append(0)
r.append(1)
color.append(i)
velocities.append(0)
def pos(k):
rand = random.uniform(0,1)
if velocities[k] < vmax:
velocities[k] += 1
if k < numCars-1 and (positions[k+1] - positions[k]) <= velocities[k]:
velocities[k] = positions[k+1] - positions[k] - 1
if k == numCars-1 and (roadLength - (positions[k] - positions[0])) <= velocities[k]:
velocities[k] = roadLength - (positions[k] - positions[0]) - 1
if rand < probability and velocities[k] > 0:
velocities[k] = velocities[k] - 1
positions[k] = positions[k] + velocities[k]
return positions[k]
def animate(frameNr):
for i in range(numCars):
theta[i] = pos(i) * 2 * math.pi / roadLength
flowr.append(velocities[i])
ax.scatter(theta, r, c=color)
return ax.scatter(theta, r, c=color),
def simulate():
anim = animation.FuncAnimation(fig, animate,
frames=numFrames, interval=10, blit=True, repeat=False)
plt.show()
simulate()
Thanks in advance!
Edit:
I fixed so that it became a point plot I am now storing velocity, position and lane in an object for a later problem. My question is:
How do I make the points have different colors?
from matplotlib import pyplot as plt
import random
import math
from matplotlib import animation
import numpy as np
from operator import attrgetter
roadLength = 100
numFrames = 200
nlanes = 1
numCars = 20
posss =[]
theta = []
r = []
color = []
probability = 0.5
vmax = 1
flowr = []
cars=[]
class Car:
def __init__(self, position, velocity, lane):
self.position = position
self.velocity = velocity
self.lane = lane
def pos(car,k):
rand = random.uniform(0,1)
if car[k].velocity < vmax:
car[k].velocity += 1
if k < numCars-1 and (car[k+1].position - car[k].position) <= car[k].velocity:
car[k].velocity = car[k+1].position - car[k].position - 1
if k == numCars-1 and (roadLength - (car[k].position - car[0].position)) <= car[k].velocity:
car[k].velocity = roadLength - (car[k].position - car[0].position) - 1
if rand < probability and car[k].velocity > 0:
car[k].velocity = car[k].velocity - 1
car[k].position = car[k].position + car[k].velocity
return car[k].position
for i in range(numCars):
cars.append(Car(i, 0, 1))
theta.append(0)
r.append(1)
color.append(i)
posss.append(i)
fig = plt.figure()
ax = fig.add_subplot(111)
point, = ax.plot(posss, r, 'o')
ax.set_xlim(-10, 1.2*numFrames)
ax.set_ylim(-2, nlanes + 3)
def animate(frameNr):
sort_cars = sorted(cars, key=attrgetter("position"))
for i in range(numCars):
pos(sort_cars,i)
for k in range(numCars):
theta[k]=cars[k].position
point.set_data(theta, r)
return point,
def simulate():
anim = animation.FuncAnimation(fig, animate,
frames=numFrames, interval=10, blit=True, repeat=False)
plt.show()
simulate()
plot() returns a Line2D object, in which all points have the same color.
If you want your points to have different colors, you should use scatter(). The logic is exactly the same, except you use set_offsets(<2D array>) instead of set_data().
See this question for more details
I want the markersize in pyplot to correspond to a diameter in my experimental data.
Can I set the markersize, so it always shows a circle with a diameter of 6*10^-8 m? My x- and y-data is of the same scale.
font_size = 20
xmin = 0
xmax = np.sqrt(4) * 1e-6
ymin = 0
ymax = np.sqrt(4) * 1e-6
n_sim = 80
print("number of simulated dots:", n_sim)
density = n_sim / (xmax * ymax)
print("density:", density * 10 ** -12)
diameter = 60e-9
dots_sim = []
def no_nearby_dots(new_dot, dots_sim, min_distance):
for dot in dots_sim:
if np.sqrt((dot[0] - new_dot[0]) ** 2 + (dot[1] - new_dot[1]) ** 2) <= min_distance:
return False
return True
new_dot = True
dots_sim.append((np.random.uniform(xmin, xmax), np.random.uniform(ymin, ymax)))
failed_attempts = 0
while new_dot:
xp = np.random.uniform(xmin, xmax)
yp = np.random.uniform(ymin, ymax)
if no_nearby_dots((xp, yp), dots_sim, diameter):
dots_sim.append((xp, yp))
failed_attempts = 0
else:
failed_attempts += 1
if len(dots_sim) == n_sim:
new_dot = False
if failed_attempts > 1000:
new_dot = False
print("ERROR...exit loop")
break
x_sim = [dot[0] for dot in dots_sim]
y_sim = [dot[1] for dot in dots_sim]
fig = plt.figure(1)
ax1 = fig.add_subplot()
plt.title('simulated pattern', fontsize=font_size)
plt.plot(x_sim, y_sim, 'b.', label='QDs', markersize=14)
plt.xlabel('x [$\mathrm{\mu m}$]', fontsize=font_size)
plt.xticks(fontsize=font_size)
plt.ylabel('y [$\mathrm{\mu m}$]', fontsize=font_size)
plt.yticks(fontsize=font_size)
plt.ticklabel_format(axis='both', style='Sci', scilimits=(-6, -6), useMathText=1)
ax1.xaxis.get_offset_text().set_visible(False)
ax1.yaxis.get_offset_text().set_visible(False)
plt.show()
This is the pattern:
SOLVED, ty Stanislav Ivanov. I didn't specify my limits on the x and y axis + I used plt.Circle() now instead of plt.plot().
fig = plt.figure(1)
ax1 = fig.add_subplot()
plt.title('simulated pattern', fontsize=font_size)
for i in range(n_sim):
plt.gca().add_artist(plt.Circle((x_sim[i], y_sim[i]), diameter/2, color='r',
fill=False, label='circumference of the QD'))
plt.xlabel('x [$\mathrm{\mu m}$]', fontsize=font_size)
plt.xticks(fontsize=font_size)
plt.ylabel('y [$\mathrm{\mu m}$]', fontsize=font_size)
plt.yticks(fontsize=font_size)
plt.xlim(xmin - 0.01*xmax, xmax + 0.01*xmax)
plt.ylim(ymin - 0.01*ymax, ymax + 0.01*ymax)
plt.ticklabel_format(axis='both', style='Sci', scilimits=(-6, -6), useMathText=1)
ax1.xaxis.get_offset_text().set_visible(False)
ax1.yaxis.get_offset_text().set_visible(False)
New Pattern
I want to make a 3d line plot animation with a point at the last point plotted, like a trajectory for a particule. In a case without blitting, I did it with a scatter plot and a line plot3d, updating both plots in the update function. I tried doing it in the same way, with blitting, but it throws the following error
TypeError: 'Path3DCollection' object is not iterable
My code is the following
import matplotlib.pyplot as m
import matplotlib.animation as animation
from mpl_toolkits.mplot3d import Axes3D
def ranstep(start,step,stop):
numberlist = int((stop - start)/step + 1)
lista = [0]*numberlist
i = 0
while i <= numberlist:
if i == 0:
lista[i] = start
i = i + 1
if i != 0 and i < numberlist - 1:
lista[i] = lista[i - 1] + step
i = i + 1
if i != 0 and i == numberlist - 1:
lista[i] = stop
break
return(lista)
k1 = 10.0
k2 = 28.0
k3 = 8.0/3.0
czeroA = 1
czeroB = 0
czeroC = 0
step = 0.001
tfinal = 100
time = list(ranstep(0,step,tfinal))
cA = [0]*len(time)
cB = [0]*len(time)
cC = [0]*len(time)
i = 0
while i <= len(time) - 1:
if i == 0:
cA[i] = czeroA
cB[i] = czeroB
cC[i] = czeroC
i = i + 1
else:
cA[i] = cA[i - 1] + k1*(cB[i - 1] - cA[i - 1])*step
cB[i] = cB[i - 1] + (cA[i - 1]*(k2 - cC[i - 1]) - cB[i-1])*step
cC[i] = cC[i - 1] + (cA[i - 1]*cB[i-1] - k3*cC[i-1])*step
i = i + 1
fig, ax = m.subplots()
ax = fig.add_subplot(111, projection='3d')
line, = ax.plot3D([], [], [], c = 'blue', lw = 1)
line, = ax.scatter([],[],[], c = 'red', s = 30)
def init():
ax.set_xlim([-20,20])
ax.set_ylim([-20,20])
ax.set_zlim([0,50])
ax.set_xlabel('$x(t)$')
ax.set_ylabel('$y(t)$')
ax.set_zlabel('$z(t)$')
return line,
def update(frame):
line.set_data(cA[0:frame*50],cB[0:frame*50])
line.set_3d_properties(cC[0:frame*50])
line.set_data(cA[frame*500,frame*500 + 1],cB[frame*500,frame*500 + 1])
line.set_3d_properties(cC[frame*500,frame*500 + 1])
return line,
ani = animation.FuncAnimation(fig, update,interval=5, frames = len(time),init_func=init, blit=True)
m.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()