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

## 60 statements

, 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

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.

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}).

20 The general idea is as follows.

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

23 .. math:

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

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

36 .. math:

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

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

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

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

53 Q = q_max

54 V = q_vol

55 N = max_points

57 # Coefs

58 a = 1.0

59 b = -3 / 2 * Q

60 c = 0.0

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

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

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)

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

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

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

80 return x

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

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

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

87 A, B, C = cell # cell vectors

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

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

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

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

98 sigma = np.arctan(tan_sigma_over_4) * 4

100 return sigma

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.

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

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

138 """

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

141 # The physicists reciprocal cell

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

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

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)

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)

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

158 # Calculate distances for pruning

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

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]

166 # Prune based on distances

167 q_points = q_points[q_distances <= q_max]

168 q_distances = q_distances[q_distances <= q_max]

170 # Pruning based on max_points

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

173 solid_angle = calculate_solid_angle(rec_cell)

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

176 q_vol = np.linalg.det(rec_cell)

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

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}')

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)

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

190 rs = np.random.RandomState(seed)

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

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

195 return q_points