路路路
************************************************
import os
os.environ['NUMERIX'] = 'numpy'
from scipy import arange, ravel, amax, amin, zeros, column_stack, sqrt, ones
from enthought.tvtk.tools import mlab
class figure3d:
聽聽聽聽""" A window for displaying 3D objects """
聽聽聽聽title = ''
聽聽聽聽xlabel = 'x'
聽聽聽聽ylabel = 'y'
聽聽聽聽zlabel = 'z'
聽聽聽聽objects = []
聽聽聽聽background = (1.0,1.0,1.0)
聽聽聽聽foreground = (0.0,0.0,0.0)
聽聽聽聽fig = None
聽聽聽聽def update(self):
聽聽聽聽聽聽聽聽"""Updates the figure"""
聽聽聽聽聽聽聽聽if self.fig:
聽聽聽聽聽聽聽聽聽聽聽聽hasTitle=0
聽聽聽聽聽聽聽聽聽聽聽聽# Update the title
聽聽聽聽聽聽聽聽聽聽聽聽for object in self.fig.objects:
聽聽聽聽聽聽聽聽聽聽聽聽聽聽聽聽if str(type(object))=="<class 'enthought.tvtk.tools.mlab.Title'>":
聽聽聽聽聽聽聽聽聽聽聽聽聽聽聽聽聽聽聽聽object.text = self.title
聽聽聽聽聽聽聽聽聽聽聽聽聽聽聽聽聽聽聽聽object.text_actor.property.color=self.foreground
聽聽聽聽聽聽聽聽聽聽聽聽聽聽聽聽聽聽聽聽hasTitle=1
聽聽聽聽聽聽聽聽聽聽聽聽if not self.title == '' and not hasTitle:
聽聽聽聽聽聽聽聽聽聽聽聽聽聽聽聽t=mlab.Title()
聽聽聽聽聽聽聽聽聽聽聽聽聽聽聽聽t.text=self.title
聽聽聽聽聽聽聽聽聽聽聽聽聽聽聽聽t.text_actor.property.color=self.foreground
聽聽聽聽聽聽聽聽聽聽聽聽聽聽聽聽self.fig.add(t)
聽聽聽聽聽聽聽聽聽聽聽聽# Update the labels
聽聽聽聽聽聽聽聽聽聽聽聽self.fig.objects[0].axis.x_label=self.xlabel
聽聽聽聽聽聽聽聽聽聽聽聽self.fig.objects[0].axis.y_label=self.ylabel
聽聽聽聽聽聽聽聽聽聽聽聽self.fig.objects[0].axis.z_label=self.zlabel
聽聽聽聽聽聽聽聽聽聽聽聽# Update the colors
聽聽聽聽聽聽聽聽聽聽聽聽self.fig.renwin.background=self.background
聽聽聽聽聽聽聽聽聽聽聽聽for object in self.fig.objects[0].actors:
聽聽聽聽聽聽聽聽聽聽聽聽聽聽聽聽object.property.color=self.foreground
聽聽聽聽聽聽聽聽聽聽聽聽self.fig.renwin.render()
聽聽聽聽
聽聽聽聽def show(self):
聽聽聽聽聽聽聽聽""" Shows the figure, if not already displayed """
聽聽聽聽聽聽聽聽if not self.fig:
聽聽聽聽聽聽聽聽聽聽聽聽self.fig = mlab.figure(browser=False)
聽聽聽聽聽聽聽聽self.fig.renwin.x_plus_view()
聽聽聽聽聽聽聽聽self.fig.renwin.camera.azimuth(-62)
聽聽聽聽聽聽聽聽self.fig.renwin.camera.elevation(19.5)
聽聽聽聽聽聽聽聽self.update()
聽聽聽聽def surf(self,z, transparency=0.0, transparent=0):
聽聽聽聽聽聽聽聽""" Plot a 3d surface from a 2D array """
聽聽聽聽聽聽聽聽y=arange(0,z.shape[0])
聽聽聽聽聽聽聽聽x=arange(0,z.shape[1])
聽聽聽聽聽聽聽聽# Flatten the matrix
聽聽聽聽聽聽聽聽z=ravel(z)
聽聽聽聽聽聽聽聽Zmax=amax(z)-amin(z)
聽聽聽聽聽聽聽聽XYmax=amax((amax(y)-amin(y),amax(x)-amin(y)))
聽聽聽聽聽聽聽聽Zscale=XYmax/float(Zmax)*0.5
聽聽聽聽聽聽聽聽s=mlab.SurfRegular(x,y,z,scale=[1.0,1.0,Zscale])
聽聽聽聽聽聽聽聽s.scalar_bar.title=''
聽聽聽聽聽聽聽聽s.show_scalar_bar=True
聽聽聽聽聽聽聽聽s.scalar_bar.orientation="horizontal"
聽聽聽聽聽聽聽聽s.scalar_bar.property.color=self.foreground
聽聽聽聽聽聽聽聽# LUT means "Look-Up Table", it give the mapping between scalar value and
聽聽聽聽聽聽聽聽# color
聽聽聽聽聽聽聽聽s.lut_type='blue-red'
聽聽聽聽聽聽聽聽transparency=1.0-transparency
聽聽聽聽聽聽聽聽s.lut.alpha=transparency
聽聽聽聽聽聽聽聽s.lut.alpha_range=(transparency,transparency)
聽聽聽聽聽聽聽聽if transparent:
聽聽聽聽聽聽聽聽聽聽聽聽s.lut.alpha_range=(0.0,1.0)
聽聽聽聽聽聽聽聽if not self.fig:
聽聽聽聽聽聽聽聽聽聽聽聽self.show()
聽聽聽聽聽聽聽聽# Scale properly the box outline
聽聽聽聽聽聽聽聽self.fig.objects[0].axis.set_ranges((amin(x),amax(x),amin(y),amax(y),amin(z),amax(z)))
聽聽聽聽聽聽聽聽self.fig.objects[0].axis.use_ranges=1
聽聽聽聽聽聽聽聽self.fig.add(s)
聽聽聽聽def plot3(self, x, y, z, color=(0,0.5,1.0), linewidth=1):
聽聽聽聽聽聽聽聽""" Plot a 3D line from 3 sets of coordinates """
聽聽聽聽聽聽聽聽assert len(x)== len(y) and len(y)==len(z), "x,y,z must have same number of coordinates"
聽聽聽聽聽聽聽聽x=ravel(x)
聽聽聽聽聽聽聽聽y=ravel(y)
聽聽聽聽聽聽聽聽z=ravel(z)
聽聽聽聽聽聽聽聽size=max((amax(x)-amin(x),amax(y)-amin(y),amax(z)-amin(z)))
聽聽聽聽聽聽聽聽pts = zeros((len(x), 3), 'd')
聽聽聽聽聽聽聽聽pts[:,0], pts[:,1], pts[:,2] = x, y, z
聽聽聽聽聽聽聽聽linewidth=0.05*linewidth/size
聽聽聽聽聽聽聽聽l = mlab.Line3(pts,color=color, radius=linewidth)
聽聽聽聽聽聽聽聽if not self.fig:
聽聽聽聽聽聽聽聽聽聽聽聽self.show()
聽聽聽聽聽聽聽聽self.fig.add(l)
聽聽聽聽def edit_traits(self):
聽聽聽聽聽聽聽聽self.fig.edit_traits()
聽聽聽聽def quiver3d(self, x, y, z, vx, vy, vz, scalars=None, color=None, autoscale=1):
聽聽聽聽聽聽聽聽""" Displays a plot of arrows located at x, y, z, and of
聽聽聽聽聽聽聽聽coordinates vx, vy, vz"""
聽聽聽聽聽聽聽聽assert len(x) == len(y) and len(y) == len(z) and len(z)== len(vx) and len(vx)==len(vy) and len(vy)==len(vz), "coordinate vectors must be of same length"
聽聽聽聽聽聽聽聽x=ravel(x)
聽聽聽聽聽聽聽聽y=ravel(y)
聽聽聽聽聽聽聽聽z=ravel(z)
聽聽聽聽聽聽聽聽vx=ravel(vx)
聽聽聽聽聽聽聽聽vy=ravel(vy)
聽聽聽聽聽聽聽聽vz=ravel(vz)
聽聽聽聽聽聽聽聽positions = column_stack((x,y,z))
聽聽聽聽聽聽聽聽norms = vx**2 + vy**2 + vz**2
聽聽聽聽聽聽聽聽norms = sqrt(norms.astype(float))
聽聽聽聽聽聽聽聽normsmax = amax(norms)
聽聽聽聽聽聽聽聽if autoscale:
聽聽聽聽聽聽聽聽聽聽聽聽vx=vx.astype(float)/normsmax
聽聽聽聽聽聽聽聽聽聽聽聽vy=vy.astype(float)/normsmax
聽聽聽聽聽聽聽聽聽聽聽聽vz=vz.astype(float)/normsmax
聽聽聽聽聽聽聽聽vectors = column_stack((vx,vy,vz))
聽聽聽聽聽聽聽聽if scalars is not None:
聽聽聽聽聽聽聽聽聽聽聽聽# scalars overrides color
聽聽聽聽聽聽聽聽聽聽聽聽color = (0.0,0.5,1.0)
聽聽聽聽聽聽聽聽聽聽聽聽scalars = ravel(scalars)
聽聽聽聽聽聽聽聽聽聽聽聽scalars = scalars + amin(scalars)
聽聽聽聽聽聽聽聽聽聽聽聽scalars = scalars.astype(float)/amax(scalars)
聽聽聽聽聽聽聽聽elif not color:
聽聽聽聽聽聽聽聽聽聽聽聽scalars = norms/normsmax
聽聽聽聽聽聽聽聽聽聽聽聽color = (0.0,0.5,1.0)
聽聽聽聽聽聽聽聽g = mlab.Arrows(positions,vectors=vectors,scalars=scalars,color=color)
聽聽聽聽聽聽聽聽g.glyph.scale_mode='scale_by_vector'
聽聽聽聽聽聽聽聽if not self.fig:
聽聽聽聽聽聽聽聽聽聽聽聽self.show()
聽聽聽聽聽聽聽聽self.fig.add(g)
聽聽聽聽def scatter3d(self, x, y, z, scalars=None, size=None, color=(0.0,0.5,1.0), autoscale=1):
聽聽聽聽聽聽聽聽""" Displays a bunch of spheres given by their 3d position.
聽聽聽聽聽聽聽聽Optional scalar matrix defines their color and optional size
聽聽聽聽聽聽聽聽matrix their radius"""
聽聽聽聽聽聽聽聽assert len(x) == len(y) and len(y) == len(z), "coordinate vectors must be of same length"
聽聽聽聽聽聽聽聽x=ravel(x)
聽聽聽聽聽聽聽聽y=ravel(y)
聽聽聽聽聽聽聽聽z=ravel(z)
聽聽聽聽聽聽聽聽positions = column_stack((x,y,z))
聽聽聽聽聽聽聽聽if scalars is not None:
聽聽聽聽聽聽聽聽聽聽聽聽scalars = ravel(scalars)
聽聽聽聽聽聽聽聽聽聽聽聽scalars = scalars + amin(scalars)
聽聽聽聽聽聽聽聽聽聽聽聽scalars = scalars.astype(float)/amax(scalars)
聽聽聽聽聽聽聽聽if size is not None:
聽聽聽聽聽聽聽聽聽聽聽聽size=ravel(size)
聽聽聽聽聽聽聽聽聽聽聽聽size = size + amin(size)
聽聽聽聽聽聽聽聽聽聽聽聽if autoscale:
聽聽聽聽聽聽聽聽聽聽聽聽聽聽聽聽size = 1.5*size.astype(float)/amax(size)
聽聽聽聽聽聽聽聽聽聽聽聽vectors = column_stack((size,size,size))
聽聽聽聽聽聽聽聽else:
聽聽聽聽聽聽聽聽聽聽聽聽vectors = None
聽聽聽聽聽聽聽聽g = mlab.Spheres(positions,vectors=vectors,scalars=scalars,color=color)
聽聽聽聽聽聽聽聽if size is not None:
聽聽聽聽聽聽聽聽聽聽聽聽g.glyph.scale_mode='scale_by_vector'
聽聽聽聽聽聽聽聽if not self.fig:
聽聽聽聽聽聽聽聽聽聽聽聽self.show()
聽聽聽聽聽聽聽聽self.fig.add(g)
聽聽聽聽def savefig(self, filename, **kw_args):
聽聽聽聽聽聽聽聽"""Saves rendered scene to one of several image formats
聽聽聽聽聽聽聽聽depending on the specified extension of the given filename.
聽聽聽聽聽聽聽聽Any extra keyword arguments are passed along to the respective
聽聽聽聽聽聽聽聽save method."""
聽聽聽聽聽聽聽聽self.fig.renwin.save(filename, **kw_args)
def surf(z,transparency=0.0,transparent=0):
聽聽聽聽""" Plot a 3d surface from a 2D array """
聽聽聽聽f = figure3d()
聽聽聽聽f.surf(z,transparency=transparency,transparent=transparent)
聽聽聽聽return f
聽聽聽聽
def plot3(x, y, z, color=(0,0.5,1.0), linewidth=0.2):
聽聽聽聽""" Plot a 3D line from 3 sets of coordinates """
聽聽聽聽f = figure3d()
聽聽聽聽f.plot3(x,y,z,color=color, linewidth=linewidth)
聽聽聽聽return f
def scatter3d(x, y, z, scalars=None, size=None, color=(0.0,0.5,1.0), autoscale=1):
聽聽聽聽""" Displays a bunch of points given by their 3d position.
聽聽聽聽Optional scalar matrix defines their color"""
聽聽聽聽f = figure3d()
聽聽聽聽f.scatter3d(x ,y ,z, scalars=scalars, size=size, color=color, autoscale=autoscale)
聽聽聽聽return f
def quiver3d(x, y, z, vx, vy, vz, scalars=None, color=None, autoscale=1):
聽聽聽聽""" Displays a plot of arrows located at x, y, z, and of
聽聽聽聽coordinates vx, vy, vz"""
聽聽聽聽f = figure3d()
聽聽聽聽f.quiver3d(x ,y ,z, vx, vy, vz, scalars=scalars, color=color, autoscale=autoscale)
聽聽聽聽return f