Рисуем цветной кубик в Mayavi

1. # Для встраивания в виджет PyQt4 нужно установить переменную ETS_TOOLKIT в значение qt4.
2. import os
3. os.environ['ETS_TOOLKIT'] = 'qt4'
4. from PySide import QtGui, QtCore
5.  
6. from enthought.mayavi import mlab
7. from enthought.tvtk.api import tvtk
8. from enthought.tvtk.pyface.api import Scene
9. from numpy import arange, nonzero, float32, min, max, median, copy, random, shape
10. from numpy.core.numeric import ravel
11.  
12. from enthought.traits.api import HasTraits, Instance, on_trait_change, \
13.     Int, Dict
14. from enthought.traits.ui.api import View, Item
15. from enthought.mayavi.core.ui.api import MayaviScene, MlabSceneModel, \
16.         SceneEditor
17.  
18. #The actual visualization
19. class Visualization(HasTraits):
20.     scene = Instance(MlabSceneModel, ())
21.     view = View(Item('scene', 
22.                      editor=SceneEditor(scene_class=MayaviScene),
23.                      # Вместо MayaviScene можно написать просто Scene
24.                      # и тогда получим чистое окошко без тулбара.
25.                      height=250, 
26.                      width=300, 
27.                      show_label=False),
28.                 resizable=True
29.                 )
30.  
31.     needUpdate = None
32.  
33.     def takePlotParametres(self, grid):
34.         self.grid = grid
35.  
36.         # Метод clf() очищает текущую сцену. Его НЕЛЬЗЯ вызывать в методе update_plot()
37.         self.scene.mlab.clf()
38.         self.needUpdate = True
39.  
40.         self.update_plot()
41.  
42.     @on_trait_change('scene.activated')
43.     def update_plot(self):
44.         # Этот метод вызывается для отрисовки.
45.         # Все, что нужно отрисовать, записывается именно сюда и больше никуда
46.  
47.         # Пока нет никаких данных, отрисуем стандартную модель
48.         if not self.needUpdate:
49.             self.scene.mlab.test_points3d()
50.         else:
51.             # Сюда пишем все свои рисовашки
52.  
53.             # Все, что отрисовывается через mlab, нужно вызывать
54.             # как self.scene.mlab.* , иначе окошко Mayavi выходит из под контроля
55.             # и начинает рисовать все в отдельном окошке.
56.             # opacity сильно влияет на производительность отрисовки!
57.             surf = self.scene.mlab.pipeline.surface(self.grid, opacity=1)
58.  
59.             # Можно так же отображать и внутреннюю сетку, раскомментировав строки
60.             #edges = mlab.pipeline.extract_edges(surf)
61.             #edgesSurf = mlab.pipeline.surface(edges)
62.             #edgesSurf.actor.property.interpolation = 'flat'
63.  
64.             # Уберем интерполяцию окраски поверхности куба. Пусть каждую ячейку будет видно отчетливо.
65.             surf.actor.property.interpolation = 'flat'
66.  
67.         self.scene.mlab.orientation_axes()
68.         self.scene.background = (0, 0, 0)
69.         self.scene.mlab.colorbar(orientation='vertical')
70.  
71.  
72. ################################################################################
73. # Виджет, в который встроим сцену. Используем его как обычный виджет PyQt.
74.  
75. class MayaviQWidget(QtGui.QWidget):
76.     def __init__(self, parent=None):
77.         QtGui.QWidget.__init__(self, parent)
78.         layout = QtGui.QVBoxLayout(self)
79.         layout.setMargin(0)
80.         layout.setSpacing(0)
81.         self.visualization = Visualization()
82.  
83.         # В идеале именно тут можно вызывать метод takePlotParametres с начальными данными, но это не обязательно
84.         # self.visualization.takePlotParametres()
85.  
86.         # Метод edit_traits генерирует окно, в котором строится сцена.
87.         # Это окошко нужно сделать Qt'шным, для чего вызываем метод .control
88.         self.ui = self.visualization.edit_traits(parent=self, 
89.                                                  kind='subpanel').control
90.         layout.addWidget(self.ui)
91.         self.ui.setParent(self)
92.  
93.     def create_cube(self):
94.         grid = self.createGrid()
95.  
96.         self.visualization.takePlotParametres(grid)
97.  
98.     def createGrid(self):
99.         x, y, z = (10, 10, 10)
100.  
101.         # Создадим сетку для нашего куба
102.         grid = tvtk.RectilinearGrid()
103.  
104.         # Важно приводить данные к float32, т.к. на 64-битной ОС numpy по умолчанию создает float64,
105.         # а Mayavi работать с ним не умеет.
106.         # Создадим рандомные данные
107.         scalars = float32(random.random((x*y*z)))
108.  
109.         # scalars - это значение, которое будет храниться в каждой ячейке нашего кубика. 
110.         # Именно по этому значению определяется окраска поверхности ячеек
111.         grid.point_data.scalars = scalars
112.         grid.point_data.scalars.name = 'scalars'
113.  
114.         grid.dimensions = (x, y, z)
115.         grid.x_coordinates = float32(arange(x))
116.         grid.y_coordinates = float32(arange(y))
117.         grid.z_coordinates = float32(arange(z))
118.  
119.         return grid
120.  
121. if __name__ == "__main__":
122.     # Важно не создавать новое приложение, а использовать то, 
123.     # которое Mayavi создает автоматически, иначе вы собьете все сигналы/слоты Traits,
124.     # а нам это как раз не нужно.
125.     # Чтобы получить ссылку, вызываем метод instance().
126.     # Не напишите по привычке QtGui.QApplication(sys.argv).instance(),
127.     # я уже наступал на эти грабли.
128.     app = QtGui.QApplication.instance()
129.  
130.     container = QtGui.QWidget()
131.     container.setWindowTitle("Hello Habrahabr!")
132.     layout = QtGui.QVBoxLayout(container)
133.  
134.     mayavi_widget = MayaviQWidget()
135.  
136.     button = QtGui.QPushButton('Create cube')
137.     button.clicked.connect(mayavi_widget.create_cube)
138.  
139.     layout.addWidget(mayavi_widget)
140.     layout.addWidget(button)
141.  
142.     container.show()
143.  
144.     app.exec_()