Coverage for ase / geometry / rdf.py: 98.73%

1"""Radial distribution function (RDF).""" 

2 

3from __future__ import annotations 

4 

5import math 

6 

7import numpy as np 

8 

9from ase import Atoms 

10from ase.cell import Cell 

11from ase.neighborlist import NeighborList 

12 

13 

14class CellTooSmall(Exception): 

15 pass 

16 

17 

18class VolumeNotDefined(Exception): 

19 pass 

20 

21 

22def get_rdf( 

23 atoms: Atoms, 

24 rmax: float, 

25 nbins: int, 

26 distance_matrix: np.ndarray | None = None, 

27 elements: list[int] | tuple | None = None, 

28 no_dists: bool = False, 

29 volume: float | None = None, 

30): 

31 """Calculate the radial distribution function (RDF) :math:`g(r)`. 

32 

33 .. versionchanged:: 3.27.0 

34 Partial RDFs are fixed to be consistent with LAMMPS ``compute rdf``. 

35 They now satisfy :math:`g_{ij}(r) = g_{ji}(r)` and 

36 :math:`g(r) = \\sum_{ij} c_i c_j g_{ij}(r)`, where `i` and `j` index 

37 the elements and `c_{i,j}` are their atomic fractions. 

38 

39 Parameters 

40 ---------- 

41 atoms : Atoms 

42 ASE ``Atoms`` object for which the RDF is computed. 

43 rmax : float 

44 The maximum distance that will contribute to the RDF. 

45 The unit cell should be large enough so that it encloses a 

46 sphere with radius rmax in the periodic directions. 

47 nbins : int 

48 Number of bins to divide the RDF into. 

49 distance_matrix : numpy.array 

50 An array of distances between atoms, typically 

51 obtained by atoms.get_all_distances(). 

52 Default None meaning that a NeighborList will be used. 

53 elements : list[int] | tuple[int, int] 

54 List of two atomic numbers. If elements is not None the partial 

55 RDF for the supplied elements will be returned. 

56 no_dists : bool 

57 If True then the second array with RDF distances will not be returned. 

58 volume : float or None 

59 Optionally specify the volume of the system. If specified, the volume 

60 will be used instead atoms.cell.volume. 

61 

62 Returns 

63 ------- 

64 rdf : np.ndarray 

65 RDFs. 

66 rr : np.ndarray 

67 Corresponding distances. 

68 

69 Notes 

70 ----- 

71 The RDF is computed following the standard solid state definition which uses 

72 the cell volume in the normalization. 

73 This may or may not be appropriate in cases where one or more directions is 

74 non-periodic. 

75 

76 """ 

77 

78 # First check whether the cell is sufficiently large 

79 vol = atoms.cell.volume if volume is None else volume 

80 if vol < 1.0e-10: 

81 raise VolumeNotDefined 

82 

83 check_cell_and_r_max(atoms, rmax) 

84 

85 natoms = len(atoms) 

86 dr = float(rmax / nbins) 

87 

88 if elements is None: 

89 i_indices = np.arange(natoms) 

90 n = natoms # number of center atoms 

91 rho = natoms / vol # average number density 

92 else: 

93 i_indices = np.where(atoms.numbers == elements[0])[0] 

94 j_indices = np.where(atoms.numbers == elements[1])[0] 

95 n = len(i_indices) # number of center atoms 

96 rho = len(j_indices) / vol # average number density 

97 

98 if distance_matrix is None: 

99 nl = NeighborList(np.ones(natoms) * rmax * 0.5, bothways=True) 

100 nl.update(atoms) 

101 

102 rdf = np.zeros(nbins + 1) 

103 for i in i_indices: 

104 j_indices, offsets = nl.get_neighbors(i) 

105 if elements is not None: 

106 mask = atoms.numbers[j_indices] == elements[1] 

107 j_indices = j_indices[mask] 

108 offsets = offsets[mask] 

109 

110 if np.count_nonzero(j_indices) == 0: 

111 continue 

112 

113 ps = atoms.positions 

114 d = ps[j_indices] + offsets @ atoms.cell - ps[i] 

115 distances = np.sqrt(np.add.reduce(d**2, axis=1)) 

116 

117 indices = np.asarray(np.ceil(distances / dr), dtype=int) 

118 rdf += np.bincount(indices, minlength=nbins + 1)[: nbins + 1] 

119 else: 

120 indices = np.asarray(np.ceil(distance_matrix / dr), dtype=int) 

121 if elements is None: 

122 x = indices.ravel() 

123 else: 

124 x = indices[i_indices][:, j_indices].ravel() 

125 rdf = np.bincount(x, minlength=nbins + 1)[: nbins + 1].astype(float) 

126 

127 rr = np.arange(dr / 2, rmax, dr) 

128 shell_volumes = 4.0 * math.pi * dr * (rr * rr + (dr * dr / 12)) 

129 rdf[1:] /= n * rho * shell_volumes 

130 

131 if no_dists: 

132 return rdf[1:] 

133 

134 return rdf[1:], rr 

135 

136 

137def check_cell_and_r_max(atoms: Atoms, rmax: float) -> None: 

138 cell = atoms.get_cell() 

139 pbc = atoms.get_pbc() 

140 

141 vol = atoms.cell.volume 

142 

143 for i in range(3): 

144 if pbc[i]: 

145 axb = np.cross(cell[(i + 1) % 3, :], cell[(i + 2) % 3, :]) 

146 h = vol / np.linalg.norm(axb) 

147 if h < 2 * rmax: 

148 recommended_r_max = get_recommended_r_max(cell, pbc) 

149 raise CellTooSmall( 

150 'The cell is not large enough in ' 

151 f'direction {i}: {h:.3f} < 2*rmax={2 * rmax: .3f}. ' 

152 f'Recommended rmax = {recommended_r_max}' 

153 ) 

154 

155 

156def get_recommended_r_max(cell: Cell, pbc: list[bool]) -> float: 

157 recommended_r_max = 5.0 

158 vol = cell.volume 

159 for i in range(3): 

160 if pbc[i]: 

161 axb = np.cross( 

162 cell[(i + 1) % 3, :], # type: ignore[index] 

163 cell[(i + 2) % 3, :], # type: ignore[index] 

164 ) 

165 h = vol / np.linalg.norm(axb) 

166 assert isinstance(h, float) # mypy 

167 recommended_r_max = min(h / 2 * 0.99, recommended_r_max) 

168 return recommended_r_max 

169 

170 

171def get_containing_cell_length(atoms: Atoms) -> np.ndarray: 

172 atom2xyz = atoms.get_positions() 

173 return np.amax(atom2xyz, axis=0) - np.amin(atom2xyz, axis=0) + 2.0 

174 

175 

176def get_volume_estimate(atoms: Atoms) -> float: 

177 return np.prod(get_containing_cell_length(atoms))