week 4 NMM pset

import numpy as np
import matplotlib.pyplot as plt
from IPython.display import set_matplotlib_formats

6.2

simple harmonic oscillator

euler's method and fixed step 4th order runge kutta

# intial conditions
y0 = 1
ydot0 = 0

# euler's method
def euler(h,y0,ydot0):
    t = np.arange(0,np.pi*100,h) #our time series
    N = t.size
    y = np.zeros(N)
    ydot = np.zeros(N)
    #set initial contitions
    y[0] = y0
    ydot[0] = ydot0
    for i in range(N-1):
        y[i+1] = y[i] + h *ydot[i]
        ydot[i+1] = ydot[i] + h * (-1*y[i])
    return y, ydot

# 4th order Runge-Kutta

def RK4f(h,y0,ydot0):
    t = np.arange(0,np.pi*100,h) #our time series
    N = t.size
    y = np.zeros(N)
    ydot = np.zeros(N)
    #set initial contitions
    y[0] = y0
    ydot[0] = ydot0
    for i in range(N-1):
        #RK parameters
        k11 = h * ydot[i]
        k12 = -h * (y[i])

        k21 = h * (ydot[i] + k12/2)
        k22 = -h * (y[i] + k11/2)

        k31 = h * (ydot[i] + k22/2)
        k32 = -h * (y[i] + k21/2)

        k41 = h * (ydot[i] + k32)
        k42 = -h * (y[i] + k31)

        y[i+1] = y[i] + k11/6 + k21/3 + k31/3 + k41/6
        ydot[i+1] = ydot[i] + k12/6 + k22/3 + k32/3 + k42/6    
    return y, ydot


# analytic solution
def yTrue(h):
    t = np.arange(0,np.pi*100,h)
    y = np.cos(t)
    ydot = -1*np.sin(t)
    return y, ydot

Some example plots of the above

# some example plots
hrange = [0.1,0.01,0.001,0.0001]

fig, axs = plt.subplots(len(hrange),1)

plt.subplots_adjust(hspace=0.75)

plt.style.use('seaborn-v0_8')

for i in range(len(hrange)):
    h = hrange[i]
    y_e, y_edot = euler(h, y0, ydot0)
    y_r, y_rdot = RK4f(h, y0, ydot0)
    y_t, y_tdot = yTrue(h)

    t = np.arange(0,np.pi*100,h)
    #str("ax"+str(i))
    #ax = plt.subplot((i+1),1,1)
    ax = axs[i]
    ax.plot(t,y_e,label='Euler')
    ax.plot(t,y_r,label='Runge-Kutta')
    ax.plot(t,y_t, 'k--', label='cos(t)')
    ax.title.set_text('step size = {}'.format(h))
    ax.legend(loc='lower left',bbox_to_anchor=(0.0, -0.62),ncols=3,prop={'size': 8})

plt.show()

errors from both appproaches

# step size range
hrange = np.logspace(-1,-5,num=12,endpoint=True, base=10)

#average local error
err_euler_avg = []
err_RK_avg = []

#final error in y
err_euler_y = []
err_euler_ydot = []

#final error in ydot
err_RK_y = []
err_RK_ydot = []

for h in hrange:
    y_e, y_edot = euler(h, y0, ydot0)
    y_r, y_rdot = RK4f(h, y0, ydot0)
    y_t, y_tdot = yTrue(h)

    t = np.arange(0,np.pi*100,h)
    N = t.size - 1

    err_euler_avg.append(np.sum(abs(y_e-y_t))/N)
    err_RK_avg.append(np.sum(abs(y_r-y_t))/N)
    
    err_euler_y.append(abs(y_e[N]-y_t[N]))
    err_euler_ydot.append(abs(y_edot[N]-y_tdot[N]))

    err_RK_y.append(abs(y_r[N]-y_t[N]))
    err_RK_ydot.append(abs(y_rdot[N]-y_tdot[N]))

# now we plot

plt.style.use('seaborn-v0_8')

fig, (ax1,ax2,ax3) = plt.subplots(3,1, constrained_layout=True)

