from typing import List, Optional, Tuple, Union
import numpy as np
from ..covalent_radii import covalentradii
from ..exceptions import NotAnElementError
__all__ = ["guess_connectivity"]
[docs]def guess_connectivity(
symbols: np.ndarray, geometry: np.ndarray, threshold: float = 1.2, default_connectivity: Optional[float] = None
) -> List[Union[Tuple[int, int], Tuple[int, int, float]]]:
r"""
Finds connected atoms based off of a covalent radii metric.
Parameters
----------
symbols
The molecular symbols (e.g., 'Zr', 'C')
geometry
The molecular geometry in Bohr
threshold
Tunes the covalent radii metric safety factor.
default_connectivity
Provides a default connectivity value
Returns
-------
~typing.List[~typing.Union[~typing.Tuple[int, int], ~typing.Tuple[int, int, float]]]
Provides a list of connected atoms, or optionally a list of
connected atoms and default connectivity if provided.
"""
geometry = np.asarray(geometry, dtype=float).reshape(-1, 3)
radii = []
for s in symbols:
try:
radii.append(covalentradii.get(s, missing=1.8))
except NotAnElementError:
radii.append(1.8)
radii = np.array(radii)
# Upper triangular
con = []
for x in range(geometry.shape[0]):
diffs = geometry[x] - geometry[x + 1 :]
dists = np.einsum("ij,ij->i", diffs, diffs)
np.sqrt(dists, out=dists)
cutoff = (radii[x] + radii[x + 1 :]) * threshold
where = np.where(dists < cutoff)[0]
where += x + 1
for atom2 in where:
con.append((x, atom2))
if default_connectivity:
con = [(x[0], x[1], default_connectivity) for x in con]
return con