I'm trying to solve the Schrödinger equation with the Numerov's method. Here is my code:
from pylab import *
from scipy.optimize import brentq
import numpy as np
l = float(input("Angular momentum l:"))
L = float(input("Width of the potential:"))
Vo = float(input("Value of the potential:"))
N = int(input("Number of steps (~10000):"))
h = float(3*L/N)
psi = np.zeros(N) #wave function
psi[0] = 0
psi[1] = h
def V(x,E):
"""
Effective potential function.
"""
if x > L:
return -2*E+l*(l+1)/x**2
else:
return -2*(Vo+E)+l*(l+1)/x**2
def Wavefunction(energy):
"""
Calculates wave function psi for the given value
of energy E and returns value at point xmax
"""
global psi
global E
E=energy
for i in range(2,N):
psi[i]=(2*(1+5*(h**2)*V(i*h,E)/12)*psi[i-1]-(1-(h**2)*V((i-1)*h,E)/12)*psi[i-2])/(1-(h**2)*V((i+1)*h,E)/12)
return psi[-1]
def find_energy_levels(x,y):
"""
Gives all zeroes in y = psi_max, x=en
"""
zeroes = []
s = np.sign(y)
for i in range(len(y)-1):
if s[i]+s[i+1] == 0: #sign change
zero = brentq(Wavefunction, x[i], x[i+1])
zeroes.append(zero)
return zeroes
def main():
energies = np.linspace(-Vo,0,int(10*Vo)) # vector of energies where we look for the stable states
psi_max = [] # vector of wave function at x = 3L for all of the energies in energies
for energy in energies:
psi_max.append(Wavefunction(energy)) # for each energy find the the psi_max at xmax
E_levels = find_energy_levels(energies,psi_max) # now find the energies where psi_max = 0
print ("Energies for the bound states are: ")
for E in E_levels:
print ("%.2f" %E)
# Plot the wavefunctions for first 4 eigenstates
x = np.linspace(0, 3*L, N)
figure()
for E in E_levels:
Wavefunction(E)
plot(x, psi, label="E = %.2f"%E)
legend(loc="upper right")
xlabel('r')
ylabel('$u(r)$', fontsize = 10)
grid()
savefig('numerov.pdf', bbox_inches='tight')
if __name__ == "__main__":
main()
Everything was working really well, this is a plot for Vo=35, l=1, but when I try whit a value of Vo=85, l=0 (is the same for Vo>50), the plot is not what I expected (the end of the plot blows up). For l=1, the error vanish. I am a novice in Python, so I do not know what would be the error. Thanks for the help.
Related
I am a student and am trying to figure out how to solve differential equations with python but I am very confused. I am also getting a syntax error on the fourth to last line, and I am not sure why. This model is supposed to show the progression of a virus.
import numpy as np
from scipy.integrate import odeint
import matplotlib.pyplot as plt
# function that returns dT/dt
# initial condition
T = 1
V = .001
I = 0
def model(t,T,I,V):
β = 10**-5
δ = 4
p = 2*(10**6)
c = 4
dTdt = -β*t*V
dIdt = (β*T*V) - (δ*I)
dVdt = (p*I) - (c* V)
return dTdt, dIdt,dVdt
# initial condition
T0 = 1
V0 = 10**-3
I0 = 0
# time points
t = np.linspace(0,25)
# solve
V = odeint(model,t,V0, args = (T0,I0)
# plot results
plt.plot(t,V)
plt.xlabel('time')
plt.ylabel('V(t)')
plt.show()
The state vector is one single argument in the ODE function. So you have to replace its first line with
def model(Y,t):
T,I,V = Y
and correct the integrator call to
Y = odeint(model,[T0,I0,V0],t)
T,I,V = Y.T
%matplotlib notebook
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from math import pi
# wave speed
c = 1
# spatial domain
xmin = 0
xmax = 1
#time domain
m=500; # num of time steps
tmin=0
T = tmin + np.arange(m+1);
tmax=500
n = 50 # num of grid points
# x grid of n points
X, dx = np.linspace(xmin, xmax, n+1, retstep=True);
X = X[:-1] # remove last point, as u(x=1,t)=u(x=0,t)
# for CFL of 0.1
CFL = 0.3
dt = CFL*dx/c
# initial conditions
def initial_u(x):
return np.sin(2*pi*x)
# each value of the U array contains the solution for all x values at each timestep
U = np.zeros((m+1,n),dtype=float)
U[0] = u = initial_u(X);
def derivatives(t,u,c,dx):
uvals = [] # u values for this time step
for j in range(len(X)):
if j == 0: # left boundary
uvals.append((-c/(2*dx))*(u[j+1]-u[n-1]))
elif j == n-1: # right boundary
uvals.append((-c/(2*dx))*(u[0]-u[j-1]))
else:
uvals.append((-c/(2*dx))*(u[j+1]-u[j-1]))
return np.asarray(uvals)
# solve for 500 time steps
for k in range(m):
t = T[k];
k1 = derivatives(t,u,c,dx)*dt;
k2 = derivatives(t+0.5*dt,u+0.5*k1,c,dx)*dt;
k3 = derivatives(t+0.5*dt,u+0.5*k2,c,dx)*dt;
k4 = derivatives(t+dt,u+k3,c,dx)*dt;
U[k+1] = u = u + (k1+2*k2+2*k3+k4)/6;
# plot solution
plt.style.use('dark_background')
fig = plt.figure()
ax1 = fig.add_subplot(1,1,1)
line, = ax1.plot(X,U[0],color='cyan')
ax1.grid(True)
ax1.set_ylim([-2,2])
ax1.set_xlim([0,1])
def animate(i):
line.set_ydata(U[i])
return line,
I want to program in Python an advection equation which is (∂u/∂t) +c (∂u/∂x) = 0. Time should be discretized with Runge-kutta 4th order. Spatial discretiziation is 2nd order finite difference. When I run my code, I get straight line which transforms into sine wave. But I gave as initial condition sine wave. Why does it start as straight line? And I want to have sine wave moving forward. Do you have any idea on how to get sine wave moving forward? I appreciate your help. Thanks in advance!
While superficially your computation steps are related to the RK4 method, they deviate from the RK4 method and the correct space discretization too much to mention it all.
The traditional way to apply ODE integration methods is to have a function derivatives(t, state, params) and then apply that to compute the Euler step or the RK4 step. In your case it would be
def derivatives(t,u,c,dx):
du = np.zeros(len(u));
p = c/(2*dx);
du[0] = p*(u[1]-u[-1]);
du[1:-1] = p*(u[2:]-u[:-2]);
du[-1] = p*(u[0]-u[-2]);
return du;
Then you can do
X, dx = np.linspace(xmin, xmax, n+1, retstep=True);
X = X[:-1] # remove last point, as u(x=1,t)=u(x=0,t)
m=500; # number of time steps
T = tmin + np.arange(m+1);
U = np.zeros((m+1,n),dtype=float)
U[0] = u = initial_u(X);
for k in range(m):
t = T[k];
k1 = derivatives(t,u,c,dx)*dt;
k2 = derivatives(t+0.5*dt,u+0.5*k1,c,dx)*dt;
k3 = derivatives(t+0.5*dt,u+0.5*k2,c,dx)*dt;
k4 = derivatives(t+dt,u+k3,c,dx)*dt;
U[k+1] = u = u + (k1+2*k2+2*k3+k4)/6;
This uses dt as computed as the primary variable in the time stepping, then constructs the arithmetic sequence from tmin with step dt. Other ways are possible, but one has to make tmax and the number of time steps compatible.
The computation up to this point should now be successful and can be used in the animation. In my understanding, you do not produce a new plot in each frame, you only draw the graph once and after that just change the line data
# animate the time data
line, = ax1.plot(X,U[0],color='cyan')
ax1.grid(True)
ax1.set_ylim([-2,2])
ax1.set_xlim([0,1])
def animate(i):
line.set_ydata(U[i])
return line,
etc.
I currently have the code and I having some trouble trying to plot it, I know that trying to plot both ymax and y won't work in this case, but how would I go about plotting just the value for y? I have plotted the function before by removing the ymax from the return, but I need to print the values and plot the solution for y.
import numpy as np
import matplotlib.pyplot as plt
def GaussElimination(A):
'''
Description: Use Gauss elimination to solve a set of simultaneous equations
Parameters: A a matrix of coefficient and constant value for the system
Return: a matrix holding the solution to the equation. This corresponds to the last n
'''
nr,nc=A.shape
B= A.copy()
# start the gauss elimination
for r in range(nr):
#pivoting
max=abs(B[r][r])
maxr = r
for rr in range(r,nr):
if max < abs(B[rr][r]):
max = abs(B[rr][r])
maxr = rr
if max == 0:
print("Singular Matrix")
return []
# swap if needed
if (maxr != r):
for c in range(nc):
temp = B[r][c]
B[r][c]=B[maxr][c]
B[maxr][c] = temp
# scale the row
scale = B[r][r]
for c in range(r,nc):
B[r][c] = B[r][c]/scale
# eliminate values in the columns
for rr in range(nr):
if rr != r:
scale = B[rr][r]
for c in range(r,nc):
B[rr][c]=B[rr][c] - scale*B[r][c]
if (nc == nr+1):
return B[:,nc-1]
else:
return B[:,(nr):nc]
def SimplySupportedBeam(n):
M = np.zeros([n+1,n+1])
C = np.array([[0],[150],[0],[0],[0],[0]])
for r in range(n-3):
M[r][r] = 1
M[r][r+1] = -4
M[r][r+2] = 6
M[r][r+3] = -4
M[r][r+4] = 1
M[n-3][1] = 1
M[n-2][n-1] = 1
M[n-1][n-5] = 1
M[n-1][n-4] = -2
M[n-1][n-3] = 1
M[n][n-2] = 1
M[n][n-1] = -2
M[n][n] = 1
A = np.concatenate((M,C), axis=1)
y0 = GaussElimination(A)
y = y0[1:n]
ymax = np.amax(abs(y))
return y, ymax
n = int(input("Index of the last node: "))
print (SimplySupportedBeam(n))
plt.figure(1)
plt.plot(SimplySupportedBeam(n))
plt.show()
How would I plot just the value I get for y from my code?
It seems like y is 1D numpy array.
If you just want to plot its values against their indices you should be able to do so using either
plt.plot(SimplySupportedBeam(n)[0])
or
y, ymax = SimplySupportedBeam(n)
plt.plot(y)
The problem was that your function returns two values, i.e. y and ymax.
(I did not
I have been working with odeint and boundary conditions. Bassically what I am trying to do is to solve the differential equations given in this figure 1
where in my code R=R, ph = Phi, al = alpha, a = a, m = m, l = l and om = omega. The initial conditions that I am trying to implement are R(0)=O(r^l); Phi(0)=O(r^{l-1}) if l/=0 and Phi(0)=O(r) if l=0; a(0) = 1 and a(inf)=1/alpha(inf) (additionally I need that R(inf)=0). I tried to applied the shooting method in order to find initial conditions for alpha that best matches with my boundary conditions. I also need to find the omega that best matches the boundary conditions for R at infinity. The code that I wrote is the following:
import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import odeint
import time
start = time.clock()
def system_DE(IC,p,r):
l = p[0]
m = p[1]
om = p[2]
R = IC[0]
ph = IC[1]
a = IC[2]
al = IC[3]
dR_dr = ph
da_dr = a*((2*l+1)*r/2*(om**2*a**2*R**2/al**2+ph**2+l*(l+1)*a**2*R**2/r**2+m**2*a**2*R**2)-(a**2-1)/(2*r))
dal_dr = al*(da_dr/a-l*(l+1)*(2*l+1)*a**2*R**2/r-(2*l+1)*m**2*a**2*r*R**2+(a**2-1)/r)
dph_dr = -2*ph/r-dal_dr*ph/al+da_dr*ph/a-om**2*a**2*R/al**2+l*(l+1)*a**2*R/r**2+m**2*a**2*R
return [dR_dr,da_dr,dal_dr,dph_dr]
def init(u,p,r):
if p==0:
return np.array([1,r,1,u])
else:
return np.array([r**l,l*r**(l-1),1,u])
l = 0
m = 1
ep = 0.3
n_om = 10
omega = np.linspace(m-ep,m+ep,n_om)
r = np.linspace(0.0001, 100, 1000)
niter = 100
u = 0
tol = 0.1
ustep = 0.01
p = np.zeros(3)
p[0] = l
p[1] = m
for j in range(len(omega)):
p[2] = omega[j]
for i in range(niter):
u += ustep
Y = odeint(system_DE(init(u,p[0],r[0]),p,r), init(u,p[0],r[0]), r)
print Y[-1,2]
print Y[-1,3]
if abs(Y[len(Y)-1,2]-1/Y[len(Y)-1,3]) < tol:
print(i,'times iterations')
print("a'(inf)) = ", Y[len(Y)-1,2])
print('y"(0) =',u)
break
if abs(Y[len(Y)-1,0]) < tol:
print(j,'times iterations in omega')
print("R'(inf)) = ", Y[len(Y)-1,0])
break
However, when I run it I am obtaining: error: The function and its Jacobian must be callable functions.
Could some one help me to understand what my mistake is?
Regards,
Luis Padilla.
To start with, the first argument to odeint is your derivative function system_DE. Just pass its name, no parentheses or arguments. Odeint with call it internally and supply arguments.
I fixed my code and now it is giving me some results. However, when I run it I am obtaining some warnings that I don't know how to solve it. Could some one help me to solve it? Basically my code is this:
import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import odeint
import time
def system_DE(IC,r,l,m,om):
R = IC[0]
ph = IC[1]
a = IC[2]
al = IC[3]
dR_dr = ph
da_dr = a*((2*l+1)*r/2*(om**2*a**2*R**2/al**2+ph**2+l*(l+1)*a**2*R**2/r**2+m**2*a**2*R**2)-(a**2-1)/(2*r))
dal_dr = al*(da_dr/a-l*(l+1)*(2*l+1)*a**2*R**2/r-(2*l+1)*m**2*a**2*r*R**2+(a**2-1)/r)
dph_dr = -2*ph/r-dal_dr*ph/al+da_dr*ph/a-om**2*a**2*R/al**2+l*(l+1)*a**2*R/r**2+m**2*a**2*R
return [dR_dr,dph_dr,da_dr,dal_dr]
def init(u,p,r):
if p==0:
return np.array([1.,r,1.,u])
else:
return np.array([r**p,l*r**(p-1),1,u])
l = 0.
m = 1.
ep = 0.2
n_om = 30
omega = np.linspace(m-ep,m+ep,n_om)
r = np.linspace(0.001, 100, 1000)
niter = 1000
tol = 0.01
ustep = 0.01
for j in range(len(omega)):
print('trying with $omega =$',omega[j])
p = (l,m,omega[j])
u = 0.001
for i in range(niter):
u += ustep
ini = init(u,p[0],r[0])
Y = odeint(system_DE, ini,r,p,mxstep=500000)
if abs(Y[len(Y)-1,2]-1/Y[len(Y)-1,3]) < tol:
break
if abs(Y[len(Y)-1,0]) < tol and abs(Y[len(Y)-1,2]-1/Y[len(Y)-1,3]) < tol:
print(j,'times iterations in omega')
print(i,'times iterations')
print("R'(inf)) = ", Y[len(Y)-1,0])
print("alpha(0)) = ", Y[0,3])
print("\omega",omega[j])
break
plt.subplot(2,1,1)
plt.plot(r,Y[:,0],'r',label = '$R$')
plt.plot(r,Y[:,1],'b',label = '$d R /dr$')
plt.xlim([0,10])
plt.legend()
plt.subplot(2,1,2)
plt.plot(r,Y[:,2],'r',label = 'a')
plt.plot(r,Y[:,3],'b', label = '$alpha$')
plt.xlim([0,10])
plt.legend()
plt.show()
But when I run it I am obtaining this:
lsoda-- warning..internal t (=r1) and h (=r2) are
such that in the machine, t + h = t on the next step
(h = step size). solver will continue anyway
in above, r1 = 0.1243782486482D+01 r2 = 0.8727680448722D-16
How could I fix the problem?
Regards,
Luis Padilla.
I am trying to deblend the emission lines of low resolution spectrum in order to get the gaussian components. This plot represents the kind of data I am using:
After searching a bit, the only option I found was the application of the gauest function from the kmpfit package (http://www.astro.rug.nl/software/kapteyn/kmpfittutorial.html#gauest). I have copied their example but I cannot make it work.
I wonder if anyone could please offer me any alternative to do this or how to correct my code:
import numpy as np
import matplotlib.pyplot as plt
from scipy import optimize
def CurveData():
x = np.array([3963.67285156, 3964.49560547, 3965.31835938, 3966.14111328, 3966.96362305,
3967.78637695, 3968.60913086, 3969.43188477, 3970.25463867, 3971.07714844,
3971.89990234, 3972.72265625, 3973.54541016, 3974.36791992, 3975.19067383])
y = np.array([1.75001533e-16, 2.15520995e-16, 2.85030769e-16, 4.10072843e-16, 7.17558032e-16,
1.27759917e-15, 1.57074192e-15, 1.40802933e-15, 1.45038722e-15, 1.55195653e-15,
1.09280316e-15, 4.96611341e-16, 2.68777266e-16, 1.87075114e-16, 1.64335999e-16])
return x, y
def FindMaxima(xval, yval):
xval = np.asarray(xval)
yval = np.asarray(yval)
sort_idx = np.argsort(xval)
yval = yval[sort_idx]
gradient = np.diff(yval)
maxima = np.diff((gradient > 0).view(np.int8))
ListIndeces = np.concatenate((([0],) if gradient[0] < 0 else ()) + (np.where(maxima == -1)[0] + 1,) + (([len(yval)-1],) if gradient[-1] > 0 else ()))
X_Maxima, Y_Maxima = [], []
for index in ListIndeces:
X_Maxima.append(xval[index])
Y_Maxima.append(yval[index])
return X_Maxima, Y_Maxima
def GaussianMixture_Model(p, x, ZeroLevel):
y = 0.0
N_Comps = int(len(p) / 3)
for i in range(N_Comps):
A, mu, sigma = p[i*3:(i+1)*3]
y += A * np.exp(-(x-mu)*(x-mu)/(2.0*sigma*sigma))
Output = y + ZeroLevel
return Output
def Residuals_GaussianMixture(p, x, y, ZeroLevel):
return GaussianMixture_Model(p, x, ZeroLevel) - y
Wave, Flux = CurveData()
Wave_Maxima, Flux_Maxima = FindMaxima(Wave, Flux)
EmLines_Number = len(Wave_Maxima)
ContinuumLevel = 1.64191e-16
# Define initial values
p_0 = []
for i in range(EmLines_Number):
p_0.append(Flux_Maxima[i])
p_0.append(Wave_Maxima[i])
p_0.append(2.0)
p1, conv = optimize.leastsq(Residuals_GaussianMixture, p_0[:],args=(Wave, Flux, ContinuumLevel))
Fig = plt.figure(figsize = (16, 10))
Axis1 = Fig.add_subplot(111)
Axis1.plot(Wave, Flux, label='Emission line')
Axis1.plot(Wave, GaussianMixture_Model(p1, Wave, ContinuumLevel), 'r', label='Fit with optimize.leastsq')
print p1
Axis1.plot(Wave, GaussianMixture_Model([p1[0],p1[1],p1[2]], Wave, ContinuumLevel), 'g:', label='Gaussian components')
Axis1.plot(Wave, GaussianMixture_Model([p1[3],p1[4],p1[5]], Wave, ContinuumLevel), 'g:')
Axis1.set_xlabel( r'Wavelength $(\AA)$',)
Axis1.set_ylabel('Flux' + r'$(erg\,cm^{-2} s^{-1} \AA^{-1})$')
plt.legend()
plt.show()
A typical simplistic way to fit:
def model(p,x):
A,x1,sig1,B,x2,sig2 = p
return A*np.exp(-(x-x1)**2/sig1**2) + B*np.exp(-(x-x2)**2/sig2**2)
def res(p,x,y):
return model(p,x) - y
from scipy import optimize
p0 = [1e-15,3968,2,1e-15,3972,2]
p1,conv = optimize.leastsq(res,p0[:],args=(x,y))
plot(x,y,'+') # data
#fitted function
plot(arange(3962,3976,0.1),model(p1,arange(3962,3976,0.1)),'-')
Where p0 is your initial guess. By the looks of things, you might want to use Lorentzian functions...
If you use full_output=True, you get all kind of info about the fitting. Also check out curve_fit and the fmin* functions in scipy.optimize. There are plenty of wrappers around these around, but often, like here, it's easier to use them directly.