Source code for


A recipe for calculating atomic coordination numbers.

.. currentmodule::
.. autosummary::

.. autofunction:: get_coordination_number
.. autofunction:: coordination_outer


from typing import Dict, Tuple, List, Optional
from itertools import combinations

import numpy as np

from scm.plams import Molecule

from nanoutils import group_by_values
from nanoCAT.bde.guess_core_dist import guess_core_core_dist

__all__ = ['get_coordination_number', 'coordination_outer']

#: A nested dictonary
NestedDict = Dict[str, Dict[int, List[int]]]

def get_indices(mol: Molecule) -> Dict[str, np.ndarray]:
    """Construct a dictionary with atomic symbols as keys and arrays of indices as values."""
    elements = [at.symbol for at in mol.atoms]
    symbol_enumerator = enumerate(elements)
    idx_dict = group_by_values(symbol_enumerator)
    for k, v in idx_dict.items():
        idx_dict[k] = np.fromiter(v, dtype=int, count=len(v))
    return idx_dict

def idx_pairs(idx_dict: Dict[str, np.ndarray]) -> Tuple[np.ndarray, np.ndarray]:
    """Construct two arrays of indice-pairs of all possible combinations in idx_dict.

    The combinations, by definition, do not contain any atom pairs where ``at1.symbol == at2.symbol``
    """  # noqa
    x, y = [], []
    symbol_combinations = combinations(idx_dict.keys(), r=2)
    for symbol1, symbol2 in symbol_combinations:
        idx1 = idx_dict[symbol1]
        idx2 = idx_dict[symbol2]
        _x, _y = np.meshgrid(idx1, idx2)
        x += _x.ravel().tolist()
        y += _y.ravel().tolist()

    x = np.fromiter(x, dtype=int, count=len(x))
    y = np.fromiter(y, dtype=int, count=len(y))
    return x, y

def upper_dist(xyz: np.ndarray, x: np.ndarray, y: np.ndarray) -> np.ndarray:
    """Construct the upper distance matrix."""
    shape = len(xyz), len(xyz)
    dist = np.full(shape, fill_value=np.nan)
    dist[x, y] = np.linalg.norm(xyz[x] - xyz[y], axis=1)
    return dist

def radius_inner(mol: Molecule, idx_dict: Dict[str, np.ndarray]) -> Dict[str, float]:
    """Search the threshold radius of the inner coordination shell for each atom type.

    This radius is calculated as the distance with the first neighbors:
    in case of heterogeneous coordinations, only the closest neighbors are considered.
    d_inner = {}
    symbol_combinations = combinations(idx_dict.keys(), r=2)
    for symbol1, symbol2 in symbol_combinations:
        d_pair = guess_core_core_dist(mol, (symbol1, symbol2))
        if d_pair < d_inner.get(symbol1, np.inf):
            d_inner[symbol1] = d_pair
        if d_pair < d_inner.get(symbol2, np.inf):
            d_inner[symbol2] = d_pair
    return d_inner

def coordination_inner(dist: np.ndarray, d_inner: Dict[str, float],
                       idx_dict: Dict[str, np.ndarray], length: int) -> NestedDict:
    """Calculate the coordination number relative to the inner shell (first neighbors)."""
    coord_inner = np.empty(length, int)
    for k, v in idx_dict.items():
        a, b = np.where(dist <= d_inner[k])
        count = np.bincount(a, minlength=length) + np.bincount(b, minlength=length)
        coord_inner[v] = count[v]
    return coord_inner

[docs]def coordination_outer(dist: np.ndarray, d_outer: float, length: int) -> np.ndarray: """Calculate the coordination number relative to the outer shell.""" a, b = np.where(dist <= d_outer) return np.bincount(a, minlength=length) + np.bincount(b, minlength=length)
def map_coordination(coord: np.ndarray, idx_dict: Dict[str, np.ndarray]) -> NestedDict: """Map atoms according to their atomic symbols and coordination number. Construct a nested dictionary ``{'atom_type1': {coord1: [indices], ...}, ...}``. """ cn_dict = {} for k, v in idx_dict.items(): cn = coord[v] mapping = {i: (v[cn == i] + 1).tolist() for i in np.unique(cn)} cn_dict[k] = mapping return cn_dict
[docs]def get_coordination_number(mol: Molecule, shell: str = 'inner', d_outer: Optional[float] = None) -> NestedDict: """Take a molecule and identify the coordination number of each atom. The function first compute the pair distance between all reference atoms in **mol**. The number of first neighbors, defined as all atoms within a threshold radius **d_inner** is then count for each atom. The threshold radius can be changed to a desired value **d_outer** (in angstrom) to obtain higher coordination numbers associated to outer coordination shells. The function finally groups the (1-based) indices of all atoms in **mol** according to their atomic symbols and coordination numbers. Parameters ---------- mol : array-like [:class:`float`], shape :math:`(n, 3)` An array-like object with the Cartesian coordinates of the molecule. shell : :class:`str` The coordination shell to be considered. Only ``'inner'`` or ``'outer'`` values are accepted. The default, ``'inner'``, refers to the first coordination shell. d_outer : :class:`float`, optional The threshold radius for defining which atoms are considered as neighbors. The default, ``None``, is accepted only if ``shell`` is ``'inner'`` Returns ------- :class:`dict` A nested dictionary ``{'Cd': {8: [0, 1, 2, 3, 4, ...], ...}, ...}`` containing lists of (1-based) indices refered to the atoms in **mol** having a given atomic symbol (*e.g.* ``'Cd'``) and coordination number (*e.g.* ``8``). Raises ------ :exc:`TypeError` Raised if no threshold radius is defined for the outer coordination shell. :exc:`ValueError` Raised if a wrong value is attributed to ``shell``. See Also -------- :func:`guess_core_core_dist()<nanoCAT.bde.guess_core_dist.guess_core_core_dist>` Estimate a value for **d_inner** based on the radial distribution function of **mol**. Can also be used to estimate **d_outer** as the distance between the atom pairs ('A', 'B'). """ # noqa # Return the Cartesian coordinates of **mol** xyz = np.asarray(mol) length = len(xyz) # Group atom indices according to atom symbols idx_dict = get_indices(mol) # Construct the upper distance matrix x, y = idx_pairs(idx_dict) dist = upper_dist(xyz, x, y) # Compute the coordination number if shell == 'inner': d_inner = radius_inner(mol, idx_dict) coord = coordination_inner(dist, d_inner, idx_dict, length) elif shell == 'outer': if d_outer is None: raise TypeError("user defined threshold radius required " "for the outer coordination shell") coord = coordination_outer(dist, d_outer, length) else: raise ValueError(f"'shell' expected to be 'inner' or 'outer'; observed value: '{shell}'") # Return the final dictionary return map_coordination(coord, idx_dict)
coordination_number = get_coordination_number