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

2''' 

3from ase.cell import Cell 

4from ase.gui.i18n import _ 

5 

6import ase.gui.ui as ui 

7import numpy as np 

8 

9 

10class CellEditor: 

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

12 def __init__(self, gui): 

13 self.gui = gui 

14 self.gui.register_vulnerable(self) 

15 

16 # Create grid control for cells 

17 # xx xy xz ||x|| pbc 

18 # yx yy yz ||y|| pbc 

19 # zx zy zz ||z|| pbc 

20 self.cell_grid = [] 

21 self.pbc = [] 

22 self.angles = [] 

23 

24 atoms = self.gui.atoms 

25 

26 cell = atoms.cell 

27 mags = cell.lengths() 

28 angles = cell.angles() 

29 pbc = atoms.pbc 

30 

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

32 row = [] 

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

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

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

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

37 rounding=7, width=9)) 

38 self.cell_grid.append(row) 

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

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

41 self.apply_angles, 

42 rounding=7, width=9)) 

43 

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

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

46 

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

48 win = self.win = ui.Window(_('Cell Editor')) 

49 

50 x, y, z = self.cell_grid 

51 

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

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

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

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

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

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

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

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

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

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

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

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

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

65 'automatically apply correctly')]) 

66 # TRANSLATORS: verb 

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

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

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

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

71 

72 def apply_center(self, *args): 

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

74 atoms.center() 

75 self.gui.new_atoms(atoms) 

76 

77 def apply_wrap(self, *args): 

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

79 atoms.wrap() 

80 self.gui.new_atoms(atoms) 

81 

82 def apply_vacuum(self, *args): 

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

84 

85 axis = [] 

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

87 if not pbc: 

88 axis.append(index) 

89 

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

91 self.gui.new_atoms(atoms) 

92 

93 def apply_vectors(self, *args): 

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

95 

96 atoms.set_cell(self.get_vectors(), 

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

98 self.gui.new_atoms(atoms) 

99 

100 def get_vectors(self): 

101 x, y, z = self.cell_grid 

102 cell = np.array( 

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

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

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

106 ) 

107 return Cell(cell) 

108 

109 def get_magnitudes(self): 

110 x, y, z = self.cell_grid 

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

112 

113 def apply_magnitudes(self, *args): 

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

115 

116 old_mags = atoms.cell.lengths() 

117 new_mags = self.get_magnitudes() 

118 

119 newcell = atoms.cell.copy() 

120 for i in range(3): 

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

122 

123 atoms.set_cell(newcell, 

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

125 

126 self.gui.new_atoms(atoms) 

127 

128 def apply_angles(self, *args): 

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

130 

131 cell_data = atoms.cell.cellpar() 

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

133 self.angles[2].value] 

134 

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

136 

137 self.gui.new_atoms(atoms) 

138 

139 def apply_pbc(self, *args): 

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

141 

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

143 atoms.set_pbc(pbc) 

144 

145 self.gui.new_atoms(atoms) 

146 

147 def notify_atoms_changed(self): 

148 atoms = self.gui.atoms 

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

150 

151 def update(self, cell, pbc): 

152 cell = Cell(cell) 

153 mags = cell.lengths() 

154 angles = cell.angles() 

155 

156 for i in range(3): 

157 for j in range(3): 

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

159 cell[i][j] = 0 

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

161 

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

163 mags[i] = 0 

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

165 

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

167 angles[i] = 0 

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

169 

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