Coverage for local_installation/dynasor/qpoints/spherical_qpoints.py: 96%

60 statements  

« prev     ^ index     » next       coverage.py v7.3.2, created at 2023-11-30 21:04 +0000

1import itertools 

2import numpy as np 

3from numpy.typing import NDArray 

4from dynasor.logging_tools import logger 

5 

6 

7def get_prune_distance( 

8 max_points: int, 

9 q_max: float, 

10 q_vol: float, 

11) -> NDArray[float]: 

12 r"""Determine distance in q-space beyond which to prune 

13 the q-point mesh to achieve near-isotropic sampling of q-space. 

14 

15 If points are selected from the full mesh with probability 

16 :math:`\min(1, (q_\mathrm{prune} / |q|)^2)`, q-space will 

17 on average be sampled with an equal number of points per radial unit 

18 (for :math:`q > q_\mathrm{prune}`). 

19 

20 The general idea is as follows. 

21 We know that the number of q-points inside a radius :math:`Q` is given by 

22 

23 .. math: 

24 

25 n = v^{-1} \int_0^Q dq 4 \pi q^2 = v^{-1} 4/3 \pi Q^3 

26 

27 where :math:`v` is the volume of one q-point. Now we want to find 

28 a distance :math:`P` such that if all points outside this radius 

29 are weighted by the function :math:`w(q)` the total number of 

30 q-points will equal the target :math:`N` (:attr:`max_points`) 

31 while the number of q-points increases linearly from :math:`P` 

32 outward. One additional constraint is that the weighting function 

33 must be 1 at :math:`P`. The weighting function which accomplishes 

34 this is :math:`w(q)=P^2/q^2` 

35 

36 .. math: 

37 

38 N = v^{-1} \left( \int_0^P 4 \pi q^2 + \int_P^Q 4 \pi q^2 P^2 / q^2 dq \right). 

39 

40 This results in a `cubic equation <https://en.wikipedia.org/wiki/Cubic_equation>`_ 

41 for :math:`P`, which is solved by this function. 

42 

43 Parameters 

44 ---------- 

45 max_points 

46 maximum number of resulting q-points; :math:`N` below 

47 max_q 

48 maximum q-value in the resulting q-point set; :math:`Q` below 

49 vol_q 

50 q-space volume for a single q-point 

51 """ 

52 

53 Q = q_max 

54 V = q_vol 

55 N = max_points 

56 

57 # Coefs 

58 a = 1.0 

59 b = -3 / 2 * Q 

60 c = 0.0 

61 d = 3 / 2 * V * N / (4 * np.pi) 

62 

63 # Eq tol solve 

64 def original_eq(x): 

65 return a * x**3 + b * x**2 + c * x + d 

66 # original_eq = lambda x: a * x**3 + b * x**2 + c * x + d 

67 

68 # Discriminant 

69 p = (3 * a * c - b**2) / (3 * a**2) 

70 q = (2 * b**3 - 9 * a * b * c + 27 * a**2 * d) / (27 * a**3) 

71 

72 D_t = - (4 * p**3 + 27 * q**2) 

73 if D_t < 0: 73 ↛ 74line 73 didn't jump to line 74, because the condition on line 73 was never true

74 return q_max 

75 

76 x = Q * (np.cos(1 / 3 * np.arccos(1 - 4 * d / Q**3) - 2 * np.pi / 3) + 0.5) 

77 

78 assert np.isclose(original_eq(x), 0), original_eq(x) 

79 

80 return x 

81 

82 

83def calculate_solid_angle(cell: NDArray[float]): 

84 # from https://en.wikipedia.org/wiki/Solid_angle#Tetrahedron 

85 assert np.linalg.det(cell) > 0 

86 

87 A, B, C = cell # cell vectors 

88 a, b, c = np.linalg.norm(cell, axis=1) # cell vector lengths 

89 

90 ac = np.arccos(A @ B / a / b) # angle a-o-b 

91 ab = np.arccos(A @ C / a / c) # angle b-o-c 

92 aa = np.arccos(B @ C / b / c) # angle b-o-c 

93 

94 a = (aa + ab + ac) / 2 

95 

96 tan_sigma_over_4 = np.sqrt(np.tan(a/2) * np.tan((a-aa)/2) * np.tan((a-ab)/2) * np.tan((a-ac)/2)) 

