Related
I have implemented a simple randomized, population-based optimization method - Grey Wolf optimizer. I am having some trouble with properly capturing the Matplotlib plots at each iteration using the camera package.
I am running GWO for the objective function f(x,y) = x^2 + y^2. I can only see the candidate solutions converging to the minima, but the contour plot doesn't show up.
Do you have any suggestions, how can I display the contour plot in the background?
GWO Algorithm implementation
%matplotlib notebook
import matplotlib.pyplot as plt
import numpy as np
from celluloid import Camera
import ffmpeg
import pillow
# X : Position vector of the initial population
# n : Initial population size
def gwo(f,max_iterations,LB,UB):
fig = plt.figure()
camera = Camera(fig)
def random_population_uniform(m,a,b):
dims = len(a)
x = [list(a + np.multiply(np.random.rand(dims),b - a)) for i in range(m)]
return np.array(x)
def search_agent_fitness(fitness):
alpha = 0
if fitness[1] < fitness[alpha]:
alpha, beta = 1, alpha
else:
beta = 1
if fitness[2] > fitness[alpha] and fitness[2] < fitness[beta]:
beta, delta = 2, beta
elif fitness[2] < fitness[alpha]:
alpha,beta,delta = 2,alpha,beta
else:
delta = 2
for i in range(3,len(fitness)):
if fitness[i] <= fitness[alpha]:
alpha, beta,delta = i, alpha, beta
elif fitness[i] > fitness[alpha] and fitness[i]<= fitness[beta]:
beta,delta = i,beta
elif fitness[i] > fitness[beta] and fitness[i]<= fitness[delta]:
delta = i
return alpha, beta, delta
def plot_search_agent_positions(f,X,alpha,beta,delta,a,b):
# Plot the positions of search agents
x = X[:,0]
y = X[:,1]
s = plt.scatter(x,y,c='gray',zorder=1)
s = plt.scatter(x[alpha],y[alpha],c='red',zorder=1)
s = plt.scatter(x[beta],y[beta],c='blue',zorder=1)
s = plt.scatter(x[delta],y[delta],c='green',zorder=1)
camera.snap()
# Initialize the position of the search agents
X = random_population_uniform(50,np.array(LB),np.array(UB))
n = len(X)
l = 1
# Plot the first image on screen
x = np.linspace(LB[0],LB[1],1000)
y = np.linspace(LB[0],UB[1],1000)
X1,X2 = np.meshgrid(x,y)
Z = f(X1,X2)
cont = plt.contour(X1,X2,Z,20,linewidths=0.75)
while (l < max_iterations):
# Take the x,y coordinates of the initial population
x = X[:,0]
y = X[:,1]
# Calculate the objective function for each search agent
fitness = list(map(f,x,y))
# Update alpha, beta and delta
alpha,beta,delta = search_agent_fitness(fitness)
# Plot search agent positions
plot_search_agent_positions(f,X,alpha,beta,delta,LB,UB)
# a decreases linearly from 2 to 0
a = 2 - l *(2 / max_iterations)
# Update the position of search agents including the Omegas
for i in range(n):
x_prey = X[alpha]
r1 = np.random.rand(2) #r1 is a random vector in [0,1] x [0,1]
r2 = np.random.rand(2) #r2 is a random vector in [0,1] x [0,1]
A1 = 2*a*r1 - a
C1 = 2*r2
D_alpha = np.abs(C1 * x_prey - X[i])
X_1 = x_prey - A1*D_alpha
x_prey = X[beta]
r1 = np.random.rand(2)
r2 = np.random.rand(2)
A2 = 2*a*r1 - a
C2 = 2*r2
D_beta = np.abs(C2 * x_prey - X[i])
X_2 = x_prey - A2*D_beta
x_prey = X[delta]
r1 = np.random.rand(2)
r2 = np.random.rand(2)
A3 = 2*a*r1 - a
C3 = 2*r2
D_delta = np.abs(C3 * x_prey - X[i])
X_3 = x_prey - A3*D_delta
X[i] = (X_1 + X_2 + X_3)/3
l = l + 1
return X[alpha],camera
Function call
# define the objective function
def f(x,y):
return x**2 + y**2
minimizer,camera = gwo(f,7,[-10,-10],[10,10])
animation = camera.animate(interval = 1000, repeat = True,
repeat_delay = 500)
Is it possible that the line x = np.linspace(LB[0],LB[1],1000) should be x = np.linspace(LB[0],UB[1],1000) instead? With your current definition of x, x is an array only filled with the value -10 which means that you are unlikely to find a contour.
Another thing that you might want to do is to move the cont = plt.contour(X1,X2,Z,20,linewidths=0.75) line inside of your plot_search_agent_positions function to ensure that the contour is plotted at each iteration of the animation.
Once you make those changes, the code looks like that:
import matplotlib.pyplot as plt
import numpy as np
from celluloid import Camera
import ffmpeg
import PIL
from matplotlib import animation, rc
from IPython.display import HTML, Image # For GIF
from scipy.interpolate import griddata
rc('animation', html='html5')
# X : Position vector of the initial population
# n : Initial population size
def gwo(f,max_iterations,LB,UB):
fig = plt.figure()
fig.gca(aspect='equal')
camera = Camera(fig)
def random_population_uniform(m,a,b):
dims = len(a)
x = [list(a + np.multiply(np.random.rand(dims),b - a)) for i in range(m)]
return np.array(x)
def search_agent_fitness(fitness):
alpha = 0
if fitness[1] < fitness[alpha]:
alpha, beta = 1, alpha
else:
beta = 1
if fitness[2] > fitness[alpha] and fitness[2] < fitness[beta]:
beta, delta = 2, beta
elif fitness[2] < fitness[alpha]:
alpha,beta,delta = 2,alpha,beta
else:
delta = 2
for i in range(3,len(fitness)):
if fitness[i] <= fitness[alpha]:
alpha, beta,delta = i, alpha, beta
elif fitness[i] > fitness[alpha] and fitness[i]<= fitness[beta]:
beta,delta = i,beta
elif fitness[i] > fitness[beta] and fitness[i]<= fitness[delta]:
delta = i
return alpha, beta, delta
def plot_search_agent_positions(f,X,alpha,beta,delta,a,b,X1,X2,Z):
# Plot the positions of search agents
x = X[:,0]
y = X[:,1]
s = plt.scatter(x,y,c='gray',zorder=1)
s = plt.scatter(x[alpha],y[alpha],c='red',zorder=1)
s = plt.scatter(x[beta],y[beta],c='blue',zorder=1)
s = plt.scatter(x[delta],y[delta],c='green',zorder=1)
Z=f(X1,X2)
cont=plt.contour(X1,X2,Z,levels=20,colors='k',norm=True)
plt.clabel(cont, cont.levels, inline=True, fontsize=10)
camera.snap()
# Initialize the position of the search agents
X = random_population_uniform(50,np.array(LB),np.array(UB))
n = len(X)
l = 1
# Plot the first image on screen
x = np.linspace(LB[0],UB[1],1000)
y = np.linspace(LB[0],UB[1],1000)
X1,X2 = np.meshgrid(x,y)
Z=f(X1,X2)
while (l < max_iterations):
# Take the x,y coordinates of the initial population
x = X[:,0]
y = X[:,1]
# Calculate the objective function for each search agent
fitness = list(map(f,x,y))
# Update alpha, beta and delta
alpha,beta,delta = search_agent_fitness(fitness)
# Plot search agent positions
plot_search_agent_positions(f,X,alpha,beta,delta,LB,UB,X1,X2,Z)
# a decreases linearly from 2 to 0
a = 2 - l *(2 / max_iterations)
# Update the position of search agents including the Omegas
for i in range(n):
x_prey = X[alpha]
r1 = np.random.rand(2) #r1 is a random vector in [0,1] x [0,1]
r2 = np.random.rand(2) #r2 is a random vector in [0,1] x [0,1]
A1 = 2*a*r1 - a
C1 = 2*r2
D_alpha = np.abs(C1 * x_prey - X[i])
X_1 = x_prey - A1*D_alpha
x_prey = X[beta]
r1 = np.random.rand(2)
r2 = np.random.rand(2)
A2 = 2*a*r1 - a
C2 = 2*r2
D_beta = np.abs(C2 * x_prey - X[i])
X_2 = x_prey - A2*D_beta
x_prey = X[delta]
r1 = np.random.rand(2)
r2 = np.random.rand(2)
A3 = 2*a*r1 - a
C3 = 2*r2
D_delta = np.abs(C3 * x_prey - X[i])
X_3 = x_prey - A3*D_delta
X[i] = (X_1 + X_2 + X_3)/3
l = l + 1
return X[alpha],camera
# define the objective function
def f(x,y):
return x**2 + y**2
minimizer,camera = gwo(f,7,[-10,-10],[10,10])
animation = camera.animate(interval = 1000, repeat = True,repeat_delay = 500)
And the output gives:
I'm trying to solve a simple minimum time optimal control problem using double integrator dynamics of the form,
dx1/dt = x2
dx2/dt = u
with the GEKKO optimization framework as follows:
from gekko import GEKKO
import numpy as np
import matplotlib.pyplot as plt
model = GEKKO(remote=False)
x1_initial = 0.0
x1_final = 10.0
x2_initial = 0.0
x2_final = 0.0
t_initial = 0.0
t_final = 25.0
num_timesteps = 1000
dt = (t_final - t_initial) / num_timesteps
x = model.Array(model.Var, (2, num_timesteps + 1))
u = model.Array(model.Var, num_timesteps + 1)
tf = model.Var()
for k in range(num_timesteps + 1):
u[k].lower = -0.4
u[k].upper = 0.4
u[k].value = 0.0
for k in range(num_timesteps + 1):
x[0, k].value = 5.0
x[1, k].value = 0.0
tf.lower = t_initial
tf.upper = t_final
tf.value = t_final
dt = (tf - t_initial) / num_timesteps
def f(x, u, k):
return np.array([x[1,k], u[k]])
for k in range(num_timesteps):
model.Equations([x[:, k + 1] == x[:, k] + (dt/2.0)*(f(x, u, k + 1) + f(x, u, k))])
# model.Equation(x[0, k + 1] == x[0, k] + (dt/2.0)*(x[1, k + 1] + x[1, k]))
# model.Equation(x[1, k + 1] == x[1, k] + (dt/2.0)*(u[k + 1] + u[k]))
model.Equation(x[0, 0] == x1_initial)
model.Equation(x[0, num_timesteps] == x1_final)
model.Equation(x[1, 0] == x2_initial)
model.Equation(x[1, num_timesteps] == x2_final)
model.Minimize(tf)
model.options.solver = 3
model.solve()
# Plotting results
t = np.linspace(t_initial, tf.value, num_timesteps + 1)
u_optimal = []
for k in range(num_timesteps + 1):
u_optimal.append(u[k].value)
x1_optimal = []
for k in range(num_timesteps + 1):
x1_optimal.append(x[0, k].value)
x2_optimal = []
for k in range(num_timesteps + 1):
x2_optimal.append(x[1, k].value)
plt.figure()
plt.plot(t, u_optimal)
plt.xlabel('time (s)')
plt.ylabel('u(t)')
plt.grid()
plt.figure()
plt.plot(t, x1_optimal)
plt.xlabel('time (s)')
plt.ylabel('x1(t)')
plt.grid()
plt.figure()
plt.plot(t, x2_optimal)
plt.xlabel('time (s)')
plt.ylabel('x2(t)')
plt.grid()
plt.show()
What I'm trying to do is to form a system of equality constraints using trapezoidal integration and then solve this system for the optimal control inputs using GEKKO. However, using the function definition,
def f(x, u, k):
return np.array([x[1,k], u[k]])
in conjunction with the system of equality constraints,
for k in range(num_timesteps):
model.Equations([x[:, k + 1] == x[:, k] + (dt/2.0)*(f(x, u, k + 1) + f(x, u, k))])
gives me the following error,
Exception: #error: Equation Definition
Equation without an equality (=) or inequality (>,<)
false
STOPPING...
I've added two commented lines of code in the above code snippet that will allow the program to run correctly, but I'm hoping to avoid having to separate each equation out, since I'd like to extend this to problems that deal with more complicated system dynamics, and to also use more sophisticated collocation methods instead of the trapezoidal approach.
I know that GEKKO has some nice features for dynamic optimization, but I'm looking to try and implement various direct collocation methods myself to understand the theory a bit better.
There are some examples of orthogonal collocation on finite elements in the Machine Learning and Dynamic Optimization course.
from __future__ import division
import numpy as np
from scipy.optimize import fsolve
from scipy.integrate import odeint
import matplotlib.pyplot as plt
# final time
tf = 1.0
# solve with ODEINT (for comparison)
def model(x,t):
u = 4.0
return (-x**2 + u)/5.0
t = np.linspace(0,tf,20)
y0 = 0
y = odeint(model,y0,t)
plt.figure(1)
plt.plot(t,y,'r-',label='ODEINT')
# ----------------------------------------------------
# Approach #1 - Write the model equations in Python
# ----------------------------------------------------
# define collocation matrices
def colloc(n):
if (n==2):
NC = np.array([[1.0]])
if (n==3):
NC = np.array([[0.75,-0.25], \
[1.00, 0.00]])
if (n==4):
NC = np.array([[0.436,-0.281, 0.121], \
[0.614, 0.064, 0.0461], \
[0.603, 0.230, 0.167]])
if (n==5):
NC = np.array([[0.278, -0.202, 0.169, -0.071], \
[0.398, 0.069, 0.064, -0.031], \
[0.387, 0.234, 0.278, -0.071], \
[0.389, 0.222, 0.389, 0.000]])
if (n==6):
NC = np.array([[0.191, -0.147, 0.139, -0.113, 0.047],
[0.276, 0.059, 0.051, -0.050, 0.022],
[0.267, 0.193, 0.252, -0.114, 0.045],
[0.269, 0.178, 0.384, 0.032, 0.019],
[0.269, 0.181, 0.374, 0.110, 0.067]])
return NC
# define collocation points from Lobatto quadrature
def tc(n):
if (n==2):
time = np.array([0.0,1.0])
if (n==3):
time = np.array([0.0,0.5,1.0])
if (n==4):
time = np.array([0.0, \
0.5-np.sqrt(5)/10.0, \
0.5+np.sqrt(5)/10.0, \
1.0])
if (n==5):
time = np.array([0.0,0.5-np.sqrt(21)/14.0, \
0.5,0.5+np.sqrt(21)/14.0, 1])
if (n==6):
time = np.array([0.0, \
0.5-np.sqrt((7.0+2.0*np.sqrt(7.0))/21.0)/2.0, \
0.5-np.sqrt((7.0-2.0*np.sqrt(7.0))/21.0)/2.0, \
0.5+np.sqrt((7.0-2.0*np.sqrt(7.0))/21.0)/2.0, \
0.5+np.sqrt((7.0+2.0*np.sqrt(7.0))/21.0)/2.0, \
1.0])
return time*tf
# solve with SciPy fsolve
def myFunction(z,*param):
n = param[0]
m = param[1]
# rename z as x and xdot variables
x = np.empty(n-1)
xdot = np.empty(n-1)
x[0:n-1] = z[0:n-1]
xdot[0:n-1] = z[n-1:m]
# initial condition (x0)
x0 = 0.0
# input parameter (u)
u = 4.0
# final time
tn = tf
# function evaluation residuals
F = np.empty(m)
# nonlinear differential equations at each node
# 5 dx/dt = -x^2 + u
F[0:n-1] = 5.0 * xdot[0:n-1] + x[0:n-1]**2 - u
# collocation equations
# tn * NC * xdot = x - x0
NC = colloc(n)
F[n-1:m] = tn * np.dot(NC,xdot) - x + x0 * np.ones(n-1)
return F
sol_py = np.empty(5) # store 5 results
for i in range(2,7):
n = i
m = (i-1)*2
zGuess = np.ones(m)
z = fsolve(myFunction,zGuess,args=(n,m))
# add to plot
yc = np.insert(z[0:n-1],0,0)
plt.plot(tc(n),yc,'o',markersize=10,label='Nodes = '+str(i))
# store just the last x[n] value
sol_py[i-2] = z[n-2]
plt.legend(loc='best')
# ----------------------------------------------------
# Approach #2 - Write model in APMonitor and let
# modeling language create the collocation equations
# ----------------------------------------------------
# load GEKKO
from gekko import GEKKO
sol_apm = np.empty(5) # store 5 results
i = 0
for nodes in range(2,7):
m = GEKKO(remote=False)
u = m.Param(value=4)
x = m.Var(value=0)
m.Equation(5*x.dt() == -x**2 + u)
m.time = [0,tf]
m.options.imode = 4
m.options.time_shift = 0
m.options.nodes = nodes
m.solve() # solve problem
sol_apm[i] = x.value[-1] # store solution (last point)
i += 1
# print the solutions
print(sol_py)
print(sol_apm)
# show plot
plt.ylabel('x(t)')
plt.xlabel('time')
plt.show()
You can define variables with the same name (such as x) or use m.Array(m.Var,n) to define the variables. One thing to look at is the model file by opening the run folder with m.open_folder() before you send the m.solve() command. Look at the .apm file in that folder with a text editor.
I am trying to make a plot of a projectile motion of a mass which is under the effect of gravitational, buoyancy and drag force. Basically, I want to show effects of the buoyancy and drag force on flight distance, flight time and velocity change on plotting.
import matplotlib.pyplot as plt
import numpy as np
V_initial = 30 # m/s
theta = np.pi/6 # 30
g = 3.711
m =1
C = 0.47
r = 0.5
S = np.pi*pow(r, 2)
ro_mars = 0.0175
t_flight = 2*(V_initial*np.sin(theta)/g)
t = np.linspace(0, t_flight, 200)
# Drag force
Ft = 0.5*C*S*ro_mars*pow(V_initial, 2)
# Buoyant Force
Fb = ro_mars*g*(4/3*np.pi*pow(r, 3))
x_loc = []
y_loc = []
for time in t:
x = V_initial*time*np.cos(theta)
y = V_initial*time*np.sin(theta) - (1/2)*g*pow(time, 2)
x_loc.append(x)
y_loc.append(y)
x_vel = []
y_vel = []
for time in t:
vx = V_initial*np.cos(theta)
vy = V_initial*np.sin(theta) - g*time
x_vel.append(vx)
y_vel.append(vy)
v_ch = [pow(i**2+ii**2, 0.5) for i in x_vel for ii in y_vel]
tau = []
for velocity in v_ch:
Ft = 0.5*C*S*ro_mars*pow(velocity, 2)
tau.append(Ft)
buoy = []
for velocity in v_ch:
Fb = ro_mars*g*(4/3*np.pi*pow(r, 3))
buoy.append(Fb)
after this point, I couldn't figure out how to plot to a projectile motion under this forces. In other words, I'm trying to compare the projectile motion of the mass under three circumstances
Mass only under the effect of gravity
Mass under the effect of gravity and air resistance
Mass under the effect of gravity, air resistance, and buoyancy
You must calculate each location based on the sum of forces at the given time. For this it is better to start from calculating the net force at any time and using this to calculate the acceleration, velocity and then position. For the following calculations, it is assumed that buoyancy and gravity are constant (which is not true in reality but the effect of their variability is negligible in this case), it is also assumed that the initial position is (0,0) though this can be trivially changed to any initial position.
F_x = tau_x
F_y = tau_y + bouyancy + gravity
Where tau_x and tau_y are the drag forces in the x and y directions, respectively. The velocities, v_x and v_y, are then given by
v_x = v_x + (F_x / (2 * m)) * dt
v_y = v_y + (F_y / (2 * m)) * dt
So the x and y positions, r_x and r_y, at any time t are given by the summation of
r_x = r_x + v_x * dt
r_y = r_y + v_y * dt
In both cases this must be evaluated from 0 to t for some dt where dt * n = t if n is the number of steps in summation.
r_x = r_x + V_i * np.cos(theta) * dt + (F_x / (2 * m)) * dt**2
r_y = r_y + V_i * np.sin(theta) * dt + (F_y / (2 * m)) * dt**2
The entire calculation can actually be done in two lines,
r_x = r_x + V_i * np.cos(theta) * dt + (tau_x / (2 * m)) * dt**2
r_y = r_y + V_i * np.sin(theta) * dt + ((tau_y + bouyancy + gravity) / (2 * m)) * dt**2
Except that v_x and v_y require updating at every time step. To loop over this and calculate the x and y positions across a range of times you can simply follow the below (edited) example.
The following code includes corrections to prevent negative y positions, since the given value of g is for the surface or Mars I assume this is appropriate - when you hit zero y and try to keep going you may end up with a rapid unscheduled disassembly, as we physicists call it.
Edit
In response to the edited question, the following example has been modified to plot all three cases requested - gravity, gravity plus drag, and gravity plus drag and buoyancy. Plot setup code has also been added
Complete example (edited)
import numpy as np
import matplotlib.pyplot as plt
def projectile(V_initial, theta, bouyancy=True, drag=True):
g = 9.81
m = 1
C = 0.47
r = 0.5
S = np.pi*pow(r, 2)
ro_mars = 0.0175
time = np.linspace(0, 100, 10000)
tof = 0.0
dt = time[1] - time[0]
bouy = ro_mars*g*(4/3*np.pi*pow(r, 3))
gravity = -g * m
V_ix = V_initial * np.cos(theta)
V_iy = V_initial * np.sin(theta)
v_x = V_ix
v_y = V_iy
r_x = 0.0
r_y = 0.0
r_xs = list()
r_ys = list()
r_xs.append(r_x)
r_ys.append(r_y)
# This gets a bit 'hand-wavy' but as dt -> 0 it approaches the analytical solution.
# Just make sure you use sufficiently small dt (dt is change in time between steps)
for t in time:
F_x = 0.0
F_y = 0.0
if (bouyancy == True):
F_y = F_y + bouy
if (drag == True):
F_y = F_y - 0.5*C*S*ro_mars*pow(v_y, 2)
F_x = F_x - 0.5*C*S*ro_mars*pow(v_x, 2) * np.sign(v_y)
F_y = F_y + gravity
r_x = r_x + v_x * dt + (F_x / (2 * m)) * dt**2
r_y = r_y + v_y * dt + (F_y / (2 * m)) * dt**2
v_x = v_x + (F_x / m) * dt
v_y = v_y + (F_y / m) * dt
if (r_y >= 0.0):
r_xs.append(r_x)
r_ys.append(r_y)
else:
tof = t
r_xs.append(r_x)
r_ys.append(r_y)
break
return r_xs, r_ys, tof
v = 30
theta = np.pi/4
fig = plt.figure(figsize=(8,4), dpi=300)
r_xs, r_ys, tof = projectile(v, theta, True, True)
plt.plot(r_xs, r_ys, 'g:', label="Gravity, Buoyancy, and Drag")
r_xs, r_ys, tof = projectile(v, theta, False, True)
plt.plot(r_xs, r_ys, 'b:', label="Gravity and Drag")
r_xs, r_ys, tof = projectile(v, theta, False, False)
plt.plot(r_xs, r_ys, 'k:', label="Gravity")
plt.title("Trajectory", fontsize=14)
plt.xlabel("Displacement in x-direction (m)")
plt.ylabel("Displacement in y-direction (m)")
plt.ylim(bottom=0.0)
plt.legend()
plt.show()
Note that this preserves and returns the time-of-flight in the variable tof.
Using vector notation, and odeint.
import numpy as np
from scipy.integrate import odeint
import scipy.constants as SPC
import matplotlib.pyplot as plt
V_initial = 30 # m/s
theta = np.pi/6 # 30
g = 3.711
m = 1 # I assume this is your mass
C = 0.47
r = 0.5
ro_mars = 0.0175
t_flight = 2*(V_initial*np.sin(theta)/g)
t = np.linspace(0, t_flight, 200)
pos0 = [0, 0]
v0 = [np.cos(theta) * V_initial, np.sin(theta) * V_initial]
def f(vector, t, C, r, ro_mars, apply_bouyancy=True, apply_resistance=True):
x, y, x_prime, y_prime = vector
# volume and surface
V = np.pi * 4/3 * r**3
S = np.pi*pow(r, 2)
# net weight bouyancy
if apply_bouyancy:
Fb = (ro_mars * V - m) * g *np.array([0,1])
else:
Fb = -m * g * np.array([0,1])
# velocity vector
v = np.array([x_prime, y_prime])
# drag force - corrected to be updated based on current velocity
# Ft = -0.5*C*S*ro_mars*pow(V_initial, 2)
if apply_resistance:
Ft = -0.5*C*S*ro_mars* v *np.linalg.norm(v)
else:
Ft = np.array([0, 0])
# resulting acceleration
x_prime2, y_prime2 = (Fb + Ft) / m
return x_prime, y_prime, x_prime2, y_prime2
sol = odeint(f, pos0 + v0 , t, args=(C, r, ro_mars))
plt.plot(sol[:,0], sol[:, 1], 'g', label='tray')
plt.legend(loc='best')
plt.xlabel('x')
plt.ylabel('y')
plt.grid()
plt.show()
Note that I corrected your drag force to use the actual (not initial) velocity, I do not know if that was your mistake or it was on purpose.
Also please check the documentation for odeint to understand better how to turn a second order ODE (like the one in your problem) to a first order vector ODE.
To remove air resistance or bouyancy, set apply_bouyancy and apply_resistance to True or False by adding them to the args=(...)
I am using pseudo-spectral method to solve for KdV equation u_t + u*u_x + u_xxx = 0. After simplification by Fourier Transform, I got two equations with two variables:
uhat = vhat * exp(-i*k^3*t)
d(vhat)/dt =-0.5i * k * exp(-i*k^3*t)*F[u^2]
where F represents Fourier Transform, uhat = F[u], vhat = F[v]
I'd like to use RK4 to solve for u by è ifft(uhat)è eventually. Now I got two variables uhat and vhat, I have 2 ideas for solving uhat and vhat:
To treat them as a system of ODEs to implement RK4.
Treat the equation 2 above as the to-be solved ODE to solve chat by
RK4, and calculate uhat = vhat * exp(-i*k^2*delta_t) after each
time step delta_t.
I got problem to implement both ideas. Here are the codes for the 2nd idea above.
import numpy as np
import math
from matplotlib import pyplot as plt
from matplotlib import animation
#----- Numerical integration of ODE via fixed-step classical Runge-Kutta -----
def RK4Step(yprime,t,y,dt):
k1 = yprime(t , y )
k2 = yprime(t+0.5*dt, y+0.5*k1*dt)
k3 = yprime(t+0.5*dt, y+0.5*k2*dt)
k4 = yprime(t+ dt, y+ k3*dt)
return y + (k1+2*k2+2*k3+k4)*dt/6.
def RK4(yprime,times,y0):
y = np.empty(times.shape+y0.shape,dtype=y0.dtype)
y[0,:] = y0 # enter initial conditions in y
steps = 4
for i in range(times.size-1):
dt = (times[i+1]-times[i])/steps
y[i+1,:] = y[i,:]
for k in range(steps):
y[i+1,:] = RK4Step(yprime, times[i]+k*dt, y[i+1,:], dt)
y[i+1,:] = y[i+1,:] * np.exp(delta_t * 1j * kx**3)
return y
#====================================================================
#----- Parameters for PDE -----
L = 100
n = 512
delta_t = 0.1
tmax = 20
c1 = 1.5 # amplitude of 1st wave
c2 = 0.5 # amplitude of 2nd wave
#----- Constructing the grid and kernel functions
x2 = np.linspace(-L/2,L/2, n+1)
x = x2[:n] # periodic B.C. #0 = #n
kx1 = np.linspace(0,n/2-1,n/2)
kx2 = np.linspace(1,n/2, n/2)
kx2 = -1*kx2[::-1]
kx = (2.* np.pi/L)*np.concatenate((kx1,kx2)); kx[0] = 1e-6
#----- Initial Condition -----
z1 = np.sqrt(c1)/2. * (x-0.1*L)
z2 = np.sqrt(c2)/2. * (x-0.4*L)
soliton = 6*(0.5 * c1 / (np.cosh(z1))**2 + 0.5 * c2/ (np.cosh(z2))**2)
uhat0 = np.fft.fft(soliton)
vhat0 = uhat0
#----- Define ODE -----
def wprime(t,vhat):
g = -0.5 * 1j* kx * np.exp(-1j * kx**3 * t)
return g * np.fft.fft(np.real(np.fft.ifft(np.exp(1j*kx**3*t)*vhat)**2))
#====================================================================
#----- Compute the numerical solution -----
TimeStart = 0.
TimeEnd = tmax+delta_t
TimeSpan = np.arange(TimeStart, TimeEnd, delta_t)
w_sol = RK4(wprime,TimeSpan, vhat0)
#----- Animate the numerical solution -----
fig = plt.figure()
ims = []
for i in TimeSpan:
w = np.real(np.fft.ifft(w_sol[i,:]))
im = plt.plot(x,w)
ims.append([im])
ani = animation.ArtistAnimation(fig, ims, interval=100, blit=False)
plt.show()
The RK4 part was from #LutzL.
In general it looks good. Better use a FuncAnimation, reduces the memory footprint and works:
#----- Animate the numerical solution -----
fig = plt.figure()
ax = plt.axes(ylim=(-5,5))
line, = ax.plot(x, soliton, '-')
NT = TimeSpan.size;
def animate(i):
i = i % (NT + 10)
if i<NT:
w = np.real(np.fft.ifft(w_sol[i,:]))
line.set_ydata(w)
return line,
anim = animation.FuncAnimation(fig, animate, interval=50, blit=True)
plt.show()
However, the first graph remains standing, I do not know why.
I have a complexed valued system from a PDE problem, the odeint() in Python cannot deal with it. I wrote a RK4 module to solve my system. It seems to work, however, the computed values are obvious incorrect. At the second time step, whole computed values are zero. Here are my code:
import matplotlib.pyplot as plt
import numpy as np
import drawnow
import time
import math
### Parameters ###
L = 20
n = 64
delta_t = 1.
tmax = 10
miu = 1e-6
x2 = np.linspace(-L/2,L/2, n+1)
x = x2[:n] # periodic B.C. #0 = #n
kx1 = np.linspace(0,n/2-1,n/2)
kx2 = np.linspace(1,n/2, n/2)
kx2 = -1*kx2[::-1]
kx = (2.*math.pi/L)*np.concatenate((kx1,kx2)); kx[0] = 1e-6
ky = kx; y = x
X, Y = np.meshgrid(x, y)
KX,KY = np.meshgrid(kx,ky)
K = KX**2 + KY**2
K2 = np.reshape(K, n**2,1)
### Initial Condition ###
vorticity = np.exp(-0.25*X**2 - 2.*Y**2)
wt = np.fft.fft2(vorticity)
wt2 = np.reshape(wt, n**2, 1) # wt2 is initial condition
### Define ODE ###
def yprime(t,rhs):
global miu, K, K2,n,KX, KY, wt2, wt
psit = -wt/ K
psix = np.real(np.fft.ifft2(1j*KX*psit))
psiy = np.real(np.fft.ifft2(1j*KY*psit))
wx = np.real(np.fft.ifft2(1j*KX*wt))
wy = np.real(np.fft.ifft2(1j*KY*wt))
rhs = -miu * K2 * wt2 + np.reshape(np.fft.fft2(wx*psiy - wy*psix), n**2,1)
return rhs
def RK4(domain,wt2,tmax):
w = np.empty((tmax+1,n**2))
w = w + 0j
t = np.empty(tmax+1) # length
w[0,:] = wt2 # enter initial conditions in y
t[0] = domain[0]
for i in range(1,tmax):
t[i+1] = t[i]+delta_t
w[i+1,:] = RK4Step(t[i], w[i,:],delta_t)
return w
def RK4Step(t,w,delta_t):
k1 = yprime(t,w)
k2 = yprime(t+0.5*delta_t, w+0.5*k1*delta_t)
k3 = yprime(t+0.5*delta_t, w+0.5*k2*delta_t)
k4 = yprime(t+delta_t, w+k3*delta_t)
return w + (k1+2*k2+2*k3+k4)*delta_t/6.
### Prediction ###
TimeStart = 0.
TimeEnd = tmax+1
TimeSpan = np.arange(TimeStart, TimeEnd, delta_t)
wt2_sol = RK4(TimeSpan, wt2, tmax)
for i in TimeSpan:
w = np.real(np.fft.ifft2(np.reshape(wt2_sol[i,:], (n, n))))
plt.pcolor(X,Y,w,shading = 'interp',cmap='jet')
drawnow
time.sleep(0.2)
plt.show()
Any idea why it doesn't work? In addition, I like to make a short video based on the solution. the function 'drawnow' and 'time.sleep() do not seem to work here.
Thank you!
My cleaned up version. Changing the number of inner steps does not change the quality of the result.
Make the Runge-Kutta solver (more) universal, input time array with (times[0],y0) being the initial point of the IVP
replace def yprime(t,rhs): with def yprime(t,wt):, since wt is your state variable, and rhs is the result. So wt is a local variable in yprime. Eliminate rhs by direct assembly in the return statement.
remove all reshape operation, act on the vector space of 2D arrays, numpy is good in treating matrices as some other kind of vector
add matplotlib.animate for the pre-generated image sequence. The tutorial for that seemed easier than the function-based animation
Played with arange to replace linspace in generation of kx. Better option is probably to use fftshift to swap the halves of the frequency array
.
import numpy as np
import math
from matplotlib import pyplot as plt
from matplotlib import animation
#----- Numerical integration of ODE via fixed-step classical Runge-Kutta -----
def RK4Step(yprime,t,y,dt):
k1 = yprime(t , y )
k2 = yprime(t+0.5*dt, y+0.5*k1*dt)
k3 = yprime(t+0.5*dt, y+0.5*k2*dt)
k4 = yprime(t+ dt, y+ k3*dt)
return y + (k1+2*k2+2*k3+k4)*dt/6.
def RK4(yprime,times,y0):
y = np.empty(times.shape+y0.shape,dtype=y0.dtype)
y[0,:] = y0 # enter initial conditions in y
steps = 4
for i in range(times.size-1):
dt = (times[i+1]-times[i])/steps
y[i+1,:] = y[i,:]
for k in range(steps):
y[i+1,:] = RK4Step(yprime, times[i]+k*dt, y[i+1,:], dt)
return y
#====================================================================
#----- Parameters for PDE -----
L = 20
n = 64
delta_t = 1.
tmax = 10
miu = 1e-6
#----- Constructing the grid and kernel functions
x2 = np.linspace(-L/2,L/2, n+1)
x = x2[:n] # periodic B.C. #0 = #n
y = x
kx = np.linspace( -n/2 , n/2-1, n)
kx = (2.*math.pi/L)*np.concatenate((np.arange(0,n/2),np.arange(-n/2,0)));
kx[0] = 1e-6
ky = kx;
X, Y = np.meshgrid(x, y)
KX,KY = np.meshgrid(kx,ky)
K = KX**2 + KY**2
#----- Initial Condition -----
vorticity = np.exp(-0.25*X**2 - 2.*Y**2)
wt0 = np.fft.fft2(vorticity)
#----- Define ODE -----
def wprime(t,wt):
global miu, K, K2,n,KX, KY
psit = -wt / K
psix = np.real(np.fft.ifft2(1j*KX*psit))
psiy = np.real(np.fft.ifft2(1j*KY*psit))
wx = np.real(np.fft.ifft2(1j*KX*wt))
wy = np.real(np.fft.ifft2(1j*KY*wt))
return -miu * K * wt + np.fft.fft2(wx*psiy - wy*psix)
#====================================================================
#----- Compute the numerical solution -----
TimeStart = 0.
TimeEnd = tmax+delta_t
TimeSpan = np.arange(TimeStart, TimeEnd, delta_t)
wt_sol = RK4(wprime,TimeSpan, wt0)
#----- Animate the numerical solution -----
fig = plt.figure()
ims = []
for i in TimeSpan:
w = np.real(np.fft.ifft2(wt_sol[i,:]))
im = plt.pcolor(X,Y,w,edgecolors='none',cmap='jet')
ims.append([im])
ani = animation.ArtistAnimation(fig, ims, interval=50, blit=True,
repeat_delay=1000)
#ani.save('PDE-animation.mp4')
plt.show()