In [101]:
import numpy as np
from matplotlib import pyplot as plt
from scipy.optimize import minimize
# variables describing the system, these would be unknown to the searcher,
# all the searcher see's is the output of pendulum_physics()
g = 9.81
l = 0.3 #m
M = 0.5 #kg
m = 0.2 #kg
b = 0.1 # coefficient of friction of the cart
I = 0.006 # mass moment inertia of a slender rod (ml^2)/15
p = I*(M+m)+M*m*l**2
# search parameters
n = 100 # number of time steps
s = 5 # number of starting positions
dt = 0.01 # size of time step
motion_limit = 15 # degrees either side of the vertical
starting_positions = []
theta_scale = 100
def populate_starting_positions(s):
for i in range(-motion_limit,motion_limit+1,int((2*motion_limit)/s)):
starting_positions.append([0,0,i*np.pi/180,0])
return np.array(starting_positions)
positions = populate_starting_positions(s)
In [102]:
def anstatt(coefficients,history):
return np.dot(coefficients,history)
def pendulum_physics(z,u):
z1 = z[1];
z2 = ((-(I + m * l**2) * b * z[1])/p) + ((m**2 * g * l**2 *z[2])/p) + (u * (I + m * l**2)/p);
z3 = z[3];
z4 = (-(m*l*b/p)*z[1]) + ((m*g*l*(M+m)/p)*z[2]) + ((m*l/p)*u);
new_z = np.array([z1,z2,z3,z4])
return new_z
def runge_kutta_4th_order(z,u):
k1 = dt*pendulum_physics(z,u)
k2 = dt*pendulum_physics(z+k1/2,u)
k3 = dt*pendulum_physics(z+k2/2,u)
k4 = dt*pendulum_physics(z+k3,u)
new_z = z + k1/6 + k2/3 + k3/3 + k4/6
return new_z
def cost_function(coefficients):
ang = []; omg = []; pos = []; vel = []
for j in range(s):
angle = []; omega = []; position = []; velocity = []
state = positions[j]
history = np.array([state[2],0,0,0,0,0])
for i in range(n):
# calculate a new input and update state
new_input = anstatt(coefficients,history)
state = runge_kutta_4th_order(state,new_input)
# update history
history = [state[2], history[0], history[1], state[0], history[3], history[4]]
# add resulting theta to the data for costing
angle.append(state[2])
omega.append(state[3])
position.append(state[0])
velocity.append(state[1])
ang.append(angle)
omg.append(omega)
pos.append(position)
vel.append(velocity)
ang = np.array(ang); omg = np.array(omg); pos = np.array(pos); vel = np.array(vel)
cost_theta = np.sum(ang**2)
cost_theta_dot = np.sum(omg**2)
total_cost = theta_scale*cost_theta + cost_theta_dot
if plot_bool == True:
return ang,omg,pos,vel
else:
return total_cost
c = 1e-3
coefficients = np.array([c,c,c,c,c,c])
plot_bool = True
angle_data, omega_data, position_data, velocity_data = cost_function(coefficients)
print(np.shape(angle_data))
In [103]:
def plot_progress(ang,omg,pos,vel):
font = {'family': 'monospace','color': '#39393d','weight': 'light','size': 24,}
fig = plt.figure(figsize=(15,18), facecolor="white")
ax1 = fig.add_subplot(411)
ax2 = fig.add_subplot(412)
ax3 = fig.add_subplot(413)
ax4 = fig.add_subplot(414)
ax1.set_title("Angle of pendulum")
ax2.set_title("Omega of pendulum")
ax3.set_title("Position of base")
ax4.set_title("Velocity of base")
t = np.arange(n)
for i in range(s):
ax1.plot(t,ang[i])
ax2.plot(t,omg[i])
ax3.plot(t,pos[i])
ax4.plot(t,vel[i])
plt.show()
plot_progress(angle_data,omega_data,position_data,velocity_data)
In [105]:
def nelder_mead(x0):
res = minimize(cost_function, x0, method='nelder-mead',options={'xtol': 1e-8, 'disp': True, 'maxfev': 5000})
return res
plot_bool = False
result = nelder_mead(coefficients)
In [106]:
plot_bool = True
print(result)
angle_data, omega_data, position_data, velocity_data = cost_function(result.x)
plot_progress(angle_data,omega_data,position_data,velocity_data)
In [ ]: