Coverage for gpaw/new/pw/bloechl_poisson.py: 91%
182 statements
« prev ^ index » next coverage.py v7.7.1, created at 2025-07-09 00:21 +0000
1"""Fast PAW Poisson-solver.
3See equations (25-28) in
4P. E. Blöchl: https://sci-hub.st/10.1103/PhysRevB.50.17953
5"""
6from __future__ import annotations
8from math import pi
10import numpy as np
11from ase.neighborlist import primitive_neighbor_list
12from gpaw.atom.radialgd import EquidistantRadialGridDescriptor as RGD
13from gpaw.atom.shapefunc import shape_functions
14from gpaw.core import PWArray, PWDesc
15from gpaw.core.atom_arrays import AtomArrays, AtomDistribution
16from gpaw.lcao.overlap import (FourierTransformer, LazySphericalHarmonics,
17 LazySphericalHarmonicsDerivative,
18 ManySiteOverlapCalculator,
19 TwoSiteOverlapCalculator)
20from gpaw.spline import Spline
21from scipy.special import erf
22from gpaw.new.pw.paw_poisson import PAWPoissonSolver
25def vg(r_r: np.ndarray, rc: float) -> np.ndarray:
26 v_lr = np.empty((3, len(r_r)))
27 x_r = r_r / rc
28 v_lr[0] = 4 * pi * erf(x_r)
29 v_lr[0, 0] = 8 * pi**0.5 / rc
30 v_lr[0, 1:] /= r_r[1:]
32 v_lr[1] = v_lr[0] / 3 - 8 * pi**0.5 / 3 * np.exp(-x_r**2) / rc
33 v_lr[1, 0] = 16 * pi**0.5 / 9 / rc**3
34 v_lr[1, 1:] /= r_r[1:]**2
36 v_lr[2] = (v_lr[0] / 5 -
37 8 * pi**0.5 / 5 * (1 + 2 * x_r**2 / 3) * np.exp(-x_r**2) / rc)
38 v_lr[2, 0] = 32 * pi**0.5 / 75 / rc**5
39 v_lr[2, 1:] /= r_r[1:]**4
41 return v_lr
44def tci(rcut, I_a, gtilde_Il, vhat_Il, ghat_Il):
45 transformer = FourierTransformer(rcut=rcut, N=2**10)
46 tsoc = TwoSiteOverlapCalculator(transformer)
48 msoc = ManySiteOverlapCalculator(tsoc, I_a, I_a)
49 gtilde_Ilq = msoc.transform(gtilde_Il)
50 vhat_Ilq = msoc.transform(vhat_Il)
51 ghat_Ilq = msoc.transform(ghat_Il)
52 l_Il = [[gtilde.l for gtilde in gtilde_l] for gtilde_l in gtilde_Il]
53 expansions1 = msoc.calculate_expansions(l_Il, gtilde_Ilq,
54 l_Il, vhat_Ilq)
55 expansions2 = msoc.calculate_expansions(l_Il, vhat_Ilq,
56 l_Il, ghat_Ilq)
57 return expansions1, expansions2
60class BloechlPAWPoissonSolver(PAWPoissonSolver):
61 def __init__(self,
62 pwg: PWDesc,
63 cutoff_a: np.ndarray,
64 poisson_solver,
65 relpos_ac: np.ndarray,
66 atomdist: AtomDistribution,
67 xp=np,
68 test=0):
69 super().__init__(poisson_solver)
70 self.xp = xp
71 self.pwg = pwg
72 self.pwg0 = pwg.new(comm=None) # not distributed
73 self.relpos_ac = relpos_ac
74 self.cutoff_a = np.asarray(cutoff_a)
75 self.r2 = self.cutoff_a.max() * 2.0
76 self.rcut = 7 * self.r2
77 d = 0.0051
78 rgd = RGD(d, int(self.rcut / d))
79 g0_lg = shape_functions(rgd, 'gauss', self.r2, lmax=2)
80 P = 25
81 if test:
82 ghat_al = []
83 for r1 in cutoff_a:
84 g_lg = shape_functions(rgd, 'gauss', r1, lmax=2)
85 ghat_al.append(
86 [rgd.spline(g_g, l=l, points=P)
87 for l, g_g in enumerate(g_lg)])
88 else:
89 ghat_l = [rgd.spline(g_g, l=l, points=P)
90 for l, g_g in enumerate(g0_lg)]
91 ghat_al = [ghat_l] * len(self.cutoff_a)
92 cache: dict[float, tuple[int, list[Spline], list[Spline]]] = {}
93 gtilde_Il = []
94 ghat_Il = []
95 vhat_Il = []
96 vhat_al = []
97 self.I_a = []
98 for r1 in cutoff_a:
99 if r1 in cache:
100 I, gtilde_l, vhat_l = cache[r1]
101 else:
102 g_lg = shape_functions(rgd, 'gauss', r1, lmax=2)
103 gtilde_l = [rgd.spline(g_g, l=l, points=P)
104 for l, g_g in enumerate(g_lg)]
105 v_lg = vg(rgd.r_g, r1) - vg(rgd.r_g, self.r2)
106 if test:
107 vhat_l = [rgd.spline(v_g * 0, l=l, points=P)
108 for l, v_g in enumerate(v_lg)]
109 else:
110 vhat_l = [rgd.spline(v_g, l=l, points=P)
111 for l, v_g in enumerate(v_lg)]
112 I = len(cache)
113 cache[r1] = I, gtilde_l, vhat_l
114 gtilde_Il.append(gtilde_l)
115 vhat_Il.append(vhat_l)
116 ghat_Il.append(ghat_l)
117 self.I_a.append(I)
118 vhat_al.append(vhat_l)
120 self.ghat_aLg = pwg.atom_centered_functions(
121 ghat_al, relpos_ac, atomdist=atomdist, xp=xp)
122 self.vhat_aLg = pwg.atom_centered_functions(
123 vhat_al, relpos_ac, atomdist=atomdist, xp=xp)
125 self.expansions = tci(self.rcut, self.I_a, gtilde_Il, vhat_Il, ghat_Il)
127 self._neighbors = None
128 self._force_av: np.ndarray | None = None
129 self._stress_vv: np.ndarray | None = None
131 def get_neighbors(self):
132 if self._neighbors is None:
133 pw = self.pwg
134 i, j, d, D = primitive_neighbor_list(
135 'ijdD', pw.pbc, pw.cell, self.relpos_ac,
136 2 * self.rcut,
137 use_scaled_positions=True,
138 self_interaction=True)
139 comm = self.pwg.comm
140 x = slice(comm.rank, None, comm.size)
141 self._neighbors = i[x], j[x], d[x], D[x]
142 return self._neighbors
144 def dipole_layer_correction(self):
145 return self.poisson_solver.dipole_layer_correction()
147 def move(self, relpos_ac, atomdist):
148 self.relpos_ac = relpos_ac
149 self.ghat_aLg.move(relpos_ac, atomdist)
150 self.vhat_aLg.move(relpos_ac, atomdist)
151 self._neighbors = None
152 self._force_av = None
153 self._stress_vv = None
155 def solve(self,
156 nt0_g: PWArray,
157 Q_aL: AtomArrays,
158 vt0_g: PWArray,
159 vHt_g: PWArray | None = None):
160 nt_g = self.pwg.empty(xp=self.xp)
161 nt_g.scatter_from(nt0_g)
162 charge_g = nt_g.copy()
163 self.ghat_aLg.add_to(charge_g, Q_aL)
164 pwg = self.pwg
165 comm = pwg.comm
167 if vHt_g is None:
168 vHt_g = pwg.empty(xp=self.xp)
170 e_coulomb1 = self.poisson_solver.solve(vHt_g, charge_g)
172 vhat_g = pwg.empty(xp=self.xp) # MYPY
173 vhat_g.data[:] = 0.0 # MYPY
175 self.vhat_aLg.add_to(vhat_g, Q_aL)
176 vhat0_g = vhat_g.gather()
177 if comm.rank == 0:
178 vt0_g.data += vhat0_g.data
179 e_coulomb2 = vhat0_g.integrate(nt0_g)
180 else:
181 e_coulomb2 = 0.0
183 vHt_g.data[0] = -vhat_g.data[0]
184 V_aL = self.ghat_aLg.integrate(vHt_g)
185 self.vhat_aLg.integrate(nt_g, V_aL, add_to=True)
187 Q_all_aL = Q_aL.gather(broadcast=True)
188 dV_all_aL = Q_all_aL.new()
189 dV_all_aL.data[:] = 0.0
191 e_coulomb3 = 0.0
192 for a1, a2, d, d_v in zip(*self.get_neighbors()):
193 rlY_lm = LazySphericalHarmonics(d_v)
194 ex1, ex2 = self.expansions
195 I1 = self.I_a[a1]
196 I2 = self.I_a[a2]
197 v_LL = self.xp.asarray(ex1.tsoe_II[I1, I2].evaluate(d, rlY_lm) +
198 ex2.tsoe_II[I1, I2].evaluate(d, rlY_lm))
199 vQ2_L = v_LL @ Q_all_aL[a2]
200 dV_all_aL[a1] += vQ2_L / 2
201 dV_all_aL[a2] += Q_all_aL[a1] @ v_LL / 2
202 e_coulomb3 -= float(Q_all_aL[a1] @ vQ2_L)
203 e_coulomb3 *= -0.5
205 comm.sum(dV_all_aL.data)
206 dV_aL = Q_aL.new()
207 dV_aL.scatter_from(dV_all_aL)
208 V_aL.data += dV_aL.data
210 vHt0_g = vHt_g.gather()
211 if comm.rank == 0:
212 vt0_g.data += vHt0_g.data
214 e_coulomb = comm.sum_scalar(e_coulomb1 / comm.size +
215 e_coulomb2 +
216 e_coulomb3)
217 return e_coulomb, vHt_g, V_aL
219 def force_contribution(self, Q_aL, vHt_g, nt_g):
220 force_av = self.xp.zeros((len(Q_aL), 3))
222 F_avL = self.ghat_aLg.derivative(vHt_g)
223 Fhat_avL = self.vhat_aLg.derivative(nt_g)
224 for a, dF_vL in F_avL.items():
225 force_av[a] += (dF_vL + Fhat_avL[a]) @ Q_aL[a]
226 pair_pot_force_av, _ = self._force_and_stress(Q_aL)
227 return force_av + pair_pot_force_av
229 def _force_and_stress(self,
230 Q_aL) -> tuple[np.ndarray, np.ndarray]:
231 if self._force_av is not None and self._stress_vv is not None:
232 return self._force_av, self._stress_vv
233 xp = self.xp
234 force_av = xp.zeros((len(Q_aL), 3))
235 stress_vv = xp.zeros((3, 3))
236 Q_aL = Q_aL.gather(broadcast=True)
237 for a1, a2, d, d_v in zip(*self.get_neighbors()):
238 if d == 0.0:
239 continue
240 rlY_lm = LazySphericalHarmonics(d_v)
241 drlYdR_lmv = LazySphericalHarmonicsDerivative(d_v)
242 ex1, ex2 = self.expansions
243 I1 = self.I_a[a1]
244 I2 = self.I_a[a2]
245 n_v = d_v / d
246 v_vLL = xp.asarray(
247 ex1.tsoe_II[I1, I2].derivative(d, n_v, rlY_lm, drlYdR_lmv) +
248 ex2.tsoe_II[I1, I2].derivative(d, n_v, rlY_lm, drlYdR_lmv))
249 f_v = (v_vLL @ Q_aL[a2]) @ Q_aL[a1] / 2
250 force_av[a1] += f_v
251 force_av[a2] -= f_v
252 stress_vv += xp.outer(xp.asarray(d_v), f_v)
253 self._force_av = force_av
254 self._stress_vv = stress_vv
255 return force_av, stress_vv
257 def stress_contribution(self, vHt_g, Q_aL):
258 _, pair_pot_stress_vv = self._force_and_stress(Q_aL)
259 return self.ghat_aLg.stress_contribution(vHt_g, Q_aL)