import sys
import scipy,pylab
import matplotlib as mpl
from mpl_toolkits.mplot3d import Axes3D

##------------------------------------------------------------------------------
def closePolys(elements):
  '''
  close the open polygon elements by adding the first node of each poly again.
  '''
  closed = []
  for nodes in elements:
    i = nodes[0]
    ii = list(nodes)
    ii.append(i)
    closed.append(ii)
  return closed

elements = [ [ (0,0), (0,1), (1,1), (1,0) ],
          [ (4,1), (2,3), (2,2), (3,1)],
          [ (0,1), (2,3), (2,2), (1,1)],
          [ (3,0), (3,1), (4,1), (4,0)]]
tmp = scipy.array([1,1.1,0.9,1])
data = [tmp, tmp*2, tmp*4]

## workaround for bug in mpl: close polys manually, no auto-closing!
polys = closePolys(elements)

##------------------------------------------------------------------------------
fig = pylab.figure()
ax = Axes3D(fig)

stacks = 3
zpos = scipy.arange(stacks)

collist = []
alldata = []
for i in range(stacks):
  col = mpl.collections.PolyCollection(polys, linewidths=0.05, closed=False)
  col.set_array(data[i])
  alldata.append(data[i])
  collist.append(col)
vmin,vmax = scipy.amin(alldata), scipy.amax(alldata)

for i,col in enumerate(collist):
  ## get number of elements per collection
  elemno = len(col.get_paths())
  ## create *iterable* for the 3D collection slices [z,z,z,...,z]
  zs = [zpos[i]]*elemno
  ## collection transparency (opaque == 1), needed for projection bug workaround
  col.set_alpha(0.3)
  ## manually set the norm
  col.norm.vmin = vmin
  col.norm.vmax = vmax
  ## thus, don't autoscale!!!
  #col.autoscale()

  ax.add_collection3d(col, zs=zs, zdir='y')
  pylab.draw()

## colorbar:

## can't set the type of colormap here:
#cb = pylab.colorbar(mappable=collist[-1], drawedges=False, shrink=0.55)#, aspect=5)
#cb.set_cmap(mpl.cm.cool)
#cb.set_label('colorbar label')

## don't know how to set a label for the colorbar here:
collist[-1].set_cmap(mpl.cm.cool)
fig.colorbar(collist[-1], shrink=0.85)#, aspect=5)

## axis limit settings:
ax.set_xlim3d(0,4)
ax.set_zlim3d(0,3)
ax.set_ylim3d(-0.1,stacks-1+0.1)

## labels:
ax.set_title('Transparency workaround')
ax.set_xlabel('x')
ax.set_ylabel('slice')
ax.set_zlabel('y')

##------------------------------------------------------------------------------
pylab.show()