1import numpy as np 

2 

3from ase.constraints import Filter, FixAtoms 

4from ase.geometry.cell import cell_to_cellpar 

5from ase.neighborlist import neighbor_list 

6 

7 

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

9 neighbor_list=neighbor_list): 

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

11 

12 Uses ase.neighborlist.neighbour_list to compute neighbors. 

13 

14 Args: 

15 atoms: ase.atoms object to calculate neighbours for 

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

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

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

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

20 ASE neighbour list with an alternative with the same call 

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

22 

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

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

25 d_list: absolute distance between the corresponding pair 

26 fixed_atoms: indices of any fixed atoms 

27 """ 

28 

29 if isinstance(atoms, Filter): 

30 atoms = atoms.atoms 

31 

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

33 

34 # filter out self-interactions (across PBC) 

35 if not self_interaction: 

36 mask = i_list != j_list 

37 i_list = i_list[mask] 

38 j_list = j_list[mask] 

39 d_list = d_list[mask] 

40 

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

42 fixed_atoms = [] 

43 for constraint in atoms.constraints: 

44 if isinstance(constraint, FixAtoms): 

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

46 

47 return i_list, j_list, d_list, fixed_atoms 

48 

49 

50def estimate_nearest_neighbour_distance(atoms, 

51 neighbor_list=neighbor_list): 

52 """ 

53 Estimate nearest neighbour distance r_NN 

54 

55 Args: 

56 atoms: Atoms object 

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

58 ASE neighbour list with an alternative with the same call 

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

60 

61 Returns: 

62 rNN: float 

63 Nearest neighbour distance 

64 """ 

65 

66 if isinstance(atoms, Filter): 

67 atoms = atoms.atoms 

68 

69 # start_time = time.time() 

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

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

72 # cutoff exceeds the size of the system 

73 r_cut = 1.0 

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

75 

76 # cell lengths and angles 

77 a, b, c, alpha, beta, gamma = cell_to_cellpar(atoms.cell) 

78 extent = [a, b, c] 

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

80 

81 while r_cut < 2.0 * max(extent): 

82 # print('estimate_nearest_neighbour_distance(): ' 

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

84 i, j, rij, fixed_atoms = get_neighbours( 

85 atoms, r_cut, self_interaction=True, 

86 neighbor_list=neighbor_list) 

87 if len(i) != 0: 

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

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

90 break 

91 r_cut *= phi 

92 else: 

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

94 'finding neighbours for all atoms. This can ' 

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

96 'setting r_cut manually') 

97 

98 # maximum of nearest neighbour distances 

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

100 r_NN = np.max(nn_distances) 

101 

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

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

104 return r_NN