Changing an ODE inside the function with ODEINT and python - python

This is my first question here on the forum, so I hope I'm doing evrything right. My problem is, I am supposed to solve an ODE with python. The ODE describes the movement of an electron inside of a "bubble" (a bubble is something that occurs in Plasmaphysics when you shoot with a laser at plasma but thats not that important for the question im asking). So my task is to build a programm that can plot the trajectory of such an electron. For that I have to basically solve 2 different ODE's since the electron inside the bubble is moving according to an ODE and if it exits the bubble it is in a field-free-area where it should move according to the classical equations of motion.
In the following I will show you the code if only 1 ODE is solved:
import numpy as np
from scipy.integrate import odeint
import matplotlib.pyplot as plt
from math import sqrt,pi
def f(r,t):
px = r[0]
py = r[1]
x = r[2]
y = r[3]
gamma = sqrt(1+px**2+py**2)
dpx_dt = -(1+V)*(x-V*t)/4 + py*y/(4*gamma)
dpy_dt = -(1+px/gamma)*y/4
dx_dt = px/gamma
dy_dt = py/gamma
return [dpx_dt, dpy_dt, dx_dt, dy_dt]
def Plot_Aufgabe5(t,p_x,p_y,xi,y):
#Plotten der Trajektorie des Elektrons
plt.plot(xi,y)
return
#Konstanten
r_b = 10
m = 9.109*10**-31
c = 2.99*10**8
roh = 9
#Anfangsbedingungen
gamma_0 = 4
V = np.sqrt(((gamma_0**2) - 1)/(gamma_0**2))
t_end = 300
N = 1000
t = np.linspace(0,t_end,N)
y = [0, 0, np.sqrt(r_b**2 - roh**2), roh]
#Array mit äquidistanten Werten erstellen
Äqui_array = np.linspace(0,2*np.pi,10)
#Array mit verschd. Anfangsimpulsen
Impuls_array = [-1,-0.5,0.5,1]
for j in Impuls_array:
plt.figure()
#Plotten eines Kreises
theta = np.linspace(0,2*np.pi,t_end)
Punkt_x = r_b*np.cos(theta)
Punkt_y = r_b*np.sin(theta)
plt.plot(Punkt_x,Punkt_y)
for i in range(10):
#Neue Anfangsbedingungen aufstellen
y = [j,j,r_b*np.cos(Äqui_array[i]),r_b*np.sin(Äqui_array[i])]
#Ausführen des Solvers
sol = odeint(f,y,t)
p_x = sol[:,0]
p_y = sol[:,1]
x = sol[:,2]
xi = x - V*t
y = sol[:,3]
#Plot
Plot_Aufgabe5(t,p_x,p_y,xi,y)
plt.axis([-15,15,-10,10])
plt.xlabel(r'$xi_{num}$ = $k_p$$\xi$')
plt.ylabel(r'$y_{num}$ = $k_p$y')
plt.title('Trajektorie des Elekrons in einer Bubble mit Impuls p = ' + str(j) + '*mc')
plt.show()
Plots for the code for 1 ODE
As you can see in the plot, the electron goes outside the bubble but it comes back in which should not happen. (Also the plot is for 10 different electrons which start on the edge of the bubble)
Therefore I made an approach to change between 2 ODE's according to the position of the electron however that did not quite work:
def f(r,t):
px = r[0] #boundary conditions for the electron
py = r[1]
x = r[2]
y = r[3]
gamma = sqrt(1+px**2+py**2)
if ((x**2 + y**2) <= r_b**2): #if electron is inside bubble
dpx_dt = -(1+V)*(x-V*t)/4 + py*y/(4*gamma)
dpy_dt = -(1+px/gamma)*y/4
dx_dt = px/gamma
dy_dt = py/gamma
else: #if it is outside of the bubble
dpx_dt = 0
dpy_dt = 0
dx_dt = px/gamma
dy_dt = py/gamma
result = [dpx_dt, dpy_dt, dx_dt, dy_dt]
return result
My Idea here was that with the if-statement i check, whether the electron is inside the bubble or not and change the ODE according to that. However it does not really work out and im not sure why, underneath you can find the plots that come with that function (also the rest of the code stayed the same as above):
Plots for the Function to change between 2 ODE's
I hope my problem gets clear and that maybe someone has an idea that could help me! Good day to evryone :)

Related

Is there a way to vary the transparency of a surface plot in matplotlib point by point?

In Matlab, I have successfully built a helpful script that projects orbital wavefunctions onto a series of concentric spherical shells and varies the color and transparency of the spherical shell based on the value of the wavefunction on the points making up the spherical shell. When trying to convert this code over to python, I was able to reproduce everything except for the variation of the transparency at the individual points, as is seems matplotlib's plot_surface function only takes alpha values that are a single number. Is there any way to perform such a task in python? Are there other surface plotting packages that achieve this that I am not aware of? It would be great for general usability if this code could be converted, as not everyone has access to Matlab.
For anyone with Matlab, I am including a functioning example of the working code. The chunk of code "'AlphaDataMapping','scaled','AlphaData',opacity.*f.^2" within the surf function does the desired trick of varying the opacity. I have not included a python example as the tests I ran were far from presentable, but plotting spherical shells with the desired color profile is not the issue.
clear
n = 3;
l = 1;
m = 0;
name = string(n) + string(l) + string(m);
Nang = 50; %defines angular grid
rmax = 10*n;
Nr = 5+30*n; %defines radial grid
opacity = 0.05*exp(-0.5*l);
r = linspace(0,rmax,Nr);
rho = 2.*r./n;
LaP = laguerreL(n-l-1,2*l+1,rho); %radial part polynomial
Rnl = exp(-rho/2).*rho.^l.*LaP;
Pnl = r.^2.*Rnl.^2./max(Rnl.^2.*r.^2);
theta = linspace(0,pi,Nang);
costheta = cos(theta);
phi = linspace(0,2*pi,2*Nang);
ThetaMatrix = zeros(2*Nang,Nang); %matrices ordered as phi, theta
PhiMatrix = zeros(2*Nang,Nang);
XMatrix = zeros(2*Nang,Nang); %ordered phi, theta
YMatrix = zeros(2*Nang,Nang);
ZMatrix = zeros(2*Nang,Nang);
YmlMatrix = zeros(2*Nang,Nang);
Psi = zeros(2*Nang,Nang); %ordered r, phi, theta
LP = legendre(l,costheta);
for i = 1:2*Nang %over the values of phi
for j = 1:Nang %over the values of theta
YmlMatrix(i,j) = LP(m+1,j)*exp((m*1i*phi(i))); %Yml(theta,phi)
end
end
maxCreal = 0;
maxCimag = 0;
for k = 1:Nr
if max(max(abs(Rnl(k)*real(YmlMatrix)))) > maxCreal
maxCreal = max(max(abs(Rnl(k)*real(YmlMatrix))));
end
if max(max(abs(Rnl(k)*imag(YmlMatrix)))) > maxCimag
maxCimag = max(max(abs(Rnl(k)*imag(YmlMatrix))));
end
end
figure(1)
hold on
a = gca;
a.FontSize = 15;
colormap jet
xlabel('x')
ylabel('y')
zlabel('z')
xlim([-rmax,rmax])
ylim([-rmax,rmax])
zlim([-rmax,rmax])
cb = colorbar;
cb.Label.String = "Re(\psi_{" + name + "})/max(\psi_{" + name + "})";
cb.FontSize = 15;
for k = 1:Nr
for i = 1:2*Nang
for j = 1:Nang
XMatrix(i,j) = r(k)*cos(phi(i))*sin(theta(j));
YMatrix(i,j) = r(k)*sin(phi(i))*sin(theta(j));
ZMatrix(i,j) = r(k)*costheta(j);
end
end
C = Rnl(k).*real(YmlMatrix); %colors the surface according to the phase
f = abs(Rnl(k)*real(YmlMatrix))/maxCreal;
h = surf(XMatrix,YMatrix,ZMatrix,C,'FaceAlpha','flat',...
'AlphaDataMapping','scaled','AlphaData',opacity.*f.^2,'EdgeAlpha',0.0);
end
set(gcf,'position',[350,0,900,900])
view([45,10])
grid on

Turning points in a squared first order ODE using Scipy.integrate

I am trying to code a radius function based on the Schwarzchild solution to a black hole given the expression:
(dr/dtau)^2= Emu^2- Veff^2
As it is a square the sign in front of the root will depend on the turning points that I have manually found and labeled tp1 and tp2. However, even though I am changing the functions sign depending on its position it behaves relatively well until it hits these turning points.
Here is the code I have so far:
(P.S: I hope this is the correct formatting and way to present a question although i have been a reader for a few years this is actually my first post).
import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import ode
tp1 = 1.08329*10**11
tp2 = 4.13115*10**11
#arbitrary initial radius
r_start = 3e11
#constants:
M = 4*10**6*(1.9891*10**30) # SMBH mass in kg
G = 6.67408*10**(-11) # Gravitational constant in N kg^-2 m^2
c = 299792458
Emu2 = 0.88*10**17
Lmu = 10**19
def odes(tau,rs):
Vef = (1-(((2*G*M)/c**2)/rs))*((c**2)+(Lmu/rs)**2)
sign = (Emu2)-Vef
signcount = 1
if sign <= 0:
if rs <= tp1:
rs = tp1+5
Vef = (1-(((2*G*M)/c**2)/rs))*((c**2)+(Lmu/rs)**2)
sign = (Emu2)-Vef
drdTau = np.sqrt(Emu2 - Vef)
signcount = 1
if rs >= tp2:
rs = tp2-5
Vef = (1-(((2*G*M)/c**2)/rs))*((c**2)+(Lmu/rs)**2)
sign = (Emu2)-Vef
drdTau = (-1)*np.sqrt((Emu2) - Vef)
signcount = 2
return [drdTau]
if sign > 0 :
if signcount == 1:
Vef = (1-(((2*G*M)/c**2)/rs))*((c**2)+(Lmu/rs)**2)
sign = (Emu2)-Vef
drdTau = np.sqrt(Emu2 - Vef)
if signcount == 2:
Vef = (1-(((2*G*M)/c**2)/rs))*((c**2)+(Lmu/rs)**2)
sign = (Emu2)-Vef
drdTau = (-1)*np.sqrt(Emu2 - Vef)
return [drdTau]
if __name__ == '__main__':
r0 = [r_start]
tspan = 5*3.154e7
# timestep
dtau = 1000
# total number of steps
n_steps = int(np.ceil(tspan/dtau))
#Initialise arrays
t = np.zeros((n_steps,1))
rs = np.zeros((n_steps,1))
ts = np.zeros((n_steps,1))
step = 1
r0 = [r_start]
t0 = [0]
t[0] = np.array(t0)
rs[0] = np.array(r0)
# initiate solver
solver = ode(odes)
solver.set_integrator('DOP83')
solver.set_initial_value(r0,0)
#propagate orbit
while solver.successful() and step<n_steps:
solver.integrate(solver.t+dtau)
ts[step] = solver.t
rs[step] = solver.y
step += 1
plt.plot(ts,rs,'s',color='#0066FF')
# axes labels
plt.xlabel('$x$')
plt.ylabel('$y$')
plt.legend('pos')
# check for and set axes limits
max_yval = np.amax(rs)
max_xval = np.amax(ts)
plt.xlim(0,max_xval)
plt.ylim(tp1 - 300000,max_yval)
plt.show()
print(rs)
I have solved it by changing the function it is integrating once it reaches tp1 and tp2 I am hoping to add a root solver to automate the turning point solving but for now a rough tratment is what I have. I'll put up the code for anyone working on this area that has come across this issue once i have refined it, but essentially instead of modifying the function I added a second one called odes2with a negative sign and solved for that once the turning points were reached.
solver = ode(odes)
solver.set_integrator('lsoda')
solver.set_initial_value(r0,0)
#propagate orbit
while solver.successful() and step<n_steps:
solver.integrate(solver.t+dtau)
ts[step] = solver.t
rs[step] = solver.y
if rs[step]>tp2:
rs[step] = tp2
solver = ode(odes2)
solver.set_initial_value(rs[step],ts[step])
if rs[step]<tp1:
rs[step] = tp1
solver = ode(odes)
solver.set_initial_value(rs[step],ts[step])
step += 1
Not perfect, but using the second order dynamic really avoids a lot of overhead
def Veff(rs): return (1-(((2*G*M)/c**2)/rs))*((c**2)+(Lmu/rs)**2)
def grad_Veff(rs):
rp=rs*(1+1e-5)-1e-5; rm = rs*(1+1e-5)+1e-5
return (Veff(rp)-Veff(rm))/(rp-rm);
def ode(t,y): return [y[1], -grad_Veff(y[0])]
t = np.arange(0, tspan, dtau)
vstart = (Emu2 - Veff(rstart))**0.5
res = solve_ivp(ode, (0,tspan), [rstart, vstart], t_eval=t, atol=1e2, rtol=1e-8)
This should integrate through the turning points, to find the turning points search for sign changes or roots in the second component. One could also define an event for this so that the solver automatically registers the turning points.

