Visualizing matrix transformations i pyplot

I was playing around with some matrix transformations and trying to visualize them using python and pyplot. Run for example as: 'ipython -pylab pyplot_matrix_transformations.py' (-pylab not really needed).

# Visualizing some vector/matrix-transformations using pylab and pyplot
# Thomas Lundqvist, 2012-01-21
# Public Domain. Use freely!

from pylab import *
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
fig = plt.figure()
# ax = fig.add_subplot(111)
ax = Axes3D(fig)

# Draws lines between all points defined as column 
# vectors in matrix a. pyplot makes each call to plot
# use a different color

def p(a):
   ax.plot(a[0].tolist()[0], a[1].tolist()[0], zs=a[2].tolist()[0])

# Original set of points

a = mat(
 [[0,0,0,1],
  [1,0,0,1],
  [1,1,0,1],
  [0,1,0,1],
  [0,0,0,1]]).T
p(a)

# Shear, new x depends on 1 * old y
# Translate z-coord +0.2
m = mat(
 [[1,1,0,  0],
  [0,1,0,  0],
  [0,0,1,0.2],
  [0,0,0,  1]])
b = m*a
p(b)

# Scale y 2 times and translate +1 in z
m2 = mat(
 [[1,0,0,0],
  [0,2,0,0],
  [0,0,1,1],
  [0,0,0,1]])
c = m2*a
p(c)

# Do both transformations (m and m2)
p(m2*b)

# Rotate the 'a' and 'c' shapes 5 degrees 20 times
# around the z-axis
for i in range(20):
   r = deg2rad((i+1)*5)
   mr = mat(
 [[cos(r), -sin(r), 0, 0],
  [sin(r),  cos(r), 0, 0],
  [     0,       0, 1, 0],
  [     0,       0, 0, 1]])
   p(mr*a)
   p(mr*c)

# Display everything
fig.show()

The result should look like this. pyplot lets me rotate the graph using the mouse: