wmaee.utils.geometry

 1import operator
 2import itertools
 3import numpy as np
 4from ase import Atoms
 5from typing import Union, Tuple, Dict, Any
 6
 7from wmaee.core.requirements import test_pmg
 8if test_pmg():
 9    from pymatgen.core import Structure
10    from pymatgen.io.ase import AseAtomsAdaptor
11
12def find_neighbors(atoms: Union[Atoms, Any],
13                   rcut: float = 3.0,
14                   mindist: float = 0.1,
15                   properties: Tuple[str, ...] = ('indices', 'distances', 'images', 'vecs')
16                   ) -> Dict[int, Tuple[np.ndarray, ...]]:
17    """
18    Find neighbors for each atom within a specified cutoff distance and minimum distance.
19    
20    This function applies the minimum image convention to consider periodic boundary conditions.
21
22    Parameters
23    ----------
24    atoms : Union[Atoms, Any]
25        The atoms object to query. It can be an Atoms object or any other compatible type.
26    rcut : float, optional
27        The cutoff distance for neighbor search, by default 3.0.
28    mindist : float, optional
29        The minimum distance for neighbor search, by default 0.1.
30    properties : Tuple[str, ...], optional
31        An iterable of properties to return for each atom. 
32        Available options are 'indices', 'distances', 'images', 'vecs', by default ('indices', 'distances', 'images', 'vecs').
33
34    Returns
35    -------
36    Dict[int, Tuple[np.ndarray, ...]]
37        A dictionary where the keys represent the atom ids, and the values are tuples containing the desired properties.
38        - 'indices': Indices of neighboring atoms.
39        - 'distances': Distances to neighboring atoms.
40        - 'images': Indices of periodic images of neighboring atoms.
41        - 'vecs': Vectors pointing from the central atom to its neighbors.
42    """
43    if test_pmg() and isinstance(atoms, Structure):
44        atoms = AseAtomsAdaptor.get_atoms(atoms)
45
46    a, b, c = atoms.cell
47    stride = len(atoms)
48    image_vector = np.zeros((stride * 27, 3), dtype=float)
49    indices = np.hstack([np.arange(len(atoms))] * 27)
50    image_indices = np.zeros_like(image_vector, dtype=int)
51
52    for i, (ta, tb, tc) in enumerate(itertools.product(*[[-1, 0, 1]] * 3)):
53        start, end = (i * stride), ((i + 1) * stride)
54        image_vector[start:end] = a * ta + b * tb + c * tc
55        image_vector[start:end] += atoms.positions
56        image_indices[start:end] = [ta, tb, tc]
57
58    result = dict()
59    position_vector = np.zeros_like(image_vector)
60
61    for i, pos in enumerate(atoms.positions):
62        position_vector[:] = pos
63        distance_vectors = position_vector - image_vector
64        distances = np.linalg.norm(distance_vectors, axis=1)
65        mask = np.logical_and(distances >= mindist, distances < rcut)
66        result[i] = dict(indices=indices[mask], distances=distances[mask], images=image_indices[mask], vecs=distance_vectors[mask, :])
67
68    mapper = operator.itemgetter(*properties)
69    return {i: mapper(d) for i, d in result.items()}