Coverage for /builds/debichem-team/python-ase/ase/visualize/mlab.py: 14.47%

1import optparse 

2 

3import numpy as np 

4 

5from ase.calculators.calculator import get_calculator_class 

6from ase.data import covalent_radii 

7from ase.data.colors import cpk_colors 

8from ase.io.cube import read_cube_data 

9 

10 

11def plot(atoms, data, contours): 

12 """Plot atoms, unit-cell and iso-surfaces using Mayavi. 

13 

14 Parameters: 

15 

16 atoms: Atoms object 

17 Positions, atomiz numbers and unit-cell. 

18 data: 3-d ndarray of float 

19 Data for iso-surfaces. 

20 countours: list of float 

21 Contour values. 

22 """ 

23 

24 # Delay slow imports: 

25 import os 

26 

27 from mayavi import mlab 

28 

29 # mayavi GUI bug fix for remote access via ssh (X11 forwarding) 

30 if "SSH_CONNECTION" in os.environ: 

31 f = mlab.gcf() 

32 f.scene._lift() 

33 

34 mlab.figure(1, bgcolor=(1, 1, 1)) # make a white figure 

35 

36 # Plot the atoms as spheres: 

37 for pos, Z in zip(atoms.positions, atoms.numbers): 

38 mlab.points3d(*pos, 

39 scale_factor=covalent_radii[Z], 

40 resolution=20, 

41 color=tuple(cpk_colors[Z])) 

42 

43 # Draw the unit cell: 

44 A = atoms.cell 

45 for i1, a in enumerate(A): 

46 i2 = (i1 + 1) % 3 

47 i3 = (i1 + 2) % 3 

48 for b in [np.zeros(3), A[i2]]: 

49 for c in [np.zeros(3), A[i3]]: 

50 p1 = b + c 

51 p2 = p1 + a 

52 mlab.plot3d([p1[0], p2[0]], 

53 [p1[1], p2[1]], 

54 [p1[2], p2[2]], 

55 tube_radius=0.1) 

56 

57 cp = mlab.contour3d(data, contours=contours, transparent=True, 

58 opacity=0.5, colormap='hot') 

59 # Do some tvtk magic in order to allow for non-orthogonal unit cells: 

60 polydata = cp.actor.actors[0].mapper.input 

61 pts = np.array(polydata.points) - 1 

62 # Transform the points to the unit cell: 

63 polydata.points = np.dot(pts, A / np.array(data.shape)[:, np.newaxis]) 

64 

65 # Apparently we need this to redraw the figure, maybe it can be done in 

66 # another way? 

67 mlab.view(azimuth=155, elevation=70, distance='auto') 

68 # Show the 3d plot: 

69 mlab.show() 

70 

71 

72def view_mlab(atoms, *args, **kwargs): 

73 return plot(atoms, *args, **kwargs) 

74 

75 

76description = """\ 

77Plot iso-surfaces from a cube-file or a wave function or an electron 

78density from a calculator-restart file.""" 

79 

80 

81def main(args=None): 

82 parser = optparse.OptionParser(usage='%prog [options] filename', 

83 description=description) 

84 add = parser.add_option 

85 add('-n', '--band-index', type=int, metavar='INDEX', 

86 help='Band index counting from zero.') 

87 add('-s', '--spin-index', type=int, metavar='SPIN', 

88 help='Spin index: zero or one.') 

89 add('-e', '--electrostatic-potential', action='store_true', 

90 help='Plot the electrostatic potential.') 

91 add('-c', '--contours', default='4', 

92 help='Use "-c 3" for 3 contours or "-c -0.5,0.5" for specific ' + 

93 'values. Default is four contours.') 

94 add('-r', '--repeat', help='Example: "-r 2,2,2".') 

95 add('-C', '--calculator-name', metavar='NAME', help='Name of calculator.') 

96 

97 opts, args = parser.parse_args(args) 

98 if len(args) != 1: 

99 parser.error('Incorrect number of arguments') 

100 

101 arg = args[0] 

102 if arg.endswith('.cube'): 

103 data, atoms = read_cube_data(arg) 

104 else: 

105 calc = get_calculator_class(opts.calculator_name)(arg, txt=None) 

106 atoms = calc.get_atoms() 

107 if opts.band_index is None: 

108 if opts.electrostatic_potential: 

109 data = calc.get_electrostatic_potential() 

110 else: 

111 data = calc.get_pseudo_density(opts.spin_index) 

112 else: 

113 data = calc.get_pseudo_wave_function(opts.band_index, 

114 opts.spin_index or 0) 

115 if data.dtype == complex: 

116 data = abs(data) 

117 

118 mn = data.min() 

119 mx = data.max() 

120 print('Min: %16.6f' % mn) 

121 print('Max: %16.6f' % mx) 

122 

123 if opts.contours.isdigit(): 

124 n = int(opts.contours) 

125 d = (mx - mn) / n 

126 contours = np.linspace(mn + d / 2, mx - d / 2, n).tolist() 

127 else: 

128 contours = [float(x) for x in opts.contours.rstrip(',').split(',')] 

129 

130 if len(contours) == 1: 

131 print('1 contour:', contours[0]) 

132 else: 

133 print('%d contours: %.6f, ..., %.6f' % 

134 (len(contours), contours[0], contours[-1])) 

135 

136 if opts.repeat: 

137 repeat = [int(r) for r in opts.repeat.split(',')] 

138 data = np.tile(data, repeat) 

139 atoms *= repeat 

140 

141 plot(atoms, data, contours) 

142 

143 

144if __name__ == '__main__': 

145 main()