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
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], 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