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))
(5, 100)
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)
Optimization terminated successfully.
         Current function value: 125.185962
         Iterations: 1388
         Function evaluations: 2252
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)
 final_simplex: (array([[-214.41480028,  218.75735307,  -37.23530629, -358.04927071,
         138.81285373,  226.06762508],
       [-214.41480028,  218.75735307,  -37.23530629, -358.04927071,
         138.81285373,  226.06762508],
       [-214.41480028,  218.75735307,  -37.23530629, -358.04927071,
         138.81285373,  226.06762508],
       [-214.41480028,  218.75735307,  -37.23530629, -358.0492707 ,
         138.81285373,  226.06762508],
       [-214.41480028,  218.75735307,  -37.23530629, -358.0492707 ,
         138.81285373,  226.06762508],
       [-214.41480028,  218.75735307,  -37.23530629, -358.04927071,
         138.81285373,  226.06762508],
       [-214.41480028,  218.75735306,  -37.23530629, -358.0492707 ,
         138.81285373,  226.06762508]]), array([ 125.18596195,  125.18596195,  125.18596195,  125.18596195,
        125.18596195,  125.18596195,  125.18596195]))
           fun: 125.18596195079743
       message: 'Optimization terminated successfully.'
          nfev: 2252
           nit: 1388
        status: 0
       success: True
             x: array([-214.41480028,  218.75735307,  -37.23530629, -358.04927071,
        138.81285373,  226.06762508])
In [ ]: