Develop Reference

Gekko optimal control. How to create multiple termination objectives/conditions?

So I am simulating plane flight. The plane flies certain distance (pathx and pathy variables) and then the simulation stops, when certain pathx and pathy values are reached. The solver is trying to minimize the fuel consumption (m.Maximize(mass*tf*final), by maximizing the mass value. The solver controls the accelerator pedal position (Tcontr).
Right now the termination condition is defined like this:
For X axis:
m.Equation(x*final<=pathx)
and
m.Minimize(final*(x-pathx)**2)
And for Y axis:
m.Equation(y*final<=pathy)
and
m.Minimize(final*(y-pathy)**2)
And right now the simulation ends when the desired X value is achieved, while desired Y values is not being achieved.
How do I force the simulation to end when both desired (X and Y) values are achieved?
My code:
import numpy as np
import matplotlib.pyplot as plt
from gekko import GEKKO
import math
#Gekko model
m = GEKKO(remote=False)
#Time points
nt = 11
tm = np.linspace(0,100,nt)
m.time = tm
# Variables
Ro = m.Var(value=1.1)#air density
g = m.Const(value=9.80665)
pressure = m.Var(value=101325)#
T = m.Var(value=281,lb=100)#temperature
T0 = m.Const(value=288)#temperature at see level
S = m.Const(value=122.6)
Cd = m.Const(value=0.1)#drag coef
Cl = m.Var(value=1)#lift couef
FuelFlow = m.Var()
D = m.Var(value=25000,lb=0)#drag
Thrmax = m.Const(value=200000)#maximum throttle
Thr = m.Var()#throttle
V = m.Var(value=100,lb=50,ub=240)#velocity
gamma = m.Var(value=0)# Flight-path angle
gammaa = gamma.value
Xi = m.Var(value=0)# Heading angle
Xii = Xi.value
x = m.Var(value=0,lb=0)#x position
y = m.Var(value=0,lb=0)#y position
h = m.Var(value=1000,lb=0)# height
mass = m.Var(value=60000,lb=10000)
pathx = m.Const(value=50000) #intended distance length
pathy = m.Const(value=50000)
L = m.Var(value=0.1)#lift
p = np.zeros(nt)
p[-1] = 1.0
final = m.Param(value=p)
m.options.MAX_ITER=10000 # iteration number
#Fixed Variable
tf = m.FV(value=1,lb=0.0001,ub=1000.0)#
tf.STATUS = 1
# Controlled parameters
Tcontr = m.MV(value=0.2,lb=0.1,ub=1)# solver controls throttle pedal position
Tcontr.STATUS = 1
Tcontr.DCOST = 0
#Mu = m.Var(value=0)
Mu = m.MV(value=0,lb=-1.5,ub=1.5)# solver controls bank angle
Mu.STATUS = 1
Mu.DCOST = 0
Muu = Mu.value
# Equations
m.Equation(Thr==Tcontr*Thrmax)
m.Equation(FuelFlow==0.75882*(1+(V/2938.5)))
m.Equation(D==0.5*Ro*(V**2)*Cd*S)
m.Equation(mass.dt()==tf*(-Thr*(FuelFlow/60000)))#
m.Equation(V.dt()==tf*((Thr-D)/mass))#
m.Equation(x.dt()==tf*(V*(math.cos(gammaa.value))*(math.cos(Xii.value))))#
m.Equation(x*final<=pathx)
#pressure and density part
m.Equation(T==T0-(0.0065*h))
m.Equation(pressure==101325*(1-(0.0065*h)/T0)**((g*0.0289652)/(8.31446*0.0065)))#
m.Equation(Ro*(8.31446*T)==(pressure*0.0289652))
#2D addition part
m.Equation(L==0.5*Ro*(V**2)*Cl*S)
m.Equation(Xi.dt()==tf*((L*math.sin(Muu.value))/(mass*V)))
m.Equation(y.dt()==tf*(V*(math.cos(gammaa.value))*(math.sin(Xii.value))))#
m.Equation(y*final<=pathy)
# Objective Function
m.Minimize(final*(x-pathx)**2) #1D part
m.Minimize(final*(y-pathy)**2) #2D part
m.Maximize(mass*tf*final) #objective function
m.options.IMODE = 6
m.options.NODES = 2 # it was 3 before
m.options.MV_TYPE = 1
m.options.SOLVER = 3
#m.open_folder() # to search for infeasibilities
m.solve()
tm = tm * tf.value[0]
fig, axs = plt.subplots(8)
fig.suptitle('Results')
axs[0].plot(tm,Tcontr,'r-',LineWidth=2,label=r'$Tcontr$')
axs[0].legend(loc='best')
axs[1].plot(tm,V.value,'b-',LineWidth=2,label=r'$V$')
axs[1].legend(loc='best')
axs[2].plot(tm,x.value,'r--',LineWidth=2,label=r'$x$')
axs[2].legend(loc='best')
axs[3].plot(tm,D.value,'g-',LineWidth=2,label=r'$D$')
axs[3].legend(loc='best')
axs[4].plot(tm,mass.value,'g:',LineWidth=2,label=r'$mass$')
axs[4].legend(loc='best')
axs[5].plot(tm,T.value,'p-',LineWidth=2,label=r'$T$')
axs[5].legend(loc='best')
axs[6].plot(tm,Mu.value,'p-',LineWidth=2,label=r'$Mu$')
axs[6].legend(loc='best')
axs[7].plot(tm,y.value,'p-',LineWidth=2,label=r'$y$')
axs[7].legend(loc='best')
plt.xlabel('Time')
#plt.ylabel('Value')
plt.show()
Important Update
It looks like termination conditions (all of the above) work as intended. Y axis values can not be reached, because it looks like the math involving it is not working properly.
Y value is being calculated by this equation:
I have turned into this: m.Equation(y.dt()==tf*(V*(math.cos(gammaa.value))*(math.sin(Xii.value))))
Where: V is true air speed (works properly);
tf is a variable that controls simulation time (works properly);
gammaa is derived from gamma, which is flight-path angle, which is 0 in this simulation (works properly);
Xii is derived from Xi, which is heading angle (the main culprit).
Xi value is defined by this equation:
I turned it into this: m.Equation(Xi.dt()==tf*((L*math.sin(Muu.value))/(mass*V)))
Where: L is lift force (works properly);
mass is mass (works properly);
V is true air speed (works properly);
Muu is derived from Mu, which is bank angle, the solver controlled variable (probably the bad apple here).
Y value calculation should be working like this: solver changes Mu, which affects The heading angle Xi, which then should affect Y value.
After some experiments it looks like the range of changes of the Mu value during the simulation is so small, that it barely affects Y value. But it should be a lot wider since Mu boundaries are quite wide (lb=-1.5,ub=1.5). I tried to remove those boundaries, but it did not affect the results of the simulation.
What could be messing up everything here?

Try adding a hard terminal constraint for both:
m.Equation((x-pathx)*final==0)
m.Equation((y-pathy)*final==0)
Alternatively, increase the weight on the terminal condition.
w = 1e3
m.Minimize(w*final*(x-pathx)**2) #1D part
m.Minimize(w*final*(y-pathy)**2) #2D part
You may need to use one or both strategies. Terminal conditions can be challenging to solve.
Response to Important Update
Use the Gekko functions with m.sin() and m.cos() instead of the math functions. Also, don't use .value in the equations. Gekko needs the full symbolic graph so that it can compile the equations into byte-code and produce function evaluations and the derivatives with automatic differentiation.
m.Equation(y.dt()==tf*(V*(m.cos(gammaa))*(m.sin(Xii))))
It also helps to avoid divide by zero by multiplying any denominator over to the other side.
m.Equation(mass*V*Xi.dt()==tf*((L*m.sin(Muu))))

Numerical solution for a pendulum

I am having a trouble in providing a graphical representation of my system, which happens to be a harmonically driven pendulum. The problem is shown below for reference.
Problem
The source code I used is shown below using the Verlet Scheme.
#Import needed modules
import numpy as np
import matplotlib.pyplot as plt
#Initialize variables (Initial conditions)
g = 9.8 #Gravitational Acceleration
L = 2.0 #Length of the Pendulum
A0 = 3.0 #Initial amplitude of the driving acceleration
v0 = 0.0 #Initial velocity
theta0 = 90*np.pi/180 #Initial Angle
drivingPeriod = 20.0 #Driving Period
#Setting time array for graph visualization
tau = 0.1 #Time Step
tStop = 10.0 #Maximum time for graph visualization derived from Kinematics
t = np.arange(0., tStop+tau, tau) #Array of time
theta = np.zeros(len(t))
v = np.zeros(len(t))
#Verlet Method
theta[0] = theta0
v[0] = v0
for i in range(len(t)-1):
accel = -((g + (A0*np.sin((2*np.pi*t) / drivingPeriod)))/L) * np.sin(theta[i])
theta[i+1] = theta[i] + tau*v[i] + 0.5*tau**2*accel[i]
v[i+1] = v[i] + 0.5*tau*(accel[i] + accel[i+1])
#Plotting and saving the resulting graph
fig, ax1 = plt.subplots(figsize=(7.5,4.5))
ax1.plot(t,theta*(180/np.pi))
ax1.set_xlabel("Time (t)")
ax1.set_ylabel("Theta")
plt.show()
A sample output is shown.
Output
The pendulum should just go back to its initial angle. How can I solve this issue? Notice that as time evolves, my angle measure (degrees) also increases. I want it to only have a domain of 0 degrees to 360 degrees.

Please change the array computation
accel = -((g + (A0*np.sin((2*np.pi*t) / drivingPeriod)))/L) * np.sin(theta[i])
theta[i+1] = theta[i] + tau*v[i] + 0.5*tau**2*accel[i]
into the proper element computation at the correct place
theta[i+1] = theta[i] + tau*v[i] + 0.5*tau**2*accel[i]
accel[i+1] = -((g + (A0*np.sin((2*np.pi*t[i+1]) / drivingPeriod)))/L) * np.sin(theta[i+1])
Note that you need to compute accel[0] separately outside the loop.
It makes the code more readable if you separate out the specifics of the physical model and declare at the start
def acceleration(t,theta):
return -((g + (A0*np.sin((2*np.pi*t) / drivingPeriod)))/L) * np.sin(theta)
so that later you just call
accel[i+1]=acceleration(t[i+1],theta[i+1])
And even then, with a forced oscillation your system is open, it is possible that the forcing action pumps energy into the pendulum, causing it to start to rotate. This is what your graph shows.
The Verlet method like any symplectic method only promises to somewhat have a constant energy if the system is closed and conservative, that is, in the most common cases, no outside influence and all forces are gradient forces.

How to do space discretization in Gekko?

The goal is to minimize time to complete the lap with Energy constraint this is why my objective is the integral of the speed over distance, but I can’t seem to figure out how to derive and integrate over distance and not time(dt).

If you don't have time in your problem then you can specify m.time as the distance points for integration. However, your differential equations are based on time such as ds/dt = v in 1D. You need to keep time as the variable because that is defined for each of the differentials.
One way to minimize the lap time is to create a new tlap=FV() and then scale all of the differentials by that new adjustable value.
tlap=FV()
m.Equation(s.dt()==v*tlap)
With this tf value, you can minimize final time to reach a final destination.
m.Minimize(tf*final)
This is similar to the rocket launch problem that minimizes final time and control action.
import numpy as np
import matplotlib.pyplot as plt
from gekko import GEKKO
# create GEKKO model
m = GEKKO()
# scale 0-1 time with tf
m.time = np.linspace(0,1,101)
# options
m.options.NODES = 6
m.options.SOLVER = 3
m.options.IMODE = 6
m.options.MAX_ITER = 500
m.options.MV_TYPE = 0
m.options.DIAGLEVEL = 0
# final time
tf = m.FV(value=1.0,lb=0.1,ub=100)
tf.STATUS = 1
# force
u = m.MV(value=0,lb=-1.1,ub=1.1)
u.STATUS = 1
u.DCOST = 1e-5
# variables
s = m.Var(value=0)
v = m.Var(value=0,lb=0,ub=1.7)
mass = m.Var(value=1,lb=0.2)
# differential equations scaled by tf
m.Equation(s.dt()==tf*v)
m.Equation(mass*v.dt()==tf*(u-0.2*v**2))
m.Equation(mass.dt()==tf*(-0.01*u**2))
# specify endpoint conditions
m.fix_final(s, 10.0)
m.fix_final(v, 0.0)
# minimize final time
m.Minimize(tf)
# Optimize launch
m.solve()
print('Optimal Solution (final time): ' + str(tf.value[0]))
# scaled time
ts = m.time * tf.value[0]
# plot results
plt.figure(1)
plt.subplot(4,1,1)
plt.plot(ts,s.value,'r-',linewidth=2)
plt.ylabel('Position')
plt.legend(['s (Position)'])
plt.subplot(4,1,2)
plt.plot(ts,v.value,'b-',linewidth=2)
plt.ylabel('Velocity')
plt.legend(['v (Velocity)'])
plt.subplot(4,1,3)
plt.plot(ts,mass.value,'k-',linewidth=2)
plt.ylabel('Mass')
plt.legend(['m (Mass)'])
plt.subplot(4,1,4)
plt.plot(ts,u.value,'g-',linewidth=2)
plt.ylabel('Force')
plt.legend(['u (Force)'])
plt.xlabel('Time')
plt.show()

There are a few problems that I fixed with your current solution:
Variables w and st are not used
The STATUS for p_s and s_s should be On (1) to be calculated by the solver
The number of time points (50000) is really long and will create a very large problem that will be hard to solve in one solution. You may consider breaking this into successive solutions that advance one cycle (m.options.TIME_SHIFT=1) or multiple (m.options.TIME_SHIFT=10) for each m.solve() command.
There may be references that can help with the problem formulation. It appears that you are taking a more physics-based approach than a data driven approach.
Switched to the APOPT solver for a successful solution.
from gekko import GEKKO
import numpy as np
import matplotlib.pyplot as plt
m = GEKKO(remote=False)
#Constants
mass = m.Const(77) #mass of the rider
g = m.Const(9.81) #gravity
H = m.Const(1.2) #height of the rider
L = m.Const(value=1.4) #lenght of the wheelbase of the bicycle
E_n = m.Const(value=22000) #Energy that can be used
c_rr = m.Const(value=0.0035) #coefficient of drag
s_max = m.Const(value=0.52) #max steer angle
W_m = m.Const(value=1800) #max power that the rider can produce
vWn = m.Const(value=50) #maximal power output variation
vSn = m.Const(value=0.52) #maximal steer output variation
kv = m.Const(value=0.13) #air drag coefficient
ws = m.Const(value=0) #wind speed
Ix = m.Const(value=77) #inertia
W_c = m.Const(value=440) #critical power(watts)
Wj1 = m.Const(value=0.01) ##weighting factors that scale the penalisation
Wj2 = m.Const(value=0.01) #weighting factors that scale the penalisation
dist = 1000 ##distance that that the rider needs to travel
nt = 100 ##calculation at every 10 meters
m.time = np.linspace(0,dist,nt)
p = np.zeros(nt)
p[-1] = 1.0
final = m.Param(value=p)
slope = np.zeros(nt) #SET THE READ CURVATURE AND SLOPE TO 0 for experimentation later we will import it from real road.
curv = np.zeros(nt) #SET THE READ CURVATURE AND SLOPE TO 0 for experimentation later we will import it from real road.
####Import Road Characterisitc####
k = m.Param(value=curv) ##road_curvature
b = m.Param(value=slope) ##slope angle
###Control Variable###
p_s = m.MV(value=1,lb=-1000,ub=1000); p_s.STATUS = 1 ##power
s_s = m.MV(value=0,lb=-100,ub=100); s_s.STATUS = 1 ##steer
###State Variable###
# Not used
#w = m.Param(value=10,lb=-10000,ub=1800) #power done by the rider (positive:pedaling, negative:braking)
#st = m.Param(value=0,lb=-30,ub=30) ##steer angle
s = m.Var(value=1,lb=1e-4,ub=100) #speed along road
v = m.Var(value=1, lb=0, ub=16) #velocity
n = m.Var(value=0,lb=-4, ub=4) ##displacement fron the center of the road upper bound and lower bound are the road width
h = m.Var(value=0,lb=-10,ub=10) #heading of the bicycle
r = m.Var(0,lb=-0.78, ub=0.78) ##roll
r_dot = m.Var(value=0,lb=-100,ub=100) ##roll_rate
W_n = m.Var(value=0.1,lb=-1, ub=1) ##normalised power
s_n = m.Var(value=0,lb=-1, ub=1) #normalised steer angle
e = m.Var(value=22000, lb=0, ub=22000) #energy remaining
####Equations####
#1 dynamics of travelling speed s(s) along the reference line
m.Equation((1-(n-k))*s.dt()==v*m.cos(h))
#2:dynamics of the longitudinal velocity of the bicycle
c1 = m.Intermediate((v*mass)/W_m,'c1')
m.Equation(c1*s*v.dt()==(W_n
-( (v/W_m) * (mass*g* (c_rr* m.cos(b)+m.sin(b))) )
-((v/W_m) * kv*(v-(ws*h))**2)
)
)
#3: dynamic of the lateral displacement
m.Equation(s*n.dt()==m.sin(k))
#4: heading of the bicycle 𝛼(s):
m.Equation((L*s)*h.dt()==(s_n*s_max)-k*(L*s))
#5&6: dynamics of the roll angle 𝜙 (rad) and its rate of change 𝜙dot(s)
m.Equation(s*r.dt()==(r_dot))
m.Equation(((h**2)*mass+Ix)*(g*L*s)*r_dot.dt()==(H*mass*g)*((v**2)*s_max*s_n+L*g*r))
#7: dynamics of the normalised power output Wn
m.Equation(s*W_n.dt()==p_s)
##8: dynamics of the normalised steering angle 𝛿n
m.Equation(s*s_n.dt()==s_s)
#9: dynamic equation describing the evolution of the anaerobic sources
# use lower bound on W_n instead of m.min2(0,W_n)
m.Equation((s*E_n)*e.dt()==(-(W_n*W_m-W_c) ))
####OBJECTIVE####
m.Minimize(m.integral( (1/s) * (1+(Wj1*((p_s/vWn)**2))+(Wj2*((s_s/vSn)**2))) )*final)
m.options.IMODE = 6 # optimal control
m.options.SOLVER = 1 # solver (APOPT)
m.options.DIAGLEVEL=0
#m.open_folder()
m.solve(disp=True, debug=True) # Solve
With this script, I get a successful solution but I haven't investigated the objective function to see if it is giving a reasonable answer.
----------------------------------------------------------------
APMonitor, Version 0.9.2
APMonitor Optimization Suite
----------------------------------------------------------------
--------- APM Model Size ------------
Each time step contains
Objects : 0
Constants : 16
Variables : 15
Intermediates: 1
Connections : 0
Equations : 12
Residuals : 11
Number of state variables: 2970
Number of total equations: - 2772
Number of slack variables: - 0
---------------------------------------
Degrees of freedom : 198
----------------------------------------------
Dynamic Control with APOPT Solver
----------------------------------------------
Iter Objective Convergence
0 2.51001E+03 1.00000E+00
1 4.36075E+04 5.66676E-01
2 3.43092E+03 5.36156E-01
3 7.36773E+03 4.16203E-01
4 2.75250E+03 9.29407E-02
5 4.12278E+03 1.93521E-02
6 5.80466E+05 7.35244E-02
7 4.99119E+04 1.27246E-01
8 2.11556E+03 4.52552E-01
9 6.32932E+03 2.14605E-01
Iter Objective Convergence
10 8.16639E+01 2.76062E-01
11 6.80002E+02 8.83214E-01
12 4.71706E+01 2.87555E-01
13 1.28152E+02 1.36994E-03
14 1.01698E+01 1.08406E+01
15 1.13082E+01 3.00869E+00
16 1.03199E+01 8.67971E+00
17 1.02638E+01 1.28697E-02
18 1.02636E+01 5.64896E-05
19 1.02636E+01 6.72710E-07
Iter Objective Convergence
20 1.02636E+01 6.72710E-07
Successful solution
---------------------------------------------------
Solver : APOPT (v1.0)
Solution time : 3.1271 sec
Objective : 10.263550885927089
Successful solution
---------------------------------------------------
You may want to create plots to make sure that the equations and solver are giving a correct solution. Here is an animation and source code that shows how to set up a model predictive controller with a finite horizon and that advances in time (or space) for each solve command.
The finite horizon approach is used commonly in industrial control to ensure that the optimizer can finish within the required cycle time and balance that with the length of the horizon to "see" future constraints and opportunities for energy or production optimization.

planetary obliquity, trace rotation axis along orbit with python

I'm trying to make a simple simulation of a planet that is being orbited by a moon. So far I have a 2 body problem that solves the planet and moon orbit. Now I would like to add a fixed rotation axis to the planet and see how it is affected by the moon. Any idea how this can be done by using python?
The two body problem can be run with the code below:
import pylab
import math
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
# Set Constants
G = 6.67e-11
AU = 1.5e11
daysec = 24.0*60*60
Ma =5.972e24 # Planet mass in Kg
Mb = 7.348e22 # Moon mass in Kg
gravconst = G*Ma*Mb
# Set up starting conditions
# Planet
xa = 0.0
ya = 0.0
za = 0.0
xva = 0.0
yva = 0.0
zva = 0.0
# Moon
xb = 384400000
yb = 0.0
zb = 0.0
xvb = 0.0
yvb = 1000.0
zvb = 0.0
# Time steps
t = 0.0
dt = 0.01*daysec
# Coordinate lists
xalist = []
yalist = []
xblist = []
yblist = []
zalist = []
zblist = []
# Loop
while t < 100.0*daysec:
# Compute Force
rx = xb-xa
ry = yb-ya
rz = zb-za
modr3 = (rx**2+ry**2+rz**2)**1.5
fx = -gravconst*rx/modr3
fy = -gravconst*ry/modr3
fz = -gravconst*rz/modr3
# Update quantities
# Moon
xvb += fx*dt/Mb
yvb += fy*dt/Mb
zvb += fz*dt/Mb
xb += xvb*dt
yb += yvb*dt
zb += zvb*dt
# Planet
xva += -fx*dt/Ma
yva += -fy*dt/Ma
zva += -fz*dt/Ma
xa += xva*dt
ya += yva*dt
za += zva*dt
t += dt
# Saving them in lists
xalist.append(xa)
yalist.append(ya)
zalist.append(za)
xblist.append(xb)
yblist.append(yb)
zblist.append(zb)
xalist[:] = [x / 1e6 for x in xalist]
yalist[:] = [x / 1e6 for x in yalist]
zalist[:] = [x / 1e6 for x in zalist]
xblist[:] = [x / 1e6 for x in xblist]
yblist[:] = [x / 1e6 for x in yblist]
zblist[:] = [x / 1e6 for x in zblist]
#Creating the point to represent the planet at the origin (not to scale),
plt.scatter(0,0,s=200,color='blue')
plt.annotate('Planet', xy=(-45,-50))
plt.scatter(xblist[0],0,s=100,color='grey')
plt.annotate('Mond', xy=(xblist[0]-45,-50))
# Plotting
pylab.plot(xalist, yalist, "-g")
pylab.plot(xblist, yblist, "-r")
plt.axhline(0, color='black')
plt.axvline(0, color='black')
pylab.axis("equal")
pylab.xlabel("X (Mio. Meter)")
pylab.ylabel("Y (Mio. Meter)")
pylab.show()

Not an answer as I am no expert in the matter but just some hints instead (was unreadable in form of comment)
The stuff you want to add is quite complicated as you would need take into account:
moving masses of both bodies
so you need "contact surfaces elevation" of body that has any moving masses (like oceans, magma, rotating core etc) so you can compute the real center of gravity for each time. Also you need to apply forces to the moving mass itself too (which is not driven just by gravity and rotation but also with resonance and mainly inertia) do not forget Earth has also core and magma around it not just oceans so you need to take into account at least 3 surfaces...
non homogenity of mass distribution of both bodies
so you can compute center of gravity, and the quadratic mass inertia for rotations in respect to actual rotation axis
planet/moon is usually at least 3 body problem not just 2
as there is the local star too affecting the moon orbit quite a lot...
Depending on the numbers some effects will be so small that can be discarded but some are not (especially with inertia+resonance in place).
The rotation equations are similar to the position/speed/acceleration like you already got. Its called Newton D'Alembert integration/physics but you would need to implement transform matrices to do this.
See few related QAs:
Rigid Body Physics Rotations
realistic n-body solar system
Pay attention to the last link for accuracy improvement of your integration as right now its very bad and the orbit will be deformed no matter how small the dt due to fact that the gravity vector changes with time but right now you affect it in wrong direction (that is correct only at the start of dt interval) for each integration iteration.
As you can see that is a lot of stuff to handle and most simulation programs I saw do not do it (mine included) ... they fake it by nutation and precession constants instead.

Simulating electron motion - differential equation with adaptive step size in python

I am trying to simulate how the oscillating electric field of an intense laser will push around an electron that is near the Coulomb potential of a +1 ion. The laser field is
E = Eo sin(wt), in the y direction.
and the Coulomb potential is
F = ke q1*q2/r^2, in the r direction.
The strong electric field causes the electron to tunnel ionize, so the initial condition of the electron is to be displaced from the atom in the y-direction. Then, the electron is pushed back and forth by the laser field and has a chance to interact with the Coulomb potential. I want to simulate how the Coulomb potential affects the flight of the electron. The simulations need to be in three dimensions, because I eventually want to include more complex laser fields that push the electron in the x and y directions and the electron can start with momentum in the z direction.
At first, I thought that this would be easy. Below is the code that I used to step through time in very small steps (1e-18 sec). When the electron is not near the ion, this works fine. However, for electrons that pass close to the ion, the results depend strongly on the time-step used in the simulations. If I make the time-step smaller, the calculations take a very long time.
So, I think in this case I am supposed to use an adaptive timestep. Also, from what I have read, the Runge-Kutta methods are supposed to be superior to the simple approach I am using. However, I don't think that scipy.odeint applies to three-dimensions. Any ideas on how to improve the accuracy and speed of these simulations?
Here is the figure showing how the time-step has a huge impact on the results (a bad thing):
And here is my code:
import numpy as np
import matplotlib.pyplot as plt
q = 1.602e-19 #Coulombs Charge of electron
h_bar = 1.054e-34 #J*s Plank's Constant div by 2Pi
c = 3.0e8 #m/s Speed of light
eo = 8.8541e-12 #C^2/(Nm^2) Permittivity of vacuum
me = 9.109e-31 #kg Mass of electron
ke = 8.985551e9 #N m^2 C-2 Coulomb's constant
def fly_trajectory(wavelength,intensity,tb=0,pulseFWHM=40.e-15,
final_time=100e-15,time_step=.001e-15,Ip=15.13,v0=(2e4,0,0)):
#Intensity is in w/cm2. Ip is in eV. Otherwise it's SI units throughout.
#The electric field of the later is in the y-direction
Ip = 15.13 * q #Calculate the ionization potential of the atom in Joules
Eo = np.sqrt(2*intensity*10**4/(c*8.85e-12)) # Electric field in V/m
w = c/wavelength * 2. * np.pi #Angular frequency of the laser
times = np.arange(tb,final_time,time_step)
Ey = Eo*np.sin(w*times) * np.exp(-times**2/(2*(pulseFWHM / 2.35482)**2))
Eb = Ey[0] #E-field at time of birth (time of tunneling)
if Eb == 0: return 0,0 #No field --> no electrons
tunnel_position = -Ip / (Eb*q)
x,y,z = 0,tunnel_position,0
vx,vy,vz = v0
y_list = np.zeros(len(times)) #store the y-values for later
for index in range(0,len(times)):
r = np.sqrt(x**2+y**2+z**2)
rx = x/r; ry = y/r; rz=z/r
Fcy = -q**2 * ke/(r**2) * ry
ay = Ey[index]*(-q)/me + Fcy/me #only y includes the laser
vy = vy + ay*time_step
y = y + vy * time_step
Fcx = -q**2 * ke/(r**2) * rx
ax = (-q)/me + Fcx/me
vx = vx + ax*time_step
x = x + vx * time_step
Fcz = -q**2 * ke/(r**2) * rz
az = (-q)/me + Fcz/me
vz = vz + az*time_step
z = z + vz * time_step
y_list[index] = y
return times,y_list
for tb in np.linspace(0.25*2.66e-15,0.5*2.66e-15,5):
print tb
times,ys = fly_trajectory(800e-9,2e14,tb=tb,time_step=.01e-15)
plt.plot(times,ys,color='r')
times,ys = fly_trajectory(800e-9,2e14,tb=tb,time_step=.001e-15)
plt.plot(times,ys,color='b')
#plot legend and labels:
plt.plot(0,0,color='r',label='10e-18 sec step')
plt.plot(0,0,color='b',label='1e-18 sec step')
plt.xlim(0,10e-15); plt.ylim(-1e-8,1e-8)
leg = plt.legend(); leg.draw_frame(False)
plt.xlabel('Time (sec)')
plt.ylabel('Y-distance (meters)')
plt.show()

As Warren Weckesser suggested, I can simply follow the Scipy cookbook for the coupled mass-spring system. First, I need to write my "right side" equations as:
x' = vx
y' = vy
z' = vz
vx' = Ac*x/r
vy' = Ac*y/r + q*E/m
vz' = Ac*z/r
where Ac=keq^2/(mr^2) is the magnitude of the acceleration due to the Coulomb potential and E is the time-dependent electric field of the laser. Then, I can use scipy.integrate.odeint to find the solutions. This is faster and more reliable than the method that I was using previously.
Here is what the electron trajectories look like with odeint. Now none of them fly away crazily:
And here is the code:
import numpy as np
import matplotlib.pyplot as plt
import scipy.integrate
q = 1.602e-19 #Coulombs Charge of electron
c = 3.0e8 #m/s Speed of light
eo = 8.8541e-12 #C^2/(Nm^2) Permittivity of vacuum
me = 9.109e-31 #kg Mass of electron
ke = 8.985551e9 #N m^2 C-2 Coulomb's constant
def tunnel_position(tb,intensity,wavelength,pulseFWHM,Ip):
Ip = 15.13 * q
Eb = E_laser(tb,intensity,wavelength,pulseFWHM)
return -Ip / (Eb*q)
def E_laser(t,intensity,wavelength,pulseFWHM):
w = c/wavelength * 2. * np.pi #Angular frequency of the laser
Eo = np.sqrt(2*intensity*10**4/(c*8.85e-12)) # Electric field in V/m
return Eo*np.sin(w*t) * np.exp(-t**2/(2*(pulseFWHM / 2.35482)**2))
def vectorfield(variables,t,params):
x,y,z,vx,vy,vz = variables
intensity,wavelength,pulseFWHM,tb = params
r = np.sqrt(x**2+y**2+z**2)
Ac = -ke*q**2/(r**2*me)
return [vx,vy,vz,
Ac*x/r,
Ac*y/r + q/me * E_laser((t-tb),intensity,wavelength,pulseFWHM),
Ac*z/r]
Ip = 15.13 # Ionization potential of Argon eV
intensity = 2e14
wavelength = 800e-9
pulseFWHM = 40e-15
period = wavelength/c
t = np.linspace(0,20*period,10000)
birth_times = np.linspace(0.01*period,0.999*period,50)
max_field = np.max(np.abs(E_laser(birth_times,intensity,wavelength,pulseFWHM)))
for tb in birth_times:
x0 = 0
y0 = tunnel_position(tb,intensity,wavelength,pulseFWHM,Ip)
z0 = 0
vx0 = 2e4
vy0 = 0
vz0 = 0
p = [intensity,wavelength,pulseFWHM,tb]
w0 = [x0,y0,z0,vx0,vy0,vz0]
solution,info = scipy.integrate.odeint(vectorfield,w0,t, args=(p,),full_output=True)
print 'Tb: %.2f fs - smallest step : %.05f attosec'%((tb*1e15),np.min(info['hu'])*1e18)
y = solution[:,1]
importance = (np.abs(E_laser(tb,intensity,wavelength,pulseFWHM))/max_field)
plt.plot(t,y,alpha=importance*0.8,lw=1)
plt.xlabel('Time (sec)')
plt.ylabel('Y-distance (meters)')
plt.show()