Coverage for /builds/kinetik161/ase/ase/visualize/mlab.py: 15.28%

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 from mayavi import mlab 

26 

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

28 

29 # Plot the atoms as spheres: 

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

31 mlab.points3d(*pos, 

32 scale_factor=covalent_radii[Z], 

33 resolution=20, 

34 color=tuple(cpk_colors[Z])) 

35 

36 # Draw the unit cell: 

37 A = atoms.cell 

38 for i1, a in enumerate(A): 

39 i2 = (i1 + 1) % 3 

40 i3 = (i1 + 2) % 3 

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

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

43 p1 = b + c 

44 p2 = p1 + a 

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

46 [p1[1], p2[1]], 

47 [p1[2], p2[2]], 

48 tube_radius=0.1) 

49 

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

51 opacity=0.5, colormap='hot') 

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

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

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

55 # Transform the points to the unit cell: 

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

57 

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

59 # another way? 

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

61 # Show the 3d plot: 

62 mlab.show() 

63 

64 

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

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

67 

68 

69description = """\ 

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

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

72 

73 

74def main(args=None): 

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

76 description=description) 

77 add = parser.add_option 

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

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

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

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

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

83 help='Plot the electrostatic potential.') 

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

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

86 'values. Default is four contours.') 

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

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

89 

90 opts, args = parser.parse_args(args) 

91 if len(args) != 1: 

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

93 

94 arg = args[0] 

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

96 data, atoms = read_cube_data(arg) 

97 else: 

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

99 atoms = calc.get_atoms() 

100 if opts.band_index is None: 

101 if opts.electrostatic_potential: 

102 data = calc.get_electrostatic_potential() 

103 else: 

104 data = calc.get_pseudo_density(opts.spin_index) 

105 else: 

106 data = calc.get_pseudo_wave_function(opts.band_index, 

107 opts.spin_index or 0) 

108 if data.dtype == complex: 

109 data = abs(data) 

110 

111 mn = data.min() 

112 mx = data.max() 

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

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

115 

116 if opts.contours.isdigit(): 

117 n = int(opts.contours) 

118 d = (mx - mn) / n 

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

120 else: 

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

122 

123 if len(contours) == 1: 

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

125 else: 

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

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

128 

129 if opts.repeat: 

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

131 data = np.tile(data, repeat) 

132 atoms *= repeat 

133 

134 plot(atoms, data, contours) 

135 

136 

137if __name__ == '__main__': 

138 main()