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

47 statements  

« prev     ^ index     » next       coverage.py v7.3.2, created at 2024-09-27 15:43 +0000

1import itertools 

2import numpy as np 

3from numpy.typing import NDArray 

4from dynasor.logging_tools import logger 

5 

6 

7def get_spherical_qpoints( 

8 cell: NDArray[float], 

9 q_max: float, 

10 max_points: int = None, 

11 seed: int = 42, 

12) -> NDArray[float]: 

13 r"""Generates all q-points on the reciprocal lattice inside a given radius 

14 :attr:`q_max`. This approach is suitable if an isotropic sampling of 

15 q-space is desired. The function returns the resulting q-points in 

16 Cartesian coordinates as an ``Nx3`` array. 

17 

18 If the number of generated q-points are large, points can be removed by 

19 specifying the :attr:`max_points`. The q-points will be randomly removed in 

20 such a way that the q-points inside are roughly uniformly distributed with 

21 respect to :math:`|q|`. If the number of q-points are binned w.r.t. their 

22 norm the function would increase quadratically up until some distance P 

23 from which point the distribution would be constant. 

24 

25 Parameters 

26 ---------- 

27 cell 

28 real cell with cell vectors as rows 

29 q_max 

30 maximum norm of generated q-points 

31 max_points 

32 Optionally limit the set to __approximately__ :attr:`max_points` points 

33 by randomly removing points from a "fully populated mesh". The points 

34 are removed in such a way that for :math:`q > q_\mathrm{prune}`, the 

35 points will be radially uniformly distributed. The value of 

36 :math:`q_\mathrm{prune}` is calculated from :attr:`max_q`, 

37 :attr:`max_points`, and the shape of the cell. 

38 seed 

39 Seed used for stochastic pruning 

40 

41 """ 

42 

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

44 # The physicists reciprocal cell 

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

46 

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

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

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

50 

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

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

53 

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

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

56 

57 # Create all q-points within a sphere 

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

59 q_points = lattice_points @ rec_cell 

60 

61 # Calculate distances for pruning 

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

63 

64 # Sort distances and q-points based on distance 

65 argsort = np.argsort(q_distances) 

66 q_distances = q_distances[argsort] 

67 q_points = q_points[argsort] 

68 

69 # Prune based on distances 

70 q_points = q_points[q_distances <= q_max] 

71 q_distances = q_distances[q_distances <= q_max] 

72 

73 # Pruning based on max_points 

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

75 

76 q_vol = np.linalg.det(rec_cell) 

77 

78 q_prune = _get_prune_distance(max_points, q_max, q_vol) 

79 

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

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

82 

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

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

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

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

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

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

89 

90 rs = np.random.RandomState(seed) 

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

92 

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

94 

95 return q_points 

96 

97 

98def _get_prune_distance( 

99 max_points: int, 

100 q_max: float, 

101 q_vol: float, 

102) -> NDArray[float]: 

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

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

105 

106 If points are selected from the full mesh with probability 

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

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

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

110 

111 The general idea is as follows. 

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

113 

114 .. math: 

115 

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

117 

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

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

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

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

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

123 outward. One additional constraint is that the weighting function 

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

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

126 

127 .. math: 

128 

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

130 

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

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

133 

134 Parameters 

135 ---------- 

136 max_points 

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

138 max_q 

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

140 vol_q 

141 q-space volume for a single q-point 

142 """ 

143 

144 Q = q_max 

145 V = q_vol 

146 N = max_points 

147 

148 # Coefs 

149 a = 1.0 

150 b = -3 / 2 * Q 

151 c = 0.0 

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

153 

154 # Eq tol solve 

155 def original_eq(x): 

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

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

158 

159 # Discriminant 

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

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

162 

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

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

165 return q_max 

166 

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

168 

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

170 

171 return x