Coverage for /builds/kinetik161/ase/ase/geometry/dimensionality/bond_generator.py: 96.43%

1import numpy as np 

2 

3from ase.data import covalent_radii 

4from ase.neighborlist import NeighborList 

5 

6 

7def get_bond_list(atoms, nl, rs): 

8 bonds = [] 

9 for i in range(len(atoms)): 

10 p = atoms.positions[i] 

11 indices, offsets = nl.get_neighbors(i) 

12 

13 for j, offset in zip(indices, offsets): 

14 q = atoms.positions[j] + np.dot(offset, atoms.get_cell()) 

15 d = np.linalg.norm(p - q) 

16 k = d / (rs[i] + rs[j]) 

17 bonds.append((k, i, j, tuple(offset))) 

18 return set(bonds) 

19 

20 

21def next_bond(atoms): 

22 """ 

23 Generates bonds (lazily) one at a time, sorted by k-value (low to high). 

24 Here, k = d_ij / (r_i + r_j), where d_ij is the bond length and r_i and r_j 

25 are the covalent radii of atoms i and j. 

26 

27 Parameters: 

28 

29 atoms: ASE atoms object 

30 

31 Returns: iterator of bonds 

32 A bond is a tuple with the following elements: 

33 

34 k: float k-value 

35 i: float index of first atom 

36 j: float index of second atom 

37 offset: tuple cell offset of second atom 

38 """ 

39 kmax = 0 

40 rs = covalent_radii[atoms.get_atomic_numbers()] 

41 seen = set() 

42 while 1: 

43 # Expand the scope of the neighbor list. 

44 kmax += 2 

45 nl = NeighborList(kmax * rs, skin=0, self_interaction=False) 

46 nl.update(atoms) 

47 

48 # Get a list of bonds, sorted by k-value. 

49 bonds = get_bond_list(atoms, nl, rs) 

50 new_bonds = bonds - seen 

51 if len(new_bonds) == 0: 

52 break 

53 

54 # Yield the bonds which we have not previously generated. 

55 seen.update(new_bonds) 

56 yield from sorted(new_bonds, key=lambda x: x[0])