Coverage for /builds/ase/ase/ase/gui/add.py: 59.38%

1# fmt: off 

2 

3import os 

4 

5import numpy as np 

6 

7import ase.gui.ui as ui 

8from ase import Atoms 

9from ase.data import atomic_numbers, chemical_symbols 

10from ase.gui.i18n import _ 

11 

12current_selection_string = _('(selection)') 

13 

14 

15class AddAtoms: 

16 def __init__(self, gui): 

17 self.gui = gui 

18 win = self.win = ui.Window(_('Add atoms'), wmtype='utility') 

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

20 

21 def choose_file(): 

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

23 filename = chooser.go() 

24 if filename is None: # No file selected 

25 return 

26 

27 self.combobox.value = filename 

28 

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

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

31 

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

33 default = current_selection_string 

34 else: 

35 default = 'H2' 

36 

37 self._filename = None 

38 self._atoms_from_file = None 

39 

40 from ase.collections import g2 

41 labels = sorted(name for name in g2.names 

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

43 values = labels 

44 

45 combobox = ui.ComboBox(labels, values) 

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

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

48 ui.bind_enter(combobox.widget, lambda e: self.add()) 

49 

50 combobox.value = default 

51 self.combobox = combobox 

52 

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

54 for __ in range(3)] 

55 

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

57 self.spinners = spinners 

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

59 'if any, else absolute.')) 

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

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

62 self.picky]) 

63 self.focus() 

64 

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

66 if filename == self._filename: 

67 # We have this file already 

68 return self._atoms_from_file 

69 

70 from ase.io import read 

71 try: 

72 atoms = read(filename) 

73 except Exception as err: 

74 ui.show_io_error(filename, err) 

75 atoms = None 

76 filename = None 

77 

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

79 self._atoms_from_file = atoms 

80 self._filename = filename 

81 return atoms 

82 

83 def get_atoms(self): 

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

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

86 

87 if val == current_selection_string: 

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

89 if selection.any(): 

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

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

92 

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

94 return Atoms(val) 

95 

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

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

98 

99 from ase.collections import g2 

100 if val in g2.names: 

101 return g2[val] 

102 

103 if os.path.exists(val): 

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

105 

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

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

108 .format(val)) 

109 

110 return None 

111 

112 def getcoords(self): 

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

114 

115 pos = self.gui.atoms.positions 

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

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

118 center = pos.mean(0) 

119 addcoords += center 

120 

121 return addcoords 

122 

123 def focus(self): 

124 self.combobox.widget.focus_set() 

125 

126 def add(self): 

127 newatoms = self.get_atoms() 

128 if newatoms is None: # Error dialog was shown 

129 return 

130 

131 newcenter = self.getcoords() 

132 

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

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

135 previous_center = newatoms.positions.mean(0) 

136 newatoms.positions += newcenter - previous_center 

137 

138 atoms = self.gui.atoms 

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

140 from ase.geometry import get_distances 

141 _disps, dists = get_distances(atoms.positions, 

142 newatoms.positions) 

143 mindist = dists.min() 

144 if mindist < 0.5: 

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

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

147 'an existing atom. To override, ' 

148 'uncheck the check positions option.')) 

149 return 

150 

151 self.gui.add_atoms_and_select(newatoms)