ax1.semilogx(hrange,err_euler_avg, label='Euler')
ax1.semilogx(hrange,err_RK_avg, label = 'Runge-Kutta')
ax1.title.set_text('average local errors')
ax1.legend()

ax2.semilogx(hrange,err_euler_y, label='Euler')
ax2.semilogx(hrange,err_RK_y, label = 'Runge-Kutta')
ax2.title.set_text('final error in y')

ax3.semilogx(hrange,err_euler_ydot, label='Euler')
ax3.semilogx(hrange,err_RK_ydot, label = 'Runge-Kutta')
ax3.title.set_text('final error in ydot')

ax3.set_xlabel('time step')

plt.show()

#let's plot again with tighter limits on y so we can see haha. 
plt.style.use('seaborn-v0_8')

fig, (ax1,ax2,ax3) = plt.subplots(3,1, constrained_layout=True)

ax1.semilogx(hrange,err_euler_avg, label='Euler')
ax1.semilogx(hrange,err_RK_avg, label = 'Runge-Kutta')
ax1.title.set_text('average local errors')
ax1.set_ylim(-0.1,0.2)
ax1.set_xlim(0,0.005)
ax1.legend()

ax2.semilogx(hrange,err_euler_y, label='Euler')
ax2.semilogx(hrange,err_RK_y, label = 'Runge-Kutta')
ax2.title.set_text('final error in y')
ax2.set_ylim(-0.01,0.2)
ax2.set_xlim(0,0.005)
ax2.legend()

ax3.semilogx(hrange,err_euler_ydot, label='Euler')
ax3.semilogx(hrange,err_RK_ydot, label = 'Runge-Kutta')
ax3.title.set_text('final error in ydot')
ax3.set_ylim(-0.01,0.2)
ax3.set_xlim(0,0.005)
ax3.legend()

ax3.set_xlabel('time step')

plt.show()

/var/folders/x2/7m_0xwh96ys74v5ztc0w94zm0000gn/T/ipykernel_24356/3699904494.py:10: UserWarning: Attempt to set non-positive xlim on a log-scaled axis will be ignored.
  ax1.set_xlim(0,0.005)
/var/folders/x2/7m_0xwh96ys74v5ztc0w94zm0000gn/T/ipykernel_24356/3699904494.py:17: UserWarning: Attempt to set non-positive xlim on a log-scaled axis will be ignored.
  ax2.set_xlim(0,0.005)
/var/folders/x2/7m_0xwh96ys74v5ztc0w94zm0000gn/T/ipykernel_24356/3699904494.py:24: UserWarning: Attempt to set non-positive xlim on a log-scaled axis will be ignored.
  ax3.set_xlim(0,0.005)

6.3 4th order adaptive runge-kutta

def RK4v(h_i,y0,ydot0,minStep, minError, maxError,lastTime,scale):
    def getRKstep(h, y, ydot):
        #RK parameters
        k11 = h * ydot
        k12 = -h * (y)

        k21 = h * (ydot + k12/2)
        k22 = -h * (y + k11/2)

        k31 = h * (ydot + k22/2)
        k32 = -h * (y + k21/2)

        k41 = h * (ydot + k32)
        k42 = -h * (y + k31)
        y = y + k11/6 + k21/3 + k31/3 + k41/6
        ydot = ydot + k12/6 + k22/3 + k32/3 + k42/6  
        return y, ydot

    #set initial contitions
    y = y0
    ydot = ydot0

    yarray = [y0]
    ydotarray = [ydot0]
    harray = []

    h = h_i
    t = 0
    #step through
    while t <= lastTime:
        #RK parameters
        yf, ydotf = getRKstep(h,y,ydot) # full step
        yh1, ydoth1 = getRKstep(h/2,y,ydot) # half step 1
        yh2,ydoth2 = getRKstep(h/2,yh1,ydoth1) # half step 2
        err = abs(yf - (yh1+yh2)) #error estimate
        if (err >= maxError) and (h/scale >= minStep):
            y = yh1
            ydot = ydoth1
            h = h/scale
        elif (err<= minError):
            y = yf
            ydot = ydotf
            h = h*scale
        else:
            y = yf
            ydot = ydotf
            h=h
        yarray.append(y)
        ydotarray.append(ydot)
        harray.append(h)
        t += h
  
    return yarray, ydotarray, harray

