parabolic reflector (animated)

Once Loop Reflect

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



fig, ax = plt.subplots(1,1)


def y(x):
    return x**2/2

def df(x,y):
    return np.array([-x, 1])


sample=10
plot_l = []
hist = []
hist_len = 40
ax.axis('off')
fig.patch.set_alpha(0.)
ax.plot(0, 0.5, 'ok', ms=10)
for i in range(sample):
    point, = ax.plot([], [], '-', ms=10)
    trace, = ax.plot([], [], '-', lw=2, c=point.get_color())
    hist_x, hist_y = deque(maxlen=hist_len), deque(maxlen=hist_len)
    plot_l.append(point)
    plot_l.append(trace)
    hist.append(hist_x)
    hist.append(hist_y)


dt = 0.02

x0 = np.linspace(-1.5, 1.5, 10)
y0 = np.array([2]*sample)
p = np.array([x0,y0])

traj_x = [np.array([])]*sample
traj_y = [np.array([])]*sample

xx = np.linspace(-2, 2, 100)
ax.plot(xx, y(xx), '-k',lw=3)
ax.text(1.7,1.2, r"$y=\frac{x^2}{2}$",fontsize=15)

for i, z in enumerate(zip(x0, y0)):
    vec = np.array([0,-1])
    for t in range(180):
        z = z + dt*vec
        x0, y0 = z
        if y0 < y(x0):
            x0 = z[0]
            y0 = y(x0)
            n = df(x0, y0)
            pn = - np.dot(vec, n)/np.dot(n, n) * n
            vec = vec + 2*pn
        if  np.all(np.abs(np.array([x0,y0]) - np.array([0,0.5])) < 0.02):
            vec = np.array([0,0])

        traj_x[i] = np.append(traj_x[i], x0)
        traj_y[i] = np.append(traj_y[i], y0)



def animate(i):
    for j in range(sample):
        thisx = traj_x[j][i]
        thisy = traj_y[j][i]

        if i == 0:
            hist[2*j].clear()
            hist[2*j+1].clear()

        hist[2*j].append(thisx)
        hist[2*j+1].append(thisy)

        plot_l[2*j].set_data(thisx, thisy)
        plot_l[2*j+1].set_data(hist[2*j], hist[2*j+1])

ani = animation.FuncAnimation(fig, animate, len(traj_x[0]),
                              interval=dt*1000,
#                              blit=True
                              )
"""
ani.save('anim_parabolic.mp4',
#         writer='pillow'
         writer="ffmpeg",
#         codec="png",
         savefig_kwargs={"transparent": True, 'facecolor': 'none'}
         )
"""
plt.show()