# -*- coding: utf-8 -*-

import scipy,pylab

##*********************************************************
##*********************************************************

## 2D collection
#import matplotlib.collections as collections
import matplotlib as mpl

vertsQuad = [ [ (0,0), (0,1), (1,1), (1,0) ],
          [ (0,1), (2,3), (2,2), (1,1) ],
          [ (2,2), (2,3), (4,1), (3,1) ],
          [ (3,0), (3,1), (4,1), (4,0) ]]
dataQuad = scipy.rand(10)

pylab.figure()
col = mpl.collections.PolyCollection(vertsQuad, linewidth=0.25)
col.set_array(dataQuad)
col.set_cmap(mpl.cm.hot)
pylab.gca().add_collection(col)
pylab.axis('equal')
pylab.savefig('example_mplCollection')

##*********************************************************
##*********************************************************

## 3D collection
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.cm as cm # colormaps

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

zpos = [0,1,2,3,4]
for i in zpos:
  poly = mpl.collections.PolyCollection(vertsQuad, linewidth=0.25)
  poly.set_alpha(0.7)# collection transparency (opaque == 1), needed for projection bug workaround!
  ## apply a colorbar
  poly.set_array(dataQuad)
  poly.set_cmap(cm.cool)
  poly.autoscale()

  ## need to have a z-value for *each* polygon = element!
  zs = [zpos[i]]*len(vertsQuad)
  ## alternatively:
  elemno = len(col.get_paths())
  zs = [zpos[i]]*elemno

  ax.add_collection3d(poly, zs=zs, zdir='y')

## axis limit settings:
ax.set_xlim3d(-0.1,4.1)
ax.set_zlim3d(-0.1,3.1)
ax.set_ylim3d(-0.1,4.1)
#ax.set_title(title)
ax.set_xlabel('x')
ax.set_ylabel('slice')
ax.set_zlabel('y')

#pylab.draw()
#pylab.show()
pylab.savefig('example_mplCollection_3D')