radius_factor = 2.0 from mayavi import mlab from mayavi.mlab import figure fig = figure(bgcolor=(0,0,0)) from tvtk.api import tvtk from misc import Spherical as sph, convert as convdir from numpy import pi, sin, cos # render semaphore class render: __count = 0 @staticmethod def down(): render.__count += 1 fig.scene.disable_render = True @staticmethod def up(): if render.__count > 0: render.__count -= 1 if render.__count == 0: fig.scene.disable_render = False def get_line(a, b): src = tvtk.LineSource(point1=a, point2=b) mapper = tvtk.PolyDataMapper(input=src.output) actor = tvtk.Actor(mapper=mapper) fig.scene.add_actor(actor) return actor def get_trunkcone(b, a): phi_base, theta_base = sph.to(a, b)[1:] quads = tvtk.CellArray() #vtk.vtkCellArray() points = tvtk.Points() #vtk.vtkPoints() Nface = 3 for i in range(Nface+1): # rotate phi, theta = convdir((i%Nface)*2*pi/Nface, pi*0.5, phi_base, theta_base) # generate new points p = tuple(sph.xyz(a[3]*0.5*radius_factor, phi, theta, a[:3])) q = tuple(sph.xyz(b[3]*0.5*radius_factor, phi, theta, b[:3])) # insert points points.append(p) points.append(q) if i >= 1: # create a face quad = tvtk.Quad() n = points.number_of_points-1 quad.point_ids.set_id(0, n-3) # p quad.point_ids.set_id(1, n-2) # q quad.point_ids.set_id(2, n) # q quad.point_ids.set_id(3, n-1) # p # insert the new face quads.insert_next_cell(quad) # create the actor polydata = tvtk.PolyData(points=points, polys=quads) mapper = tvtk.PolyDataMapper(input=polydata) actor = tvtk.Actor(mapper=mapper) fig.scene.add_actor(actor) return actor def get_cone(base, radius, v): if type(base) != tuple: base = tuple(base) if type(v) != tuple: v = tuple(v) src = tvtk.ConeSource(center=base, radius=radius*radius_factor, height=radius, direction=v, resolution=20) mapper = tvtk.PolyDataMapper(input=src.output) actor = tvtk.Actor(mapper=mapper) fig.scene.add_actor(actor) return actor def get_sphere(p, radius, res=8): if type(p) != tuple: p = tuple(p) src = tvtk.SphereSource(center=p, radius=radius, phi_resolution=res, theta_resolution=res) mapper = tvtk.PolyDataMapper(input=src.output) actor = tvtk.Actor(mapper=mapper) fig.scene.add_actor(actor) return actor def get_many_spheres(pts, radius, res=32): actors = [] render.down() for p in pts: actors.append(get_sphere(p, radius, resolution=res)) render.up() return actors def get_ellipsoid(center, axis, res=8): dphi = 2*pi/res dtheta = pi/res [phi, theta] = numpy.mgrid[0:2*pi+dphi:dphi, 0:pi+dtheta:dtheta] x = axis[0]*0.5*cos(phi)*sin(theta)+center[0] y = axis[1]*0.5*sin(phi)*sin(theta)+center[1] z = axis[2]*0.5*cos(theta)+center[2] return mlab.mesh(x, y, z, color=(0,0,0)).actor def get_mesh(x, y, z): return mlab.mesh(x, y, z, color=(0,0,0)).actor def remove_actor(actor): fig.scene.remove_actor(actor) def start(): mlab.show()