Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
Drawing cylinders
#cyl_example.py
from enthought.tvtk.tools import visual
def draw_cyl(coords):
x_ax = coords['x1'] -coords['x0']
y_ax = coords['y1'] -coords['y0']
z_ax = coords['z1'] -coords['z0']
cyl = visual.Cylinder(pos=(coords['x0'],coords['y0'],coords['z0']),
axis=(x_ax,y_ax,z_ax), radius=1, length=10)
return cyl
if __name__ =='__main__':
cyl_a_coords = {'y1': 0.0, 'y0': 0.0, 'x0': 0.0, 'x1': 30.0, 'z0': 0.0, 'z1': 0.0}
cyl_b_coords = {'y1': 0.0, 'y0': 0.0, 'x0': 30.0, 'x1': 630.0, 'z0': 0.0, 'z1': 0.0}
cyl_c_coords = {'y1': -77.88, 'y0': 0.0, 'x0': 0.0, 'x1': -184.21, 'z0': 0.0, 'z1': 0.0}
cyl_d_coords = {'y1': 389.41, 'y0': 0.0, 'x0': 0.0, 'x1': -921.06, 'z0': 0.0, 'z1': 0.0}
cyl_a = {'color': (1.0, 0.0, 0.0), 'coords' : cyl_a_coords}
cyl_b = {'color': (1.0, 0.63, 0.04), 'coords' : cyl_b_coords}
cyl_c = {'color': (1.0, 0.8, 0.0), 'coords' : cyl_c_coords}
cyl_d = {'color': (1.0, 1.0, 0.0), 'coords' : cyl_d_coords}
cyls_obj = []
for cyl in [cyl_a, cyl_b, cyl_c, cyl_d]:
drawn_cyl = draw_cyl(cyl['coords'])
drawn_cyl.color = cyl['color']
cyls_obj.append(drawn_cyl)
import numpy as np
from enthought.mayavi import mlab
from enthought.tvtk.api import tvtk
from bisect import bisect_left
class CylSelector(object):
def __init__(self):
self.outline = mlab.outline(line_width=1, color=(1.0, 1.0, 1.0))
self.outline.outline_mode = 'cornered'
self.outline.visible = False
self.picker = None
def select_cylinder(self, cyl):
y0 = cyl.center[1] - cyl.height/2. - 0.1
x0 = cyl.center[0] - cyl.radius
z0 = cyl.center[2] - cyl.radius
y1 = cyl.center[1] + cyl.height/2. + 0.1
x1 = cyl.center[0] + cyl.radius
z1 = cyl.center[2] + cyl.radius
self.outline.bounds = (x0, x1, y0, y1, z0, z1)
self.outline.visible = True
def picker_callback(self, picker):
""" Picker callback: this get called when on pick events.
"""
self.outline.visible = False
for cyl in cylinders:
x_b = [cyl.center[0] - cyl.radius, cyl.center[0] + cyl.radius]
z_b = [cyl.center[2] - cyl.radius, cyl.center[2] + cyl.radius]
y_b = [cyl.center[1] - cyl.height/2., cyl.center[1] + cyl.height/2.]
if bisect_left(x_b, self.picker.pick_position[0]) == 1:
if bisect_left(y_b, self.picker.pick_position[1]) == 1:
if bisect_left(z_b, self.picker.pick_position[2]) == 1:
self.select_cylinder(cyl)
break
num_cyl = 20
x, y, z, r, h = np.random.normal(size=(5, num_cyl))
h = .1+ np.abs(h)
r = .01+ .1*np.abs(r)
cylinders = [tvtk.CylinderSource(center=(this_x, this_y, this_z),
radius=this_r,
height=this_h,
)
for this_x, this_y, this_z, this_r, this_h
in zip(x, y, z, r, h)]
combined_source = tvtk.AppendPolyData(input=cylinders[0].output)
for cylinder in cylinders:
combined_source.add_input(cylinder.output)
combination = combined_source.output
figure = mlab.figure(1, size=(500, 500))
mlab.clf()
# Scalars
voltage = np.linspace(-90, 60, num=num_cyl)
surf = mlab.pipeline.surface(combination, vmin=voltage.min(),
vmax=voltage.max())
combination.point_data.scalars = np.repeat(voltage, 24) # 24 points per cylinder
combination.point_data.scalars.name = 'Voltage'
# Selector and Picker
cyl_sel = CylSelector()
cyl_sel.picker = figure.on_mouse_pick(cyl_sel.picker_callback, type='cell')
mlab.scalarbar(surf, orientation='vertical')
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.