1from typing import List, Sequence, Set, Dict, Union, Iterator 

2import warnings 

3import collections.abc 

4 

5import numpy as np 

6 

7from ase.data import atomic_numbers, chemical_symbols 

8from ase.formula import Formula 

9 

10 

11Integers = Union[Sequence[int], np.ndarray] 

12 

13 

14def string2symbols(s: str) -> List[str]: 

15 """Convert string to list of chemical symbols.""" 

16 return list(Formula(s)) 

17 

18 

19def symbols2numbers(symbols) -> List[int]: 

20 if isinstance(symbols, str): 

21 symbols = string2symbols(symbols) 

22 numbers = [] 

23 for s in symbols: 

24 if isinstance(s, str): 

25 numbers.append(atomic_numbers[s]) 

26 else: 

27 numbers.append(int(s)) 

28 return numbers 

29 

30 

31class Symbols(collections.abc.Sequence): 

32 """A sequence of chemical symbols. 

33 

34 ``atoms.symbols`` is a :class:`ase.symbols.Symbols` object. This 

35 object works like an editable view of ``atoms.numbers``, except 

36 its elements are manipulated as strings. 

37 

38 Examples: 

39 

40 >>> from ase.build import molecule 

41 >>> atoms = molecule('CH3CH2OH') 

42 >>> atoms.symbols 

43 Symbols('C2OH6') 

44 >>> atoms.symbols[:3] 

45 Symbols('C2O') 

46 >>> atoms.symbols == 'H' 

47 array([False, False, False, True, True, True, True, True, True], \ 

48dtype=bool) 

49 >>> atoms.symbols[-3:] = 'Pu' 

50 >>> atoms.symbols 

51 Symbols('C2OH3Pu3') 

52 >>> atoms.symbols[3:6] = 'Mo2U' 

53 >>> atoms.symbols 

54 Symbols('C2OMo2UPu3') 

55 >>> atoms.symbols.formula 

56 Formula('C2OMo2UPu3') 

57 

58 The :class:`ase.formula.Formula` object is useful for extended 

59 formatting options and analysis. 

60 

61 """ 

62 

63 def __init__(self, numbers) -> None: 

64 self.numbers = np.asarray(numbers, int) 

65 

66 @classmethod 

67 def fromsymbols(cls, symbols) -> 'Symbols': 

68 numbers = symbols2numbers(symbols) 

69 return cls(np.array(numbers)) 

70 

71 @property 

72 def formula(self) -> Formula: 

73 """Formula object.""" 

74 string = Formula.from_list(self).format('reduce') 

75 return Formula(string) 

76 

77 def __getitem__(self, key) -> Union['Symbols', str]: 

78 num = self.numbers[key] 

79 if np.isscalar(num): 

80 return chemical_symbols[num] 

81 return Symbols(num) 

82 

83 def __iter__(self) -> Iterator[str]: 

84 for num in self.numbers: 

85 yield chemical_symbols[num] 

86 

87 def __setitem__(self, key, value) -> None: 

88 numbers = symbols2numbers(value) 

89 if len(numbers) == 1: 

90 self.numbers[key] = numbers[0] 

91 else: 

92 self.numbers[key] = numbers 

93 

94 def __len__(self) -> int: 

95 return len(self.numbers) 

96 

97 def __str__(self) -> str: 

98 return self.get_chemical_formula('reduce') 

99 

100 def __repr__(self) -> str: 

101 return 'Symbols(\'{}\')'.format(self) 

102 

103 def __eq__(self, obj) -> bool: 

104 if not hasattr(obj, '__len__'): 

105 return False 

106 

107 try: 

108 symbols = Symbols.fromsymbols(obj) 

109 except Exception: 

110 # Typically this would happen if obj cannot be converged to 

111 # atomic numbers. 

112 return False 

113 return self.numbers == symbols.numbers 

114 

115 def get_chemical_formula( 

116 self, 

117 mode: str = 'hill', 

118 empirical: bool = False, 

119 ) -> str: 

120 """Get chemical formula. 

121 

122 See documentation of ase.atoms.Atoms.get_chemical_formula().""" 

123 # XXX Delegate the work to the Formula object! 

124 if mode in ('reduce', 'all') and empirical: 

125 warnings.warn("Empirical chemical formula not available " 

126 "for mode '{}'".format(mode)) 

127 

128 if len(self) == 0: 

129 return '' 

130 

131 numbers = self.numbers 

132 

133 if mode == 'reduce': 

134 n = len(numbers) 

135 changes = np.concatenate(([0], np.arange(1, n)[numbers[1:] != 

136 numbers[:-1]])) 

137 symbols = [chemical_symbols[e] for e in numbers[changes]] 

138 counts = np.append(changes[1:], n) - changes 

139 

140 tokens = [] 

141 for s, c in zip(symbols, counts): 

142 tokens.append(s) 

143 if c > 1: 

144 tokens.append(str(c)) 

145 formula = ''.join(tokens) 

146 elif mode == 'all': 

147 formula = ''.join([chemical_symbols[n] for n in numbers]) 

148 else: 

149 symbols = [chemical_symbols[Z] for Z in numbers] 

150 f = Formula('', _tree=[(symbols, 1)]) 

151 if empirical: 

152 f, _ = f.reduce() 

153 if mode in {'hill', 'metal'}: 

154 formula = f.format(mode) 

155 else: 

156 raise ValueError( 

157 "Use mode = 'all', 'reduce', 'hill' or 'metal'.") 

158 

159 return formula 

160 

161 def search(self, symbols) -> Integers: 

162 """Return the indices of elements with given symbol or symbols.""" 

163 numbers = set(symbols2numbers(symbols)) 

164 indices = [i for i, number in enumerate(self.numbers) 

165 if number in numbers] 

166 return np.array(indices, int) 

167 

168 def species(self) -> Set[str]: 

169 """Return unique symbols as a set.""" 

170 return set(self) 

171 

172 def indices(self) -> Dict[str, Integers]: 

173 """Return dictionary mapping each unique symbol to indices. 

174 

175 >>> from ase.build import molecule 

176 >>> atoms = molecule('CH3CH2OH') 

177 >>> atoms.symbols.indices() 

178 {'C': array([0, 1]), 'O': array([2]), 'H': array([3, 4, 5, 6, 7, 8])} 

179 

180 """ 

181 dct: Dict[str, List[int]] = {} 

182 for i, symbol in enumerate(self): 

183 dct.setdefault(symbol, []).append(i) 

184 return {key: np.array(value, int) for key, value in dct.items()} 

185 

186 def species_indices(self) -> Sequence[int]: 

187 """Return the indices of each atom within their individual species. 

188 

189 >>> from ase import Atoms 

190 >>> atoms = Atoms('CH3CH2OH') 

191 >>> atoms.symbols.species_indices() 

192 [0, 0, 1, 2, 1, 3, 4, 0, 5] 

193 

194 ^ ^ ^ ^ ^ ^ ^ ^ ^ 

195 C H H H C H H O H 

196 

197 """ 

198 

199 counts: Dict[str, int] = {} 

200 result = [] 

201 for i, n in enumerate(self.numbers): 

202 counts[n] = counts.get(n, -1) + 1 

203 result.append(counts[n]) 

204 

205 return result