DISPATCH
mayavi_utils.py
1 # Pythn 2/3 compatibility
2 from __future__ import print_function
3 
4 import matplotlib.mlab as mml
5 import mayavi.mlab as ml
6 import numpy as np
7 import os
8 
9 def mayavi_utils():
10  '''
11  spin(): Spin the view interactively
12  zoom(): Zoom the view a factor exp(e), in n steps
13  dolly(): Dolly the view a factor exp(e), in n steps
14  spin_zoom(): Spin + zoom the view interactively
15  spin_movie(): Spin one turn, saving images in dir
16  cube(): Construct a cube
17  torus(): Construct a torus of cubes
18 
19  Type e.g. spin?<RETURN> for iPython help, or spin( for
20  pop-up help in Canopy.
21  '''
22  print("Type mayavi_utils?<RETURN> for help text, or \
23 mayavi_utils( for pop-up help in Canopy")
24 
25 @ml.animate(delay=30)
26 def spin():
27  '''
28  Spin the view interactively
29  '''
30  f = ml.gcf()
31  while 1:
32  f.scene.camera.azimuth(1)
33  f.scene.render()
34  yield
35 
36 @ml.animate(delay=10,ui=None)
37 def zoom(e=1.0,n=200):
38  '''
39  Zoom the view a factor exp(e), in n steps
40  '''
41  f=np.exp(e)**(1.0/n)
42  g=ml.gcf()
43  for i in range(n):
44  g.scene.camera.zoom(f)
45  g.scene.render()
46  yield
47 
48 @ml.animate(delay=10)
49 def spin_zoom(factor=2.,n=360,phi=360,name='im',save=0):
50  '''
51  Spin + zoom the view interactively
52  '''
53  f = ml.gcf()
54  eps=factor/n
55  dphi=phi/n
56  for i in range(n):
57  f.scene.camera.azimuth(dphi)
58  f.scene.camera.zoom(1.0+eps)
59  f.scene.render()
60  if save:
61  file=name+'_{:04d}.png'.format(i)
62  ml.savefig(file)
63  yield
64  for i in range(n):
65  f.scene.camera.azimuth(dphi)
66  f.scene.camera.zoom(1.0-eps)
67  f.scene.render()
68  if save:
69  file=name+'_{:04d}.png'.format(i+n)
70  ml.savefig(file)
71  yield
72 
73 @ml.animate(delay=10,ui=None)
74 def dolly(e=1.0,n=200):
75  '''
76  Dolly the view a factor exp(e), in n steps
77  '''
78  f=np.exp(e)**(1.0/n)
79  g=ml.gcf()
80  for i in range(n):
81  g.scene.camera.dolly(f)
82  g.scene.render()
83  yield
84 
85 @ml.animate(delay=10)
86 def spin_movie(dir='anim',n=360):
87  '''
88  Spin the view 360 degrees, in n steps, saving images in dir
89  '''
90  os.system('mkdir -p '+dir)
91  g = ml.gcf()
92  dphi=360.0/n
93  for i in range(n):
94  g.scene.camera.azimuth(dphi)
95  g.scene.render()
96  ml.show()
97  file=dir+'/im_{:03d}.png'.format(i)
98  print(file)
99  yield
100  #ml.savefig(file)
101 
102 def cube(x=0,y=0,z=0,size=1,phi=30):
103  '''
104  Construct a 3D cube
105  '''
106  c=np.cos(phi+np.pi/4)
107  s=np.sin(phi+np.pi/4)
108  x1=0.5*size*c*np.sqrt(2.)
109  y1=0.5*size*s*np.sqrt(2.)
110  z1=0.5*size
111  x2=x+np.array([+x1,-y1,-x1,+y1,+x1])
112  y2=y+np.array([+y1,+x1,-y1,-x1,+y1])
113  z2=z+np.array([+z1,+z1,+z1,+z1,+z1])
114  lw=0.1
115  ml.plot3d(x2,y2,+z2,tube_radius=None,line_width=lw)
116  ml.plot3d(x2,y2,-z2,tube_radius=None,line_width=lw)
117  for i in range(4):
118  ml.plot3d([x2[i],x2[i]],[y2[i],y2[i]],[-z1,z1],tube_radius=None,line_width=lw)
119 
120 def torus(size=1.,levels=1,nphi=31,n=31,azimuth=90.):
121  '''
122  Construct a torus consisting of cubes with given size
123  '''
124  dphi=2.*np.pi/nphi
125  if levels>1:
126  size=2.2*0.9**(1./3.)*3**levels
127  radius=size/dphi
128  for i in range(n):
129  phi=i*dphi+np.pi/2.
130  x=radius*np.cos(phi)
131  y=radius*(np.sin(phi)-1.0)
132  z=0.0
133  cube(x,y,z,size,phi)
134  ml.draw()
135  ml.view(focalpoint=[0.,0.,0.],azimuth=azimuth,elevation=75.,
136  distance=4*radius)