Coverage for gpaw/atom/atompaw.py: 98%
252 statements
« prev ^ index » next coverage.py v7.7.1, created at 2025-07-20 00:19 +0000
1from math import pi, sqrt
3import numpy as np
4from ase.atoms import Atoms
5from scipy.linalg import eigh
7from gpaw.calculator import GPAW
8from gpaw.wavefunctions.base import WaveFunctions
9from gpaw.atom.radialgd import EquidistantRadialGridDescriptor
10from gpaw.utilities import unpack_hermitian
11from gpaw.occupations import OccupationNumberCalculator
12import gpaw.mpi as mpi
15class MakeWaveFunctions:
16 name = 'atompaw'
17 interpolation = 9
18 force_complex_dtype = False
20 def __init__(self, gd):
21 self.gd = gd
23 def __call__(self, paw, gd, *args, **kwargs):
24 return AtomWaveFunctions(self.gd, *args, **kwargs)
27class AtomWaveFunctionsArray:
28 def __init__(self, psit_nG):
29 self.array = psit_nG
32class AtomWaveFunctions(WaveFunctions):
33 mode = 'atompaw'
35 def initialize(self, density, hamiltonian, spos_ac):
36 setup = self.setups[0]
37 bf = AtomBasisFunctions(self.gd, setup.basis_functions_J)
38 density.initialize_from_atomic_densities(bf)
39 hamiltonian.update(density)
40 return 0, 0
42 def add_to_density_from_k_point(self, nt_sG, kpt):
43 nt_sG[kpt.s] += np.dot(kpt.f_n / 4 / pi, kpt.psit_nG**2)
45 def summary(self, log):
46 log('Mode: Spherically symmetric atomic solver')
49class AtomPoissonSolver:
50 def get_description(self):
51 return 'Radial equidistant'
53 def set_grid_descriptor(self, gd):
54 self.gd = gd
55 self.relax_method = 0
56 self.nn = 1
58 def initialize(self):
59 pass
61 def get_stencil(self):
62 return 'Exact'
64 def solve(self, vHt_g, rhot_g, charge=0, timer=None):
65 r = self.gd.r_g
66 dp = rhot_g * r * self.gd.dr_g
67 dq = dp * r
68 p = np.add.accumulate(dp[::-1])[::-1]
69 q = np.add.accumulate(dq[::-1])[::-1]
70 vHt_g[:] = 4 * pi * (p - 0.5 * dp - (q - 0.5 * dq - q[0]) / r)
71 return 1
74class AtomEigensolver:
75 def __init__(self, gd, f_sln):
76 self.gd = gd
77 self.f_sln = f_sln
78 self.error = 0.0
79 self.initialized = False
81 def reset(self):
82 self.initialized = False
84 def initialize(self, wfs):
85 r = self.gd.r_g
86 h = r[0]
87 N = len(r)
88 lmax = len(self.f_sln[0]) - 1
90 self.T_l = [np.eye(N) * (1.0 / h**2)]
91 self.T_l[0].flat[1::N + 1] = -0.5 / h**2
92 self.T_l[0].flat[N::N + 1] = -0.5 / h**2
93 for l in range(1, lmax + 1):
94 self.T_l.append(self.T_l[0] + np.diag(l * (l + 1) / 2.0 / r**2))
96 self.S_l = [np.eye(N) for l in range(lmax + 1)]
97 setup = wfs.setups[0]
98 self.pt_j = np.array([[pt(x) * x**l for x in r]
99 for pt, l in zip(setup.pt_j, setup.l_j)])
101 dS_ii = setup.dO_ii
102 i1 = 0
103 for pt1, l1 in zip(self.pt_j, setup.l_j):
104 i2 = 0
105 for pt2, l2 in zip(self.pt_j, setup.l_j):
106 if l1 == l2 and l1 <= lmax:
107 self.S_l[l1] += (np.outer(pt1 * r, pt2 * r) *
108 h * dS_ii[i1, i2])
109 i2 += 2 * l2 + 1
110 i1 += 2 * l1 + 1
112 for kpt in wfs.kpt_u:
113 kpt.eps_n = np.empty(wfs.bd.nbands)
114 kpt.psit = AtomWaveFunctionsArray(self.gd.empty(wfs.bd.nbands))
115 kpt.projections = {0: np.zeros((wfs.bd.nbands, len(dS_ii)))}
117 self.initialized = True
119 def iterate(self, hamiltonian, wfs):
120 if not self.initialized:
121 self.initialize(wfs)
123 r = self.gd.r_g
124 h = r[0]
125 N = len(r)
126 lmax = len(self.f_sln[0]) - 1
127 setup = wfs.setups[0]
129 e_n = np.zeros(N)
131 for s in range(wfs.nspins):
132 dH_ii = unpack_hermitian(hamiltonian.dH_asp[0][s])
133 kpt = wfs.kpt_u[s]
134 N1 = 0
135 for l in range(lmax + 1):
136 H = self.T_l[l] + np.diag(hamiltonian.vt_sg[s])
137 i1 = 0
138 for pt1, l1 in zip(self.pt_j, setup.l_j):
139 i2 = 0
140 for pt2, l2 in zip(self.pt_j, setup.l_j):
141 if l1 == l2 == l:
142 H += (h * dH_ii[i1, i2] *
143 np.outer(pt1 * r, pt2 * r))
144 i2 += 2 * l2 + 1
145 i1 += 2 * l1 + 1
146 e_n, H = eigh(H, self.S_l[l].copy())
148 for n in range(len(self.f_sln[s][l])):
149 N2 = N1 + 2 * l + 1
150 kpt.eps_n[N1:N2] = e_n[n]
151 kpt.psit_nG[N1:N2] = H[:, n] / r / sqrt(h)
152 i1 = 0
153 for pt, ll in zip(self.pt_j, setup.l_j):
154 i2 = i1 + 2 * ll + 1
155 if ll == l:
156 P = np.dot(kpt.psit_nG[N1], pt * r**2) * h
157 kpt.P_ani[0][N1:N2, i1:i2] = P * np.eye(2 * l + 1)
158 i1 = i2
159 N1 = N2
162class AtomLocalizedFunctionsCollection:
163 def __init__(self, gd, spline_aj):
164 self.gd = gd
165 spline = spline_aj[0][0]
166 self.b_g = np.array([spline(r) for r in gd.r_g]) / sqrt(4 * pi)
167 self.nfunctions = sum(2 * spline.get_angular_momentum_number() + 1
168 for spline in spline_aj[0])
170 def set_positions(self, spos_ac, atom_partition=None):
171 pass
173 def add(self, a_xG, c_axi=1.0, q=-1):
174 assert q == -1
175 if isinstance(c_axi, float):
176 a_xG += c_axi * self.b_g
177 else:
178 a_xG += c_axi[0][0] * self.b_g
180 def integrate(self, a_g, c_ai, q=-1):
181 assert a_g.ndim == 1
182 assert q == -1
183 c_ai[0][0] = self.gd.integrate(a_g, self.b_g)
184 c_ai[0][1:] = 0.0
186 def dict(self):
187 return {0: np.empty(self.nfunctions)}
190class AtomBasisFunctions:
191 def __init__(self, gd, bfs_J):
192 self.gd = gd
193 self.bl_J = []
194 self.Mmax = 0
196 for bf in bfs_J:
197 l = bf.get_angular_momentum_number()
198 self.bl_J.append((np.array([bf(x) * x**l for x in gd.r_g]), l))
199 self.Mmax += 2 * l + 1
200 self.atom_indices = [0]
201 self.my_atom_indices = [0]
203 def set_positions(self, spos_ac):
204 pass
206 def add_to_density(self, nt_sG, f_asi):
207 i = 0
208 for b_g, l in self.bl_J:
209 nt_sG += f_asi[0][:, i:i + 1] * (2 * l + 1) / 4 / pi * b_g**2
210 i += 2 * l + 1
213class AtomGridDescriptor(EquidistantRadialGridDescriptor):
214 def __init__(self, h, rcut):
215 ng = int(float(rcut) / h + 0.5) - 1
216 rcut = ng * h
217 super().__init__(h, ng, h0=h)
218 self.sdisp_cd = np.empty((3, 2))
219 self.comm = mpi.serial_comm
220 self.pbc_c = np.zeros(3, bool)
221 self.cell_cv = np.eye(3) * rcut
222 self.N_c = np.ones(3, dtype=int) * 2 * ng
223 self.h_cv = self.cell_cv / self.N_c
224 self.dv = (rcut / 2 / ng)**3
225 self.orthogonal = False
226 self.parsize_c = (1, 1, 1)
228 def get_ranks_from_positions(self, spos_ac):
229 return np.array([0])
231 def refine(self):
232 return self
234 def get_lfc(self, gd, spline_aj):
235 return AtomLocalizedFunctionsCollection(gd, spline_aj)
237 def integrate(self, a_xg, b_xg=None, global_integral=True):
238 """Integrate function(s) in array over domain."""
239 if b_xg is None:
240 return np.dot(a_xg, self.dv_g)
241 else:
242 return np.dot(a_xg * b_xg, self.dv_g)
244 def calculate_dipole_moment(self, rhot_g):
245 return np.zeros(3)
247 def symmetrize(self, a_g, op_scc, ft_sc=None):
248 pass
250 def get_grid_spacings(self):
251 return self.h_cv.diagonal()
253 def get_size_of_global_array(self, pad=None):
254 return np.array(len(self.N_c))
256 def new_descriptor(self, *args, **kwargs):
257 return self
259 def get_processor_position_from_rank(self, rank):
260 return (0, 0, 0)
263class AtomOccupations(OccupationNumberCalculator):
264 extrapolate_factor = 0.0
266 def __init__(self, f_sln):
267 self.f_sln = f_sln
268 super().__init__()
270 def _calculate(self,
271 nelectrons,
272 eig_qn,
273 weight_q,
274 f_qn,
275 fermi_level_guess,
276 fix_fermi_level=False):
277 for s, f_n in enumerate(f_qn):
278 n1 = 0
279 for l, f0_n in enumerate(self.f_sln[s]):
280 for f in f0_n:
281 n2 = n1 + 2 * l + 1
282 f_n[n1:n2] = f / (2 * l + 1) / 2
283 n1 = n2
285 return np.inf, 0.0
288class AtomPAW(GPAW):
289 def __init__(self, symbol, f_sln, h=0.05, rcut=10.0, **kwargs):
290 assert len(f_sln) in [1, 2]
291 self.symbol = symbol
293 gd = AtomGridDescriptor(h, rcut)
294 super().__init__(mode=MakeWaveFunctions(gd),
295 eigensolver=AtomEigensolver(gd, f_sln),
296 poissonsolver=AtomPoissonSolver(),
297 nbands=sum([(2 * l + 1) * len(f_n)
298 for l, f_n in enumerate(f_sln[0])]),
299 communicator=mpi.serial_comm,
300 parallel=dict(augment_grids=False),
301 occupations=AtomOccupations(f_sln),
302 **kwargs)
303 # Initialize function will raise an error unless we set a (bogus) cell
304 self.initialize(Atoms(symbol, calculator=self,
305 cell=np.eye(3)))
306 self.density.charge_eps = 1e-3
307 self.calculate(system_changes=['positions'])
309 def dry_run(self):
310 pass
312 def state_iter(self):
313 """Yield the tuples (l, n, f, eps, psit_G) of states.
315 Skips degenerate states."""
316 f_sln = self.wfs.occupations.f_sln
317 assert len(f_sln) == 1, 'Not yet implemented with more spins'
318 f_ln = f_sln[0]
319 kpt = self.wfs.kpt_u[0]
321 band = 0
322 for l, f_n in enumerate(f_ln):
323 for n, f in enumerate(f_n):
324 psit_G = kpt.psit_nG[band]
325 eps = kpt.eps_n[band]
326 yield l, n, f, eps, psit_G
327 band += 2 * l + 1
329 def extract_basis_functions(self, n_j, l_j, basis_name='atompaw.sz'):
330 """Create BasisFunctions object with pseudo wave functions."""
331 from gpaw.basis_data import Basis, BasisFunction
332 assert self.wfs.nspins == 1
333 d = self.wfs.gd.r_g[0]
334 ng = self.wfs.gd.N + 1
335 rgd = EquidistantRadialGridDescriptor(d, ng)
337 bf_j = []
338 for l, n, f, eps, psit_G in self.state_iter():
339 n = [N for N, L in zip(n_j, l_j) if L == l][n]
340 bf_g = rgd.empty()
341 bf_g[0] = 0.0
342 bf_g[1:] = psit_G
343 bf_g *= np.sign(psit_G[-1])
345 # If there's no node at zero, we shouldn't set bf_g to zero
346 # We'll make an ugly hack
347 if abs(bf_g[1]) > 3.0 * abs(bf_g[2] - bf_g[1]):
348 bf_g[0] = bf_g[1]
349 bf = BasisFunction(n, l, self.wfs.gd.r_g[-1], bf_g,
350 f'{n}{"spdfgh"[l]} e={eps:.3f} f={f:.3f}')
351 bf_j.append(bf)
353 return Basis(self.symbol, basis_name, rgd=rgd,
354 generatordata='AtomPAW', bf_j=bf_j)