1import numpy as np 

2from ase.neighborlist import NeighborList 

3from ase.data import covalent_radii 

4 

5 

6def get_bond_list(atoms, nl, rs): 

7 bonds = [] 

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

9 p = atoms.positions[i] 

10 indices, offsets = nl.get_neighbors(i) 

11 

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

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

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

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

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

17 return set(bonds) 

18 

19 

20def next_bond(atoms): 

21 """ 

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

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

24 are the covalent radii of atoms i and j. 

25 

26 Parameters: 

27 

28 atoms: ASE atoms object 

29 

30 Returns: iterator of bonds 

31 A bond is a tuple with the following elements: 

32 

33 k: float k-value 

34 i: float index of first atom 

35 j: float index of second atom 

36 offset: tuple cell offset of second atom 

37 """ 

38 kmax = 0 

39 rs = covalent_radii[atoms.get_atomic_numbers()] 

40 seen = set() 

41 while 1: 

42 # Expand the scope of the neighbor list. 

43 kmax += 2 

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

45 nl.update(atoms) 

46 

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

48 bonds = get_bond_list(atoms, nl, rs) 

49 new_bonds = bonds - seen 

50 if len(new_bonds) == 0: 

51 break 

52 

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

54 seen.update(new_bonds) 

55 for b in sorted(new_bonds, key=lambda x: x[0]): 

56 yield b