Large angle pendulum plot did not show as expected

I am trying to plot the relationship between period and amplitude for an undamped and undriven pendulum for when small angle approximation breaks down, however, my code did not do what I expected...
I think I should be expecting a strictly increasing graph as shown in this video: https://www.youtube.com/watch?v=34zcw_nNFGU
Here is my code, I used zero crossing method to calculate period:
import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import solve_ivp
from itertools import chain
# Second order differential equation to be solved:
# d^2 theta/dt^2 = - (g/l)*sin(theta) - q* (d theta/dt) + F*sin(omega*t)
# set g = l and omega = 2/3 rad per second
# Let y[0] = theta, y[1] = d(theta)/dt
def derivatives(t,y,q,F):
return [y[1], -np.sin(y[0])-q*y[1]+F*np.sin((2/3)*t)]
t = np.linspace(0.0, 100, 10000)
#initial conditions:theta0, omega0
theta0 = np.linspace(0.0,np.pi,100)
q = 0.0 #alpha / (mass*g), resistive term
F = 0.0 #G*np.sin(2*t/3)
value = []
amp = []
period = []
for i in range (len(theta0)):
sol = solve_ivp(derivatives, (0.0,100.0), (theta0[i], 0.0), method = 'RK45', t_eval = t,args = (q,F))
velocity = sol.y[1]
time = sol.t
zero_cross = 0
for k in range (len(velocity)-1):
if (velocity[k+1]*velocity[k]) < 0:
zero_cross += 1
value.append(k)
else:
zero_cross += 0
if zero_cross != 0:
amp.append(theta0[i])
# period calculated using the time evolved between the first and last zero-crossing detected
period.append((2*(time[value[zero_cross - 1]] - time[value[0]])) / (zero_cross -1))
plt.plot(amp,period)
plt.title('Period of oscillation of an undamped, undriven pendulum \nwith varying initial angular displacemnet')
plt.xlabel('Initial Displacement')
plt.ylabel('Period/s')
plt.show()
enter image description here
You can use the event mechanism of solve_ivp for such tasks, it is designed for such "simple" situations
def halfperiod(t,y): return y[1]
halfperiod.terminal=True # stop when root found
halfperiod.direction=1 # find sign changes from negative to positive
for i in range (1,len(theta0)): # amp==0 gives no useful result
sol = solve_ivp(derivatives, (0.0,100.0), (theta0[i], 0.0), method = 'RK45', events =(halfperiod,) )
if sol.success and len(sol.t_events[-1])>0:
period.append(2*sol.t_events[-1][0]) # the full period is twice the event time
amp.append(theta0[i])
This results in the plot

Poincare Section of a system of second order odes

