#214 Animated 3D random walk

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
# Fixing random state for reproducibility
np.random.seed(19680801)

def random_walk(num_steps, max_step=0.05):
    """Return a 3D random walk as (num_steps, 3) array."""
    start_pos = np.random.random(3)
    steps = np.random.uniform(-max_step, max_step, size=(num_steps, 3))
    walk = start_pos + np.cumsum(steps, axis=0)
    return walk

def update_lines(num, walks, lines):
    for line, walk in zip(lines, walks):
        # NOTE: there is no .set_data() for 3 dim data...
        line.set_data(walk[:num, :2].T)
        line.set_3d_properties(walk[:num, 2])
    return lines

# Data: 20 random walks as (num_steps, 3) arrays
num_steps = 100
walks = [random_walk(num_steps) for index in range(20)]
# Attaching 3D axis to the figure
fig = plt.figure()
ax = fig.add_subplot(projection="3d")
# Create lines initially without data
lines = [ax.plot([], [], [],lw=0.5)[0] for _ in walks]
# Setting the axes properties
ax.set(xlim3d=(0, 1), xlabel='X')
ax.set(ylim3d=(0, 1), ylabel='Y')
ax.set(zlim3d=(0, 1), zlabel='Z')
# Creating the Animation object
ani = animation.FuncAnimation(fig, update_lines, num_steps, fargs=(walks, lines), interval=100,repeat=False)
plt.savefig('ex214.png', dpi=72)
plt.show()