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

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

2 

3from __future__ import annotations 

4 

5from collections.abc import Iterable 

6from math import pi 

7 

8import numpy as np 

9 

10from ase import Atoms 

11from ase.cell import Cell 

12from ase.neighborlist import NeighborList 

13from ase.symbols import symbols2numbers 

14 

15 

16class CellTooSmall(Exception): 

17 pass 

18 

19 

20class VolumeNotDefined(Exception): 

21 pass 

22 

23 

24def get_rdf( 

25 atoms: Atoms | Iterable[Atoms], 

26 rmax: float, 

27 nbins: int, 

28 distance_matrix: np.ndarray | None = None, 

29 elements: Iterable[int | str] | None = None, 

30 no_dists: bool = False, 

31 volume: float | None = None, 

32): 

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

34 

35 .. versionchanged:: 3.27.0 

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

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

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

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

40 

41 Parameters 

42 ---------- 

43 atoms : Atoms | Iterable[Atoms] 

44 :class:`~ase.Atoms` or their iterable for which the RDF is computed. 

45 

46 .. versionadded:: 3.28.0 

47 This can accept an iterable of :class:`~ase.Atoms` 

48 to compute the average RDF over the images. 

49 

50 rmax : float 

51 The maximum distance that will contribute to the RDF. 

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

53 sphere with radius rmax in the periodic directions. 

54 nbins : int 

55 Number of bins to divide the RDF into. 

56 distance_matrix : numpy.array 

57 An array of distances between atoms, typically 

58 obtained by atoms.get_all_distances(). 

59 Default None meaning that a NeighborList will be used. 

60 elements : Iterable[int | str] | None, default: None 

61 Two atomic numbers of chemical symbols. If not None, the partial 

62 RDF for the supplied elements will be returned. 

63 

64 .. versionadded:: 3.28.0 

65 Each item of ``elements`` can be a chemical element symbol. 

66 

67 no_dists : bool 

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

69 volume : float or None 

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

71 will be used instead atoms.cell.volume. 

72 

73 Returns 

74 ------- 

75 rdf : np.ndarray 

76 RDFs. 

77 rr : np.ndarray 

78 Corresponding distances. 

79 

80 Notes 

81 ----- 

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

83 the cell volume in the normalization. 

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

85 non-periodic. 

86 

87 """ 

88 images = [atoms] if isinstance(atoms, Atoms) else atoms 

89 rdfs = [] 

90 bin_edges = np.linspace(0.0, rmax, nbins + 1) 

91 for atoms in images: 

92 rdf = _get_rdf_per_frame( 

93 atoms, 

94 bin_edges, 

95 distance_matrix=distance_matrix, 

96 elements=elements, 

97 volume=volume, 

98 ) 

99 rdfs.append(rdf) 

100 rdf = np.mean(rdfs, axis=0) 

101 rr = 0.5 * (bin_edges[:-1] + bin_edges[1:]) # mid point of each bin 

102 return rdf if no_dists else (rdf, rr) 

103 

104 

105def _get_rdf_per_frame( 

106 atoms: Atoms, 

107 bin_edges: np.ndarray, 

108 *, 

109 elements: Iterable[int | str] | None = None, 

110 distance_matrix: np.ndarray | None = None, 

111 volume: float | None = None, 

112) -> np.ndarray: 

113 # First check whether the cell is sufficiently large 

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

115 if vol < 1.0e-10: 

116 raise VolumeNotDefined 

117 

118 rmax = bin_edges.max() 

119 nbins = bin_edges.size - 1 

120 

121 check_cell_and_r_max(atoms, rmax) 

122 

123 natoms = len(atoms) 

124 dr = float(rmax / nbins) 

125 

126 atomic_numbers = None if elements is None else symbols2numbers(elements) 

127 

128 if atomic_numbers is None: 

129 i_indices = np.arange(natoms) 

130 n = natoms # number of center atoms 

131 rho = natoms / vol # average number density 

132 else: 

133 i_indices = np.where(atoms.numbers == atomic_numbers[0])[0] 

134 j_indices = np.where(atoms.numbers == atomic_numbers[1])[0] 

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

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

137 

138 if distance_matrix is None: 

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

140 nl.update(atoms) 

141 

142 rdf = np.zeros(nbins + 1) 

143 for i in i_indices: 

144 j_indices, offsets = nl.get_neighbors(i) 

145 if atomic_numbers is not None: 

146 mask = atoms.numbers[j_indices] == atomic_numbers[1] 

147 j_indices = j_indices[mask] 

148 offsets = offsets[mask] 

149 

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

151 continue 

152 

153 ps = atoms.positions 

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

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

156 

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

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

159 else: 

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

161 if atomic_numbers is None: 

162 x = indices.ravel() 

163 else: 

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

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

166 

167 shell_vols = 4.0 * pi * dr * (bin_edges**2 + bin_edges * dr + dr**2 / 3.0) 

168 rdf[1:] /= n * rho * shell_vols[:-1] 

169 

170 return rdf[1:] 

171 

172 

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

174 cell = atoms.get_cell() 

175 pbc = atoms.get_pbc() 

176 

177 vol = atoms.cell.volume 

178 

179 for i in range(3): 

180 if pbc[i]: 

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

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

183 if h < 2 * rmax: 

184 recommended_r_max = get_recommended_r_max(cell, pbc) 

185 raise CellTooSmall( 

186 'The cell is not large enough in ' 

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

188 f'Recommended rmax = {recommended_r_max}' 

189 ) 

190 

191 

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

193 recommended_r_max = 5.0 

194 vol = cell.volume 

195 for i in range(3): 

196 if pbc[i]: 

197 axb = np.cross( 

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

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

200 ) 

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

202 assert isinstance(h, float) # mypy 

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

204 return recommended_r_max 

205 

206 

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

208 atom2xyz = atoms.get_positions() 

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

210 

211 

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

213 return np.prod(get_containing_cell_length(atoms))