Naive Random Tree Animation

naive_random_tree.py

import matplotlib
import matplotlib.pyplot as plt
from matplotlib import cm
import numpy as np
from matplotlib.animation import FuncAnimation

N = 100
a_max = 0.1
s = 1

rng = np.random.RandomState(0)

fig = plt.figure()

q_start = [0, 0]
plt.scatter(*q_start, c='r', s=16, zorder=2)
q_goal = [0.5, 0.5]
plt.scatter(*q_goal, c='g', s=16, zorder=2)
nodes = [q_start]


def propagate(q, a):
    q_new = [q[0] + a[0], q[1] + a[1]]
    return q_new


def update(t):
    print(t, N)
    rand_idx = rng.randint(0, t + 1)
    c = cm.jet(t / N)
    q_start = nodes[rand_idx]

    x_rand = rng.uniform(-a_max, a_max)
    y_rand = rng.uniform(-a_max, a_max)
    a_rand = [x_rand, y_rand]
    q_new = propagate(q_start, a_rand)
    if -s < q_new[0] < s and -s < q_new[1] < s:
        nodes.append(q_new)
        plt.plot([q_start[0], q_new[0]], [q_start[1], q_new[1]], color=c, linewidth=1, zorder=0)
        plt.scatter(*q_new, color=c, s=7, zorder=1)


plt.axis("equal")
plt.xlim([-s, s])
plt.ylim([-s, s])
anim = FuncAnimation(fig, update, frames=N, interval=1)
anim.save("naive_random_tree.gif", writer='imagemagick', fps=120)
plt.show()