source: inundation/pyvolution/vtk_realtime_visualiser.py @ 2194

# VTK-Based visualiser
# intended to replace realtime_visualisation_new.py

import threading
from Queue import Queue, Empty
import Tkinter
import vtk
from vtk.tk.vtkTkRenderWidget import vtkTkRenderWidget

class Visualiser(threading.Thread):

    """A VTK-powered visualiser, designed to replace the VPython one.
    Intended to be run in its own thread
    """

    def __init__(self, domain, default_scale_z=1.0, rect=None, title='Test'):
        threading.Thread.__init__(self)
        # Initialise data structures. coloring, setup and updating are maps
        # quantity name -> boolean, setup means to render it,
        # update means to update with time.
        # qcolor maps quantity name -> (float, float, float): the colour (r, g, b)
        # to render a quantity.
        self.setup = {}
        self.updating = {}
        self.coloring = {}
        self.qcolor = {}
        self.scale_z = {}
        self.default_scale_z = default_scale_z
        self.domain = domain
        self.vertices = domain.vertex_coordinates

        self.messages = Queue()

        # Internal use - storage of vtk objects
        self.grids = {}
        self.actors = {}
        self.polydata = {}
        self.mappers = {}

        # Default options
        for x in self.domain.quantities:
            self.setup[x] = False
            self.updating[x] = False
            self.qcolor[x] = (0.0, 0.0, 0.0)
            self.coloring[x] = False

        # Bounding box.
        if rect is None:
            self.max_x = max(max(self.vertices[:,0]),max(self.vertices[:,2]),max(self.vertices[:,4]))
            self.min_x = min(min(self.vertices[:,0]),min(self.vertices[:,2]),min(self.vertices[:,4]))
            self.max_y = max(max(self.vertices[:,1]),max(self.vertices[:,3]),max(self.vertices[:,5]))
            self.min_y = min(min(self.vertices[:,1]),min(self.vertices[:,3]),min(self.vertices[:,5]))
        else:
            self.max_x = rect[2]
            self.min_x = rect[0]
            self.max_y = rect[3]
            self.min_y = rect[1]

        self.range_x = self.max_x - self.min_x
        self.range_y = self.max_y - self.min_y
        self.range_xy = max(self.range_x, self.range_y)
        
    def run(self):
        self.initialise_gui()
        self.add_axes()
        self.setup_all()
        self.root.mainloop()

    def initialise_gui(self):
        """Prepare the GUI for the Visualiser, and set up the
        renderer.
        
        """
        # Using the TK VTK widget allows extendability
        # should a gui need adding.
        self.root = Tkinter.Tk()

        # Message handling with after
        self.root.after(100, self.handle_messages)
        
        self.renderWidget = vtkTkRenderWidget(self.root, width=400, height=400)
        self.renderWidget.pack(expand='true', fill='both')
        self.renderWindow = self.renderWidget.GetRenderWindow()
        self.renderer = vtk.vtkRenderer()
        self.renderWindow.AddRenderer(self.renderer)

    def add_axes(self):
        """Add axes to this Visualiser
        """
        pass

    def setup_all(self):
        """Draw in the data that is specified by setup or update
        """

        self.N_tri = len(self.domain.triangles)
        self.cells = vtk.vtkCellArray()
        self.vertices = self.domain.get_vertex_coordinates()
        self.vert_index = {}

        for n in range(len(self.vertices)):
            for i in range(3):
                self.vert_index[self.domain.triangles[n][i]] = self.vertices[n][i*2:i*2+2]

        self.N_vert = len(self.vert_index.keys())

        # Prepare the list of cells
        for t in self.domain.triangles:
            self.cells.InsertNextCell(3)
            for i in range(3):
                self.cells.InsertCellPoint(t[i])

        # Set up the rendering of each quantity
        for q in self.domain.quantities:
            if self.setup[q] | self.updating[q]:
                self.draw_quantity(q)

    def draw_quantity(self, q):
        qty_index = {}
        for n in range(self.N_tri):
            for v in range(3):
                qty_index[self.domain.triangles[n][v]] = self.domain.quantities[q].vertex_values[n][v]

        self.grids[q] = vtk.vtkPoints()
        grid = self.grids[q]

        if self.scale_z.has_key(q):
            scale = self.scale_z[q]
        else:
            scale = self.default_scale_z
                
        for v in range(self.N_vert):
            grid.InsertNextPoint(self.vert_index[v][0],
                                 self.vert_index[v][1],
                                 qty_index[v] * scale)

        # Can't recycle vtkPolyData objects: Apparently they behave
        # unusually if the points (i.e. vertex data) is set after
        # the polys (i.e. triangle data)
        self.polydata[q] = vtk.vtkPolyData()
        polydata = self.polydata[q]
        
        polydata.SetPoints(grid)
        polydata.SetPolys(self.cells)

        if self.mappers.has_key(q):
            mapper = self.mappers[q]
            mapper.SetInput(polydata)
            mapper.Update()
        else:
            self.mappers[q] = vtk.vtkPolyDataMapper()
            mapper = self.mappers[q]
            mapper.SetInput(polydata)

        if not self.actors.has_key(q):
            self.actors[q] = vtk.vtkActor()
            actor = self.actors[q]
            actor.SetMapper(mapper)

            if self.coloring[q] is True:
                actor.GetProperty().SetColor(self.qcolor[q])
            else:
                actor.GetProperty().SetColor(0.5, 0.5, 0.5)

            self.renderer.AddActor(actor)

    def update_all(self):
        pass

    def update_timer(self):
        pass

    def handle_messages(self):
        # Expandable to include more kinds of messages if need be
        # This lookup table is copied because it is modified later
        handlers = {'redraw': self.message_redraw}.copy()
        try:
            while True:
                message = self.messages.get_nowait()
                try:
                    handlers[message]()
                    # Guard against redrawing more than once per
                    # call to handle_messages.
                    if message == 'redraw':
                        del handlers[message]
                except KeyError:
                    pass
        except Empty:
            pass
        self.root.after(100, self.handle_messages)

    def message_redraw(self):
        """Re-render and repaint the rendered image.
        This function should only be called from the message handler.
        Use request_redraw to ask for a redraw"""

        for q in self.domain.quantities:
            if self.updating[q]:
                self.draw_quantity(q)
        
        self.renderWindow.Render()
        self.root.update_idletasks()

    def request_redraw(self):
        self.messages.put('redraw')