source: inundation/pyvolution/vtk_realtime_visualiser.py @ 2202

# 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.

    Customisation options for the visualiser are as follows:

    setup: Dictionary mapping quantity name -> boolean.
    if setup[q] is true, draw the quantity when the visualiser starts.
    
    updating: Dictionary mapping quantity name -> boolean.
    if updating[q] is true, update the rendering of this quantity each
    timestep.
    
    qcolor: Dictionary mapping quantity name -> (float, float, float)
    if the name of a quantity is found in qcolor, the colour specified (in
    (r, g, b) from 0.0 to 1.0) is used for display of that quantity instead of
    (0.5, 0.5, 0.5) (the default)
        
    scale_z: Dictionary mapping quantity name -> float
    Override the z scaling of this quantity with the float specified.

    default_scale_z: float
    multiply all z coordinates by this amount unless an overriding value
    exists in scale_z.
    """

    def __init__(self, domain, default_scale_z=1.0, rect=None, title='Test'):
        threading.Thread.__init__(self)
        # Initialise data structures. 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)
        self.setup = {}
        self.updating = {}
        self.qcolor = {}
        self.scale_z = {}
        self.default_scale_z = default_scale_z
        self.domain = domain
        self.vertices = domain.vertex_coordinates

        # Message queue
        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

        # 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.qcolor.has_key(q):
                actor.GetProperty().SetColor(self.qcolor[q])
            else:
                actor.GetProperty().SetColor(0.5, 0.5, 0.5)

            self.renderer.AddActor(actor)

    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')