Inspired by this post by Allen Downey, let's model simple sand piles. The model is the following:
%matplotlib inline
from pylab import *
def generate_grid(n):
return random_integers(0, 10, size=n*n).reshape(n, n)
grid = generate_grid(4)
grid
array([[ 9, 7, 2, 7], [ 8, 4, 3, 7], [ 5, 2, 7, 3], [ 8, 9, 10, 1]])
This sort of grid can easily be visualized with the help of the Axes3D
toolkit from Matplotlib.
from mpl_toolkits.mplot3d import Axes3D
def plot_grid(grid, ax):
n = grid.shape[0]
xpos, ypos = meshgrid(arange(n), arange(n))
xpos = xpos.flatten()
ypos = ypos.flatten()
zpos = 0.01 + zeros(n * n) # the little offset helps avoiding division by zero in certain cases
dx = 0.5 * ones_like(zpos)
dy = dx.copy()
dz = grid.flatten()
ax.bar3d(xpos, ypos, zpos, dx, dy, dz)
fig = figure()
ax = fig.add_subplot(111, projection='3d')
plot_grid(grid, ax)
From here, we can implement the iteration procedure.
def iterate_timestep(old_grid):
n = old_grid.shape[0]
new_grid = zeros(old_grid.shape)
for i in range(n):
for j in range(n):
if old_grid[i, j] > 4:
new_grid[i, j] += (old_grid[i, j] - 4)
if i != 0:
new_grid[i - 1, j] += 1
if i != n - 1:
new_grid[i + 1, j] += 1
if j != 0:
new_grid[i, j - 1] += 1
if j != n - 1:
new_grid[i, j + 1] += 1
else:
new_grid[i, j] += old_grid[i, j]
return new_grid
This being written, let's visualize one iteration on a given figure.
# simulation
n = 15
grid = generate_grid(n)
new_grid = iterate_timestep(grid)
# plotting
fig = plt.figure(figsize=(10, 4))
ax = fig.add_subplot(121, projection='3d')
plot_grid(grid, ax)
ax = fig.add_subplot(122, projection='3d')
plot_grid(new_grid, ax)
C:\Python27\lib\site-packages\mpl_toolkits\mplot3d\axes3d.py:1673: RuntimeWarning: invalid value encountered in divide for n in normals])
Let's again apply the algorithm to see the outcome of this.
# simulation
old_grid = new_grid
new_grid = iterate_timestep(old_grid)
# plotting
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
plot_grid(new_grid, ax)
We can make an animation with this code and display it as an HTML video. First, let's use the boilerplate code given by Jake Vanderplas on his blog.
from tempfile import NamedTemporaryFile
VIDEO_TAG = """<video controls>
<source src="data:video/x-m4v;base64,{0}" type="video/mp4">
Your browser does not support the video tag.
</video>"""
def anim_to_html(anim):
if not hasattr(anim, '_encoded_video'):
f = NamedTemporaryFile(suffix='.mp4', delete=False)
anim.save(f.name, fps=4, writer='ffmpeg', extra_args=['-vcodec', 'libx264', '-pix_fmt', 'yuv420p'])
video = open(f.name, "rb").read()
f.close()
anim._encoded_video = video.encode("base64")
return VIDEO_TAG.format(anim._encoded_video)
from IPython.display import HTML
def display_animation(anim):
close(anim._fig)
return HTML(anim_to_html(anim))
from matplotlib import animation
# First set up the figure, the axis, and the plot element we want to animate
fig = figure(figsize=(10, 10))
ax = fig.add_subplot(111, projection='3d')
grid = generate_grid(10)
maxi = grid.max()
# animation function. This is called sequentially
def animate(i):
global grid
grid = iterate_timestep(grid)
ax.cla()
plot_grid(grid, ax)
ax.set_zlim(0, maxi)
# call the animator.
anim = animation.FuncAnimation(fig,
animate,
frames=50)
# call our new function to display the animation
display_animation(anim)