billiard trajectory in a circle (animated)

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from collections import deque

def df(x):
    return np.array([2*x[0], 2*x[1]])

def intersection(vec0, coord):
    a = vec0[1]/vec0[0]
    if np.abs(a) > 1e7:
        z = np.array([x1, -y0])

    x0 = coord[0]
    y0 = coord[1]
    sg = np.sign(vec0[0])
    D = 1 - a**2*(-1+x0**2)+ 2*a*x0*y0 -y0**2
    x1 = - ( -a**2*x0 + a*y0 - sg * np.sqrt(D) ) / (1 + a**2)

    y1 = (y0 -a * ( x0 - sg* np.sqrt(D) ) )/ (1 + a**2)
    z = np.array([x1, y1])
    return z

fig = plt.figure(figsize=(5, 5))
ax = fig.add_subplot(autoscale_on=False)
ax.set_aspect('equal')
ax.axis('off')

theta = np.linspace(-np.pi, np.pi, 100)
z = np.exp(1j*theta)
ax.plot(z.real, z.imag, '-k',lw=3)
ax.set_xlim(-1.1, 1.1)
ax.set_ylim(-1.1, 1.1)

point0, = ax.plot([], [], 'o')
trace0, = ax.plot([], [], '-', lw=2, ms=2, c=point0.get_color())


dt = 0.2

history_len = 50
history_x0, history_y0 = deque(maxlen=history_len), deque(maxlen=history_len)

x0 = np.array([0.1, 0.4])
vec0 = np.array([0.1, 0.02])


traj0 = [np.array([x0[0]]), np.array([x0[1]])]

for i in range(300):
    x1 = x0 + vec0*dt

    if x1[0]**2 + x1[1]**2 > 1:
        x1 = intersection(vec0, x0)
        n = -df(x1)
        pn = -np.dot(vec0, n)/np.dot(n, n) * n
        vec0 = vec0 + 2*pn

    traj0[0] = np.append(traj0[0], x1[0])
    traj0[1] = np.append(traj0[1], x1[1])
    x0 = x1.copy()


def animate(i):
    traj_x0,traj_y0 = traj0[0], traj0[1]

    thisx0 = traj_x0[i]
    thisy0 = traj_y0[i]


    if i == 0:
        history_x0.clear()
        history_y0.clear()


    history_x0.appendleft(thisx0)
    history_y0.appendleft(thisy0)

    point0.set_data(thisx0, thisy0)
    trace0.set_data(history_x0, history_y0)

    return point0, trace0

ani = animation.FuncAnimation(
    fig, animate, len(traj0[0]), interval=200*dt, blit=True) #, len(y1), interval=dt*1000, blit=True)

plt.show()

Total running time of the script: ( 0 minutes 7.123 seconds)

Gallery generated by Sphinx-Gallery