I have to give a short presentation at SOSP next week, and for it, I needed to have some nice pictures representing a distributed array. After trying out several tools for trying to create these, I began to lament and cry over the state of Linux drawing software. But that’s a different story. I ended up writing a simple matplotlib script to generate the pictures I needed, and since it worked out pretty well, I thought I’d share it here.
Here’s the kind of picture I’m referring to:
It turns out this is pretty straightforward using matplotlib. Here’s the basic function:
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 | 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 |
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 axThe first part of this just turns off the various lines for the axes. We then iterate through the elements of the array and create a Rectangle() for each one; each “layer” (z-axis) is shifted off to the right a little bit from the previous, to give our illusion of depth. (We don’t want a normal perspective projection, as it would hide too much of the deeper layers).
The “sigh” comment is where I’m using a hack to determine the size of the text we’re going to put in so I can center it in the array cell. I couldn’t find an easier way to do this, and no, I don’t know why I have to divide the result by 30.
The input array has 3 fields which specify how to render each rectangle:
1 | dtype=([('color', 'f,f,f'), ('name', 'i'), ('transparency', 'f')])) |
dtype=([('color', 'f,f,f'), ('name', 'i'), ('transparency', 'f')]))Now we can construct an arbitrary array and feed it into our function:
1 2 3 4 5 6 7 8 | shape = (3,3,5) a = np.ndarray(shape, dtype=([('color', 'f,f,f'), ('name', 'i'), ('transparency', 'f')])) a['name'] = np.arange(np.prod(shape)).reshape(shape) a['transparency'] = 1.0 a['color'] = (1,1,1) return a draw_array(a, target='array.pdf') |
shape = (3,3,5)
a = np.ndarray(shape, dtype=([('color', 'f,f,f'), ('name', 'i'), ('transparency', 'f')]))
a['name'] = np.arange(np.prod(shape)).reshape(shape)
a['transparency'] = 1.0
a['color'] = (1,1,1)
return a
draw_array(a, target='array.pdf')Once we have the basics out of the way, we can do some fancy rendering really easily. First, let’s make a little helper class to draw slices:
1 2 3 4 5 6 7 8 9 10 | 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) |
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)We can wrap an array in draw_slice() to make it easy to construct pictures of slices:
draw_slice(a)[:,:,1]
We can be fancier if we like too, drawing the results of a filter operation:,
draw_slice(a)[a['name'] <= 1]
If you are interested, the full code for creating these figures is here: https://gist.github.com/rjpower/7249729<. All you need is matplotlib and numpy.