Quick matplotlib script to create pictures for array visualizations.

array-viz.py

import pylab
import pandas as P
import numpy as np

def draw_array(a, target=None):
    fig = pylab.gcf()
    fig.frameon = False
    
    ax = fig.gca()
    #ax.set_axis_off()
    
    ax.patch.set_facecolor('white')
    ax.set_aspect('equal', 'box')
    ax.xaxis.set_major_locator(plt.NullLocator())
    ax.yaxis.set_major_locator(plt.NullLocator())
    
    size = 1.0
    z_scale = 1.4
    i = 0
    for z in reversed(range(a.shape[2])):
        for (x,y),v in np.ndenumerate(a[:, :, z]):
            i += 2
            alpha = a['transparency'][x,y,z]
            color = tuple(a['color'][x,y,z])
            off_x = 0.01 + x + size + z / z_scale
            off_y = y + size + z / z_scale
            
            rect = pylab.Rectangle([off_x, off_y], size, size,
                                   facecolor=color, edgecolor=(0,0,0),
                                   zorder = i, alpha = alpha)
            ax.add_patch(rect)
            
            cx = off_x + size/2
            cy = off_y + size/2
            
            # sigh
            label = str(a['name'][x,y,z])
            w, h = pylab.matplotlib.text.TextPath((0,0), label).get_extents().size / 30
            
            #print w, h
            
            text = pylab.Text(cx - w / 2, cy - h / 2, label, zorder=i+1)
            ax.add_artist(text)
    
    ax.autoscale_view()
    if target is not None:
        pylab.savefig(target)
    return ax

class draw_slice(object):
    def __init__(self, a, target=None):
        self.a = a
        self.target = target
    
    def __getitem__(self, slc):
        slice_z = np.copy(self.a)
        slice_z['color'][slc] = (0.9, 0.5, 0.3)
        slice_z['transparency'] = 0.9
        draw_array(slice_z, self.target)

        
default_size = (3, 3, 5)
np.random.seed(1)
names = np.random.randint(0, 5, np.prod(size)).reshape(size)

def make_array(shape=default_size):
    a = np.ndarray(shape, dtype=([('color', 'f,f,f'),
                              ('name', 'i'),
                              ('transparency', 'f')]))
    a['transparency'] = 1.0
    a['color'] = (1,1,1)
    return a

default = make_array()
default['name'] = names
ax = draw_array(default, target='array.pdf')


mapped = np.copy(default)
mapped['name'] += 100
draw_array(mapped, 'array+100.pdf')

draw_slice(default, 'slice-x.pdf')[1,:,:]
draw_slice(default, 'slice-y.pdf')[:,1,:]
draw_slice(default, 'slice-z.pdf')[:,:,1]
draw_slice(default, 'slice-box.pdf')[1:3,1:3,1:3]

mapped = default.copy()
mapped['name'][1:3,1:3,1:3] += 100
draw_slice(mapped, 'slice-add.pdf')[1:3,1:3,1:3]


draw_slice(default, 'filter-small.pdf')[default['name'] <= 1]

mapped = np.copy(default)
filter = mapped['name'] <= 1
mapped['name'][filter] += 100
draw_slice(mapped, 'filter-add.pdf')[filter]

reduced = make_array(shape=(3,3,1))
reduced['name'][:,:,0] = np.sum(default['name'], axis=2)
draw_array(reduced, 'reduced.pdf')