Coverage for gpaw/sphere/integrate.py: 90%
122 statements
« prev ^ index » next coverage.py v7.7.1, created at 2025-07-20 00:19 +0000
1import numpy as np
3from scipy.optimize import minimize
5from gpaw.spline import Spline
6from gpaw.lfc import LocalizedFunctionsCollection
7from gpaw.sphere.lebedev import weight_n
10def integrate_lebedev(f_nx):
11 """Integrate the function f(r) on the angular Lebedev quadrature.
13 Here, n is the quadrature index for the angular dependence of the function
14 defined on a spherical grid, while x are some arbitrary extra dimensions.
15 """
16 return 4. * np.pi * np.tensordot(weight_n, f_nx, axes=([0], [0]))
19def integrate_radial_grid(f_xg, r_g, rcut=None):
20 """Integrate the function f(r) on the radial grid.
22 Computes the integral
24 /
25 | r^2 dr f(r)
26 /
28 for the range of values r on the grid r_g (up to rcut, if specified).
29 """
30 if rcut is not None:
31 f_xg, r_g = truncate_radial_grid(f_xg, r_g, rcut)
33 # Perform actual integration using the radial trapezoidal rule
34 f_x = radial_trapz(f_xg, r_g)
36 return f_x
39def truncate_radial_grid(f_xg, r_g, rcut):
40 """Truncate the radial grid representation of a function f(r) at r=rcut.
42 If rcut is not part of the original grid, it will be added as a grid point,
43 with f(rcut) determined by linear interpolation."""
44 assert rcut > 0.
45 assert np.any(r_g >= rcut)
46 if rcut not in r_g:
47 f_gx = np.moveaxis(f_xg, -1, 0)
48 # Find the two points closest to rcut and interpolate between them
49 # to get the value at rcut
50 g1, g2 = find_two_closest_grid_points(r_g, rcut)
51 r1 = r_g[g1]
52 r2 = r_g[g2]
53 lambd = (rcut - r1) / (r2 - r1)
54 f_interpolated_x = (1 - lambd) * f_gx[g1] + lambd * f_gx[g2]
55 # Add rcut as a grid point
56 r_g = np.append(r_g, np.array([rcut]))
57 f_gx = np.append(f_gx, np.array([f_interpolated_x]), axis=0)
58 f_xg = np.moveaxis(f_gx, 0, -1)
59 # Pick out the grid points inside rcut
60 mask_g = r_g <= rcut
61 r_g = r_g[mask_g]
62 f_xg = f_xg[..., mask_g]
64 return f_xg, r_g
67def find_two_closest_grid_points(r_g, rcut):
68 """Find the two closest grid point to a specified rcut."""
69 # Find the two smallest absolute differences
70 abs_diff_g = abs(r_g - rcut)
71 ad1, ad2 = np.partition(abs_diff_g, 1)[:2]
73 # Identify the corresponding indices
74 g1 = np.where(abs_diff_g == ad1)[0][0]
75 g2 = np.where(abs_diff_g == ad2)[0][0]
77 return g1, g2
80def radial_trapz(f_xg, r_g):
81 r"""Integrate the function f(r) using the radial trapezoidal rule.
83 Linearly interpolating,
85 r - r0
86 f(r) ≃ f(r0) + ‾‾‾‾‾‾‾ (f(r1) - f(r0)) for r0 <= r <= r1
87 r1 - r0
89 the integral
91 /
92 | r^2 dr f(r)
93 /
95 can be constructed in a piecewise manner from each discretized interval
96 r_(n-1) <= r <= r_n, using:
98 r1
99 / 1
100 | r^2 dr f(r) ≃ ‾ (r1^3 f(r1) - r0^3 f(r0))
101 / 4
102 r0 r1^3 - r0^3
103 + ‾‾‾‾‾‾‾‾‾‾‾ (r1 f(r0) - r0 f(r1))
104 12(r1 - r0)
105 """
106 assert np.all(r_g >= 0.)
107 assert f_xg.shape[-1] == len(r_g)
109 # Start and end of each discretized interval
110 r0_g = r_g[:-1]
111 r1_g = r_g[1:]
112 f0_xg = f_xg[..., :-1]
113 f1_xg = f_xg[..., 1:]
114 assert np.all(r1_g - r0_g > 0.), \
115 'Please give the radial grid in ascending order'
117 # Linearly interpolate f(r) between r0 and r1 and integrate r^2 f(r)
118 # in this area
119 integrand_xg = (r1_g**3. * f1_xg - r0_g**3. * f0_xg) / 4.
120 integrand_xg += (r1_g**3. - r0_g**3.) * (r1_g * f0_xg - r0_g * f1_xg)\
121 / (12. * (r1_g - r0_g))
123 # Sum over the discretized integration intervals
124 return np.sum(integrand_xg, axis=-1)
127def radial_truncation_function(r_g, rcut, drcut=None, lambd=None):
128 r"""Generate smooth radial truncation function θ(r<rc).
130 The function is generated to interpolate smoothly between the values
132 ( 1 for r <= rc - Δrc/2
133 θ(r) = < λ for r = rc
134 ( 0 for r >= rc + Δrc/2
136 In the interpolation region, rc - Δrc/2 < r < rc + Δrc/2, the nonanalytic
137 smooth function
139 ( exp(-1/x) for x > 0
140 f(x) = <
141 ( 0 for x <= 0
143 is used to define θ(r), in order for all derivatives to be continous:
145 f(1/2-[r-rc]/Δrc)
146 θ(r) = ‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾
147 f(1/2-[r-rc]/Δrc) + (1-λ)f(1/2+[r-rc]/Δrc)/λ
149 Unless given as an input, we choose 0 < λ < 1 to conserve the spherical
150 integration volume, 4π rc^3/3.
151 """
152 assert np.all(r_g >= 0.)
153 if drcut is None:
154 # As a default, define Δrc to match twice the grid sampling around the
155 # cutoff
156 g1, g2 = find_two_closest_grid_points(r_g, rcut)
157 drcut = 2 * abs(r_g[g2] - r_g[g1])
158 assert rcut > 0. and drcut > 0. and rcut - drcut / 2. >= 0.
159 if lambd is None:
160 lambd = find_volume_conserving_lambd(rcut, drcut, r_g)
161 assert 0. < lambd and lambd < 1.
162 assert np.any(r_g >= rcut + drcut / 2.)
164 def f(x):
165 out = np.zeros_like(x)
166 out[x > 0] = np.exp(-1 / x[x > 0])
167 return out
169 # Create array of ones inside rc + Δrc/2
170 theta_g = np.ones_like(r_g)
171 theta_g[r_g >= rcut + drcut / 2.] = 0.
173 # Add smooth truncation
174 gmask = np.logical_and(rcut - drcut / 2. < r_g, r_g < rcut + drcut / 2.)
175 tr_g = r_g[gmask]
176 theta_g[gmask] = f(1 / 2. - (tr_g - rcut) / drcut) \
177 / (f(1 / 2. - (tr_g - rcut) / drcut)
178 + (1 - lambd) * f(1 / 2. + (tr_g - rcut) / drcut) / lambd)
180 return theta_g
183def find_volume_conserving_lambd(rcut, drcut, r_g=None):
184 r"""Determine the scaling factor λ to conserve the spherical volume.
186 For a given rc and drc, λ is determined to make θ(r) numerically satisfy:
187 ∞
188 /
189 4π | r^2 dr θ(r) = 4π rc^3/3
190 /
191 0
192 """
193 if r_g is None:
194 r_g = _uniform_radial_grid(rcut, drcut)
195 ref = 4 * np.pi * rcut**3. / 3.
197 def integration_volume_error(lambd):
198 theta_g = radial_truncation_function(r_g, rcut, drcut, lambd)
199 vol = 4 * np.pi * radial_trapz(theta_g, r_g)
200 return (vol - ref)**2.
202 opt_result = minimize(integration_volume_error,
203 1 / 2., # start guess
204 bounds=[(1e-8, 1 - 1e-8)],
205 method='L-BFGS-B',
206 options={'ftol': 1e-12})
207 if opt_result.success:
208 lambd = opt_result.x[0]
209 else:
210 raise Exception('Could not find an appropriate truncation scaling λ',
211 opt_result.message)
213 return lambd
216def periodic_truncation_function(gd, spos_c, rcut, drcut=None, lambd=None):
217 r"""Generate periodic images of the spherical truncation function θ.
219 The smooth radial truncation function θ(r<rc) is used to define a
220 smoothly truncated sphere of radius rcut, centered at the scaled
221 position spos_c. The sphere is periodically repeated on the real-space grid
222 described by gd.
223 """
224 # Generate a spherical truncation function collection with a single
225 # truncation function
226 stfc = spherical_truncation_function_collection(
227 gd, spos_ac=[spos_c],
228 rcut_aj=[[rcut]], drcut=drcut, lambd_aj=[[lambd]])
230 # Evaluate the spherical truncation function (with its periodic images) on
231 # the real-space grid
232 theta_R = gd.zeros(dtype=float)
233 stfc.add(theta_R)
235 return theta_R
238def spherical_truncation_function_collection(gd, spos_ac, rcut_aj,
239 drcut=None, lambd_aj=None,
240 kd=None, dtype=float):
241 """Generate a collection of spherical truncation functions θ(|r-r_a|<rc_aj)
243 Generates a LocalizedFunctionsCollection with radial truncation functions
244 θ(r<rc_aj), centered at the scaled positions spos_ac.
246 See radial_truncation_function() for the functional form of θ(r<rc).
247 """
248 if drcut is None:
249 drcut = default_spherical_drcut(gd)
250 if lambd_aj is None:
251 # Match the nested list structure of rcut_aj and determine the actual
252 # lambda values later
253 lambd_aj = [[None] * len(rcut_j) for rcut_j in rcut_aj]
255 # Generate splines for each atomic site and radial truncation function
256 spline_aj = []
257 for spos_c, rcut_j, lambd_j in zip(spos_ac, rcut_aj, lambd_aj):
258 spline_j = []
259 for rcut, lambd in zip(rcut_j, lambd_j):
260 spline_j.append(radial_truncation_function_spline(
261 rcut, drcut, lambd))
262 spline_aj.append(spline_j)
264 # Generate the spherical truncation function collection (stfc)
265 stfc = LocalizedFunctionsCollection(gd, spline_aj, kd=kd, dtype=dtype)
266 stfc.set_positions(spos_ac)
268 return stfc
271def radial_truncation_function_spline(rcut, drcut, lambd=None):
272 """Generate spline representation of the radial truncation function θ(r<rc)
273 """
274 if lambd is None:
275 lambd = find_volume_conserving_lambd(rcut, drcut)
277 # Lay out truncation function on a radial grid and generate spline
278 r_g = _uniform_radial_grid(rcut, drcut)
279 theta_g = radial_truncation_function(r_g, rcut, drcut, lambd)
280 spline = Spline.from_data(
281 l=0, rmax=max(r_g),
282 # The input f_g is the expansion coefficient of the requested
283 # spherical harmonic for the function in question.
284 # For l=0, Y = 1/sqrt(4π):
285 f_g=np.sqrt(4 * np.pi) * theta_g,
286 )
287 return spline
290def default_spherical_drcut(gd):
291 """Define default width for the spherical truncation function."""
292 # Find the side-length corresponding to a cubic grid volume element
293 length = gd.dv**(1. / 3.)
294 # Use this side-length as the default
295 return length
298def _uniform_radial_grid(rcut, drcut):
299 return np.linspace(0., rcut + 2 * drcut, 251)
302if __name__ == '__main__':
303 import matplotlib.pyplot as plt
305 r_g = np.linspace(0., 4., 200)
307 plt.subplot(1, 2, 1)
308 plt.plot(r_g, radial_truncation_function(r_g, 1.0, 1.0))
309 plt.plot(r_g, radial_truncation_function(r_g, 2.0, 1.0))
310 plt.plot(r_g, radial_truncation_function(r_g, 3.0, 1.0))
311 plt.subplot(1, 2, 2)
312 plt.plot(r_g, radial_truncation_function(r_g, 1.0, 2.0))
313 plt.plot(r_g, radial_truncation_function(r_g, 2.0, 2.0))
314 plt.plot(r_g, radial_truncation_function(r_g, 3.0, 2.0))
315 plt.show()