matplotlibでランダムウォーク

ランダムウォークをmatplotlibで表示してみた。
各粒子の動作は乱数で上下左右に動くように設定。
マシンが非力なせいか、粒子数を増やすとgif作成にやたら時間がかかる。

#-*- coding:utf-8 -*-
import matplotlib.animation as anm
import matplotlib.pyplot as plt
import numpy as np
import random

def pstep(xy):
   nnum = 0
   pnum = 0
   while nnum < len(xy):
      while pnum < 2:
         rr = random.random()
         if rr > 0.5:
            xy[nnum][pnum] =  xy[nnum][pnum] + 0.05
         else:
            xy[nnum][pnum] =  xy[nnum][pnum] - 0.05
         if xy[nnum][pnum] < 0:
            xy[nnum][pnum] += 10
         if xy[nnum][pnum] > 10:
            xy[nnum][pnum] -= 10

         pnum = pnum + 1
      pnum = 0
      nnum = nnum + 1

def update(i, fig_title, A, xy, total_step):
    if i != 0:
        plt.cla()
    ax.set_xlim(0,10)
    ax.set_ylim(0,10)

    pstep(xy)
    nnum = 0

    while nnum < len(xy):
       plt.plot(xy[nnum][0], xy[nnum][1], "o")
       nnum = nnum + 1

    plt.title(fig_title + ' step = ' + str(i) + '/' + str(total_step))


if __name__ == "__main__":
    fig = plt.figure(figsize = (10, 6))
    fig.set_size_inches(10,10)
    ax = fig.add_subplot(111)
    perticle_num = 50
    xy = [[0 for i in range(2)] for j in range(perticle_num)]

    nnum = 0
    pnum = 0
    while nnum < len(xy):
        while pnum < 2:
            xy[nnum][pnum] = random.random() * 10
            pnum = pnum + 1
        pnum = 0
        nnum = nnum + 1

    total_step =200
    ani = anm.FuncAnimation(fig, \
                        update, fargs = ('Random Work', 2.0, xy, total_step), \
                        interval = 100, frames =  total_step)

    ##plt.show()
    ani.save("Sample_randomwork.gif", writer = 'imagemagick', dpi=50)

実行した結果が下記。

random work sample by python and matplotlib.