Plotting code for EE364b problem 8.3

traj_avoid_data.py

import matplotlib.pyplot as plt
import numpy as np

from matplotlib.animation import FuncAnimation
from matplotlib.patches import Circle

n = 2; N = 3; T = 30; D = 0.4

p1start = np.array([-1.0, 0.0]).reshape([2,1])
p1final = np.array([ 1.0, 0.0]).reshape([2,1])
p2start = np.array([-1.0, 0.5]).reshape([2,1])
p2final = np.array([ 1.0,-0.5]).reshape([2,1])
p3start = np.array([-1.0, 1.0]).reshape([2,1])
p3final = np.array([ 1.0,-1.0]).reshape([2,1])

steps = np.linspace(0, 1, T)
p1 = np.tile(p1start, (1,T)) + (p1final-p1start)*steps
p2 = np.tile(p2start, (1,T)) + (p2final-p2start)*steps 
p3 = np.tile(p3start, (1,T)) + (p3final-p3start)*steps 

fig, ax = plt.subplots(figsize=(5, 5))
ax.set(xlim=(-1.5, 1.5), ylim=(-1.5, 1.5))

ax.plot(*p1, 'bd')
ax.plot(*p2, 'rd')
ax.plot(*p3, 'kd')

c1 = Circle([0, 0], D/2, color='b', lw=1.5, fill=False)
c2 = Circle([0, 0], D/2, color='r', lw=1.5, fill=False)
c3 = Circle([0, 0], D/2, color='k', lw=1.5, fill=False)

def init():
    ax.add_artist(c1)
    ax.add_artist(c2)
    ax.add_artist(c3)
    return (c1, c2, c3)

def animate(t):
    ax.set_title('Time ' + str(t))
    c1.center = p1[0, t], p1[1, t]
    c2.center = p2[0, t], p2[1, t]
    c3.center = p3[0, t], p3[1, t]
    return

anim = FuncAnimation(fig, animate, init_func=init, interval=500, frames=T)
anim.save('traj_avoid_plot.mp4')

# Uncomment for display inside of a ipython notebook
#from IPython.display import HTML
#HTML(anim.to_html5_video())