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

28 statements  

« prev     ^ index     » next       coverage.py v7.13.5, created at 2026-03-30 08:22 +0000

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 

30 atoms: ASE atoms object 

31 

32 Returns: iterator of bonds 

33 A bond is a tuple with the following elements: 

34 

35 k: float k-value 

36 i: float index of first atom 

37 j: float index of second atom 

38 offset: tuple cell offset of second atom 

39 """ 

40 kmax = 0 

41 rs = covalent_radii[atoms.get_atomic_numbers()] 

42 seen = set() 

43 while 1: 

44 # Expand the scope of the neighbor list. 

45 kmax += 2 

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

47 nl.update(atoms) 

48 

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

50 bonds = get_bond_list(atoms, nl, rs) 

51 new_bonds = bonds - seen 

52 if len(new_bonds) == 0: 

53 break 

54 

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

56 seen.update(new_bonds) 

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