#quick example plot
h_i = 0.01
minStep = 10**-3
minError = 10**-3
maxError = 10**-2
lastTime = 100*np.pi
scale = 1.1

y,ydot,h_array = RK4v(h_i,y0,ydot0,minStep, minError, maxError,lastTime,scale)

t = np.arange(0,lastTime,lastTime/len(y))
#print(t)
y_t, y_tdot = yTrue(lastTime/len(y))

plt.plot(t,y,label='RK')
plt.plot(t,y_t,'k--',label='True')
plt.title("max error: {}".format(maxError))
plt.legend(bbox_to_anchor=(0.0, 0.0))
plt.show()

#find relationship between desired error and step size
h_i = 0.01
minStep = 10**-4.5
minError = 10**-4.5
maxError = np.logspace(-1,-4,num=4,endpoint=True, base=10)
lastTime = 100*np.pi
scale = 1.1

y0 = 1
ydot0 = 0

avg_step = []

for i in range(len(maxError)):
    #ax = axs[i]
    e = maxError[i]
    #print(e)
    y,ydot,h_array = RK4v(h_i,y0,ydot0,minStep, minError, e,lastTime,scale)
    N = len(h_array)
    print(N)
    avg_step.append(np.sum(h_array)/N)

9517464
9557655
9562105
9562550

plt.semilogx(maxError,avg_step)
plt.xlabel('desired maximum error')
plt.ylabel('average step size')
plt.show()

6.4

solving a pendulum on a moving base (base moves vertically tho!)

# we will do this with drake. b/c the question says to use a numerical method. it does not specify that we had to use one of our numerical methods.
from pydrake.all import (AddMultibodyPlantSceneGraph, ControllabilityMatrix,
                         DiagramBuilder, Linearize, LinearQuadraticRegulator,
                         MeshcatVisualizer, Parser, Saturation, SceneGraph,
                         Simulator, StartMeshcat, WrapToSystem, ConstantVectorSource, 
                         TrajectorySource, Trajectory, VectorSystem, DirectCollocation, 
                         MultibodyPlant, LogVectorOutput, PiecewisePolynomial, PlanarSceneGraphVisualizer,
                         MultibodyPositionToGeometryPose)
from pydrake.solvers import MathematicalProgram, Solve

# Start the visualizer (run this cell only once, each instance consumes a port)
meshcat = StartMeshcat()

INFO:drake:Meshcat listening for connections at http://localhost:7002

class Utraj(VectorSystem):

        def __init__(self):
            # 4 inputs: state: x, theta, xdot, thetadot
            # 1 output: spring force
            VectorSystem.__init__(self, 4,  1)
            
            # spring stiffness
            self.k = meshcat.GetSliderValue('w')**2#200#get_hopper_parameters()[0]
            self.A = meshcat.GetSliderValue('A')#0.3#

        # note that this function is called at each time step
        def DoCalcVectorOutput(self, context, state_delta, unused, spring_force):
            
            # set the output of the spring block to
            # the value of the elastic force
            l = state_delta[0]
            delta = meshcat.GetSliderValue('A')#self.A #state_delta[4]
            k = meshcat.GetSliderValue('w')**2
            spring_force[:] = [- k * (l + delta)]

