# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

from IPython.display import SVG
SVG(filename='n-pendulum-with-cart.svg')

import sympy as sp
import sympy.physics.mechanics as mech

n = 5

# +
q = mech.dynamicsymbols('q:{}'.format(n + 1))  # Generalized coordinates
qd = mech.dynamicsymbols('q:{}'.format(n + 1),1)  # Generalized coordinates
f = mech.dynamicsymbols('f')                   # Force applied to the cart
    
m = sp.symbols('m:{}'.format(n + 1))         # Mass of each bob
l = sp.symbols('l:{}'.format(n))             # Length of each link
g, t = sp.symbols('g t')                     # Gravity and time
# -

N = mech.ReferenceFrame('N')  # Inertial reference frame
pointN = mech.Point('N*')           # Origin point
pointN.set_vel(N, 0)             # Origin's velocity is zero

points = []
particles = []
forces = []

# +
points = [pointN.locatenew('point_0', l[0]*(sp.sin(q[0])*N.x-sp.cos(q[0])*N.y))]
points[0].set_vel(N, points[0].pos_from(pointN).dt(N))
particles = [mech.Particle('particle_0',points[0],m[0])]  
forces = [(points[0],-particles[0].mass*g*N.y)]

for i in range(1,n):
    points.append(points[i-1].locatenew('point_'+str(i), l[i]*(sp.sin(q[i])*N.x-sp.cos(q[i])*N.y)))
    points[i].set_vel(N, points[i].pos_from(pointN).dt(N))
    particles.append(mech.Particle('particle_'+str(i),pointi,m[i]))
    forces.append((points[i],-particles[i].mass*g*N.y))
# -

# Construct the Lagrangian
L = mech.Lagrangian(N, *particles)
L

# Create the LagrangesMethod object
LM = mech.LagrangesMethod(L, points, hol_coneqs=None, forcelist=forces, frame=N)

LM.form_lagranges_equations()

# Form Lagranges Equations
ELeqns = LM.form_lagranges_equations()

sp.trigsimp(LM.mass_matrix)

mech.find_dynamicsymbols(LM.mass_matrix)

sp.trigsimp(LM.forcing)

mech.find_dynamicsymbols(LM.forcing)

# Simulation
# ==========
#
# Now that the symbolic equations of motion are available we can simulate the pendulum's motion. We will need some more SymPy functionality and several NumPy functions, and most importantly the integration function from SciPy, `odeint`.

import numpy as np
from numpy.linalg import solve
from scipy.integrate import odeint

# First, define some numeric values for all of the constant parameters in the problem.

arm_length = 1. / n                          # The maximum length of the pendulum is 1 meter
bob_mass = 0.01 / n                          # The maximum mass of the bobs is 10 grams
parameters = [g, m[0]]                       # Parameter definitions starting with gravity and the first bob
parameter_vals = [9.81, 0.01 / n]            # Numerical values for the first two
for i in range(n):                           # Then each mass and length
    parameters += [l[i], m[i + 1]]            
    parameter_vals += [arm_length, bob_mass]

# Mathematica has a really nice `NDSolve` function for quickly integrating their symbolic differential equations. We make use of SymPy's lambdify function to do something similar, i.e. to create functions that will evaluate the "full" mass matrix, $M$, and "full" forcing vector, $f$ from $M\dot{x} = f(x, r, t)$ as a NumPy function.

# +
dynamic = q + qd    # Make a list of the states
dynamic.append(f)  # Add the input force

M_func = sp.lambdify(dynamic + parameters, LM.mass_matrix_full) # Create a callable function to evaluate the mass matrix 
f_func = sp.lambdify(dynamic + parameters, LM.forcing_full)     # Create a callable function to evaluate the forcing vector 


# -

# To integrate the ODE's we need to define a function that returns the derivatives of the states given the current state and time.

def right_hand_side(x, t, args):
    """Returns the derivatives of the states.

    Parameters
    ----------
    x : ndarray, shape(2 * (n + 1))
        The current state vector.
    t : float
        The current time.
    args : ndarray
        The constants.

    Returns
    -------
    dx : ndarray, shape(2 * (n + 1))
        The derivative of the state.
    
    """
    r = 0.0                                  # The input force is always zero     
    arguments = np.hstack((x, r, args))      # States, input, and parameters
    dx = np.array(solve(M_func(*arguments),  # Solving for the derivatives
                  f_func(*arguments))).T[0]
    
    return dx


