Coverage for ase / symbols.py: 96.08%

1# fmt: off 

2 

3import collections.abc 

4import numbers 

5import warnings 

6from collections.abc import Iterator, Sequence 

7 

8import numpy as np 

9 

10from ase.data import atomic_numbers, chemical_symbols 

11from ase.formula import Formula 

12 

13Integers = Sequence[int] | np.ndarray 

14 

15 

16def string2symbols(s: str) -> list[str]: 

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

18 return list(Formula(s)) 

19 

20 

21def symbols2numbers(symbols) -> list[int]: 

22 if isinstance(symbols, str): 

23 symbols = string2symbols(symbols) 

24 numbers = [] 

25 for s in symbols: 

26 if isinstance(s, str): 

27 numbers.append(atomic_numbers[s]) 

28 else: 

29 numbers.append(int(s)) 

30 return numbers 

31 

32 

33class Symbols(collections.abc.Sequence): 

34 """A sequence of chemical symbols. 

35 

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

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

38 its elements are manipulated as strings. 

39 

40 Examples 

41 -------- 

42 

43 >>> from ase.build import molecule 

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

45 >>> atoms.symbols 

46 Symbols('C2OH6') 

47 >>> atoms.symbols[:3] 

48 Symbols('C2O') 

49 >>> atoms.symbols == 'H' # doctest: +ELLIPSIS 

50 array([False, False, False, True, True, True, True, True, True]...) 

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

52 >>> atoms.symbols 

53 Symbols('C2OH3Pu3') 

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

55 >>> atoms.symbols 

56 Symbols('C2OMo2UPu3') 

57 >>> atoms.symbols.formula 

58 Formula('C2OMo2UPu3') 

59 

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

61 formatting options and analysis. 

62 

63 """ 

64 

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

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

67 

68 @classmethod 

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

70 numbers = symbols2numbers(symbols) 

71 return cls(np.array(numbers)) 

72 

73 @property 

74 def formula(self) -> Formula: 

75 """Formula object.""" 

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

77 return Formula(string) 

78 

79 def __getitem__(self, key) -> 'Symbols | str': 

80 num = self.numbers[key] 

81 if isinstance(num, numbers.Integral): 

82 return chemical_symbols[num] 

83 return Symbols(num) 

84 

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

86 for num in self.numbers: 

87 yield chemical_symbols[num] 

88 

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

90 numbers = symbols2numbers(value) 

91 if len(numbers) == 1: 

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

93 else: 

94 self.numbers[key] = numbers 

95 

96 def __len__(self) -> int: 

97 return len(self.numbers) 

98 

99 def __str__(self) -> str: 

100 return self.get_chemical_formula('reduce') 

101 

102 def __repr__(self) -> str: 

103 return f'Symbols(\'{self}\')' 

104 

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

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

107 return False 

108 

109 try: 

110 symbols = Symbols.fromsymbols(obj) 

111 except Exception: 

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

113 # atomic numbers. 

114 return False 

115 return self.numbers == symbols.numbers 

116 

117 def get_chemical_formula( 

118 self, 

119 mode: str = 'hill', 

120 empirical: bool = False, 

121 ) -> str: 

122 """Get chemical formula. 

123 

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

125 # XXX Delegate the work to the Formula object! 

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

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

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

129 

130 if len(self) == 0: 

131 return '' 

132 

133 numbers = self.numbers 

134 

135 if mode == 'reduce': 

136 n = len(numbers) 

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

138 numbers[:-1]])) 

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

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

141 

142 tokens = [] 

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

144 tokens.append(s) 

145 if c > 1: 

146 tokens.append(str(c)) 

147 formula = ''.join(tokens) 

148 elif mode == 'all': 

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

150 else: 

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

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

153 if empirical: 

154 f, _ = f.reduce() 

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

156 formula = f.format(mode) 

157 else: 

158 raise ValueError( 

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

160 

161 return formula 

162 

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

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

165 numbers = set(symbols2numbers(symbols)) 

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

167 if number in numbers] 

168 return np.array(indices, int) 

169 

170 def species(self) -> set[str]: 

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

172 return set(self) 

173 

174 def indices(self) -> dict[str, Integers]: 

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

176 

177 >>> from ase.build import molecule 

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

179 >>> atoms.symbols.indices() 

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

181 

182 """ 

183 dct: dict[str, list[int]] = {} 

184 for i, symbol in enumerate(self): 

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

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

187 

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

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

190 

191 >>> from ase import Atoms 

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

193 >>> atoms.symbols.species_indices() 

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

195 

196 ^ ^ ^ ^ ^ ^ ^ ^ ^ 

197 C H H H C H H O H 

198 

199 """ 

200 

201 counts: dict[str, int] = {} 

202 result = [] 

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

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

205 result.append(counts[n]) 

206 

207 return result