It is the first time I am trying to write a Poincare section code at Python.
I borrowed the piece of code from here:
https://github.com/williamgilpin/rk4/blob/master/rk4_demo.py
and I have tried to run it for my system of second order coupled odes. The problem is that I do not see what I was expecting to. Actually, I need the Poincare section when x=0 and px>0.
I believe that my implementation is not the best out there. I would like to:
Improve the way that the initial conditions are chosen.
Apply the correct conditions (x=0 and px>0) in order to acquire the correct Poincare section.
Create one plot with all the collected poincare section data, not four separate ones.
I would appreciate any help.
This is the code:
from matplotlib.pyplot import *
from scipy import *
from numpy import *
# a simple Runge-Kutta integrator for multiple dependent variables and one independent variable
def rungekutta4(yprime, time, y0):
# yprime is a list of functions, y0 is a list of initial values of y
# time is a list of t-values at which solutions are computed
#
# Dependency: numpy
N = len(time)
y = array([thing*ones(N) for thing in y0]).T
for ii in xrange(N-1):
dt = time[ii+1] - time[ii]
k1 = dt*yprime(y[ii], time[ii])
k2 = dt*yprime(y[ii] + 0.5*k1, time[ii] + 0.5*dt)
k3 = dt*yprime(y[ii] + 0.5*k2, time[ii] + 0.5*dt)
k4 = dt*yprime(y[ii] + k3, time[ii+1])
y[ii+1] = y[ii] + (k1 + 2.0*(k2 + k3) + k4)/6.0
return y
# Miscellaneous functions
n= 1.0/3.0
kappa1 = 0.1
kappa2 = 0.1
kappa3 = 0.1
def total_energy(valpair):
(x, y, px, py) = tuple(valpair)
return .5*(px**2 + py**2) + (1.0/(1.0*(n+1)))*(kappa1*np.absolute(x)**(n+1)+kappa2*np.absolute(y-x)**(n+1)+kappa3*np.absolute(y)**(n+1))
def pqdot(valpair, tval):
# input: [x, y, px, py], t
# takes a pair of x and y values and returns \dot{p} according to the Hamiltonian
(x, y, px, py) = tuple(valpair)
return np.array([px, py, -kappa1*np.sign(x)*np.absolute(x)**n+kappa2*np.sign(y-x)*np.absolute(y-x)**n, kappa2*np.sign(y-x)*np.absolute(y-x)**n-kappa3*np.sign(y)*np.absolute(y)**n]).T
def findcrossings(data, data1):
# returns indices in 1D data set where the data crossed zero. Useful for generating Poincare map at 0
prb = list()
for ii in xrange(len(data)-1):
if (((data[ii] > 0) and (data[ii+1] < 0)) or ((data[ii] < 0) and (data[ii+1] > 0))) and data1[ii] > 0:
prb.append(ii)
return array(prb)
t = linspace(0, 1000.0, 100000)
print ("step size is " + str(t[1]-t[0]))
# Representative initial conditions for E=1
E = 1
x0=0
y0=0
init_cons = [[x0, y0, np.sqrt(2*E-(1.0*i/10.0)*(1.0*i/10.0)-2.0/(n+1)*(kappa1*np.absolute(x0)**(n+1)+kappa2*np.absolute(y0-x0)**(n+1)+kappa3*np.absolute(y0)**(n+1))), 1.0*i/10.0] for i in range(-10,11)]
outs = list()
for con in init_cons:
outs.append( rungekutta4(pqdot, t, con) )
# plot the results
fig1 = figure(1)
for ii in xrange(4):
subplot(2, 2, ii+1)
plot(outs[ii][:,1],outs[ii][:,3])
ylabel("py")
xlabel("y")
title("Full trajectory projected onto the plane")
fig1.suptitle('Full trajectories E = 1', fontsize=10)
# Plot Poincare sections at x=0 and px>0
fig2 = figure(2)
for ii in xrange(4):
subplot(2, 2, ii+1)
xcrossings = findcrossings(outs[ii][:,0], outs[ii][:,3])
yints = [.5*(outs[ii][cross, 1] + outs[ii][cross+1, 1]) for cross in xcrossings]
pyints = [.5*(outs[ii][cross, 3] + outs[ii][cross+1, 3]) for cross in xcrossings]
plot(yints, pyints,'.')
ylabel("py")
xlabel("y")
title("Poincare section x = 0")
fig2.suptitle('Poincare Sections E = 1', fontsize=10)
show()
You need to compute the derivatives of the Hamiltonian correctly. The derivative of |y-x|^n for x is
n*(x-y)*|x-y|^(n-2)=n*sign(x-y)*|x-y|^(n-1)
and the derivative for y is almost, but not exactly (as in your code), the same,
n*(y-x)*|x-y|^(n-2)=n*sign(y-x)*|x-y|^(n-1),
note the sign difference. With this correction you can take larger time steps, with correct linear interpolation probably even larger ones, to obtain the images
I changed the integration of the ODE to
t = linspace(0, 1000.0, 2000+1)
...
E_kin = E-total_energy([x0,y0,0,0])
init_cons = [[x0, y0, (2*E_kin-py**2)**0.5, py] for py in np.linspace(-10,10,8)]
outs = [ odeint(pqdot, con, t, atol=1e-9, rtol=1e-8) ) for con in init_cons[:8] ]
Obviously the number and parametrization of initial conditions may change.
The computation and display of the zero-crossings was changed to
def refine_crossing(a,b):
tf = -a[0]/a[2]
while abs(b[0])>1e-6:
b = odeint(pqdot, a, [0,tf], atol=1e-8, rtol=1e-6)[-1];
# Newton step using that b[0]=x(tf) and b[2]=x'(tf)
tf -= b[0]/b[2]
return [ b[1], b[3] ]
# Plot Poincare sections at x=0 and px>0
fig2 = figure(2)
for ii in xrange(8):
#subplot(4, 2, ii+1)
xcrossings = findcrossings(outs[ii][:,0], outs[ii][:,3])
ycrossings = [ refine_crossing(outs[ii][cross], outs[ii][cross+1]) for cross in xcrossings]
yints, pyints = array(ycrossings).T
plot(yints, pyints,'.')
ylabel("py")
xlabel("y")
title("Poincare section x = 0")
and evaluating the result of a longer integration interval

