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 

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

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

13 return list(Formula(s)) 

14 

15 

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

17 if isinstance(symbols, str): 

18 symbols = string2symbols(symbols) 

19 numbers = [] 

20 for s in symbols: 

21 if isinstance(s, str): 

22 numbers.append(atomic_numbers[s]) 

23 else: 

24 numbers.append(int(s)) 

25 return numbers 

26 

27 

28class Symbols(collections.abc.Sequence): 

29 """A sequence of chemical symbols. 

30 

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

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

33 its elements are manipulated as strings. 

34 

35 Examples: 

36 

37 >>> from ase.build import molecule 

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

39 >>> atoms.symbols 

40 Symbols('C2OH6') 

41 >>> atoms.symbols[:3] 

42 Symbols('C2O') 

43 >>> atoms.symbols == 'H' 

44 array([False, False, False, True, True, True, True, True, True], dtype=bool) 

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

46 >>> atoms.symbols 

47 Symbols('C2OH3Pu3') 

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

49 >>> atoms.symbols 

50 Symbols('C2OMo2UPu3') 

51 >>> atoms.symbols.formula 

52 Formula('C2OMo2UPu3') 

53 

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

55 formatting options and analysis. 

56 

57 """ 

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

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

60 

61 @classmethod 

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

63 numbers = symbols2numbers(symbols) 

64 return cls(np.array(numbers)) 

65 

66 @property 

67 def formula(self) -> Formula: 

68 """Formula object.""" 

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

70 return Formula(string) 

71 

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

73 num = self.numbers[key] 

74 if np.isscalar(num): 

75 return chemical_symbols[num] 

76 return Symbols(num) 

77 

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

79 for num in self.numbers: 

80 yield chemical_symbols[num] 

81 

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

83 numbers = symbols2numbers(value) 

84 if len(numbers) == 1: 

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

86 else: 

87 self.numbers[key] = numbers 

88 

89 def __len__(self) -> int: 

90 return len(self.numbers) 

91 

92 def __str__(self) -> str: 

93 return self.get_chemical_formula('reduce') 

94 

95 def __repr__(self) -> str: 

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

97 

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

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

100 return False 

101 

102 try: 

103 symbols = Symbols.fromsymbols(obj) 

104 except Exception: 

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

106 # atomic numbers. 

107 return False 

108 return self.numbers == symbols.numbers 

109 

110 def get_chemical_formula( 

111 self, 

112 mode: str = 'hill', 

113 empirical: bool = False, 

114 ) -> str: 

115 """Get chemical formula. 

116 

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

118 # XXX Delegate the work to the Formula object! 

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

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

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

122 

123 if len(self) == 0: 

124 return '' 

125 

126 numbers = self.numbers 

127 

128 if mode == 'reduce': 

129 n = len(numbers) 

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

131 numbers[:-1]])) 

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

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

134 

135 tokens = [] 

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

137 tokens.append(s) 

138 if c > 1: 

139 tokens.append(str(c)) 

140 formula = ''.join(tokens) 

141 elif mode == 'all': 

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

143 else: 

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

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

146 if empirical: 

147 f, _ = f.reduce() 

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

149 formula = f.format(mode) 

150 else: 

151 raise ValueError( 

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

153 

154 return formula 

155 

156 def search(self, symbols) -> Sequence[int]: 

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

158 numbers = set(symbols2numbers(symbols)) 

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

160 if number in numbers] 

161 return np.array(indices, int) 

162 

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

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

165 return set(self) 

166 

167 def indices(self) -> Dict[str, Sequence[int]]: 

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

169 

170 >>> from ase.build import molecule 

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

172 >>> atoms.symbols.indices() 

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

174 

175 """ 

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

177 for i, symbol in enumerate(self): 

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

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