def cartpole_demo():
    builder = DiagramBuilder()
    plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=0.0)
    file_name = "models/cartpolev.urdf"
    Parser(plant).AddModelFromFile(file_name)
    plant.Finalize()

    # reset meshcat visualization
    meshcat.Delete()
    meshcat.DeleteAddedControls()
    # configure for 2D
    meshcat.Set2dRenderMode(xmin=-20.5, xmax=10.5, ymin=-10.5, ymax=10.5)
    MeshcatVisualizer.AddToBuilder(builder, scene_graph, meshcat)

    # Setup slider input
    meshcat.AddSlider('w', min=5., max=40, step=10, value=10.0)
    meshcat.AddSlider('A', min=0., max=5, step=0.1, value=0.5)
    #a = meshcat.GetSliderValue('k')
    #print(a)

    #force_system = builder.AddSystem(ConstantVectorSource([200]))
    # omega = 1
    # A = 200
    # t = np.arange(0,2,0.1)

    # Add the periodic motion forcer to the base
    force_system = builder.AddSystem(Utraj())
    builder.Connect(plant.get_state_output_port(), force_system.get_input_port())
    builder.Connect(force_system.get_output_port(),
                    plant.get_actuation_input_port())

    diagram = builder.Build()

    # Set up a simulator to run this diagram
    simulator = Simulator(diagram)
    context = simulator.get_mutable_context()

    # Set the initial conditions (x, theta, xdot, thetadot)
    context.SetContinuousState([0, 1, 0, 0])

    simulator.set_target_realtime_rate(1.0)

    print('Use the slider in the MeshCat controls to change w and A')
    print("Press 'Stop Simulation' in MeshCat to continue.")
    meshcat.AddButton('Stop Simulation')
    while meshcat.GetButtonClicks('Stop Simulation') < 1:
        simulator.AdvanceTo(simulator.get_context().get_time() + 1.0)

    # simulator.set_target_realtime_rate(1.0)
    # simulator.AdvanceTo(1.0)

cartpole_demo()

Press 'Stop Simulation' in MeshCat to continue.

def dircol_cartpole(initial_state,final_state):
    builder = DiagramBuilder()
    plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=0.0)
    file_name = "models/cartpolev.urdf"
    Parser(plant).AddModelFromFile(file_name)
    plant.Finalize()

    context = plant.CreateDefaultContext()

    dircol = DirectCollocation(
        plant,
        context,
        num_time_samples=21,
        minimum_timestep=0.05,
        maximum_timestep=0.4,
        input_port_index=plant.get_actuation_input_port().get_index())
    
    prog = dircol.prog()

    # dircol.AddEqualTimeIntervalsConstraints()

    #initial_state = [0., 1.0, 0., 0.]
    prog.AddBoundingBoxConstraint(initial_state, initial_state,
                                  dircol.initial_state())

    #final_state = [0,np.pi, 0., 0.]
    #prog.AddBoundingBoxConstraint(np.pi-0.05, np.pi+0.05, dircol.final_state()[1])
    prog.AddBoundingBoxConstraint(final_state[1:], final_state[1:],
                                  dircol.final_state()[1:])
    prog.AddLinearConstraint(dircol.final_state()[0]<=5.)
    prog.AddLinearConstraint(dircol.final_state()[0]>=-5.)
    
    #x = dircol.state()
    #dircol.AddRunningCost(100*(x[0])**4)
    #dircol.AddRunningCost(x[2]**2)
    #dircol.AddRunningCost((x[3])**2)
    # prog.AddLinearConstraint(dircol.final_state() == final_state)
    # dircol.AddFinalCost(100*(dircol.final_state()[0])**2)

    R = 10  # Cost on input "effort".
    u = dircol.input()
    dircol.AddRunningCost(R * u[0]**2)

    # Add a final cost equal to the total duration.
    #dircol.AddFinalCost(dircol.time())
    #dircol.AddFinalCost(100*(dircol.final_state()[0])**2)
    #dircol.AddFinalCost(dircol.final_state()[2])

    #initial guess: interpolation between initial and final state
    initial_x_trajectory = PiecewisePolynomial.FirstOrderHold(
        [0., 4.], np.column_stack((initial_state, final_state))) 
    dircol.SetInitialTrajectory(PiecewisePolynomial(), initial_x_trajectory)

    result = Solve(prog)
    print(result.is_success())

    u_trajectory = dircol.ReconstructInputTrajectory(result)
    x_trajectory = dircol.ReconstructStateTrajectory(result)
    return u_trajectory, x_trajectory

initial_state = [0., 1.0, 0., 0.]
final_state = [0,np.pi, 0., 0.]

u_trajectory, x_trajectory = dircol_cartpole(initial_state, final_state)

True

times = np.linspace(u_trajectory.start_time(), u_trajectory.end_time(), 100)
u_lookup = np.vectorize(u_trajectory.value)
u_values = u_lookup(times)
x_values = x_trajectory.vector_values(times)

