Coverage for gpaw/new/pw/paw_poisson.py: 89%
102 statements
« prev ^ index » next coverage.py v7.7.1, created at 2025-07-20 00:19 +0000
1"""PAW Poisson-solvers.
3Adds smooth compensation charges to the pseudo density.
4"""
5from __future__ import annotations
7import numpy as np
8from gpaw.atom.radialgd import EquidistantRadialGridDescriptor as RGD
9from gpaw.atom.shapefunc import shape_functions
10from gpaw.core import PWArray, PWDesc
11from gpaw.core.atom_arrays import AtomArrays, AtomDistribution
12from gpaw.gpu import cupy as cp
13from gpaw.setup import Setups
14from gpaw.spline import Spline
17class PAWPoissonSolver:
18 def __init__(self,
19 poisson_solver):
20 self.poisson_solver = poisson_solver
22 def dipole_layer_correction(self) -> None:
23 return self.poisson_solver.dipole_layer_correction()
25 def move(self,
26 relpos_ac: np.ndarray,
27 atomdist: AtomDistribution) -> None:
28 raise NotImplementedError
30 def solve(self,
31 nt_g: PWArray,
32 Q_aL: AtomArrays,
33 vt0_g: PWArray,
34 vHt_h: PWArray | None = None) -> tuple[float,
35 PWArray,
36 AtomArrays]:
37 raise NotImplementedError
40class SlowPAWPoissonSolver(PAWPoissonSolver):
41 """Solve Poisson-equation on very fine grid."""
42 def __init__(self,
43 pwg: PWDesc,
44 # cutoff_a,
45 setups: Setups,
46 poisson_solver,
47 relpos_ac: np.ndarray,
48 atomdist: AtomDistribution,
49 xp=np):
50 super().__init__(poisson_solver)
51 self.xp = xp
52 self.pwg = pwg
53 self.pwg0 = pwg.new(comm=None) # not distributed
54 self.pwh = poisson_solver.pw
55 self.ghat_aLh = setups.create_compensation_charges(
56 self.pwh, relpos_ac, atomdist, xp)
57 self.h_g, self.g_r = self.pwh.map_indices(self.pwg0)
58 if xp is cp:
59 self.h_g = cp.asarray(self.h_g)
60 self.g_r = [cp.asarray(g) for g in self.g_r]
62 def move(self,
63 relpos_ac: np.ndarray,
64 atomdist: AtomDistribution) -> None:
65 self.ghat_aLh.move(relpos_ac, atomdist)
67 def solve(self,
68 nt0_g: PWArray,
69 Q_aL: AtomArrays,
70 vt0_g: PWArray,
71 vHt_h: PWArray | None = None) -> tuple[float,
72 PWArray,
73 AtomArrays]:
74 charge_h = self.pwh.zeros(xp=self.xp)
75 self.ghat_aLh.add_to(charge_h, Q_aL)
76 pwg = self.pwg
78 if pwg.comm.rank == 0:
79 for rank, g in enumerate(self.g_r):
80 if rank == 0:
81 charge_h.data[self.h_g] += nt0_g.data[g]
82 else:
83 pwg.comm.send(nt0_g.data[g], rank)
84 else:
85 data = self.xp.empty(len(self.h_g), complex)
86 pwg.comm.receive(data, 0)
87 charge_h.data[self.h_g] += data
89 if vHt_h is None:
90 vHt_h = self.pwh.zeros(xp=self.xp)
92 e_coulomb = self.poisson_solver.solve(vHt_h, charge_h)
94 if pwg.comm.rank == 0:
95 for rank, g in enumerate(self.g_r):
96 if rank == 0:
97 vt0_g.data[g] += vHt_h.data[self.h_g]
98 else:
99 data = self.xp.empty(len(g), complex)
100 pwg.comm.receive(data, rank)
101 vt0_g.data[g] += data
102 else:
103 pwg.comm.send(vHt_h.data[self.h_g], 0)
105 V_aL = self.ghat_aLh.integrate(vHt_h)
107 return e_coulomb, vHt_h, V_aL
109 def force_contribution(self, Q_aL, vHt_h, nt_g):
110 force_av = self.xp.zeros((len(Q_aL), 3))
112 F_avL = self.ghat_aLh.derivative(vHt_h)
113 for a, dF_vL in F_avL.items():
114 force_av[a] += dF_vL @ Q_aL[a]
116 return force_av
118 def stress_contribution(self, vHt_h, Q_aL):
119 return self.ghat_aLh.stress_contribution(vHt_h, Q_aL)
122class SimplePAWPoissonSolver(PAWPoissonSolver):
123 """For testing only!"""
124 def __init__(self,
125 pwg: PWDesc,
126 cutoff_a,
127 poisson_solver,
128 relpos_ac: np.ndarray,
129 atomdist: AtomDistribution,
130 xp=np):
131 self.xp = xp
132 self.pwg = pwg
133 self.pwg0 = pwg.new(comm=None) # not distributed
134 self.poisson_solver = poisson_solver
135 d = 0.005
136 rgd = RGD(d, int(10.0 / d))
137 cache: dict[float, list[Spline]] = {}
138 ghat_al = []
139 for rc in cutoff_a:
140 if rc in cache:
141 ghat_l = cache[rc]
142 else:
143 g_lg = shape_functions(rgd, 'gauss', rc, lmax=2)
144 ghat_l = [rgd.spline(g_g, l=l) for l, g_g in enumerate(g_lg)]
145 cache[rc] = ghat_l
146 ghat_al.append(ghat_l)
148 self.ghat_aLg = pwg.atom_centered_functions(
149 ghat_al, relpos_ac, atomdist=atomdist, xp=xp)
151 def solve(self,
152 nt0_g: PWArray,
153 Q_aL: AtomArrays,
154 vt0_g: PWArray,
155 vHt_g: PWArray | None = None):
157 charge_g = self.pwg.empty(xp=self.xp)
158 charge_g.scatter_from(nt0_g)
159 self.ghat_aLg.add_to(charge_g, Q_aL)
160 pwg = self.pwg
161 if vHt_g is None:
162 vHt_g = pwg.empty(xp=self.xp)
163 e_coulomb = self.poisson_solver.solve(vHt_g, charge_g)
164 vHt0_g = vHt_g.gather()
165 if pwg.comm.rank == 0:
166 vt0_g.data += vHt0_g.data
167 V_aL = self.ghat_aLg.integrate(vHt_g)
168 return e_coulomb, vHt_g, V_aL
170 def force_contribution(self, Q_aL, vHt_g, nt_g):
171 force_av = self.xp.zeros((len(Q_aL), 3))
173 F_avL = self.ghat_aLg.derivative(vHt_g)
174 for a, dF_vL in F_avL.items():
175 force_av[a] += dF_vL @ Q_aL[a]
176 return force_av