97 

98 sigma = np.arctan(tan_sigma_over_4) * 4 

99 

100 return sigma 

101 

102 

103def get_spherical_qpoints( 

104 cell: NDArray[float], 

105 q_max: float, 

106 max_points: int = None, 

107 seed: int = 42, 

108) -> NDArray[float]: 

109 r"""Generates all q-points in the first octant of the reciprocal 

110 lattice inside a given radius :attr:`q_max`. This approach is 

111 suitable if an isotropic sampling of q-space is desired. The 

112 function returns the resulting q-points in Cartesian coordinates 

113 as an ``Nx3`` array. 

114 

115 If the number of q-points generated exceeds :attr:`max_points`, 

116 points will be randomly pruned. If the number of generated 

117 q-points exceeds :attr:`max_points`, q-points will be randomly 

118 removed in such a way that the q-points inside are roughly 

119 uniformly distributed with respect to :math:`|q|`. 

120 

121 Parameters 

122 ---------- 

123 cell 

124 real cell with cell vectors as rows 

125 q_max 

126 maximum norm of generated q-points 

127 max_points 

128 Optionally limit the set to approximately :attr:`max_points` 

129 points by randomly removing points from a "fully populated 

130 mesh". The points are removed in such a way that for 

131 :math:`q > q_\mathrm{prune}`, the points will be radially uniformly 

132 distributed. The value of :math:`q_\mathrm{prune}` is 

133 calculated from :attr:`max_q`, :attr:`max_points`, and the 

134 shape of the cell. 

135 seed 

136 Seed used for stochastic pruning 

137 

138 """ 

139 

140 # inv(A.T) == inv(A).T 

141 # The physicists reciprocal cell 

142 rec_cell = np.linalg.inv(cell.T) * 2 * np.pi 

143 

144 # We want to find all points on the lattice defined by the reciprocal cell 

145 # such that all points within max_q are in this set 

146 inv_rec_cell = np.linalg.inv(rec_cell.T) # cell / 2pi 

147 

148 # h is the height of the rec_cell perpendicular to the other two vectors 

149 h = 1 / np.linalg.norm(inv_rec_cell, axis=1) 

150 

151 # If a q_point has a coordinate larger than this number it must be further away than q_max 

152 N = np.ceil(q_max / h).astype(int) 

153 

154 # Create q-points in the first octant 

155 lattice_points = list(itertools.product(*[range(n+1) for n in N])) 

156 q_points = lattice_points @ rec_cell 

157 

158 # Calculate distances for pruning 

159 q_distances = np.linalg.norm(q_points, axis=1) # Find distances 

160 

161 # Sort distances and q-points based on distance 

162 argsort = np.argsort(q_distances) 

163 q_distances = q_distances[argsort] 

164 q_points = q_points[argsort] 

165 

166 # Prune based on distances 

167 q_points = q_points[q_distances <= q_max] 

168 q_distances = q_distances[q_distances <= q_max] 

169 

170 # Pruning based on max_points 

171 if max_points is not None and max_points < len(q_points): 

172 

173 solid_angle = calculate_solid_angle(rec_cell) 

174 angle_factor = solid_angle / (4 * np.pi) 

175 

176 q_vol = np.linalg.det(rec_cell) 

177 

178 q_prune = get_prune_distance(max_points / angle_factor, q_max, q_vol) 

179 

180 if q_prune < q_max: 180 ↛ 195line 180 didn't jump to line 195, because the condition on line 180 was never false

181 logger.info(f'Pruning q-points from the range {q_prune:.3} < |q| < {q_max}') 

182 

183 # Keep point with probability min(1, (q_prune/|q|)^2) -> 

184 # aim for an equal number of points per equally thick "onion peel" 

185 # to get equal number of points per radial unit. 

186 p = np.ones(len(q_points)) 

187 assert np.isclose(q_distances[0], 0) 

188 p[1:] = (q_prune / q_distances[1:]) ** 2 

189 

190 rs = np.random.RandomState(seed) 

191 q_points = q_points[p > rs.rand(len(q_points))] 

192 

193 logger.info(f'Pruned from {len(q_distances)} q-points to {len(q_points)}') 

194 

195 return q_points