1import os 

2import numpy as np 

3 

4from ase.gui.i18n import _ 

5from ase import Atoms 

6import ase.gui.ui as ui 

7from ase.data import atomic_numbers, chemical_symbols 

8 

9 

10current_selection_string = _('(selection)') 

11 

12 

13class AddAtoms: 

14 def __init__(self, gui): 

15 self.gui = gui 

16 win = self.win = ui.Window(_('Add atoms')) 

17 win.add(_('Specify chemical symbol, formula, or filename.')) 

18 

19 def choose_file(): 

20 chooser = ui.ASEFileChooser(self.win.win) 

21 filename = chooser.go() 

22 if filename is None: # No file selected 

23 return 

24 

25 self.combobox.value = filename 

26 

27 # Load the file immediately, so we can warn now in case of error 

28 self.readfile(filename, format=chooser.format) 

29 

30 if self.gui.images.selected.any(): 

31 default = current_selection_string 

32 else: 

33 default = 'H2' 

34 

35 self._filename = None 

36 self._atoms_from_file = None 

37 

38 from ase.collections import g2 

39 labels = list(sorted(name for name in g2.names 

40 if len(g2[name]) > 1)) 

41 values = labels 

42 

43 combobox = ui.ComboBox(labels, values) 

44 win.add([_('Add:'), combobox, 

45 ui.Button(_('File ...'), callback=choose_file)]) 

46 combobox.widget.bind('<Return>', lambda e: self.add()) 

47 

48 combobox.value = default 

49 self.combobox = combobox 

50 

51 spinners = [ui.SpinBox(0.0, -1e3, 1e3, 0.1, rounding=2, width=3) 

52 for __ in range(3)] 

53 

54 win.add([_('Coordinates:')] + spinners) 

55 self.spinners = spinners 

56 win.add(_('Coordinates are relative to the center of the selection, ' 

57 'if any, else absolute.')) 

58 self.picky = ui.CheckButton(_('Check positions'), True) 

59 win.add([ui.Button(_('Add'), self.add), 

60 self.picky]) 

61 self.focus() 

62 

63 def readfile(self, filename, format=None): 

64 if filename == self._filename: 

65 # We have this file already 

66 return self._atoms_from_file 

67 

68 from ase.io import read 

69 try: 

70 atoms = read(filename) 

71 except Exception as err: 

72 ui.show_io_error(filename, err) 

73 atoms = None 

74 filename = None 

75 

76 # Cache selected Atoms/filename (or None) for future calls 

77 self._atoms_from_file = atoms 

78 self._filename = filename 

79 return atoms 

80 

81 def get_atoms(self): 

82 # Get the text, whether it's a combobox item or not 

83 val = self.combobox.widget.get() 

84 

85 if val == current_selection_string: 

86 selection = self.gui.images.selected.copy() 

87 if selection.any(): 

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

89 return atoms[selection[:len(self.gui.atoms)]] 

90 

91 if val in atomic_numbers: # Note: This means val is a symbol! 

92 return Atoms(val) 

93 

94 if val.isdigit() and int(val) < len(chemical_symbols): 

95 return Atoms(numbers=[int(val)]) 

96 

97 from ase.collections import g2 

98 if val in g2.names: 

99 return g2[val] 

100 

101 if os.path.exists(val): 

102 return self.readfile(val) # May show UI error 

103 

104 ui.showerror(_('Cannot add atoms'), 

105 _('{} is neither atom, molecule, nor file') 

106 .format(val)) 

107 

108 return None 

109 

110 def getcoords(self): 

111 addcoords = np.array([spinner.value for spinner in self.spinners]) 

112 

113 pos = self.gui.atoms.positions 

114 if self.gui.images.selected[:len(pos)].any(): 

115 pos = pos[self.gui.images.selected[:len(pos)]] 

116 center = pos.mean(0) 

117 addcoords += center 

118 

119 return addcoords 

120 

121 def focus(self): 

122 self.combobox.widget.focus_set() 

123 

124 def add(self): 

125 newatoms = self.get_atoms() 

126 if newatoms is None: # Error dialog was shown 

127 return 

128 

129 newcenter = self.getcoords() 

130 

131 # Not newatoms.center() because we want the same centering method 

132 # used for adding atoms relative to selections (mean). 

133 previous_center = newatoms.positions.mean(0) 

134 newatoms.positions += newcenter - previous_center 

135 

136 atoms = self.gui.atoms 

137 if len(atoms) and self.picky.value: 

138 from ase.geometry import get_distances 

139 disps, dists = get_distances(atoms.positions, 

140 newatoms.positions) 

141 mindist = dists.min() 

142 if mindist < 0.5: 

143 ui.showerror(_('Bad positions'), 

144 _('Atom would be less than 0.5 Å from ' 

145 'an existing atom. To override, ' 

146 'uncheck the check positions option.')) 

147 return 

148 

149 self.gui.add_atoms_and_select(newatoms)