# Now that we have the right hand side function, the initial conditions are set such that the pendulum is in the vertical equilibrium and a slight initial rate is set for each speed to ensure the pendulum falls. The equations can then be integrated with SciPy's `odeint` function given a time series.

x0 = np.hstack((np.pi / 2 * np.ones(len(q)),  # q0...qn+1
                1e-3 * np.ones(len(qd))))          # qd0...un+1

t = np.linspace(0.0, 10.0, num=500)  # Time vector

x = odeint(right_hand_side, x0, t, args=(parameter_vals,))  # Numerical integration

# Plotting
# ========
#
# The results of the simulation can be plotted with matplotlib. First, load the plotting functionality.

import matplotlib.pyplot as plt

# %matplotlib inline

from IPython.core.pylabtools import figsize
figsize(8.0, 6.0)

# The coordinate trajectories are plotted below.

lines = plt.plot(t, x[:, :x.shape[1] // 2])
lab = plt.xlabel('Time [sec]')
leg = plt.legend(dynamic[:x.shape[1] // 2])

# And the generalized speed trajectories.

lines = plt.plot(t, x[:, x.shape[1] // 2:])
lab = plt.xlabel('Time [sec]')
leg = plt.legend(dynamic[x.shape[1] // 2:])

# Animation
# =========
#
# matplotlib now includes very nice animation functions for animating matplotlib plots. First we import the necessary functions for creating the animation.

from matplotlib import animation


def animate_pendulum(t, states, length, filename=None):
    """Animates the n-pendulum and optionally saves it to file.

    Parameters
    ----------
    t : ndarray, shape(m)
        Time array.
    states: ndarray, shape(m,p)
        State time history.
    length: float
        The length of the pendulum links.
    filename: string or None, optional
        If true a movie file will be saved of the animation. This may take some time.

    Returns
    -------
    fig : matplotlib.Figure
        The figure.
    anim : matplotlib.FuncAnimation
        The animation.

    """
    # the number of pendulum bobs
    numpoints = states.shape[1] // 2

    # first set up the figure, the axis, and the plot elements we want to animate
    fig = plt.figure()
    
    # set the limits based on the motion
    xmin = -numpoints*length
    xmax = numpoints*length
    
    # create the axes
    ax = plt.axes(xlim=(xmin, xmax), ylim=(-1.1, 1.1), aspect='equal')
    
    # display the current time
    time_text = ax.text(0.04, 0.9, '', transform=ax.transAxes)
        
    # blank line for the pendulum
    line, = ax.plot([], [], lw=2, marker='o', markersize=6)

    # initialization function: plot the background of each frame
    def init():
        time_text.set_text('')
        line.set_data([], [])
        return time_text, line,

    # animation function: update the objects
    def animate(i):
        time_text.set_text('time = {:2.2f}'.format(t[i]))
        x = np.hstack((states[i, 0], np.zeros((numpoints))))
        y = np.hstack((states[i, 0], np.zeros((numpoints))))
        
        for j in np.arange(1, numpoints):
            x[j] = x[j - 1] + length * np.cos(states[i, j])
            y[j] = y[j - 1] + length * np.sin(states[i, j])
        line.set_data(x, y)
        return time_text, line,

    # call the animator function
    anim = animation.FuncAnimation(fig, animate, frames=len(t), init_func=init,
            interval=t[-1] / len(t) * 1000, blit=True, repeat=False)
    
    # save the animation if a filename is given
    if filename is not None:
        anim.save(filename, fps=30, codec='libx264')


animate_pendulum(t, x, arm_length, filename="open-loop.mp4")

from IPython.display import HTML
html = \
"""
<video width="640" height="480" controls>
  <source src="open-loop.mp4" type="video/mp4">
Your browser does not support the video tag, check out the YouTube version instead: http://youtu.be/Nj3_npq7MZI.
</video>
"""
HTML(html)



StarcoderdataJupyterScriptsDedupFiltered