Source

bubble-economy / plots.py

#!/usr/bin/env python
# encoding: utf-8
"""
plosimstate.py

various handy graphs to show what's going on 

Created by dan mackinlay on 2010-10-27.
Copyright (c) 2010 __MyCompanyName__. All righsimstate reserved.
"""

import numpy as np
from matplotlib import pyplot as PLT
from matplotlib import cm
import os.path
import subprocess
from path import path
from string import Template

def tri_plot(A):
    #from http://stackoverflow.com/questions/2318529/plotting-only-upper-lower-triangle-of-a-heatmap
    # A = plt.random.randint(10, 100, 100).reshape(10, 10)
    # create an upper triangular 'matrix' from A
    A2 = np.triu(A)
    fig = plt.figure()
    ax1 = fig.add_subplot(111)
    # use dir(matplotlib.cm) to get a list of the installed colormaps
    # the "_r" means "reversed" and accounsimstate for why zero values are plotted as white
    cmap = cm.get_cmap('gray_r', 10)
    ax1.imshow(A2, interpolation="nearest", cmap=cmap)
    ax1.grid(True)
    plt.show()

def flock_quiver_plot(simstate, title=None):
    """plot arrow vectors for each particle"""
    vels = simstate.vels
    locs = simstate.locs
    
    X = locs[:,0]
    Y = locs[:,1]
    U = vels[:,0]
    V = vels[:,1]
    
    PLT.clf()
    PLT.quiver(X,Y,U,V)
    ax = PLT.gca()
    ax.set_xbound(0,1,)
    ax.set_ybound(0,1,)
    if title is None:
        title = Template('$dt dt, $num_agents agents, $radius radius, $noise noise').substitute(simstate.params)
    PLT.suptitle(title)

def animated_plot(simstate, 
        base_path,
        png_basename="2d_frame",
        movie_name="2d_anim.avi",
        fps=25,
        limit=500):
    if limit is None:
        limit = sys.maxint
    base_path = path(base_path)
    movie_path = base_path/movie_name
    png_partial_path = base_path/png_basename
    png_glob = base_path/png_basename + "*.png"
    png_subs = png_partial_path+ "%03d.png"
    try:
      os.remove(movie_path)
    except OSError:
      pass

    for i in range(min(simstate.keep_steps, limit)):
        PLT.clf()
        flock_quiver_plot(simstate, i)
        PLT.savefig(png_subs % i)
    movie_command = [
      'ffmpeg', 
      '-i',
      str(png_subs),
      '-r',
      '%d' % fps,
      str(movie_path)
    ]
    print ' '.join(movie_command)
    subprocess.call(movie_command)