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

1# fmt: off 

2 

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

4''' 

5import numpy as np 

6 

7import ase.gui.ui as ui 

8from ase.cell import Cell 

9from ase.gui.i18n import _ 

10 

11 

12class CellEditor: 

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

14 

15 def __init__(self, gui): 

16 self.gui = gui 

17 self.gui.obs.set_atoms.register(self.notify_atoms_changed) 

18 

19 # Create grid control for cells 

20 # xx xy xz ||x|| pbc 

21 # yx yy yz ||y|| pbc 

22 # zx zy zz ||z|| pbc 

23 self.cell_grid = [] 

24 self.pbc = [] 

25 self.angles = [] 

26 

27 atoms = self.gui.atoms 

28 

29 cell = atoms.cell 

30 mags = cell.lengths() 

31 angles = cell.angles() 

32 pbc = atoms.pbc 

33 

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

35 row = [] 

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

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

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

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

40 rounding=7, width=9)) 

41 self.cell_grid.append(row) 

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

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

44 self.apply_angles, 

45 rounding=7, width=9)) 

46 

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

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

49 

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

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

52 

53 x, y, z = self.cell_grid 

54 

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

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

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

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

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

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

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

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

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

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

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

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

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

68 'automatically apply correctly')]) 

69 # TRANSLATORS: verb 

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

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

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

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

74 

75 def apply_center(self, *args): 

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

77 atoms.center() 

78 self.gui.new_atoms(atoms) 

79 

80 def apply_wrap(self, *args): 

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

82 atoms.wrap() 

83 self.gui.new_atoms(atoms) 

84 

85 def apply_vacuum(self, *args): 

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

87 

88 axis = [] 

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

90 if not pbc: 

91 axis.append(index) 

92 

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

94 self.gui.new_atoms(atoms) 

95 

96 def apply_vectors(self, *args): 

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

98 

99 atoms.set_cell(self.get_vectors(), 

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

101 self.gui.new_atoms(atoms) 

102 

103 def get_vectors(self): 

104 x, y, z = self.cell_grid 

105 cell = np.array( 

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

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

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

109 ) 

110 return Cell(cell) 

111 

112 def get_magnitudes(self): 

113 x, y, z = self.cell_grid 

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

115 

116 def apply_magnitudes(self, *args): 

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

118 

119 old_mags = atoms.cell.lengths() 

120 new_mags = self.get_magnitudes() 

121 

122 newcell = atoms.cell.copy() 

123 for i in range(3): 

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

125 

126 atoms.set_cell(newcell, 

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

128 

129 self.gui.new_atoms(atoms) 

130 

131 def apply_angles(self, *args): 

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

133 

134 cell_data = atoms.cell.cellpar() 

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

136 self.angles[2].value] 

137 

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

139 

140 self.gui.new_atoms(atoms) 

141 

142 def apply_pbc(self, *args): 

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

144 

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

146 atoms.set_pbc(pbc) 

147 

148 self.gui.new_atoms(atoms) 

149 

150 def notify_atoms_changed(self): 

151 atoms = self.gui.atoms 

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

153 

154 def update(self, cell, pbc): 

155 cell = Cell(cell) 

156 mags = cell.lengths() 

157 angles = cell.angles() 

158 

159 for i in range(3): 

160 for j in range(3): 

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

162 cell[i][j] = 0 

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

164 

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

166 mags[i] = 0 

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

168 

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

170 angles[i] = 0 

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

172 

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