index | modules | gitlab

# Source code for ase.formula

    import re
    from functools import lru_cache
    from math import gcd
    from typing import Dict, List, Sequence, Tuple, Union

    from ase.data import atomic_numbers, chemical_symbols

    # For type hints (A, A2, A+B):
    Tree = Union[str, Tuple['Tree', int], List['Tree']]

    class Formula:
        def __init__(self,
                     formula: Union[str, 'Formula'] = '',
                     *,
                     strict: bool = False,
                     format: str = '',
                     _tree: Tree = None,
                     _count: Dict[str, int] = None):
            """Chemical formula object.

            Parameters
            ----------
            formula: str
                Text string representation of formula.  Examples: ``'6CO2'``,
                ``'30Cu+2CO'``, ``'Pt(CO)6'``.
            strict: bool
                Only allow real chemical symbols.
            format: str
                Reorder according to *format*.  Must be one of hill, metal,
                ab2, a2b, periodic or reduce.

            Examples
            --------
            >>> from ase.formula import Formula
            >>> w = Formula('H2O')
            >>> w.count()
            {'H': 2, 'O': 1}
            >>> 'H' in w
            True
            >>> w == 'HOH'
            True
            >>> f'{w:latex}'
            'H$_{2}$O'
            >>> w.format('latex')
            'H$_{2}$O'
            >>> divmod(6 * w + 'Cu', w)
            (6, Formula('Cu'))

            Raises
            ------
            ValueError
                on malformed formula
            """

            # Be sure that Formula(x) works the same whether x is string or Formula
            assert isinstance(formula, (str, Formula))
            formula = str(formula)

            if format:
                assert _tree is None and _count is None
                if format not in {'hill', 'metal', 'abc', 'reduce', 'ab2', 'a2b',
                                  'periodic'}:
                    raise ValueError(f'Illegal format: {format}')
                formula = Formula(formula).format(format)

            self._formula = formula

            self._tree = _tree or parse(formula)
            self._count = _count or count_tree(self._tree)
            if strict:
                for symbol in self._count:
                    if symbol not in atomic_numbers:
                        raise ValueError('Unknown chemical symbol: ' + symbol)

        def convert(self, fmt: str) -> 'Formula':
            """Reformat this formula as a new Formula.

            Same formatting rules as Formula(format=...) keyword.
            """
            return Formula(self._formula, format=fmt)

        def count(self) -> Dict[str, int]:
            """Return dictionary mapping chemical symbol to number of atoms.

            Example
            -------
            >>> Formula('H2O').count()
            {'H': 2, 'O': 1}
            """
            return self._count.copy()

        def reduce(self) -> Tuple['Formula', int]:
            """Reduce formula.

            Returns
            -------
            formula: Formula
                Reduced formula.
            n: int
                Number of reduced formula units.

            Example
            -------
            >>> Formula('2H2O').reduce()
            (Formula('H2O'), 2)
            """
            dct, N = self._reduce()
            return self.from_dict(dct), N

        def stoichiometry(self) -> Tuple['Formula', 'Formula', int]:
            """Reduce to unique stoichiometry using "chemical symbols" A, B, C, ...

            Examples
            --------
            >>> Formula('CO2').stoichiometry()
            (Formula('AB2'), Formula('CO2'), 1)
            >>> Formula('(H2O)4').stoichiometry()
            (Formula('AB2'), Formula('OH2'), 4)
            """
            count1, N = self._reduce()
            c = ord('A')
            count2 = {}
            count3 = {}
            for n, symb in sorted((n, symb)
                                  for symb, n in count1.items()):
                count2[chr(c)] = n
                count3[symb] = n
                c += 1
            return self.from_dict(count2), self.from_dict(count3), N

        def format(self, fmt: str = '') -> str:
            """Format formula as string.

            Formats:

            * ``'hill'``: alphabetically ordered with C and H first
            * ``'metal'``: alphabetically ordered with metals first
            * ``'ab2'``: count-ordered first then alphabetically ordered
            * ``'abc'``: old name for ``'ab2'``
            * ``'a2b'``: reverse count-ordered first then alphabetically ordered
            * ``'periodic'``: periodic-table ordered: period first then group
            * ``'reduce'``: Reduce and keep order (ABBBC -> AB3C)
            * ``'latex'``: LaTeX representation
            * ``'html'``: HTML representation
            * ``'rest'``: reStructuredText representation

            Example
            -------
            >>> Formula('H2O').format('html')
            'H<sub>2</sub>O'
            """
            return format(self, fmt)

        def __format__(self, fmt: str) -> str:
            """Format Formula as str.

            Possible formats: ``'hill'``, ``'metal'``, ``'abc'``, ``'reduce'``,
            ``'latex'``, ``'html'``, ``'rest'``.

            Example
            -------
            >>> f = Formula('OH2')
            >>> '{f}, {f:hill}, {f:latex}'.format(f=f)
            'OH2, H2O, OH$_{2}$'
            """

            if fmt == 'hill':
                count = self.count()
                count2 = {symb: count.pop(symb) for symb in 'CH' if symb in count}
                for symb, n in sorted(count.items()):
                    count2[symb] = n
                return dict2str(count2)

            if fmt == 'metal':
                count = self.count()
                result2 = [(s, count.pop(s)) for s in non_metals if s in count]
                result = [(s, count[s]) for s in sorted(count)]
                result += sorted(result2)
                return dict2str(dict(result))

            if fmt == 'abc' or fmt == 'ab2':
                _, f, N = self.stoichiometry()
                return dict2str({symb: n * N for symb, n in f._count.items()})

            if fmt == 'a2b':
                _, f, N = self.stoichiometry()
                return dict2str({symb: -n * N
                                 for n, symb
                                 in sorted([(-n, symb) for symb, n
                                            in f._count.items()])})

            if fmt == 'periodic':
                count = self.count()
                order = periodic_table_order()
                items = sorted(count.items(),
                               key=lambda item: order.get(item[0], 0))
                return ''.join(symb + (str(n) if n > 1 else '')
                               for symb, n in items)

            if fmt == 'reduce':
                symbols = list(self)
                nsymb = len(symbols)
                parts = []
                i1 = 0
                for i2, symbol in enumerate(symbols):
                    if i2 == nsymb - 1 or symbol != symbols[i2 + 1]:
                        parts.append(symbol)
                        m = i2 + 1 - i1
                        if m > 1:
                            parts.append(str(m))
                        i1 = i2 + 1
                return ''.join(parts)

            if fmt == 'latex':
                return self._tostr('$_{', '}$')

            if fmt == 'html':
                return self._tostr('<sub>', '</sub>')

            if fmt == 'rest':
                return self._tostr(r'\ :sub:`', r'`\ ')

            if fmt == '':
                return self._formula

            raise ValueError('Invalid format specifier')

        @staticmethod
        def from_dict(dct: Dict[str, int]) -> 'Formula':
            """Convert dict to Formula.

            >>> Formula.from_dict({'H': 2})
            Formula('H2')
            """
            dct2 = {}
            for symb, n in dct.items():
                if not (isinstance(symb, str) and isinstance(n, int) and n >= 0):
                    raise ValueError(f'Bad dictionary: {dct}')
                if n > 0:  # filter out n=0 symbols
                    dct2[symb] = n
            return Formula(dict2str(dct2),
                           _tree=[([(symb, n) for symb, n in dct2.items()], 1)],
                           _count=dct2)

        @staticmethod
        def from_list(symbols: Sequence[str]) -> 'Formula':
            """Convert list of chemical symbols to Formula."""
            return Formula(''.join(symbols),
                           _tree=[(symbols[:], 1)])  # type: ignore[list-item]

        def __len__(self) -> int:
            """Number of atoms."""
            return sum(self._count.values())

        def __getitem__(self, symb: str) -> int:
            """Number of atoms with chemical symbol *symb*."""
            return self._count.get(symb, 0)

        def __contains__(self, f: Union[str, 'Formula']) -> bool:
            """Check if formula contains chemical symbols in *f*.

            Type of *f* must be str or Formula.

            Examples
            --------
            >>> 'OH' in Formula('H2O')
            True
            >>> 'O2' in Formula('H2O')
            False
            """
            if isinstance(f, str):
                f = Formula(f)
            for symb, n in f._count.items():
                if self[symb] < n:
                    return False
            return True

        def __eq__(self, other) -> bool:
            """Equality check.

            Note that order is not important.

            Example
            -------
            >>> Formula('CO') == Formula('OC')
            True
            """
            if isinstance(other, str):
                other = Formula(other)
            elif not isinstance(other, Formula):
                return False
            return self._count == other._count

        def __add__(self, other: Union[str, 'Formula']) -> 'Formula':
            """Add two formulas."""
            if not isinstance(other, str):
                other = other._formula
            return Formula(self._formula + '+' + other)

        def __radd__(self, other: str):  # -> Formula
            return Formula(other) + self

        def __mul__(self, N: int) -> 'Formula':
            """Repeat formula `N` times."""
            if N == 0:
                return Formula('')
            return self.from_dict({symb: n * N
                                   for symb, n in self._count.items()})

        def __rmul__(self, N: int):  # -> Formula
            return self * N

        def __divmod__(self,
                       other: Union['Formula', str]) -> Tuple[int, 'Formula']:
            """Return the tuple (self // other, self % other).

            Invariant::

                div, mod = divmod(self, other)
                div * other + mod == self

            Example
            -------
            >>> divmod(Formula('H2O'), 'H')
            (2, Formula('O'))
            """
            if isinstance(other, str):
                other = Formula(other)
            N = min(self[symb] // n for symb, n in other._count.items())
            dct = self.count()
            if N:
                for symb, n in other._count.items():
                    dct[symb] -= n * N
                    if dct[symb] == 0:
                        del dct[symb]
            return N, self.from_dict(dct)

        def __rdivmod__(self, other):
            return divmod(Formula(other), self)

        def __mod__(self, other):
            return divmod(self, other)[1]

        def __rmod__(self, other):
            return Formula(other) % self

        def __floordiv__(self, other):
            return divmod(self, other)[0]

        def __rfloordiv__(self, other):
            return Formula(other) // self

        def __iter__(self):
            return self._tree_iter()

        def _tree_iter(self, tree=None):
            if tree is None:
                tree = self._tree
            if isinstance(tree, str):
                yield tree
            elif isinstance(tree, tuple):
                tree, N = tree
                for _ in range(N):
                    yield from self._tree_iter(tree)
            else:
                for tree in tree:
                    yield from self._tree_iter(tree)

        def __str__(self):
            return self._formula

        def __repr__(self):
            return f'Formula({self._formula!r})'

        def _reduce(self):
            N = 0
            for n in self._count.values():
                if N == 0:
                    N = n
                else:
                    N = gcd(n, N)
            dct = {symb: n // N for symb, n in self._count.items()}
            return dct, N

        def _tostr(self, sub1, sub2):
            parts = []
            for tree, n in self._tree:
                s = tree2str(tree, sub1, sub2)
                if s[0] == '(' and s[-1] == ')':
                    s = s[1:-1]
                if n > 1:
                    s = str(n) + s
                parts.append(s)
            return '+'.join(parts)

    def dict2str(dct: Dict[str, int]) -> str:
        """Convert symbol-to-number dict to str.

        >>> dict2str({'A': 1, 'B': 2})
        'AB2'
        """
        return ''.join(symb + (str(n) if n > 1 else '')
                       for symb, n in dct.items())

    def parse(f: str) -> Tree:
        """Convert formula string to tree structure.

        >>> parse('2A+BC2')
        [('A', 2), (['B', ('C', 2)], 1)]
        """
        if not f:
            return []
        parts = f.split('+')
        result = []
        for part in parts:
            n, f = strip_number(part)
            result.append((parse2(f), n))
        return result  # type: ignore[return-value]

    def parse2(f: str) -> Tree:
        """Convert formula string to tree structure (no "+" symbols).

        >>> parse('10(H2O)')
        [(([('H', 2), 'O'], 1), 10)]
        """
        units = []
        while f:
            unit: Union[str, Tuple[str, int], Tree]
            if f[0] == '(':
                level = 0
                for i, c in enumerate(f[1:], 1):
                    if c == '(':
                        level += 1
                    elif c == ')':
                        if level == 0:
                            break
                        level -= 1
                else:
                    raise ValueError
                f2 = f[1:i]
                n, f = strip_number(f[i + 1:])
                unit = (parse2(f2), n)
            else:
                m = re.match('([A-Z][a-z]?)([0-9]*)', f)
                if m is None:
                    raise ValueError
                symb = m.group(1)
                number = m.group(2)
                if number:
                    unit = (symb, int(number))
                else:
                    unit = symb
                f = f[m.end():]
            units.append(unit)
        if len(units) == 1:
            return unit
        return units

    def strip_number(s: str) -> Tuple[int, str]:
        """Strip leading nuimber.

        >>> strip_number('10AB2')
        (10, 'AB2')
        >>> strip_number('AB2')
        (1, 'AB2')
        """
        m = re.match('[0-9]*', s)
        assert m is not None
        return int(m.group() or 1), s[m.end():]

    def tree2str(tree: Tree,
                 sub1: str, sub2: str) -> str:
        """Helper function for html, latex and rest formats."""
        if isinstance(tree, str):
            return tree
        if isinstance(tree, tuple):
            tree, N = tree
            s = tree2str(tree, sub1, sub2)
            if N == 1:
                if s[0] == '(' and s[-1] == ')':
                    return s[1:-1]
                return s
            return s + sub1 + str(N) + sub2
        return '(' + ''.join(tree2str(tree, sub1, sub2) for tree in tree) + ')'

    def count_tree(tree: Tree) -> Dict[str, int]:
        if isinstance(tree, str):
            return {tree: 1}
        if isinstance(tree, tuple):
            tree, N = tree
            return {symb: n * N for symb, n in count_tree(tree).items()}
        dct = {}  # type: Dict[str, int]
        for tree in tree:
            for symb, n in count_tree(tree).items():
                m = dct.get(symb, 0)
                dct[symb] = m + n
        return dct

    # non metals, half-metals/metalloid, halogen, noble gas:
    non_metals = ['H', 'He', 'B', 'C', 'N', 'O', 'F', 'Ne',
                  'Si', 'P', 'S', 'Cl', 'Ar',
                  'Ge', 'As', 'Se', 'Br', 'Kr',
                  'Sb', 'Te', 'I', 'Xe',
                  'Po', 'At', 'Rn']

    @lru_cache
    def periodic_table_order() -> Dict[str, int]:
        """Create dict for sorting after period first then row."""
        return {symbol: n for n, symbol in enumerate(chemical_symbols[87:] +
                                                     chemical_symbols[55:87] +
                                                     chemical_symbols[37:55] +
                                                     chemical_symbols[19:37] +
                                                     chemical_symbols[11:19] +
                                                     chemical_symbols[3:11] +
                                                     chemical_symbols[1:3])}

    # Backwards compatibility:

    def formula_hill(numbers, empirical=False):
        """Convert list of atomic numbers to a chemical formula as a string.

        Elements are alphabetically ordered with C and H first.

        If argument `empirical`, element counts will be divided by greatest common
        divisor to yield an empirical formula"""
        symbols = [chemical_symbols[Z] for Z in numbers]
        f = Formula('', _tree=[(symbols, 1)])
        if empirical:
            f, _ = f.reduce()
        return f.format('hill')

    # Backwards compatibility:

    def formula_metal(numbers, empirical=False):
        """Convert list of atomic numbers to a chemical formula as a string.

        Elements are alphabetically ordered with metals first.

        If argument `empirical`, element counts will be divided by greatest common
        divisor to yield an empirical formula"""
        symbols = [chemical_symbols[Z] for Z in numbers]
        f = Formula('', _tree=[(symbols, 1)])
        if empirical:
            f, _ = f.reduce()
        return f.format('metal')