DynamicalSystems.jl logo animation using Python

double_pendulum_pyplot.py

# -*- coding: utf-8 -*-

"""
===========================
The double pendulum problem
===========================

This animation illustrates the double pendulum problem.
"""

# Double pendulum formula translated from the C code at
# http://www.physics.usyd.edu.au/~wheat/dpend_html/solve_dpend.c

from numpy import sin, cos
import numpy as np
import matplotlib.pyplot as plt
import scipy.integrate as integrate
import matplotlib.animation as animation


G = 9.8  # acceleration due to gravity, in m/s^2
L1 = 0.85  # length of pendulum 1 in m
L2 = 0.85  # length of pendulum 2 in m
M1 = 1.0  # mass of pendulum 1 in kg
M2 = 1.0  # mass of pendulum 2 in kg


def derivs(state, t):

    dydx = np.zeros_like(state)
    dydx[0] = state[1]

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

    dydx[2] = state[3]

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

    return dydx

# create a time array
dt = 0.01
t = np.arange(0.0, 30, 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

th1 = 120.0
w1 = 60.0
th2 = -10.0
w2 = 60.0

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

# integrate your ODE using scipy.integrate.
y = integrate.odeint(derivs, state, t)

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

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

fig = plt.figure()
ax = fig.add_subplot(111, autoscale_on=False, xlim=(-2.2, 2.2), ylim=(-2.2, 2.2))
ax.axis("off")

# Julia Colors:
darker_blue = (0.251, 0.388, 0.847)
lighter_blue = (0.4, 0.51, 0.878)
darker_purple = (0.584, 0.345, 0.698)
lighter_purple  = (0.667, 0.475, 0.757)
darker_green  = (0.22, 0.596, 0.149)
lighter_green  = (0.376, 0.678, 0.318)
darker_red  = (0.796, 0.235, 0.2)
lighter_red  = (0.835, 0.388, 0.361)

MS = 70 #marker size
MEW = 6 #marker edge width
TRACE_LEN = 500 #length (in datapoints) of the trace of last ball
PENDULUM_WIDTH = 5 #linewidth of pendulum
TRACE_WIDTH = 3#width of the trace line

traces = []

for i in range(TRACE_LEN):
    alphas = np.linspace(1.0, 0.0, num=TRACE_LEN)
    trace, = ax.plot([],[], color =darker_blue, lw = TRACE_WIDTH, alpha = alphas[i],\
                     zorder =-1, solid_capstyle="butt")
    traces.append(trace)

line, = ax.plot([], [], color = "black", lw=PENDULUM_WIDTH, zorder = -1)

cp, = ax.plot([0], [0], lw=0, markersize = MS, marker="o", mfc = lighter_green,\
              mec=darker_green, mew = MEW, zorder = 2)

p1, = ax.plot([], [], lw=0, markersize = MS, marker="o", mfc = lighter_red,\
              mec=darker_red, mew = MEW, zorder=2)
p2, = ax.plot([], [],  lw=0, markersize = MS, marker="o",\
              mew = MEW, mec = darker_purple, mfc = lighter_purple, zorder =2 )


def init():
    line.set_data([], [])
    cp.set_data([0], [0])
    p1.set_data([], [])
    p2.set_data([], [])
    for i in range(TRACE_LEN):
        traces[i].set_data([], [])

    return (line, p1, p2, cp) + tuple(traces)


def animate(i):
    
    realj = min(i, TRACE_LEN) # maximum amount of segments that will be created
    for k in range(TRACE_LEN): # k iteratates over trace segments
        if k < realj:
            traces[k].set_data([x2[i-k-2:i-k]], [y2[i-k-2:i-k]])
    
    thisx = [0, x1[i], x2[i]]
    thisy = [0, y1[i], y2[i]]
    line.set_data(thisx, thisy)       
  
    cp.set_data([0], [0])
    p1.set_data(x1[i], y1[i])
    p2.set_data(x2[i], y2[i])
        
    return (line, p1, p2, cp) + tuple(traces)


# plt.ioff() #this doesn't show the animation (but still saves it)

ani = animation.FuncAnimation(fig, animate, np.arange(1, len(y)),
                              interval=1, blit=True, init_func=init,repeat = False)

ani.save('double_pendulum_500_fast_hq.mp4', fps=90)
plt.show()

In case you ever want a pure-Julia solution, here's one. :)