DoublePendulum

DoublePendulum.py

start_time = 0
end_time = 60
from manimlib.imports import *
from manimlib import constants
from numpy import sin, cos
import numpy as np
import scipy.integrate as integrate

# Stolen from https://matplotlib.org/3.1.1/gallery/animation/double_pendulum_sgskip.html

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
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]

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

    dydx[2] = state[3]

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

    return dydx
fps = 60
# create a time array from 0..100 sampled at 0.05 second steps
dt = (1 / fps)
t = np.arange(start_time, end_time, dt)
print(t.shape)

# 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 your ODE using scipy.integrate.
y = integrate.odeint(derivs, state, t)

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

ap = np.asarray((x1, y1, np.zeros_like(x1)))

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

bp = np.asarray((x2, y2, np.zeros_like(x1)))

#print(x1, y1)
constants.FRAME_WIDTH = 7
constants.FRAME_HEIGHT = constants.FRAME_WIDTH * constants.DEFAULT_PIXEL_HEIGHT / constants.DEFAULT_PIXEL_WIDTH
constants.FRAME_Y_RADIUS = constants.FRAME_HEIGHT / 2
constants.FRAME_X_RADIUS = constants.FRAME_WIDTH / 2
class DoublePendulum(MovingCameraScene):
    def construct(self):
        self.camera.set_frame_height(constants.FRAME_HEIGHT)
        self.camera.set_frame_width(constants.FRAME_WIDTH)

        grid = NumberPlane()
        self.add(grid)

        c2a = Line(stroke_width=6)
        a2b = Line(stroke_width=6)
        rrr = Elbow()
        self.add(c2a)
        self.add(a2b)

        c = Dot(point=(np.asarray((0,0,0))))
        a = Dot(point=(np.asarray((0,-1,0))), color=COLOR_MAP["RED_E"])
        b = Dot(point=(np.asarray((0,-2,0))))
        # aline = CubicBezier([],color=COLOR_MAP['RED_C'])
        # bline = CubicBezier([],color=COLOR_MAP['GREEN_C'])
        # self.add(aline)
        # self.add(bline)
        ashades = []
        bshades = []
        for i in range(0, 2 * fps):
            ashade = Line(color=COLOR_MAP['RED_C'], fill_opacity=1 - i / (2*fps))
            bshade = Line(color=COLOR_MAP['GREEN_C'], fill_opacity=1 - i / (2*fps))
            ashade.set_opacity(i / (2*fps))
            bshade.set_opacity(i / (2*fps))

            self.add(ashade)
            self.add(bshade)
            ashades.append(ashade)
            bshades.append(bshade)
        self.add(a)
        self.add(b)
        self.add(c)
        for i in range(0, x1.shape[0]):
            a.move_to(np.array((x1[i],y1[i], 0.)))
            b.move_to(np.array((x2[i],y2[i], 0.)))
            c2a.put_start_and_end_on(np.array((0., 0., 0.)), np.array((x1[i],y1[i], 0.)))
            a2b.put_start_and_end_on(np.array((x1[i],y1[i], 0.)), np.array((x2[i],y2[i], 0.)))
            # aline.set_points(ap[ min(0, i - 2 * fps):i ])
            # bline.set_points(bp[ min(0, i - 2 * fps):i ])
            for tt in range(0, 2 * fps):
                pos = max(0, i - 2 * fps + tt)
                ashades[tt].put_start_and_end_on(np.array((x1[pos], y1[pos], 0.)),
                                                 np.array((x1[pos + 1], y1[pos + 1], 0.)))
                bshades[tt].put_start_and_end_on(np.array((x2[pos], y2[pos], 0.)),
                                                 np.array((x2[pos + 1], y2[pos + 1], 0.)))
            self.wait(dt)