Coverage for /builds/ase/ase/ase/optimize/precon/neighbors.py: 97.30%

37 statements  

« prev     ^ index     » next       coverage.py v7.5.3, created at 2025-08-02 00:12 +0000

1# fmt: off 

2 

3import numpy as np 

4 

5from ase.constraints import FixAtoms 

6from ase.filters import Filter 

7from ase.geometry.cell import cell_to_cellpar 

8from ase.neighborlist import neighbor_list 

9 

10 

11def get_neighbours(atoms, r_cut, self_interaction=False, 

12 neighbor_list=neighbor_list): 

13 """Return a list of pairs of atoms within a given distance of each other. 

14 

15 Uses ase.neighborlist.neighbour_list to compute neighbors. 

16 

17 Args: 

18 atoms: ase.atoms object to calculate neighbours for 

19 r_cut: cutoff radius (float). Pairs of atoms are considered neighbours 

20 if they are within a distance r_cut of each other (note that this 

21 is double the parameter used in the ASE's neighborlist module) 

22 neighbor_list: function (optional). Optionally replace the built-in 

23 ASE neighbour list with an alternative with the same call 

24 signature, e.g. `matscipy.neighbours.neighbour_list`. 

25 

26 Returns: a tuple (i_list, j_list, d_list, fixed_atoms): 

27 i_list, j_list: i and j indices of each neighbour pair 

28 d_list: absolute distance between the corresponding pair 

29 fixed_atoms: indices of any fixed atoms 

30 """ 

31 

32 if isinstance(atoms, Filter): 

33 atoms = atoms.atoms 

34 

35 i_list, j_list, d_list = neighbor_list('ijd', atoms, r_cut) 

36 

37 # filter out self-interactions (across PBC) 

38 if not self_interaction: 

39 mask = i_list != j_list 

40 i_list = i_list[mask] 

41 j_list = j_list[mask] 

42 d_list = d_list[mask] 

43 

44 # filter out bonds where 1st atom (i) in pair is fixed 

45 fixed_atoms = [] 

46 for constraint in atoms.constraints: 

47 if isinstance(constraint, FixAtoms): 

48 fixed_atoms.extend(list(constraint.index)) 

49 

50 return i_list, j_list, d_list, fixed_atoms 

51 

52 

53def estimate_nearest_neighbour_distance(atoms, 

54 neighbor_list=neighbor_list): 

55 """ 

56 Estimate nearest neighbour distance r_NN 

57 

58 Args: 

59 atoms: Atoms object 

60 neighbor_list: function (optional). Optionally replace the built-in 

61 ASE neighbour list with an alternative with the same call 

62 signature, e.g. `matscipy.neighbours.neighbour_list`. 

63 

64 Returns: 

65 rNN: float 

66 Nearest neighbour distance 

67 """ 

68 

69 if isinstance(atoms, Filter): 

70 atoms = atoms.atoms 

71 

72 # start_time = time.time() 

73 # compute number of neighbours of each atom. If any atom doesn't 

74 # have a neighbour we increase the cutoff and try again, until our 

75 # cutoff exceeds the size of the system 

76 r_cut = 1.0 

77 phi = (1.0 + np.sqrt(5.0)) / 2.0 # Golden ratio 

78 

79 # cell lengths and angles 

80 a, b, c, _alpha, _beta, _gamma = cell_to_cellpar(atoms.cell) 

81 extent = [a, b, c] 

82 # print('estimate_nearest_neighbour_distance(): extent=%r' % extent) 

83 

84 while r_cut < 2.0 * max(extent): 

85 # print('estimate_nearest_neighbour_distance(): ' 

86 # 'calling neighbour_list with r_cut=%.2f A' % r_cut) 

87 i, _j, rij, _fixed_atoms = get_neighbours( 

88 atoms, r_cut, self_interaction=True, 

89 neighbor_list=neighbor_list) 

90 if len(i) != 0: 

91 nn_i = np.bincount(i, minlength=len(atoms)) 

92 if (nn_i != 0).all(): 

93 break 

94 r_cut *= phi 

95 else: 

96 raise RuntimeError('increased r_cut to twice system extent without ' 

97 'finding neighbours for all atoms. This can ' 

98 'happen if your system is too small; try ' 

99 'setting r_cut manually') 

100 

101 # maximum of nearest neighbour distances 

102 nn_distances = [np.min(rij[i == I]) for I in range(len(atoms))] 

103 r_NN = np.max(nn_distances) 

104 

105 # print('estimate_nearest_neighbour_distance(): got r_NN=%.3f in %s s' % 

106 # (r_NN, time.time() - start_time)) 

107 return r_NN