code used to create this video: https://youtu.be/7YmBSqWw7PY

dancing_flower_1.py

"""
Copyright (C) 2018 Torin Cooper-Bennun

    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <https://www.gnu.org/licenses/>.
    
simulate a double pendulum's motion by numerically integrating the equations of motion as derived from the system's
Lagrangian.
"""
import time

import numpy as np
import scipy.integrate as integrate
import matplotlib.pyplot as plt
from matplotlib.animation import FFMpegWriter
from matplotlib.collections import LineCollection


def f(t, y):
    """
    equations of motion for the double pendulum. motion is defined in terms of the angle of each limb from the vertical
    (these are the 2 generalised coordinates) and their conjugate momenta. this function calculates the first time
    derivative of each of these quantities, and returns them as a vector. see
    <https://en.wikipedia.org/wiki/Double_pendulum#Lagrangian> for the equations.

    args:
        t : time (s) [scalar]
        y : variables of motion: (theta_1, theta_2, p_theta_1, p_theta_2) [np.ndarray, shape (1,)]

    the limb lengths and masses are taken to be equal to 1m and 1kg respectively.
    """
    m = 1
    l = 1

    # gravitational acceleration
    g = 9.81

    theta_1, theta_2, p_theta_1, p_theta_2 = y

    A = 3*np.cos(theta_1 - theta_2)
    B = 6/(m*l*l) / (16 - A*A)

    theta_1_dt = B * (2*p_theta_1 - A*p_theta_2)
    theta_2_dt = B * (8*p_theta_2 - A*p_theta_1)

    C = -m*l*l/2
    D = theta_1_dt*theta_2_dt*np.sin(theta_1 - theta_2)

    p_theta_1_dt = C*(D + 3*g/l * np.sin(theta_1))
    p_theta_2_dt = C*(-D + g/l * np.sin(theta_2))

    return np.array([theta_1_dt, theta_2_dt, p_theta_1_dt, p_theta_2_dt])


def double_pendulum_trajectory(theta_1_initial, theta_2_initial, t_0, t_f, N):
    """
    calculate the trajectory of the double pendulum, given the initial angles of the limbs (from the vertical), between
    times t_0 and t_f (N points in time).
    """
    t_eval = np.linspace(t_0, t_f, N, dtype=np.float64)
    y_0 = np.array([theta_1_initial, theta_2_initial, 0, 0], dtype=np.float64)

    return integrate.solve_ivp(f, (t_0, t_f), y_0, t_eval=t_eval)


def coordinates_from_trajectory(y_out):
    """
    given integration return value y, which has the trajectory in terms of the angles, convert to x,y and return
    (assumes l = 1m, m = 1kg, for both limbs)
    """
    l = 1
    m = 1

    theta_1 = y_out[0]
    theta_2 = y_out[1]

    # calculate end position from angles
    x = l*(np.sin(theta_1) + np.sin(theta_2))
    y = -l*(np.cos(theta_1) + np.cos(theta_2))

    return x, y


def animate_trajectory_coloured(t, y, img_width=1920, img_height=1080, basic=False, fps=30):
    """
    use matplotlib's FFmpeg writer class to create an animation of the double pendulum's trajectory
    """
    if basic:
        # override fps in basic mode
        fps = 10

    # initialise animation writer
    metadata = dict(title='double pendulum', artist='tcbennun')
    writer = FFMpegWriter(fps=fps, metadata=metadata)

    # pendulum things
    l = 1
    m = 1
    theta_1 = y[0]
    theta_2 = y[1]

    # get x, y coordinates, and then convert them into an array of line segments
    x, y = coordinates_from_trajectory(y)

    # coordinates of the limbs
    limbs_x = np.array([np.zeros_like(theta_1), l*np.sin(theta_1), x])
    limbs_y = np.array([np.zeros_like(theta_1), -l*np.cos(theta_1), y])

    # initialise figure
    dpi = 100
    fig, ax = plt.subplots(figsize=(img_width/dpi, img_height/dpi), dpi=dpi)

    if basic:
        traj_line, = plt.plot([], [], linewidth=2, color="white")
    else:
        # convert coordinates into array of segments (each containing 2 points)
        points = np.array([x, y]).T.reshape(-1, 1, 2)
        segments = np.concatenate([points[:-1], points[1:]], axis=1)
        # use time array for colouring
        color_arr = 0.5*(t[1:] + t[:-1])
        norm = plt.Normalize(color_arr.min(), color_arr.max())

    # line object used to render the limbs
    pend_line, = ax.plot([], [], linewidth=10, color="xkcd:bright blue")

    # aesthetic display options
    ax.tick_params(which="both", left=False, bottom=False, labelleft=False, labelbottom=False)
    ax.set_aspect(1)
    ax.set_facecolor("black")

    # ensure the subplot fills the figure, and the aspect ratio matches
    height = 4.4
    width = img_width/img_height * height
    ax.set_xlim(-width/2, width/2)
    ax.set_ylim(-height/2, height/2)
    ax.set_position([0, 0, 1, 1])

    # animation fiddling
    length = t[-1] - t[0]
    frames = int(length * fps)
    points = len(t)
    points_per_frame = int(np.ceil(points/frames))

    print(f"{frames} frames to process. Recording...")
    t0 = time.perf_counter()

    # produce the animation
    prev_i = 0
    with writer.saving(fig, "writer_test.mp4", dpi):
        for i in np.arange(0, points, points_per_frame):
            print(f"{t[i]:.3f} seconds of {length} ({t[i]/length*100:.1f}%)...          ", end="\r")

            if basic:
                traj_line.set_data(x[:i], y[:i])
            else:
                # need to add a new bunch of segments upon each frame
                lc = LineCollection(segments[prev_i:i], cmap="RdYlBu", norm=norm)
                lc.set_array(color_arr[prev_i:i])
                lc.set_linewidth(2)
                ax.add_collection(lc)

            pend_line.set_data(limbs_x[:, i], limbs_y[:, i])
            writer.grab_frame()

            prev_i = i

    t1 = time.perf_counter()
    print(f"\nDone in {t1-t0:.3f} s.")


def dancing_flower_1():
    t_max = 30
    time_res = 1/1000
    ret = double_pendulum_trajectory(np.pi - 0.1, 0.1, 0, t_max, int(np.ceil(t_max/time_res)))
    # visualise_aesthetic(ret.t, ret.y, 1080, 1080)
    print("Animating...")
    animate_trajectory_coloured(ret.t, ret.y, 3840, 2160, fps=60)


if __name__ == "__main__":
    dancing_flower_1()