Coverage for dynasor / qpoints / spherical_qpoints.py: 86%
62 statements
« prev ^ index » next coverage.py v7.13.4, created at 2026-03-16 12:31 +0000
1from typing import Optional
2import itertools
3import numpy as np
4from numpy.typing import NDArray
5from dynasor.logging_tools import logger
8def get_spherical_qpoints(
9 cell: NDArray[float],
10 q_max: float,
11 q_min: Optional[float] = 0.0,
12 max_points: Optional[int] = None,
13 seed: Optional[int] = 42,
14 chunk_size: Optional[int] = 100000,
15) -> NDArray[float]:
16 r"""Generates all q-points on the reciprocal lattice inside a given radius
17 :attr:`q_max`. This approach is suitable if an isotropic sampling of
18 q-space is desired. The function returns the resulting q-points in
19 Cartesian coordinates as an ``Nx3`` array.
21 If the number of generated q-points are large, points can be removed by
22 specifying :attr:`max_points`. The q-points will be randomly removed in
23 such a way that the q-points inside are roughly uniformly distributed with
24 respect to :math:`|q|`. If the number of q-points are binned w.r.t. their
25 norm the function would increase quadratically up until some distance P
26 from which point the distribution would be constant.
28 Parameters
29 ----------
30 cell
31 Real cell with cell vectors as rows.
32 q_max
33 Maximum norm of generated q-points
34 (in units of rad/Å, i.e. including factor of :math:`2\pi`).
35 q_min
36 Minimum norm of generated q-points
37 (in units of rad/Å, i.e. including factor of :math:`2\pi`).
38 max_points
39 Optionally limit the set to __approximately__ :attr:`max_points` points
40 by randomly removing points from a "fully populated mesh". The points
41 are removed in such a way that for :math:`q > q_\mathrm{prune}`, the
42 points will be radially uniformly distributed. The value of
43 :math:`q_\mathrm{prune}` is calculated from :attr:`q_max`,
44 :attr:`max_points`, and the shape of the cell.
45 seed
46 Seed used for stochastic pruning.
47 chunk_size
48 Number of q-points to process at once (controls memory use).
49 """
51 if q_min >= q_max: 51 ↛ 52line 51 didn't jump to line 52 because the condition on line 51 was never true
52 raise ValueError('q_min must be smaller than q_max.')
53 if q_min < 0: 53 ↛ 54line 53 didn't jump to line 54 because the condition on line 53 was never true
54 raise ValueError('q_min can not be negative.')
55 if q_max < 0: 55 ↛ 56line 55 didn't jump to line 56 because the condition on line 55 was never true
56 raise ValueError('q_max can not be negative.')
58 # inv(A.T) == inv(A).T
59 # The physicists reciprocal cell
60 cell = cell.astype(np.float64)
61 rec_cell = np.linalg.inv(cell.T) * 2 * np.pi
63 # We want to find all points on the lattice defined by the reciprocal cell
64 # such that all points within q_max are in this set
65 inv_rec_cell = np.linalg.inv(rec_cell.T) # cell / 2pi
67 # h is the height of the rec_cell perpendicular to the other two vectors
68 h = 1 / np.linalg.norm(inv_rec_cell, axis=1)
70 # If a q_point has a coordinate larger than this number it must be further away than q_max
71 N = np.ceil(q_max / h).astype(int)
73 # Generate q-points on the grid with in chunks to keep memory consumption low
74 iterator = itertools.product(*[range(-n, n+1) for n in N])
75 q_points = []
77 while True:
79 # Take a chunk of points
80 lattice_points_chunk = list(itertools.islice(iterator, chunk_size))
81 if not lattice_points_chunk:
82 break
83 lattice_points_chunk = np.array(lattice_points_chunk)
84 q_points_chunk = lattice_points_chunk @ rec_cell # (chunk, 3)
86 # Filter by norm
87 q_distances = np.linalg.norm(q_points_chunk, axis=1)
88 mask = np.logical_and(q_distances >= q_min, q_distances <= q_max)
89 if np.any(mask):
90 q_points.append(q_points_chunk[mask])
92 if len(q_points) == 0: 92 ↛ 93line 92 didn't jump to line 93 because the condition on line 92 was never true
93 return np.empty((0, 3))
94 q_points = np.vstack(q_points)
96 # Pruning based on max_points
97 if max_points is not None and max_points < len(q_points):
99 q_vol = np.linalg.det(rec_cell)
100 q_prune = _get_prune_distance(max_points, q_min, q_max, q_vol)
102 if q_prune < q_max: 102 ↛ 118line 102 didn't jump to line 118 because the condition on line 102 was always true
103 logger.info(f'Pruning q-points from the range {q_prune:.3} < |q| < {q_max}')
105 # Keep point with probability min(1, (q_prune/|q|)^2) ->
106 # aim for an equal number of points per equally thick "onion peel"
107 # to get equal number of points per radial unit.
108 q_distances = np.linalg.norm(q_points, axis=1)
109 p = np.ones(len(q_points))
110 p = np.divide(q_prune**2, q_distances**2, out=np.ones_like(q_distances),
111 where=np.logical_not(np.isclose(q_distances, 0)))
113 rs = np.random.RandomState(seed)
114 q_points = q_points[p > rs.rand(len(q_points))]
116 logger.info(f'Pruned from {len(q_distances)} q-points to {len(q_points)}')
118 return q_points
121def _get_prune_distance(
122 max_points: int,
123 q_min: float,
124 q_max: float,
125 q_vol: float,
126) -> NDArray[float]:
127 r"""Determine distance in q-space beyond which to prune
128 the q-point mesh to achieve near-isotropic sampling of q-space.
130 If points are selected from the full mesh with probability
131 :math:`\min(1, (q_\mathrm{prune} / |q|)^2)`, q-space will
132 on average be sampled with an equal number of points per radial unit
133 (for :math:`q > q_\mathrm{prune}`).
135 The general idea is as follows.
136 We know that the number of q-points inside a radius :math:`Q` is given by
138 .. math:
140 n = v^{-1} \int_0^Q dq 4 \pi q^2 = v^{-1} 4/3 \pi Q^3
142 where :math:`v` is the volume of one q-point. Now we want to find
143 a distance :math:`P` such that if all points outside this radius
144 are weighted by the function :math:`w(q)` the total number of
145 q-points will equal the target :math:`N` (:attr:`max_points`)
146 while the number of q-points increases linearly from :math:`P`
147 outward. One additional constraint is that the weighting function
148 must be 1 at :math:`P`. The weighting function which accomplishes
149 this is :math:`w(q)=P^2/q^2`
151 .. math:
153 N = v^{-1} \left( \int_0^P 4 \pi q^2 + \int_P^Q 4 \pi q^2 P^2 / q^2 dq \right).
155 This results in a `cubic equation <https://en.wikipedia.org/wiki/Cubic_equation>`_
156 for :math:`P`, which is solved by this function.
158 Parameters
159 ----------
160 max_points
161 Maximum number of resulting q-points; :math:`N` below.
162 q_min
163 Minimum q-value in the resulting q-point set.
164 q_max
165 Maximum q-value in the resulting q-point set; :math:`Q` below.
166 vol_q
167 q-space volume for a single q-point.
168 """
170 Q = q_max
171 V = q_vol
172 N = max_points
174 # Coefs
175 a = 1.0
176 b = -3 / 2 * Q
177 c = 0.0
178 d = 3 / 2 * V * N / (4 * np.pi) + 0.5 * q_min**3
180 # Eq tol solve
181 def original_eq(x):
182 return a * x**3 + b * x**2 + c * x + d
183 # original_eq = lambda x: a * x**3 + b * x**2 + c * x + d
185 # Discriminant
186 p = (3 * a * c - b**2) / (3 * a**2)
187 q = (2 * b**3 - 9 * a * b * c + 27 * a**2 * d) / (27 * a**3)
189 D_t = - (4 * p**3 + 27 * q**2)
190 if D_t < 0: 190 ↛ 191line 190 didn't jump to line 191 because the condition on line 190 was never true
191 return q_max
193 x = Q * (np.cos(1 / 3 * np.arccos(1 - 4 * d / Q**3) - 2 * np.pi / 3) + 0.5)
195 assert np.isclose(original_eq(x), 0), original_eq(x)
197 return x