source: inundation/pyvolution/realtime_visualisation_new.py @ 2334

Last change on this file since 2334 was 2334, checked in by steve, 18 years ago

Adding coloring to vtk visualisation

  • Property svn:executable set to *
File size: 20.0 KB
Line 
1
2from visual import *
3
4elevation_color = (0.3,0.4,0.5)
5stage_color     = (0.1,0.4,0.99)
6friction_color  = (0.1,0.99,0.1)
7other_color     = (0.99,0.4,0.1)
8
9
10class Visualiser:
11
12    def __init__(self,domain,scale_z=1.0,rect=None,title='Test'):
13        """Create visualisation of domain
14        """
15
16        self.scene = display(title=title)
17        self.scene.center = (0.5,0.5,0.0)
18        self.scene.forward = vector(-0.0, 0.5, -0.8)
19        self.scene.ambient = 0.4
20        self.scene.lights = [(0.6, 0.3, 0.2), (0.1, -0.5, 0.4)]
21
22        self.frame  = frame()
23
24        self.domain          = domain
25        self.default_scale_z = scale_z
26        self.vertices        = domain.vertex_coordinates
27
28
29        # models for each quantity
30        self.vpython_z_models = {}
31        #print keys
32        for key in self.domain.quantities:
33            self.vpython_z_models[key] = faces(frame=self.frame)
34
35        #print self.z_models
36
37        if rect is None:
38            self.max_x = max(max(self.vertices[:,0]),max(self.vertices[:,2]),max(self.vertices[:,4]))
39            self.min_x = min(min(self.vertices[:,0]),min(self.vertices[:,2]),min(self.vertices[:,4]))
40            self.max_y = max(max(self.vertices[:,1]),max(self.vertices[:,3]),max(self.vertices[:,5]))
41            self.min_y = min(min(self.vertices[:,1]),min(self.vertices[:,3]),min(self.vertices[:,5]))
42        else:
43            self.max_x = rect[2]
44            self.min_x = rect[0]
45            self.max_y = rect[3]
46            self.min_y = rect[1]
47
48
49        self.range_x = self.max_x - self.min_x
50        self.range_y = self.max_y - self.min_y
51        self.range_xy = max(self.range_x, self.range_y)
52
53
54#        print 'min_x=',self.min_x
55#        print 'max_x=',self.max_x
56#        print 'range_x=',self.range_x
57#        print 'min_y=',self.min_y
58#        print 'max_y=',self.max_y
59#        print 'range_y',self.range_y
60
61        self.stage_color     = stage_color
62        self.elevation_color = elevation_color
63        self.friction_color  = friction_color
64        self.other_color     = other_color
65
66        self.qcolor = {}
67        self.scale_z = {}
68        self.coloring = {}
69        self.updating = {}
70        self.setup = {}
71
72        self.qcolor['stage']  = self.stage_color
73        self.scale_z['stage']  = 1.0
74        self.coloring['stage'] = False
75        self.updating['stage'] = True
76        self.setup['stage'] = True
77
78        self.qcolor['elevation']  = self.elevation_color
79        self.scale_z['elevation']  = 1.0
80        self.coloring['elevation'] = False
81        self.updating['elevation'] = False
82        self.setup['elevation'] = False
83
84        #print self.max_x, self.min_x, self.max_y, self.min_y,
85        #self.min_bed print self.max_bed, self.range_x, self.range_y,
86        #self.range_xy, self.range_bed
87        #print 'Calculating triangles'
88
89        self.pos     = zeros( (6*len(domain),3), Float)
90        self.colour  = zeros( (6*len(domain),3), Float)
91        self.normals = zeros( (6*len(domain),3), Float)
92
93        #print 'keys',self.domain.quantities.keys()
94        #print 'shape of stage',shape(self.stage)
95
96        self.border_model = curve(frame = self.frame,
97                                  pos=[(0,0),(0,1),(1,1),(1,0),(0,0)])
98        self.timer=label(pos=(0.75,0.75,0.5),text='Time=%10.5e'%self.domain.time,
99                         visible=False)
100
101
102        scene.autoscale=0
103        #self.update_all()
104        #scene.uniform=0
105        self.range_z  = 1.0
106        if domain.visualise_range_z == None:
107            self.setup_range_z()
108        else:
109            self.range_z = domain.visualise_range_z
110
111        self.max_z = self.range_z/2.0
112        self.min_z = -self.range_z/2.0
113        self.range_z = max(self.max_z - self.min_z, 1.0e-10)
114
115#        print self.range_z
116#        print self.max_z
117#        print self.min_z
118
119    def setup_range_z(self,qname1='elevation',qname2='stage'):
120
121        #print qname1
122        #print qname2
123        range_z = 1.0e-10
124        try:
125            q1 = self.domain.quantities[qname1].vertex_values
126            max_z = max(max(q1))
127            min_z = min(min(q1))
128            print max_z, min_z
129            range_z = max(range_z, max_z - min_z)
130        except:
131            #print 'Visualisation: could not find range of '+qname1
132            pass
133
134        try:
135            q2 = self.domain.quantities[qname2].vertex_values
136            max_z = max(max(q2))
137            min_z = min(min(q2))
138            print max_z, min_z
139            range_z = max(range_z, max_z - min_z)
140        except:
141            #print 'Visualisation: could not find range of '+qname2
142            pass
143
144        self.range_z = max(range_z,self.range_z)
145
146
147    def setup_all(self):
148
149        self.scene.select()
150
151        for key in ['elevation','stage']:
152            if self.setup[key] | self.updating[key]:
153                self.update_quantity(key)
154
155
156
157    def update_all(self):
158
159        self.scene.select()
160
161        for key in ['elevation','stage']:
162            if self.updating[key]:
163                self.update_quantity(key)
164
165
166    def update_timer(self):
167
168        self.scene.select()
169
170        if self.domain.visualise_timer == True:
171            self.timer.visible = True
172            self.timer.text='Time=%10.5e'%self.domain.time
173
174
175    def update_quantity(self,qname):
176
177        #print 'update '+qname+' arrays'
178        try:
179            qcolor   = self.qcolor[qname]
180            scale_z  = self.scale_z[qname]
181            coloring = self.coloring[qname]
182            updating = self.updating[qname]
183        except:
184            qcolor   = None
185            scale_z  = None
186            coloring = False
187            updating = True
188
189
190        if qcolor is None:
191            qcolor = self.other_color
192
193            if qname=='elevation':
194                qcolor = self.elevation_color
195
196            if qname=='stage':
197                qcolor = self.stage_color
198
199            if qname=='friction':
200                qcolor = self.friction_color
201
202        if scale_z is None:
203            scale_z = self.default_scale_z
204
205        #try:
206        if coloring:
207            self.update_arrays_color(self.domain.quantities[qname].vertex_values, qcolor, scale_z)
208        else:
209            self.update_arrays(self.domain.quantities[qname].vertex_values, qcolor, scale_z)
210
211        #print 'update bed image'
212        if qname=='elevation':
213            self.pos[:,2] = self.pos[:,2]+1.0e-3
214
215        self.vpython_z_models[qname].pos    = self.pos
216        self.vpython_z_models[qname].color  = self.colour
217        self.vpython_z_models[qname].normal = self.normals
218        #except:
219        #    print 'Visualisation: Could not update quantity '+qname
220        #    pass
221
222    def update_quantity_color(self,qname,qcolor=None,scale_z=None):
223
224        #print 'update '+qname+' arrays'
225
226        if qcolor is None:
227            if qname=='elevation':
228                qcolor = self.elevation_color
229
230            if qname=='stage':
231                qcolor = self.stage_color
232
233            if qname=='friction':
234                qcolor = self.friction_color
235
236        if scale_z is None:
237            scale_z = self.scale_z
238
239        #try:
240        self.update_arrays_color(self.domain.quantities[qname].vertex_values, qcolor, scale_z)
241
242            #print 'update bed image'
243        self.vpython_z_models[qname].pos    = self.pos
244        self.vpython_z_models[qname].color  = self.colour
245        self.vpython_z_models[qname].normal = self.normals
246        #except:
247        #    print 'Visualisation: Could not update quantity '+qname
248        #    pass
249
250
251    def update_arrays(self,quantity,qcolor,scale_z):
252
253        col = asarray(qcolor)
254
255        N = len(self.domain)
256
257        min_x = self.min_x
258        min_y = self.min_y
259        range_xy = self.range_xy
260        range_z = self.range_z
261
262        vertices = self.vertices
263        pos      = self.pos
264        normals  = self.normals
265        colour   = self.colour
266
267        try:
268            update_arrays_weave(vertices,quantity,col,pos,normals,colour,N,
269                                min_x,min_y,range_xy,range_z,scale_z)
270        except:
271            update_arrays_python(vertices,quantity,col,pos,normals,colour,N,
272                                 min_x,min_y,range_xy,range_z,scale_z)
273
274
275
276    def update_arrays_color(self,quantity,qcolor,scale_z):
277
278        col = asarray(qcolor)
279
280        N = len(self.domain)
281
282        min_x = self.min_x
283        min_y = self.min_y
284        range_xy = self.range_xy
285        range_z = self.range_z
286
287        vertices = self.vertices
288        pos      = self.pos
289        normals  = self.normals
290        colour   = self.colour
291
292        #try:
293        update_arrays_color_python(vertices,quantity,col,pos,normals,colour,N,
294                                min_x,min_y,range_xy,range_z,scale_z)
295        #except:
296        #    update_arrays_color_python(vertices,quantity,col,pos,normals,colour,N,
297        #                         min_x,min_y,range_xy,range_z,scale_z)
298
299
300#==================================================================================
301
302def  update_arrays_python(vertices,quantity,col,pos,normals,colour,N,
303                          min_x,min_y,range_xy,range_z,scale_z):
304
305    from math import sqrt
306    normal = zeros( 3, Float)
307    v = zeros( (3,3), Float)
308    s = 1.0
309
310    for i in range( N ):
311
312        for j in range(3):
313            v[j,0] = (vertices[i,2*]-min_x)/range_xy
314            v[j,1] = (vertices[i,2*j+1]-min_y)/range_xy
315            v[j,2] = quantity[i,j]/range_z*scale_z*0.5
316
317        v10 = v[1,:]-v[0,:]
318        v20 = v[2,:]-v[0,:]
319
320        normal[0] = v10[1]*v20[2] - v20[1]*v10[2]
321        normal[1] = v10[2]*v20[0] - v20[2]*v10[0]
322        normal[2] = v10[0]*v20[1] - v20[0]*v10[1]
323
324        norm = sqrt( normal[0]**2 + normal[1]**2 + normal[2]**2)
325
326        normal[0] = normal[0]/norm
327        normal[1] = normal[1]/norm
328        normal[2] = normal[2]/norm
329
330        pos[6*,:] = v[0,:]
331        pos[6*i+1,:] = v[1,:]
332        pos[6*i+2,:] = v[2,:]
333        pos[6*i+3,:] = v[0,:]
334        pos[6*i+4,:] = v[2,:]
335        pos[6*i+5,:] = v[1,:]
336
337
338
339        colour[6*,:] = s*col
340        colour[6*i+1,:] = s*col
341        colour[6*i+2,:] = s*col
342        colour[6*i+3,:] = s*col
343        colour[6*i+4,:] = s*col
344        colour[6*i+5,:] = s*col
345
346        s =  15.0/8.0 - s
347
348        normals[6*,:] = normal
349        normals[6*i+1,:] = normal
350        normals[6*i+2,:] = normal
351        normals[6*i+3,:] = -normal
352        normals[6*i+4,:] = -normal
353        normals[6*i+5,:] = -normal
354
355
356
357def  update_arrays_weave(vertices,quantity,col,pos,normals,colour,
358                         N,min_x,min_y,range_xy,range_z,scale_z):
359
360    import weave
361    from weave import converters
362
363    from math import sqrt
364    normal = zeros( 3, Float)
365    v = zeros( (3,3), Float)
366    v10 = zeros( 3, Float)
367    v20 = zeros( 3, Float)
368
369    code1 = """
370
371        double s = 1.0;
372
373        for (int i=0; i<N ; i++ ) {
374            for (int j=0; j<3 ; j++) {
375                v(j,0) = (vertices(i,2*j  )-min_x)/range_xy;
376                v(j,1) = (vertices(i,2*j+1)-min_y)/range_xy;
377                v(j,2) = quantity(i,j)/range_z*scale_z*0.5;
378            }
379
380
381            for (int j=0; j<3; j++) {
382                v10(j) = v(1,j)-v(0,j);
383                v20(j) = v(2,j)-v(0,j);
384            }
385
386            normal(0) = v10(1)*v20(2) - v20(1)*v10(2);
387            normal(1) = v10(2)*v20(0) - v20(2)*v10(0);
388            normal(2) = v10(0)*v20(1) - v20(0)*v10(1);
389
390            double norm =  sqrt(normal(0)*normal(0) + normal(1)*normal(1) + normal(2)*normal(2));
391
392            normal(0) = normal(0)/norm;
393            normal(1) = normal(1)/norm;
394            normal(2) = normal(2)/norm;
395
396
397            for (int j=0; j<3; j++) {
398                pos(6*i  ,j) = v(0,j);
399                pos(6*i+1,j) = v(1,j);
400                pos(6*i+2,j) = v(2,j);
401                pos(6*i+3,j) = v(0,j);
402                pos(6*i+4,j) = v(2,j);
403                pos(6*i+5,j) = v(1,j);
404
405                colour(6*i  ,j) = s*col(j);
406                colour(6*i+1,j) = s*col(j);
407                colour(6*i+2,j) = s*col(j);
408                colour(6*i+3,j) = s*col(j);
409                colour(6*i+4,j) = s*col(j);
410                colour(6*i+5,j) = s*col(j);
411
412                normals(6*i  ,j) = normal(j);
413                normals(6*i+1,j) = normal(j);
414                normals(6*i+2,j) = normal(j);
415                normals(6*i+3,j) = -normal(j);
416                normals(6*i+4,j) = -normal(j);
417                normals(6*i+5,j) = -normal(j);
418                }
419
420            s =  15.0/8.0 - s;
421        }
422
423        """
424
425    weave.inline(code1, ['vertices','quantity','col','pos','normals','colour','N',
426                         'min_x','min_y','range_xy','range_z','scale_z','v','normal','v10','v20'],
427                 type_converters = converters.blitz, compiler='gcc');
428
429def  update_arrays_color_python(vertices,quantity,col,pos,normals,colour,N,
430                          min_x,min_y,range_xy,range_z,scale_z):
431
432    from math import sqrt
433    normal = zeros( 3, Float)
434    v = zeros( (3,3), Float)
435    s = 1.0
436
437    for i in range( N ):
438
439        for j in range(3):
440            v[j,0] = (vertices[i,2*]-min_x)/range_xy
441            v[j,1] = (vertices[i,2*j+1]-min_y)/range_xy
442            v[j,2] = quantity[i,j]/range_z*scale_z*0.5
443
444        v10 = v[1,:]-v[0,:]
445        v20 = v[2,:]-v[0,:]
446
447        normal[0] = v10[1]*v20[2] - v20[1]*v10[2]
448        normal[1] = v10[2]*v20[0] - v20[2]*v10[0]
449        normal[2] = v10[0]*v20[1] - v20[0]*v10[1]
450
451        norm = sqrt( normal[0]**2 + normal[1]**2 + normal[2]**2)
452
453        normal[0] = normal[0]/norm
454        normal[1] = normal[1]/norm
455        normal[2] = normal[2]/norm
456
457        pos[6*,:] = v[0,:]
458        pos[6*i+1,:] = v[1,:]
459        pos[6*i+2,:] = v[2,:]
460        pos[6*i+3,:] = v[0,:]
461        pos[6*i+4,:] = v[2,:]
462        pos[6*i+5,:] = v[1,:]
463
464        q0 = quantity[i,0]/0.4
465        q1 = quantity[i,1]/0.4
466        q2 = quantity[i,2]/0.4
467
468        q0r = min(q0,0.0)
469        q1r = min(q1,0.0)
470        q2r = min(q2,0.0)
471        q0b = max(q0,0.0)
472        q1b = max(q1,0.0)
473        q2b = max(q2,0.0)
474
475
476#        colour[6*i  ,:] = s*array([0.3 - q0r, 0.3 - q0, 0.3 + q0b])
477#        colour[6*i+1,:] = s*array([0.3 - q1r, 0.3 - q1, 0.3 + q1b])
478#        colour[6*i+2,:] = s*array([0.3 - q2r, 0.3 - q2, 0.3 + q2b])
479#        colour[6*i+3,:] = s*array([0.3 - q0r, 0.3 - q0, 0.3 + q0b])
480#        colour[6*i+4,:] = s*array([0.3 - q2r, 0.3 - q2, 0.3 + q2b])
481#        colour[6*i+5,:] = s*array([0.3 - q1r, 0.3 - q1, 0.3 + q1b])
482
483        s = 0.1
484        colour[6*,:] = s*array([0.3 - q0r, 0.3, 0.3 + q0b])
485        colour[6*i+1,:] = s*array([0.3 - q1r, 0.3, 0.3 + q1b])
486        colour[6*i+2,:] = s*array([0.3 - q2r, 0.3, 0.3 + q2b])
487        colour[6*i+3,:] = s*array([0.3 - q0r, 0.3, 0.3 + q0b])
488        colour[6*i+4,:] = s*array([0.3 - q2r, 0.3, 0.3 + q2b])
489        colour[6*i+5,:] = s*array([0.3 - q1r, 0.3, 0.3 + q1b])
490
491
492        s =  15.0/8.0 - s
493
494        normals[6*,:] = normal
495        normals[6*i+1,:] = normal
496        normals[6*i+2,:] = normal
497        normals[6*i+3,:] = -normal
498        normals[6*i+4,:] = -normal
499        normals[6*i+5,:] = -normal
500
501
502
503def  update_arrays_color_weave(vertices,quantity,col,pos,normals,colour,
504                         N,min_x,min_y,range_xy,range_z,scale_z):
505
506    import weave
507    from weave import converters
508
509    from math import sqrt
510    normal = zeros( 3, Float)
511    v = zeros( (3,3), Float)
512    v10 = zeros( 3, Float)
513    v20 = zeros( 3, Float)
514
515    code1 = """
516        double s = 1.0;
517
518        for (int i=0; i<N ; i++ ) {
519            for (int j=0; j<3 ; j++) {
520                v(j,0) = (vertices(i,2*j  )-min_x)/range_xy;
521                v(j,1) = (vertices(i,2*j+1)-min_y)/range_xy;
522                v(j,2) = quantity(i,j)/range_z*scale_z*0.5;
523            }
524
525
526            for (int j=0; j<3; j++) {
527                v10(j) = v(1,j)-v(0,j);
528                v20(j) = v(2,j)-v(0,j);
529            }
530
531            normal(0) = v10(1)*v20(2) - v20(1)*v10(2);
532            normal(1) = v10(2)*v20(0) - v20(2)*v10(0);
533            normal(2) = v10(0)*v20(1) - v20(0)*v10(1);
534
535            double norm =  sqrt(normal(0)*normal(0) + normal(1)*normal(1) + normal(2)*normal(2));
536
537            normal(0) = normal(0)/norm;
538            normal(1) = normal(1)/norm;
539            normal(2) = normal(2)/norm;
540
541
542
543            for (int j=0; j<3; j++) {
544                pos(6*i  ,j) = v(0,j);
545                pos(6*i+1,j) = v(1,j);
546                pos(6*i+2,j) = v(2,j);
547                pos(6*i+3,j) = v(0,j);
548                pos(6*i+4,j) = v(2,j);
549                pos(6*i+5,j) = v(1,j);
550
551                normals(6*i  ,j) = normal(j);
552                normals(6*i+1,j) = normal(j);
553                normals(6*i+2,j) = normal(j);
554                normals(6*i+3,j) = -normal(j);
555                normals(6*i+4,j) = -normal(j);
556                normals(6*i+5,j) = -normal(j);
557                }
558
559            s = 0.2+fabs((v(0,2)+v(1,2)+v(2,2))/3.0);
560            for (int j=0; j<3; j++) {
561                colour(6*i  ,j) = s*col(j);
562                colour(6*i+1,j) = s*col(j);
563                colour(6*i+2,j) = s*col(j);
564                colour(6*i+3,j) = s*col(j);
565                colour(6*i+4,j) = s*col(j);
566                colour(6*i+5,j) = s*col(j);
567                }
568            s =  15.0/8.0 - s;
569        }
570
571        """
572
573    weave.inline(code1, ['vertices','quantity','col','pos','normals','colour','N',
574                         'min_x','min_y','range_xy','range_z','scale_z','v','normal','v10','v20'],
575                 type_converters = converters.blitz, compiler='gcc');
576
577
578def setup_scene():
579
580    #scene.width = 1000
581    #scene.height = 800
582
583    #Original
584    scene.center = (0.5,0.5,0.0)
585    scene.forward = vector(-0.0, 0.5, -0.8)
586
587    #Temporary (for bedslope)
588    #scene.forward = vector(0.0006, 0.7, -0.03)
589
590
591    #Temporary for hackett - begin
592    #scene.autoscale = 0
593    #scene.scale = (0.002, 0.002, 0.01) #Scale z so that countours stand out more
594    #scene.center = (300.0,500.0,-10)
595    #Temporary for hackett - end
596
597
598    scene.ambient = 0.4
599    #scene.lights = [(0.6, 0.3, 0.2), (0.1, -0.5, 0.4), (-0.1, 0.1, -0.4),
600    #               (-0.2, 0.2, 0.1)]
601
602    scene.lights = [(0.6, 0.3, 0.2), (0.1, -0.5, 0.4)]
603
604
605
606def create_surface(domain,scale_z=1.0):
607
608    domain.initialise_visualiser(scale_z)
609
610    domain.visualiser.update_quantity('elevation')
611    domain.visualiser.update_quantity('stage')
612    domain.visualiser.update_timer()
613
614def update(domain):
615
616    if domain.visualise_color_stage:
617        domain.visualiser.update_quantity_color('stage')
618    else:
619        domain.visualiser.update_quantity('stage')
620
621    domain.visualiser.update_timer()
622
623
624
625if __name__ == '__main__':
626
627    from advection import Domain as A_Domain
628    from shallow_water import Domain as SW_Domain
629    from shallow_water import Constant_height, Constant_stage
630    from mesh_factory import rectangular
631
632    points, vertices, boundary = rectangular(60, 60, len1=200, len2=200)
633    a_domain  = A_Domain(points, vertices, boundary)
634    print 'No of Elements' , a_domain.number_of_elements
635    #sw_domain = SW_Domain(points, vertices, boundary)
636    #sw_domain.set_quantity('elevation', Constant_stage(0.0))
637    #sw_domain.set_quantity('stage', Constant_stage(0.0))
638
639    a_domain.initialise_visualiser()
640    a_domain.visualiser.setup_all()
641
642
643    #sw_domain.initialise_visualiser()
644    #sw_domain.visualiser.setup['elevation']=True
645    #sw_domain.visualiser.setup_all()
646
647    time = 0.0
648    while (True):
649        time = time+0.001
650        #print 'Time = ',time
651        #sw_domain.set_quantity('stage', time)
652        a_domain.set_quantity('stage',time)
653
654        #sw_domain.visualiser.update_quantity('stage')
655        a_domain.visualiser.update_quantity('stage')
656
657        if time>1.0 :
658            break
659
660    #a_v = Visualiser(domain=a_domain, title='advection')
661    #a_v.update_all()
662
663
664    #sw_v = Visualiser(domain=sw_domain, title='shallow water')
665    #sw_v.update_all()
Note: See TracBrowser for help on using the repository browser.