source: trunk/anuga_core/source/anuga/visualiser_new/visualiser.py @ 8968

Last change on this file since 8968 was 4629, checked in by jack, 17 years ago

Fixed broken pause button.
File size: 9.8 KB
Line 
1from os.path import splitext
2from Tkinter import Tk, Button, Frame, Label, Scale
3from Tkinter import N, E, S, W, HORIZONTAL, VERTICAL
4from vtk import vtkCellArray, vtkRenderer, vtkWindowToImageFilter
5from vtk import vtkJPEGWriter, vtkPNGWriter, vtkPNMWriter
6from vtk.tk.vtkTkRenderWidget import vtkTkRenderWidget
7
8class Visualiser(object):
9    '''
10    Generic Offline Visualiser. Subclasses need to be used that specify
11    how to handle a data source.
12    '''
13   
14    ### Public functions ###
15
16    def __init__(self,
17                 title="Visualisation",
18                 width=400,
19                 height=400,
20                 recording=False,
21                 recordPattern=None,
22                 paused=False,
23                 source=None):
24        '''
25        Constructor.
26
27        Params:
28        title: string - Title of the visualisation window
29        width: int - Width of the visualisation window
30        height: int - Height of the visualisation window
31        recording: boolean - Start with recording enabled?
32        recordPattern: string - Pattern for recorded images,
33          e.g., cylinders%05g.png
34        paused: boolean - Start with playback paused?
35        source:- The data source to read.
36          What is required here varies by visualiser.
37        '''
38
39        # Visualisation options
40        self.vis_features = []
41        self.vis_frame = 0
42        self.vis_frameStep = 1
43        self.vis_jumping = False
44        self.vis_recording = recording
45        self.vis_recordPattern = recordPattern
46        self.vis_paused = paused
47        self.vis_source = source
48
49        # VTK structures
50        self.vtk_cells = vtkCellArray()
51        self.vtk_renderer = vtkRenderer()
52
53        # Tk structures
54        self.tk_root = Tk()
55        self.tk_root.title(title)
56        self.tk_root.grid_rowconfigure(0, weight=1)
57        self.tk_root.grid_columnconfigure(0, weight=3)
58        self.tk_root.bind('<Destroy>', self.destroyed)
59        if not self.vis_paused: self.tk_root.after(100, self.animate)           
60
61        self.tk_renderWidget = vtkTkRenderWidget(self.tk_root,
62                                                 width=width,
63                                                 height=height)
64        self.tk_renderWidget.grid(row=0, column=0, sticky=N+S+E+W)
65        self.tk_renderWidget.GetRenderWindow().AddRenderer(self.vtk_renderer)
66
67        self.tk_featureFrame = Frame(self.tk_root)
68        self.tk_featureFrame.grid(row=0, column=1, rowspan=2)
69        Label(self.tk_featureFrame, text='Features:').grid(row=0, column=0)
70
71        self.tk_controlFrame = Frame(self.tk_root)
72        self.tk_controlFrame.grid(row=1, column=0)
73
74        self.tk_quit = Button(self.tk_controlFrame, text="Quit",
75                              command=self.shutdown)
76        self.tk_quit.grid(row=0, column=0, columnspan=2)
77
78        def pause():
79            if self.vis_paused:
80                self.tk_pause.config(text='Pause')
81                self.tk_root.after(100, self.animate)
82            else:
83                self.tk_pause.config(text='Resume')
84            self.vis_paused ^= True
85        self.tk_pause = Button(self.tk_controlFrame, text="Pause", command=pause)
86        self.tk_pause.grid(row=0, column=2, columnspan=2)
87
88        if self.vis_recordPattern is not None:
89            def record():
90                if self.vis_recording:
91                    self.tk_record.config(text='Start Recording')
92                else:
93                    self.tk_record.config(text='Stop Recording')
94                self.vis_recording ^= True
95            self.tk_record = Button(self.tk_controlFrame, text="Start Recording", command=record)
96            self.tk_record.grid(row=0, column=4, columnspan=2)
97            if self.vis_recording:
98                self.tk_record.config(text="Stop Recording")
99
100        def make_seek_button(label, column, frame):
101            def jump():
102                self.jumpTo(frame)
103            b = Button(self.tk_controlFrame,
104                       text=label,
105                       command=jump)
106            b.grid(row=1, column=column, sticky=W+E)
107            return b
108        self.tk_seek_start = make_seek_button("|<", 0, 0)
109        self.tk_seek_back10 = make_seek_button("<<", 1,
110                                               lambda: self.vis_frame - 10)
111        self.tk_seek_back1 = make_seek_button("<", 2,
112                                              lambda: self.vis_frame - 1)
113        self.tk_seek_forward1 = make_seek_button(">", 3,
114                                                 lambda: self.vis_frame + 1)
115        self.tk_seek_forward10 = make_seek_button(">>", 4,
116                                                  lambda: self.vis_frame + 10)
117        self.tk_seek_end = make_seek_button(">|", 5,
118                                            self.getMaxFrameNumber)
119
120        Label(self.tk_controlFrame, text='Frame').grid(row=2, column=0,
121                                                       sticky=W+E)
122        def changeFrame(frame):
123            if not self.vis_jumping:
124                self.vis_jumping = True
125                self.jumpTo(self.tk_frame.get())
126                self.vis_jumping = False
127        self.tk_frame = Scale(self.tk_controlFrame, command=changeFrame,
128                              from_=0, to=0, orient=HORIZONTAL)
129        self.tk_frame.grid(row=2, column=1, columnspan=2, sticky=W+E)
130
131        Label(self.tk_controlFrame, text='Step').grid(row=2, column=3,
132                                                     sticky=W+E)
133        def changeFrameStep(step):
134            self.vis_frameStep = int(step)
135        self.tk_frameStep = Scale(self.tk_controlFrame, command=changeFrameStep,
136                                  from_=1, to=1, orient=HORIZONTAL)
137        self.tk_frameStep.grid(row=2, column=4, columnspan=2, sticky=W+E)
138
139        self.setupGrid()
140
141    def add_feature(self, feature):
142        '''Add a feature to this visualiser'''
143        self.vis_features.append(feature)
144        feature.button(self.tk_featureFrame).grid(row=len(self.vis_features),
145                                                  column=0, sticky=W+E)
146        feature.visualiser = self
147
148    def run(self):
149        '''Start the visualiser'''
150        self.redraw()
151        self.tk_root.mainloop()
152
153    ### Private funcitions ###
154
155    def resetSliders(self):
156        '''
157        Recalculate the upper bound on the frame and frameStep
158        sliders.
159        '''
160        maxFrame = self.getMaxFrameNumber()
161        self.tk_frame.config(to=maxFrame)
162        self.tk_frameStep.config(to=maxFrame)
163
164    def jumpTo(self, frame):
165        '''
166        Jump to a given frame. If frame is a function, jump to the
167        return value of frame().
168        '''
169        oldFrame = self.vis_frame
170        if hasattr(frame, '__call__'):
171            self.vis_frame = frame()
172        else:
173            self.vis_frame = frame
174
175        maxFrame = self.getMaxFrameNumber()
176
177        if self.vis_frame < 0:
178            self.vis_frame = 0
179        elif self.vis_frame > maxFrame:
180            self.vis_frame = maxFrame
181            self.vis_paused = True
182            self.tk_pause.config(text='Resume')
183
184        self.tk_frame.set(self.vis_frame)
185        self.redraw(oldFrame != self.vis_frame)
186
187    def redraw(self, update=False):
188        self.resetSliders()
189        for feature in [ f for f in self.vis_features if not f.dynamic ]:
190            feature.draw(self.vtk_renderer)
191        if update:
192            for feature in [ f for f in self.vis_features if f.dynamic]:
193                f.redraw(self.vtk_renderer)
194        self.tk_renderWidget.GetRenderWindow().Render()
195        self.tk_root.update_idletasks()
196
197    ### Gui events ###
198
199    def destroyed(self, event):
200        if event.widget == self.tk_root:
201            self.shutdown()
202
203    def shutdown(self):
204        self.tk_root.withdraw()
205        self.tk_root.destroy()
206
207    def animate(self):
208        if not self.vis_paused:
209            self.jumpTo(self.vis_frame + self.vis_frameStep)
210            if self.vis_recording and self.vis_recordPattern is not None:
211                self.save_image()
212            self.tk_root.after(100, self.animate)
213
214    def save_image(self):
215        extmap = {'.jpg' : vtkJPEGWriter,
216                  '.jpeg' : vtkJPEGWriter,
217                  '.png' : vtkPNGWriter,
218                  '.pnm' : vtkPNMWriter}
219        _, ext = splitext(self.vis_recordPattern)
220        try: writer = extmap[ext.lower()]()
221        except KeyError:
222            print 'ERROR: Can\'t handle %s extension. Recording disabled.' % ext
223            self.vis_recordPattern = None
224            return
225
226        win = self.vtk_renderer.GetRenderWindow()
227        w2i = vtkWindowToImageFilter()
228        w2i.SetInput(win)
229        w2i.Update()
230        writer.SetInput(w2i.GetOutput())
231        writer.SetFileName(self.vis_recordPattern % self.vis_frame)
232        win.Render()
233        writer.Write()
234
235    ### Things subclasses need to override ###
236
237    def setupGrid(self):
238        '''
239        Populate the vtkCellArray instance at
240        self.vtk_cells. Subclasses are required to override this
241        function to read from their source as appropriate.
242        '''
243        raise NotImplementedError('Subclass needs to override Visualiser::setupGrid!')
244
245    def getMaxFrameNumber(self):
246        '''
247        Return the maximum frame number. This will need to be defined
248        by a subclass.
249        '''
250        raise NotImplementedError('Subclass needs to override Visualiser::getMaxFrameNumber!')
251
252    def getQuantityPoints(self, quantityName, dynamic=True, frameNumber=0):
253        '''
254        Return the points of a quantity at a given frame as a list
255        [float]. Subclasses need to override this.
256        '''
257        raise NotImplementedError('Subclass needs to override Visualiser::getQuantityPoints!')
258
259    def getQuantityDict(self):
260        '''
261        Return the values of all quantities at a given time as a
262        dictionary. Sublclasses need to override this.
263        '''
264        raise NotImplementedError('Subclass needs to override Visualiser::getQuantityDict!')
Note: See TracBrowser for help on using the repository browser.