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

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 ------- 

66 rNN: float 

67 Nearest neighbour distance 

68 """ 

69 

70 if isinstance(atoms, Filter): 

71 atoms = atoms.atoms 

72 

73 # start_time = time.time() 

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

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

76 # cutoff exceeds the size of the system 

77 r_cut = 1.0 

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

79 

80 # cell lengths and angles 

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

82 extent = [a, b, c] 

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

84 

85 while r_cut < 2.0 * max(extent): 

86 # print('estimate_nearest_neighbour_distance(): ' 

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

88 i, _j, rij, _fixed_atoms = get_neighbours( 

89 atoms, r_cut, self_interaction=True, 

90 neighbor_list=neighbor_list) 

91 if len(i) != 0: 

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

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

94 break 

95 r_cut *= phi 

96 else: 

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

98 'finding neighbours for all atoms. This can ' 

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

100 'setting r_cut manually') 

101 

102 # maximum of nearest neighbour distances 

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

104 r_NN = np.max(nn_distances) 

105 

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

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

108 return r_NN