coupled_oscillations.py

#!/usr/bin/env python3

from scipy.integrate import odeint
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import animation

# phase space points y = (X, phi, X', phi')
# length in units of l
# time in units of 1 / \omega_P

# inverse mass ratio: (m + M)/m
mu_inv = 4

# squared ratio of uncoupled frequencies: \omega_F^2 / \omega_P^2
w2 = 4**2


def f(y, t):
    pos, vel = y[:2], y[2:]
    f = np.zeros(len(y))
    f[:2] = vel
    acc = f[2:]
    sin, cos = np.sin(pos[1]), np.cos(pos[1])
    acc[0] = ((cos + vel[1]**2)*sin - w2*pos[0]) / (mu_inv - cos**2)
    acc[1] = -sin - cos*acc[0]
    return f

phi0 = np.pi * 0.8
X0 = phi0/2

y0 = (X0, phi0, 0, 0)
dt = 0.01
ts = np.arange(0.0, 50, dt)
result = odeint(f, y0, ts)


colors = ['black']*2  # mpl.rcParams["axes.color_cycle"]
fig = plt.figure()
ax = fig.add_subplot(111, autoscale_on=False, aspect='equal')
ax.set_xlim(-2.5, 2.5)
ax.set_ylim(-1.2, 1.2)
ax.set_xlabel(r'$x\;[\ell]$')
ax.set_ylabel(r'$y\;[\ell]$')
ax.grid()
spring_line = ax.plot([], [], color=colors[0])[0]
spring_mass = ax.plot([], [], 'o', color=colors[0])[0]
pendulum_line = ax.plot([], [], color=colors[1])[0]
pendulum_mass = ax.plot([], [], 'o', color=colors[1])[0]
time_template = r'$t = %.1f / \omega_P$'
time_text = ax.text(0.05, 0.9, '', transform=ax.transAxes,
                    bbox=dict(facecolor='white', edgecolor='black'))


def set_plot_data(y):
    X, phi = y[0:2]
    spring_line.set_data([-2.5, X], [0, 0])
    spring_mass.set_data([X], [0])
    sin, cos = np.sin(phi), np.cos(phi)
    pendulum_line.set_data([X, X + sin], [0, -cos])
    pendulum_mass.set_data(X + sin, -cos)


def init():
    set_plot_data(y0)


animation_step_size = 3
frames = result.shape[0]//animation_step_size


def animate(i):
    i *= animation_step_size
    set_plot_data(result[i])
    time_text.set_text(time_template % (i * dt))


ani = animation.FuncAnimation(fig, animate, frames=frames, blit=False,
                              repeat=True)
ani.save('coupled_oscillations.mp4', fps=60, extra_args=['-vcodec', 'libx264'])

fig.clear()
plt.plot(ts, result[:, :2])
plt.show()