Coverage for local_installation/dynasor/qpoints/spherical_qpoints.py: 95%
47 statements
« prev ^ index » next coverage.py v7.3.2, created at 2025-04-16 06:13 +0000
1import itertools
2import numpy as np
3from numpy.typing import NDArray
4from dynasor.logging_tools import logger
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.
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.
25 Parameters
26 ----------
27 cell
28 real cell with cell vectors as rows
29 q_max
30 maximum norm of generated q-points (in units of rad/Å, i.e. including factor of 2pi)
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
41 """
43 # inv(A.T) == inv(A).T
44 # The physicists reciprocal cell
45 rec_cell = np.linalg.inv(cell.T) * 2 * np.pi
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
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)
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)
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
61 # Calculate distances for pruning
62 q_distances = np.linalg.norm(q_points, axis=1) # Find distances
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]
69 # Prune based on distances
70 q_points = q_points[q_distances <= q_max]
71 q_distances = q_distances[q_distances <= q_max]
73 # Pruning based on max_points
74 if max_points is not None and max_points < len(q_points):
76 q_vol = np.linalg.det(rec_cell)
78 q_prune = _get_prune_distance(max_points, q_max, q_vol)
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}')
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
90 rs = np.random.RandomState(seed)
91 q_points = q_points[p > rs.rand(len(q_points))]
93 logger.info(f'Pruned from {len(q_distances)} q-points to {len(q_points)}')
95 return q_points
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.
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}`).
111 The general idea is as follows.
112 We know that the number of q-points inside a radius :math:`Q` is given by
114 .. math:
116 n = v^{-1} \int_0^Q dq 4 \pi q^2 = v^{-1} 4/3 \pi Q^3
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`
127 .. math:
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).
131 This results in a `cubic equation <https://en.wikipedia.org/wiki/Cubic_equation>`_
132 for :math:`P`, which is solved by this function.
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 """
144 Q = q_max
145 V = q_vol
146 N = max_points
148 # Coefs
149 a = 1.0
150 b = -3 / 2 * Q
151 c = 0.0
152 d = 3 / 2 * V * N / (4 * np.pi)
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
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)
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
167 x = Q * (np.cos(1 / 3 * np.arccos(1 - 4 * d / Q**3) - 2 * np.pi / 3) + 0.5)
169 assert np.isclose(original_eq(x), 0), original_eq(x)
171 return x