fig, (ax1,ax2) = plt.subplots(2,1, constrained_layout=True)
ax1.plot(times, u_values)
ax1.set_xlabel("time (seconds)")
ax1.set_ylabel("force (Newtons)")

ax2.plot(times,x_values[0,:],label='z')
ax2.plot(times,x_values[1,:],label='theta')
ax2.set_xlabel("time (seconds)")
ax2.set_ylabel("position")
ax2.legend()
plt.show()

def dircol_vis(x_trajectory):
    builder = DiagramBuilder()
    #plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=0.0)
    plant = MultibodyPlant(time_step=0.0)
    scene_graph = SceneGraph()
    plant.RegisterAsSourceForSceneGraph(scene_graph)
    file_name = "models/cartpolev.urdf"
    Parser(plant).AddModelFromFile(file_name)
    plant.Finalize()

    #force_system = builder.AddSystem(TrajectorySource(u_trajectory))
    #builder.Connect(plant.get_state_output_port(), force_system.get_input_port())
    # builder.Connect(force_system.get_output_port(),
    #                 plant.get_actuation_input_port())

    #input_logger = LogVectorOutput(force_system.get_output_port(), builder)
    #state_logger = LogVectorOutput(plant.get_state_output_port(), builder)
    ######

    source = builder.AddSystem(TrajectorySource(x_trajectory))
    builder.AddSystem(scene_graph)
    pos_to_pose = builder.AddSystem(MultibodyPositionToGeometryPose(plant, input_multibody_state=True))

    builder.Connect(source.get_output_port(0), pos_to_pose.get_input_port())
    builder.Connect(pos_to_pose.get_output_port(), scene_graph.get_source_pose_port(plant.get_source_id()))

    meshcat.Set2dRenderMode(xmin=-20.5, xmax=10.5, ymin=-10.5, ymax=10.5)
    MeshcatVisualizer.AddToBuilder(builder, scene_graph, meshcat)

    meshcat.Delete()
    meshcat.DeleteAddedControls()


    diagram = builder.Build()

    # Set up a simulator to run this diagram
    simulator = Simulator(diagram)

    context = simulator.get_mutable_context()

    # Set the initial conditions (x, theta, xdot, thetadot)
    # context.SetContinuousState([0, 1, 0, 0])

    #simulator = Simulator(builder.Build())
    simulator.set_target_realtime_rate(0.5)
    simulator.AdvanceTo(u_trajectory.end_time())

dircol_vis(x_trajectory)

#let's go again with a different angle
initial_state = [0., 1.0, 0., 0.]
final_state = [0,np.pi/2, 0., 0.]

u2_trajectory, x2_trajectory = dircol_cartpole(initial_state, final_state)

True

#run simulation
dircol_vis(x2_trajectory)

times = np.linspace(u2_trajectory.start_time(), u2_trajectory.end_time(), 100)
u_lookup = np.vectorize(u2_trajectory.value)
u_values = u_lookup(times)
x_values = x2_trajectory.vector_values(times)
# x_lookup = np.vectorize(x_trajectory.value)
# x_values = x_lookup(times)

# input_log = input_logger.FindLog(simulator.get_context())
# state_log = state_logger.FindLog(simulator.get_context())

fig, (ax1,ax2) = plt.subplots(2,1, constrained_layout=True)
ax1.plot(times, u_values, label='u')
# ax1.plot(input_log.sample_times(), input_log.data()[0,:].T, 'k--', label = 'u actual')
ax1.set_xlabel("time (seconds)")
ax1.set_ylabel("force (Newtons)")
ax1.legend()

ax2.plot(times,x_values[0,:],label='z')
# ax2.plot(state_log.sample_times(), state_log.data()[0,:].T,'k--',label='z actual')
ax2.plot(times,x_values[1,:],label='theta')
# ax2.plot(state_log.sample_times(), state_log.data()[1,:].T,'r--',label='theta actual')
ax2.set_xlabel("time (seconds)")
ax2.legend()
plt.show()

print('hi')

hi