Coverage for /builds/kinetik161/ase/ase/gui/celleditor.py: 84.31%

1'''celleditor.py - Window for editing the cell of an atoms object 

2''' 

3import numpy as np 

4 

5import ase.gui.ui as ui 

6from ase.cell import Cell 

7from ase.gui.i18n import _ 

8 

9 

10class CellEditor: 

11 '''Window for editing the cell of an atoms object.''' 

12 

13 def __init__(self, gui): 

14 self.gui = gui 

15 self.gui.register_vulnerable(self) 

16 

17 # Create grid control for cells 

18 # xx xy xz ||x|| pbc 

19 # yx yy yz ||y|| pbc 

20 # zx zy zz ||z|| pbc 

21 self.cell_grid = [] 

22 self.pbc = [] 

23 self.angles = [] 

24 

25 atoms = self.gui.atoms 

26 

27 cell = atoms.cell 

28 mags = cell.lengths() 

29 angles = cell.angles() 

30 pbc = atoms.pbc 

31 

32 for i in [0, 1, 2]: # x_ y_ z_ 

33 row = [] 

34 for j in [0, 1, 2]: # _x _y _z 

35 row.append(ui.SpinBox(cell[i][j], -30, 30, 0.1, 

36 self.apply_vectors, rounding=7, width=9)) 

37 row.append(ui.SpinBox(mags[i], -30, 30, 0.1, self.apply_magnitudes, 

38 rounding=7, width=9)) 

39 self.cell_grid.append(row) 

40 self.pbc.append(ui.CheckButton('', bool(pbc[i]), self.apply_pbc)) 

41 self.angles.append(ui.SpinBox(angles[i], -360, 360, 15, 

42 self.apply_angles, 

43 rounding=7, width=9)) 

44 

45 self.scale_atoms = ui.CheckButton('', False) 

46 self.vacuum = ui.SpinBox(5, 0, 15, 0.1, self.apply_vacuum) 

47 

48 # TRANSLATORS: This is a title of a window. 

49 win = self.win = ui.Window(_('Cell Editor'), wmtype='utility') 

50 

51 x, y, z = self.cell_grid 

52 

53 win.add([_('A:'), x[0], x[1], x[2], _('||A||:'), x[3], 

54 _('periodic:'), self.pbc[0]]) 

55 win.add([_('B:'), y[0], y[1], y[2], _('||B||:'), y[3], 

56 _('periodic:'), self.pbc[1]]) 

57 win.add([_('C:'), z[0], z[1], z[2], _('||C||:'), z[3], 

58 _('periodic:'), self.pbc[2]]) 

59 win.add([_('∠BC:'), self.angles[0], _('∠AC:'), self.angles[1], 

60 _('∠AB:'), self.angles[2]]) 

61 win.add([_('Scale atoms with cell:'), self.scale_atoms]) 

62 win.add([ui.Button(_('Apply Vectors'), self.apply_vectors), 

63 ui.Button(_('Apply Magnitudes'), self.apply_magnitudes), 

64 ui.Button(_('Apply Angles'), self.apply_angles)]) 

65 win.add([_('Pressing 〈Enter〉 as you enter values will ' 

66 'automatically apply correctly')]) 

67 # TRANSLATORS: verb 

68 win.add([ui.Button(_('Center'), self.apply_center), 

69 ui.Button(_('Wrap'), self.apply_wrap), 

70 _('Vacuum:'), self.vacuum, 

71 ui.Button(_('Apply Vacuum'), self.apply_vacuum)]) 

72 

73 def apply_center(self, *args): 

74 atoms = self.gui.atoms.copy() 

75 atoms.center() 

76 self.gui.new_atoms(atoms) 

77 

78 def apply_wrap(self, *args): 

79 atoms = self.gui.atoms.copy() 

80 atoms.wrap() 

81 self.gui.new_atoms(atoms) 

82 

83 def apply_vacuum(self, *args): 

84 atoms = self.gui.atoms.copy() 

85 

86 axis = [] 

87 for index, pbc in enumerate(atoms.pbc): 

88 if not pbc: 

89 axis.append(index) 

90 

91 atoms.center(vacuum=self.vacuum.value, axis=axis) 

92 self.gui.new_atoms(atoms) 

93 

94 def apply_vectors(self, *args): 

95 atoms = self.gui.atoms.copy() 

96 

97 atoms.set_cell(self.get_vectors(), 

98 scale_atoms=self.scale_atoms.var.get()) 

99 self.gui.new_atoms(atoms) 

100 

101 def get_vectors(self): 

102 x, y, z = self.cell_grid 

103 cell = np.array( 

104 [[x[0].value, x[1].value, x[2].value], 

105 [y[0].value, y[1].value, y[2].value], 

106 [z[0].value, z[1].value, z[2].value]] 

107 ) 

108 return Cell(cell) 

109 

110 def get_magnitudes(self): 

111 x, y, z = self.cell_grid 

112 return np.array([x[3].value, y[3].value, z[3].value]) 

113 

114 def apply_magnitudes(self, *args): 

115 atoms = self.gui.atoms.copy() 

116 

117 old_mags = atoms.cell.lengths() 

118 new_mags = self.get_magnitudes() 

119 

120 newcell = atoms.cell.copy() 

121 for i in range(3): 

122 newcell[i] *= new_mags[i] / old_mags[i] 

123 

124 atoms.set_cell(newcell, 

125 scale_atoms=self.scale_atoms.var.get()) 

126 

127 self.gui.new_atoms(atoms) 

128 

129 def apply_angles(self, *args): 

130 atoms = self.gui.atoms.copy() 

131 

132 cell_data = atoms.cell.cellpar() 

133 cell_data[3:7] = [self.angles[0].value, self.angles[1].value, 

134 self.angles[2].value] 

135 

136 atoms.set_cell(cell_data, scale_atoms=self.scale_atoms.var.get()) 

137 

138 self.gui.new_atoms(atoms) 

139 

140 def apply_pbc(self, *args): 

141 atoms = self.gui.atoms.copy() 

142 

143 pbc = [pbc.var.get() for pbc in self.pbc] 

144 atoms.set_pbc(pbc) 

145 

146 self.gui.new_atoms(atoms) 

147 

148 def notify_atoms_changed(self): 

149 atoms = self.gui.atoms 

150 self.update(atoms.cell, atoms.pbc) 

151 

152 def update(self, cell, pbc): 

153 cell = Cell(cell) 

154 mags = cell.lengths() 

155 angles = cell.angles() 

156 

157 for i in range(3): 

158 for j in range(3): 

159 if np.isnan(cell[i][j]): 

160 cell[i][j] = 0 

161 self.cell_grid[i][j].value = cell[i][j] 

162 

163 if np.isnan(mags[i]): 

164 mags[i] = 0 

165 self.cell_grid[i][3].value = mags[i] 

166 

167 if np.isnan(angles[i]): 

168 angles[i] = 0 

169 self.angles[i].value = angles[i] 

170 

171 self.pbc[i].var.set(bool(pbc[i]))