Double pendulum

This example uses the sequence() and subdraw() animations to animate plot plot objects.

This example is a Diplotocus implementation of this animation from the matplotlib docs.

We first define a function get_pendulum_positions() that returns the double pendulum nodes positions for a range of times.

Show code cell source

Hide code cell source

import numpy as np

def get_pendulum_positions(G,L1,L2,M1,M2,t_stop):
    def derivs(state):
        dydx = np.zeros_like(state)

        dydx[0] = state[1]

        delta = state[2] - state[0]
        den1 = (M1+M2) * L1 - M2 * L1 * np.cos(delta) * np.cos(delta)
        dydx[1] = ((M2 * L1 * state[1] * state[1] * np.sin(delta) * np.cos(delta)
                    + M2 * G * np.sin(state[2]) * np.cos(delta)
                    + M2 * L2 * state[3] * state[3] * np.sin(delta)
                    - (M1+M2) * G * np.sin(state[0]))
                / den1)

        dydx[2] = state[3]

        den2 = (L2/L1) * den1
        dydx[3] = ((- M2 * L2 * state[3] * state[3] * np.sin(delta) * np.cos(delta)
                    + (M1+M2) * G * np.sin(state[0]) * np.cos(delta)
                    - (M1+M2) * L1 * state[1] * state[1] * np.sin(delta)
                    - (M1+M2) * G * np.sin(state[2]))
                / den2)

        return dydx

    # create a time array from 0..t_stop sampled at 0.01 second steps
    dt = 0.01
    t = np.arange(0, t_stop, dt)

    # th1 and th2 are the initial angles (degrees)
    # w10 and w20 are the initial angular velocities (degrees per second)
    th1 = 120.0
    w1 = 0.0
    th2 = -10.0
    w2 = 0.0

    # initial state
    state = np.radians([th1, w1, th2, w2])

    # integrate the ODE using Euler's method
    y = np.empty((len(t), 4))
    y[0] = state
    for i in range(1, len(t)):
        y[i] = y[i - 1] + derivs(y[i - 1]) * dt

    x0 = np.zeros(len(t))
    y0 = np.zeros(len(t))

    x1 = L1*np.sin(y[:, 0])
    y1 = -L1*np.cos(y[:, 0])

    x2 = L2*np.sin(y[:, 2]) + x1
    y2 = -L2*np.cos(y[:, 2]) + y1

    return x0,x1,x2,y0,y1,y2

Then we generate the positions, create a plot() object p that displays the double pendulum and another plot() object trace that displays the last 5 positions of its end node:

import diplotocus as dpl

G = 9.8  # acceleration due to gravity, in m/s^2
L1 = 1.0  # length of pendulum 1 in m
L2 = 1.0  # length of pendulum 2 in m
L = L1 + L2 # maximal length of the combined pendulum
M1 = 1.0  # mass of pendulum 1 in kg
M2 = 1.0  # mass of pendulum 2 in kg
t_stop = 10  # how many seconds to simulate

x0,x1,x2,y0,y1,y2 = get_pendulum_positions(G,L1,L2,M1,M2,t_stop)

xs = np.stack((x0,x1,x2),axis=-1)
ys = np.stack((y0,y1,y2),axis=-1)

tl = dpl.Timeline(xlim=(-L,L),ylim=(-L,1))

p = dpl.plot(xs,ys,marker='o',lw=2,ls='-',zorder=2)
p.sequence(600)

trace = dpl.plot(x2,y2,lw=4,ls='-',zorder=1)
trace.subdraw(-5,0,600)

tl.animate((p,trace))
tl.save_video('../../_static/examples/pendulum.mp4')