axes3d.py~

#!/usr/bin/env python
from misc import rotation3D
from mpl_toolkits.mplot3d import Axes3D as _Axes3D
from matplotlib import cm
import numpy as np
import tcmap

class Axes3D(_Axes3D):
  def __init__(self,*args,**kwargs):
    if (len(args) > 1) | kwargs.has_key('rect'):
      _Axes3D.__init__(self,*args,**kwargs)

    else:
      rect = (0.1,0.1,0.8,0.8)
      _Axes3D.__init__(self,*args,rect=rect,**kwargs)
    
  def cross_section(self,func,anchor,
                    zrot,yrot,xrot,
                    length,width,
                    Nl=20,Nw=20,
                    cmap=tcmap.slip,
                    clim=None):

    x = np.linspace(0,length,Nl)
    y = np.linspace(0,-width,Nw)
    x,y = np.meshgrid(x,y)
    z = 0.0*x
    R = rotation3D(zrot,yrot,xrot)
    p = np.concatenate((x[None,:,:],
                        y[None,:,:],
                        z[None,:,:]),
                        axis=0)
    p = np.einsum('ij,jkl->kli',R,p)
    p += anchor
    x = p[:,:,0]
    y = p[:,:,1]
    z = p[:,:,2]
    c = f(x,y,z)
    if clim is None:
      clim = (np.min(c),np.max(c))

    cnorm = (c - clim[0])/(clim[1] -clim[0])
    s = self.plot_surface(x,y,z,
                          shade=False,
                          facecolors=cmap(cnorm),
                          rstride=1,cstride=1)
     
    m = cm.ScalarMappable(cmap=cmap)
    m.set_array(c)

    idx1 = np.array([[0,-1],[0,-1]])
    idx2 = np.array([[0,0],[-1,-1]])

    w = self.plot_wireframe(x[idx1,idx2],
                            y[idx1,idx2],
                            z[idx1,idx2],
                            color='k')
    return m