Coverage for /builds/ase/ase/ase/formula.py : 91.45%
Hot-keys on this page
r m x p toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1import re
2from functools import lru_cache
3from math import gcd
4from typing import Dict, List, Sequence, Tuple, Union
6from ase.data import atomic_numbers, chemical_symbols
8# For type hints (A, A2, A+B):
9Tree = Union[str, Tuple['Tree', int], List['Tree']] # type: ignore
12class Formula:
13 def __init__(self,
14 formula: Union[str, 'Formula'] = '',
15 *,
16 strict: bool = False,
17 format: str = '',
18 _tree: Tree = None,
19 _count: Dict[str, int] = None):
20 """Chemical formula object.
22 Parameters
23 ----------
24 formula: str
25 Text string representation of formula. Examples: ``'6CO2'``,
26 ``'30Cu+2CO'``, ``'Pt(CO)6'``.
27 strict: bool
28 Only allow real chemical symbols.
29 format: str
30 Reorder according to *format*. Must be one of hill, metal,
31 ab2, a2b, periodic or reduce.
33 Examples
34 --------
35 >>> from ase.formula import Formula
36 >>> w = Formula('H2O')
37 >>> w.count()
38 {'H': 2, 'O': 1}
39 >>> 'H' in w
40 True
41 >>> w == 'HOH'
42 True
43 >>> f'{w:latex}'
44 'H$_{2}$O'
45 >>> w.format('latex')
46 'H$_{2}$O'
47 >>> divmod(6 * w + 'Cu', w)
48 (6, Formula('Cu'))
50 Raises
51 ------
52 ValueError
53 on malformed formula
54 """
56 # Be sure that Formula(x) works the same whether x is string or Formula
57 assert isinstance(formula, (str, Formula))
58 formula = str(formula)
60 if format:
61 assert _tree is None and _count is None
62 if format not in {'hill', 'metal', 'abc', 'reduce', 'ab2', 'a2b',
63 'periodic'}:
64 raise ValueError(f'Illegal format: {format}')
65 formula = Formula(formula).format(format)
67 self._formula = formula
69 self._tree = _tree or parse(formula)
70 self._count = _count or count_tree(self._tree)
71 if strict:
72 for symbol in self._count:
73 if symbol not in atomic_numbers:
74 raise ValueError('Unknown chemical symbol: ' + symbol)
76 def convert(self, fmt: str) -> 'Formula':
77 """Reformat this formula as a new Formula.
79 Same formatting rules as Formula(format=...) keyword.
80 """
81 return Formula(self._formula, format=fmt)
83 def count(self) -> Dict[str, int]:
84 """Return dictionary mapping chemical symbol to number of atoms.
86 Example
87 -------
88 >>> Formula('H2O').count()
89 {'H': 2, 'O': 1}
90 """
91 return self._count.copy()
93 def reduce(self) -> Tuple['Formula', int]:
94 """Reduce formula.
96 Returns
97 -------
98 formula: Formula
99 Reduced formula.
100 n: int
101 Number of reduced formula units.
103 Example
104 -------
105 >>> Formula('2H2O').reduce()
106 (Formula('H2O'), 2)
107 """
108 dct, N = self._reduce()
109 return self.from_dict(dct), N
111 def stoichiometry(self) -> Tuple['Formula', 'Formula', int]:
112 """Reduce to unique stoichiometry using "chemical symbols" A, B, C, ...
114 Examples
115 --------
116 >>> Formula('CO2').stoichiometry()
117 (Formula('AB2'), Formula('CO2'), 1)
118 >>> Formula('(H2O)4').stoichiometry()
119 (Formula('AB2'), Formula('OH2'), 4)
120 """
121 count1, N = self._reduce()
122 c = ord('A')
123 count2 = {}
124 count3 = {}
125 for n, symb in sorted((n, symb)
126 for symb, n in count1.items()):
127 count2[chr(c)] = n
128 count3[symb] = n
129 c += 1
130 return self.from_dict(count2), self.from_dict(count3), N
132 def format(self, fmt: str = '') -> str:
133 """Format formula as string.
135 Formats:
137 * ``'hill'``: alphabetically ordered with C and H first
138 * ``'metal'``: alphabetically ordered with metals first
139 * ``'ab2'``: count-ordered first then alphabetically ordered
140 * ``'abc'``: old name for ``'ab2'``
141 * ``'a2b'``: reverse count-ordered first then alphabetically ordered
142 * ``'periodic'``: periodic-table ordered: period first then group
143 * ``'reduce'``: Reduce and keep order (ABBBC -> AB3C)
144 * ``'latex'``: LaTeX representation
145 * ``'html'``: HTML representation
146 * ``'rest'``: reStructuredText representation
148 Example
149 -------
150 >>> Formula('H2O').format('html')
151 'H<sub>2</sub>O'
152 """
153 return format(self, fmt)
155 def __format__(self, fmt: str) -> str:
156 """Format Formula as str.
158 Possible formats: ``'hill'``, ``'metal'``, ``'abc'``, ``'reduce'``,
159 ``'latex'``, ``'html'``, ``'rest'``.
161 Example
162 -------
163 >>> f = Formula('OH2')
164 >>> '{f}, {f:hill}, {f:latex}'.format(f=f)
165 'OH2, H2O, OH$_{2}$'
166 """
168 if fmt == 'hill':
169 count = self.count()
170 count2 = {}
171 for symb in 'CH':
172 if symb in count:
173 count2[symb] = count.pop(symb)
174 for symb, n in sorted(count.items()):
175 count2[symb] = n
176 return dict2str(count2)
178 if fmt == 'metal':
179 count = self.count()
180 result2 = [(s, count.pop(s)) for s in non_metals if s in count]
181 result = [(s, count[s]) for s in sorted(count)]
182 result += sorted(result2)
183 return dict2str(dict(result))
185 if fmt == 'abc' or fmt == 'ab2':
186 _, f, N = self.stoichiometry()
187 return dict2str({symb: n * N for symb, n in f._count.items()})
189 if fmt == 'a2b':
190 _, f, N = self.stoichiometry()
191 return dict2str({symb: -n * N
192 for n, symb
193 in sorted([(-n, symb) for symb, n
194 in f._count.items()])})
196 if fmt == 'periodic':
197 count = self.count()
198 order = periodic_table_order()
199 items = sorted(count.items(),
200 key=lambda item: order.get(item[0], 0))
201 return ''.join(symb + (str(n) if n > 1 else '')
202 for symb, n in items)
204 if fmt == 'reduce':
205 symbols = list(self)
206 nsymb = len(symbols)
207 parts = []
208 i1 = 0
209 for i2, symbol in enumerate(symbols):
210 if i2 == nsymb - 1 or symbol != symbols[i2 + 1]:
211 parts.append(symbol)
212 m = i2 + 1 - i1
213 if m > 1:
214 parts.append(str(m))
215 i1 = i2 + 1
216 return ''.join(parts)
218 if fmt == 'latex':
219 return self._tostr('$_{', '}$')
221 if fmt == 'html':
222 return self._tostr('<sub>', '</sub>')
224 if fmt == 'rest':
225 return self._tostr(r'\ :sub:`', r'`\ ')
227 if fmt == '':
228 return self._formula
230 raise ValueError('Invalid format specifier')
232 @staticmethod
233 def from_dict(dct: Dict[str, int]) -> 'Formula':
234 """Convert dict to Formula.
236 >>> Formula.from_dict({'H': 2})
237 Formula('H2')
238 """
239 dct2 = {}
240 for symb, n in dct.items():
241 if not (isinstance(symb, str) and isinstance(n, int) and n >= 0):
242 raise ValueError('Bad dictionary: {dct}'.format(dct=dct))
243 if n > 0: # filter out n=0 symbols
244 dct2[symb] = n
245 return Formula(dict2str(dct2),
246 _tree=[([(symb, n) for symb, n in dct2.items()], 1)],
247 _count=dct2)
249 @staticmethod
250 def from_list(symbols: Sequence[str]) -> 'Formula':
251 """Convert list of chemical symbols to Formula."""
252 return Formula(''.join(symbols),
253 _tree=[(symbols[:], 1)])
255 def __len__(self) -> int:
256 """Number of atoms."""
257 return sum(self._count.values())
259 def __getitem__(self, symb: str) -> int:
260 """Number of atoms with chemical symbol *symb*."""
261 return self._count.get(symb, 0)
263 def __contains__(self, f: Union[str, 'Formula']) -> bool:
264 """Check if formula contains chemical symbols in *f*.
266 Type of *f* must be str or Formula.
268 Examples
269 --------
270 >>> 'OH' in Formula('H2O')
271 True
272 >>> 'O2' in Formula('H2O')
273 False
274 """
275 if isinstance(f, str):
276 f = Formula(f)
277 for symb, n in f._count.items():
278 if self[symb] < n:
279 return False
280 return True
282 def __eq__(self, other) -> bool:
283 """Equality check.
285 Note that order is not important.
287 Example
288 -------
289 >>> Formula('CO') == Formula('OC')
290 True
291 """
292 if isinstance(other, str):
293 other = Formula(other)
294 elif not isinstance(other, Formula):
295 return False
296 return self._count == other._count
298 def __add__(self, other: Union[str, 'Formula']) -> 'Formula':
299 """Add two formulas."""
300 if not isinstance(other, str):
301 other = other._formula
302 return Formula(self._formula + '+' + other)
304 def __radd__(self, other: str): # -> Formula
305 return Formula(other) + self
307 def __mul__(self, N: int) -> 'Formula':
308 """Repeat formula `N` times."""
309 if N == 0:
310 return Formula('')
311 return self.from_dict({symb: n * N
312 for symb, n in self._count.items()})
314 def __rmul__(self, N: int): # -> Formula
315 return self * N
317 def __divmod__(self,
318 other: Union['Formula', str]) -> Tuple[int, 'Formula']:
319 """Return the tuple (self // other, self % other).
321 Invariant::
323 div, mod = divmod(self, other)
324 div * other + mod == self
326 Example
327 -------
328 >>> divmod(Formula('H2O'), 'H')
329 (2, Formula('O'))
330 """
331 if isinstance(other, str):
332 other = Formula(other)
333 N = min(self[symb] // n for symb, n in other._count.items())
334 dct = self.count()
335 if N:
336 for symb, n in other._count.items():
337 dct[symb] -= n * N
338 if dct[symb] == 0:
339 del dct[symb]
340 return N, self.from_dict(dct)
342 def __rdivmod__(self, other):
343 return divmod(Formula(other), self)
345 def __mod__(self, other):
346 return divmod(self, other)[1]
348 def __rmod__(self, other):
349 return Formula(other) % self
351 def __floordiv__(self, other):
352 return divmod(self, other)[0]
354 def __rfloordiv__(self, other):
355 return Formula(other) // self
357 def __iter__(self, tree=None):
358 if tree is None:
359 tree = self._tree
360 if isinstance(tree, str):
361 yield tree
362 elif isinstance(tree, tuple):
363 tree, N = tree
364 for _ in range(N):
365 yield from self.__iter__(tree)
366 else:
367 for tree in tree:
368 yield from self.__iter__(tree)
370 def __str__(self):
371 return self._formula
373 def __repr__(self):
374 return 'Formula({!r})'.format(self._formula)
376 def _reduce(self):
377 N = 0
378 for n in self._count.values():
379 if N == 0:
380 N = n
381 else:
382 N = gcd(n, N)
383 dct = {symb: n // N for symb, n in self._count.items()}
384 return dct, N
386 def _tostr(self, sub1, sub2):
387 parts = []
388 for tree, n in self._tree:
389 s = tree2str(tree, sub1, sub2)
390 if s[0] == '(' and s[-1] == ')':
391 s = s[1:-1]
392 if n > 1:
393 s = str(n) + s
394 parts.append(s)
395 return '+'.join(parts)
398def dict2str(dct: Dict[str, int]) -> str:
399 """Convert symbol-to-number dict to str.
401 >>> dict2str({'A': 1, 'B': 2})
402 'AB2'
403 """
404 return ''.join(symb + (str(n) if n > 1 else '')
405 for symb, n in dct.items())
408def parse(f: str) -> Tree:
409 """Convert formula string to tree structure.
411 >>> parse('2A+BC2')
412 [('A', 2), (['B', ('C', 2)], 1)]
413 """
414 if not f:
415 return []
416 parts = f.split('+')
417 result = []
418 for part in parts:
419 n, f = strip_number(part)
420 result.append((parse2(f), n))
421 return result
424def parse2(f: str) -> Tree:
425 """Convert formula string to tree structure (no "+" symbols).
427 >>> parse('10(H2O)')
428 [(([('H', 2), 'O'], 1), 10)]
429 """
430 units = []
431 while f:
432 unit: Union[str, Tuple[str, int], Tree]
433 if f[0] == '(':
434 level = 0
435 for i, c in enumerate(f[1:], 1):
436 if c == '(':
437 level += 1
438 elif c == ')':
439 if level == 0:
440 break
441 level -= 1
442 else:
443 raise ValueError
444 f2 = f[1:i]
445 n, f = strip_number(f[i + 1:])
446 unit = (parse2(f2), n)
447 else:
448 m = re.match('([A-Z][a-z]?)([0-9]*)', f)
449 if m is None:
450 raise ValueError
451 symb = m.group(1)
452 number = m.group(2)
453 if number:
454 unit = (symb, int(number))
455 else:
456 unit = symb
457 f = f[m.end():]
458 units.append(unit)
459 if len(units) == 1:
460 return unit
461 return units
464def strip_number(s: str) -> Tuple[int, str]:
465 """Strip leading nuimber.
467 >>> strip_number('10AB2')
468 (10, 'AB2')
469 >>> strip_number('AB2')
470 (1, 'AB2')
471 """
472 m = re.match('[0-9]*', s)
473 assert m is not None
474 return int(m.group() or 1), s[m.end():]
477def tree2str(tree: Tree,
478 sub1: str, sub2: str) -> str:
479 """Helper function for html, latex and rest formats."""
480 if isinstance(tree, str):
481 return tree
482 if isinstance(tree, tuple):
483 tree, N = tree
484 s = tree2str(tree, sub1, sub2)
485 if N == 1:
486 if s[0] == '(' and s[-1] == ')':
487 return s[1:-1]
488 return s
489 return s + sub1 + str(N) + sub2
490 return '(' + ''.join(tree2str(tree, sub1, sub2) for tree in tree) + ')'
493def count_tree(tree: Tree) -> Dict[str, int]:
494 if isinstance(tree, str):
495 return {tree: 1}
496 if isinstance(tree, tuple):
497 tree, N = tree
498 return {symb: n * N for symb, n in count_tree(tree).items()}
499 dct = {} # type: Dict[str, int]
500 for tree in tree:
501 for symb, n in count_tree(tree).items():
502 m = dct.get(symb, 0)
503 dct[symb] = m + n
504 return dct
507# non metals, half-metals/metalloid, halogen, noble gas:
508non_metals = ['H', 'He', 'B', 'C', 'N', 'O', 'F', 'Ne',
509 'Si', 'P', 'S', 'Cl', 'Ar',
510 'Ge', 'As', 'Se', 'Br', 'Kr',
511 'Sb', 'Te', 'I', 'Xe',
512 'Po', 'At', 'Rn']
515@lru_cache()
516def periodic_table_order() -> Dict[str, int]:
517 """Create dict for sorting after period first then row."""
518 return {symbol: n for n, symbol in enumerate(chemical_symbols[87:] +
519 chemical_symbols[55:87] +
520 chemical_symbols[37:55] +
521 chemical_symbols[19:37] +
522 chemical_symbols[11:19] +
523 chemical_symbols[3:11] +
524 chemical_symbols[1:3])}
527# Backwards compatibility:
528def formula_hill(numbers, empirical=False):
529 """Convert list of atomic numbers to a chemical formula as a string.
531 Elements are alphabetically ordered with C and H first.
533 If argument `empirical`, element counts will be divided by greatest common
534 divisor to yield an empirical formula"""
535 symbols = [chemical_symbols[Z] for Z in numbers]
536 f = Formula('', _tree=[(symbols, 1)])
537 if empirical:
538 f, _ = f.reduce()
539 return f.format('hill')
542# Backwards compatibility:
543def formula_metal(numbers, empirical=False):
544 """Convert list of atomic numbers to a chemical formula as a string.
546 Elements are alphabetically ordered with metals first.
548 If argument `empirical`, element counts will be divided by greatest common
549 divisor to yield an empirical formula"""
550 symbols = [chemical_symbols[Z] for Z in numbers]
551 f = Formula('', _tree=[(symbols, 1)])
552 if empirical:
553 f, _ = f.reduce()
554 return f.format('metal')