How could I use scipy library to get the distance?

I just started to learn python, and my teacher asks me to simulate the model rocket trajectory knowing the thrust of the engine.
I have already got the speed and the acceleration of the rocket using odeint function. However, I don't know how to use the speed and time.
I've already got to get the distance the rocket traveled since the speed is solved by odeint function.
here's the code I wrote to get the speed:
def getforce(t):
if 0<=t<0.15:
F = 40*t
elif 0.15<=t<0.7:
F = -9.09*t+7.36
elif 0.7<=t<1.25:
F = 1
elif 1.25<=t<1.65:
F = 7.5*t-8.375
elif 1.65<=t<1.8:
F = -26.6*t+48
else:
F = 0
return F
def getspeed(x,t):
Ft = getforce(t)
y0,y1 = x
dy0 = y1
dy1 = (Ft-0.0001277422*y1**2*np.sign(y1)-0.174)/0.0177
return dy0,dy1
t = np.linspace(0,10,100)
sol = si.odeint(getspeed,(0,0),t)
plt.plot(t,sol[:,0])
plt.show()
Assuming that everything else is correct, you just integrate the speed by hand.
(The reason I can't check on the overall correctness easily is that you are using an unorthodox (given) expression in the speed section, rather than solving the rocked equation that includes the mass loss F= m a --> d(m*v)/dt = dm/dt * v + m * dv/dt.)
import numpy as np
import matplotlib.pyplot as plt
import scipy.integrate as si
%matplotlib inline
def getforce(t):
if 0<=t<0.15:
F = 40*t
elif 0.15<=t<0.7:
F = -9.09*t+7.36
elif 0.7<=t<1.25:
F = 1
elif 1.25<=t<1.65:
F = 7.5*t-8.375
elif 1.65<=t<1.8:
F = -26.6*t+48
else:
F = 0
return F
def getspeed(x,t):
Ft = getforce(t)
y0,y1 = x
dy0 = y1
dy1 = (Ft-0.0001277422*y1**2*np.sign(y1)-0.174)/0.0177
return dy0,dy1
t = np.linspace(0,10,100)
sol = si.odeint(getspeed,(0,0),t)
v=sol[:,0]
x=0
xs=[]
dt=t[1]-t[0] # use linspace with 101 to get the sample distance you'd normally expect
for i in range(len(v)):
x=x+v[i]*dt
xs.append(x)
plt.subplot(121)
plt.plot(t,v)
plt.subplot(122)
plt.plot(t,xs)
plt.show()
I didn't use numpy or lambda expressions to integrate to keep it easily readable and because the speed of execution is immaterial for this case.

Categories

Resources