Coverage for gpaw/atom/optimize.py: 13%
293 statements
« prev ^ index » next coverage.py v7.7.1, created at 2025-07-14 00:18 +0000
1import multiprocessing as mp
2import os
3import re
4import sys
5import time
6from collections import defaultdict
8import numpy as np
9from scipy.optimize import differential_evolution as DE
10from ase import Atoms
11from ase.data import covalent_radii, atomic_numbers
12from ase.units import Bohr, Ha
14from gpaw import GPAW, PW, setup_paths, Mixer, ConvergenceError
15from gpaw.atom.generator2 import generate # , DatasetGenerationError
16from gpaw.atom.aeatom import AllElectronAtom
17from gpaw.atom.atompaw import AtomPAW
18from gpaw.setup import create_setup
21my_covalent_radii = covalent_radii.copy()
22for e in ['Bk', 'Cf', 'Es', 'Fm', 'Md', 'No', 'Lr']: # missing radii
23 my_covalent_radii[atomic_numbers[e]] = 1.7
25my_radii = {
26 1: 0.41, 2: 1.46, 3: 1.47, 4: 1.08, 5: 0.83, 6: 0.75, 7: 0.67,
27 8: 0.68, 9: 0.7, 10: 1.63, 11: 1.7, 12: 1.47, 13: 1.4, 14: 1.18,
28 15: 1.07, 16: 1.15, 17: 1.06, 19: 2.09, 20: 1.71, 21: 1.58, 22: 1.44,
29 23: 1.32, 24: 1.28, 25: 1.31, 26: 1.28, 27: 1.28, 28: 1.28, 29: 1.29,
30 30: 1.28, 31: 1.26, 32: 1.28, 33: 1.26, 34: 1.18, 35: 1.19, 37: 2.24,
31 38: 1.94, 39: 1.73, 40: 1.6, 41: 1.45, 42: 1.39, 43: 1.31, 44: 1.32,
32 45: 1.34, 46: 1.39, 47: 1.52, 48: 1.53, 49: 1.57, 50: 1.45, 51: 1.43,
33 52: 1.45, 53: 1.42, 56: 1.95, 57: 1.89, 58: 1.84, 59: 1.91, 60: 1.88,
34 61: 1.61, 62: 1.85, 63: 1.83, 64: 1.82, 65: 1.76, 66: 1.74, 67: 1.73,
35 68: 1.72, 69: 1.71, 70: 1.7, 71: 1.69, 72: 1.57, 73: 1.45, 74: 1.39,
36 75: 1.34, 76: 1.35, 77: 1.35, 78: 1.39, 79: 1.38, 80: 1.7, 81: 2.27,
37 82: 1.81, 83: 1.59, 84: 1.69, 88: 2.06, 89: 1.95, 90: 1.92}
40class PAWDataError(Exception):
41 """Error in PAW-data generation."""
44class DatasetOptimizer:
45 tolerances = np.array([0.2, # radii
46 0.3, # log. derivs.
47 40, # iterations
48 1.2 * 2 / 3 * 300 * 0.1**0.25, # convergence
49 0.0005, # eggbox error
50 0.05]) # IP
52 def __init__(self, symbol='H', nc=False):
53 self.old = False
55 self.symbol = symbol
56 self.nc = nc
58 line = Path('start.txt').read_text()
59 words = line.split()
60 assert words[1] == symbol
61 projectors = words[3]
62 radii = [float(f) for f in words[5].split(',')]
63 r0 = float(words[7].split(',')[1])
65 self.Z = atomic_numbers[symbol]
66 rc = self.rc = my_radii[self.Z] / Bohr
68 # Parse projectors string:
69 pattern = r'(-?\d+\.\d)'
70 energies = []
71 for m in re.finditer(pattern, projectors):
72 energies.append(float(projectors[m.start():m.end()]))
73 self.projectors = re.sub(pattern, '{:.1f}', projectors)
74 self.nenergies = len(energies)
76 self.x = energies + radii + [r0]
77 self.bounds = ([(-1.0, 4.0)] * self.nenergies +
78 [(rc * 0.7, rc * 1.0) for r in radii] +
79 [(0.3, rc)])
81 self.ecut1 = 450.0
82 self.ecut2 = 800.0
84 setup_paths[:0] = ['.']
86 self.logfile = None
87 self.tflush = time.time() + 60
89 def run(self):
90 print(self.symbol, self.rc / Bohr, self.projectors)
91 print(self.x)
92 print(self.bounds)
93 init = 'latinhypercube'
94 if 0: # os.path.isfile('data.csv'):
95 n = len(self.x)
96 data = self.read()[:15 * n]
97 if np.isfinite(data[:, n]).all() and len(data) == 15 * n:
98 init = data[:, :n]
100 DE(self, self.bounds, init=init, workers=8, updating='deferred')
102 def generate(self, fd, projectors, radii, r0, xc,
103 scalar_relativistic=True, tag=None, logderivs=True):
105 if projectors[-1].isupper():
106 nderiv0 = 5
107 else:
108 nderiv0 = 2
110 type = 'poly'
111 if self.nc:
112 type = 'nc'
114 try:
115 gen = generate(self.symbol, projectors, radii, r0, nderiv0,
116 xc, scalar_relativistic, (type, 4), output=fd)
117 except np.linalg.LinAlgError:
118 raise PAWDataError('LinAlgError')
120 if not scalar_relativistic:
121 if not gen.check_all():
122 raise PAWDataError('dataset check failed')
124 if tag is not None:
125 gen.make_paw_setup(tag or None).write_xml()
127 r = 1.1 * gen.rcmax
129 lmax = 2
130 if 'f' in projectors:
131 lmax = 3
133 error = 0.0
134 if logderivs:
135 for l in range(lmax + 1):
136 emin = -1.5
137 emax = 2.0
138 n0 = gen.number_of_core_states(l)
139 if n0 > 0:
140 e0_n = gen.aea.channels[l].e_n
141 emin = max(emin, e0_n[n0 - 1] + 0.1)
142 energies = np.linspace(emin, emax, 100)
143 de = energies[1] - energies[0]
144 ld1 = gen.aea.logarithmic_derivative(l, energies, r)
145 ld2 = gen.logarithmic_derivative(l, energies, r)
146 error += abs(ld1 - ld2).sum() * de
148 return error
150 def parameters(self, x):
151 energies = x[:self.nenergies]
152 radii = x[self.nenergies:-1]
153 r0 = x[-1]
154 projectors = self.projectors.format(*energies)
155 return energies, radii, r0, projectors
157 def __call__(self, x):
158 id = mp.current_process().name[-1]
159 self.setup = 'de' + id
160 energies, radii, r0, projectors = self.parameters(x)
162 self.logfile = open(f'data-{id}.csv', 'a')
164 with open(f'out-{id}.txt', 'w') as fd:
165 errors, msg, convenergies, eggenergies, ips = \
166 self.test(fd, projectors, radii, r0)
167 error = ((errors / self.tolerances)**2).sum()
169 if msg:
170 print(msg, x, error, errors, convenergies, eggenergies, ips,
171 file=sys.stderr)
173 convenergies += [0] * (7 - len(convenergies))
175 print(', '.join(repr(number) for number in
176 list(x) + [error] + errors +
177 convenergies + eggenergies + ips),
178 file=self.logfile)
180 if time.time() > self.tflush:
181 self.logfile.flush()
182 self.tflush = time.time() + 60
184 return error
186 def test_old_paw_data(self):
187 with open('old.txt', 'w') as fd:
188 self._test_old_paw_data(fd)
190 def _test_old_paw_data(self, fd):
191 area, niter, convenergies = self.convergence(fd, 'paw')
192 eggenergies = self.eggbox(fd, 'paw')
193 print('RESULTS:',
194 ', '.join(repr(number) for number in
195 [area, niter, max(eggenergies)] +
196 convenergies + eggenergies),
197 file=fd)
199 def test(self, fd, projectors, radii, r0):
200 errors = [np.inf] * 6
201 energies = []
202 eggenergies = [0, 0, 0]
203 ip = 0.0
204 ip0 = 0.0
205 msg = ''
207 try:
208 if any(r < r0 for r in radii):
209 raise PAWDataError('Core radius too large')
211 rc = self.rc
212 errors[0] = sum(r - rc for r in radii if r > rc)
214 error = 0.0
215 for kwargs in [dict(xc='PBE', tag=self.setup),
216 dict(xc='PBE', scalar_relativistic=False),
217 dict(xc='LDA', tag=self.setup),
218 dict(xc='PBEsol'),
219 dict(xc='RPBE'),
220 dict(xc='PW91')]:
221 error += self.generate(fd, projectors, radii, r0, **kwargs)
222 errors[1] = error
224 area, niter, energies = self.convergence(fd)
225 errors[2] = niter
226 errors[3] = area
228 eggenergies = self.eggbox(fd)
229 errors[4] = max(eggenergies)
231 ip, ip0 = self.ip(fd)
232 errors[5] = ip - ip0
234 except (ConvergenceError, PAWDataError, RuntimeError,
235 np.linalg.LinAlgError) as e:
236 msg = str(e)
238 return errors, msg, energies, eggenergies, [ip, ip0]
240 def eggbox(self, fd, setup='de'):
241 energies = []
242 for h in [0.16, 0.18, 0.2]:
243 a0 = 16 * h
244 atoms = Atoms(self.symbol, cell=(a0, a0, 2 * a0), pbc=True)
245 if 58 <= self.Z <= 70 or 90 <= self.Z <= 102:
246 M = 999
247 mixer = {'mixer': Mixer(0.01, 5)}
248 else:
249 M = 333
250 mixer = {}
251 atoms.calc = GPAW(h=h,
252 xc='PBE',
253 symmetry='off',
254 setups=self.setup,
255 maxiter=M,
256 txt=fd,
257 **mixer)
258 atoms.positions += h / 2 # start with broken symmetry
259 e0 = atoms.get_potential_energy()
260 atoms.positions -= h / 6
261 e1 = atoms.get_potential_energy()
262 atoms.positions -= h / 6
263 e2 = atoms.get_potential_energy()
264 atoms.positions -= h / 6
265 e3 = atoms.get_potential_energy()
266 energies.append(np.ptp([e0, e1, e2, e3]))
267 # print(energies)
268 return energies
270 def convergence(self, fd, setup='de'):
271 a = 3.0
272 atoms = Atoms(self.symbol, cell=(a, a, a), pbc=True)
273 if 58 <= self.Z <= 70 or 90 <= self.Z <= 102:
274 M = 999
275 mixer = {'mixer': Mixer(0.01, 5)}
276 else:
277 M = 333
278 mixer = {}
279 atoms.calc = GPAW(mode=PW(1500),
280 xc='PBE',
281 setups=self.setup,
282 symmetry='off',
283 maxiter=M,
284 txt=fd,
285 **mixer)
286 e0 = atoms.get_potential_energy()
287 energies = [e0]
288 iters = atoms.calc.get_number_of_iterations()
289 oldfde = None
290 area = 0.0
292 def f(x):
293 return x**0.25
295 for ec in range(800, 200, -100):
296 atoms.calc.set(mode=PW(ec))
297 atoms.calc.set(eigensolver='rmm-diis')
298 de = atoms.get_potential_energy() - e0
299 energies.append(de)
300 # print(ec, de)
301 fde = f(abs(de))
302 if fde > f(0.1):
303 if oldfde is None:
304 return np.inf, iters, energies
305 ec0 = ec + (fde - f(0.1)) / (fde - oldfde) * 100
306 area += ((ec + 100) - ec0) * (f(0.1) + oldfde) / 2
307 break
309 if oldfde is not None:
310 area += 100 * (fde + oldfde) / 2
311 oldfde = fde
313 return area, iters, energies
315 def ip(self, fd):
316 IP, IP0 = ip(self.symbol, fd, self.setup)
317 return IP, IP0
319 def read(self):
320 data = []
321 for i in '12345678s':
322 try:
323 d = np.loadtxt(f'data-{i}.csv', delimiter=',')
324 except OSError:
325 if i == 's':
326 pass
327 if i == '1':
328 data = np.loadtxt('data.csv', delimiter=',')
329 break
330 data += d.tolist()
331 data = np.array(data)
332 return data[data[:, len(self.x)].argsort()]
334 def summary(self, N=10):
335 n = len(self.x)
336 for x in self.read()[:N]:
337 print('{:3} {:2} {:9.1f} ({}) ({}, {})'
338 .format(self.Z,
339 self.symbol,
340 x[n],
341 ', '.join(f'{e:4.1f}' + s
342 for e, s
343 in zip(x[n + 1:n + 7] / self.tolerances,
344 'rlciex')),
345 ', '.join(f'{e:+.2f}'
346 for e in x[:self.nenergies]),
347 ', '.join(f'{r:.2f}'
348 for r in x[self.nenergies:n])))
350 def best(self):
351 n = len(self.x)
352 a = self.read()[0]
353 x = a[:n]
354 error = a[n]
355 energies, radii, r0, projectors = self.parameters(x)
356 if 1:
357 if projectors[-1].isupper():
358 nderiv0 = 5
359 else:
360 nderiv0 = 2
361 fmt = '{:3} {:2} -P {:31} -r {:20} -0 {},{:.2f} # {:10.3f}'
362 print(fmt.format(self.Z,
363 self.symbol,
364 projectors,
365 ','.join(f'{r:.2f}' for r in radii),
366 nderiv0,
367 r0,
368 error))
369 if 1 and error != np.inf and error != np.nan:
370 self.generate(None, projectors, radii, r0, 'PBE', True, 'a7o',
371 logderivs=False)
374def ip(symbol, fd, setup):
375 xc = 'LDA'
376 aea = AllElectronAtom(symbol, log=fd, scalar_relativistic=True)
377 aea.initialize()
378 aea.run()
379 aea.refine()
380 energy = aea.ekin + aea.eH + aea.eZ + aea.exc
381 eigs = []
382 for l, channel in enumerate(aea.channels):
383 n = l + 1
384 for e, f in zip(channel.e_n, channel.f_n):
385 if f == 0:
386 break
387 eigs.append((e, n, l))
388 n += 1
389 e0, n0, l0 = max(eigs)
390 aea = AllElectronAtom(symbol, log=fd, scalar_relativistic=True)
391 aea.add(n0, l0, -1)
392 aea.initialize()
393 aea.run()
394 aea.refine()
395 IP = aea.ekin + aea.eH + aea.eZ + aea.exc - energy
396 IP *= Ha
398 s = create_setup(symbol, type=setup, xc=xc)
399 f_ln = defaultdict(list)
400 for l, f in zip(s.l_j, s. f_j):
401 if f:
402 f_ln[l].append(f)
404 f_sln = [[f_ln[l] for l in range(1 + max(f_ln))]]
405 calc = AtomPAW(symbol, f_sln, xc=xc, txt=fd, setups=setup)
406 energy = calc.results['energy']
407 # eps_n = calc.wfs.kpt_u[0].eps_n
409 f_sln[0][l0][-1] -= 1
410 calc = AtomPAW(symbol, f_sln, xc=xc, charge=1, txt=fd, setups=setup)
411 IP2 = calc.results['energy'] - energy
412 return IP, IP2
415if __name__ == '__main__':
416 import argparse
417 from pathlib import Path
418 parser = argparse.ArgumentParser(usage='python -m gpaw.atom.optimize '
419 '[options] [folder folder ...]',
420 description='Optimize PAW data')
421 parser.add_argument('-s', '--summary', type=int)
422 parser.add_argument('-b', '--best', action='store_true')
423 parser.add_argument('-n', '--norm-conserving', action='store_true')
424 parser.add_argument('-o', '--old-setups', action='store_true')
425 parser.add_argument('folder', nargs='*')
426 args = parser.parse_args()
427 folders = [Path(folder) for folder in args.folder or ['.']]
428 home = Path.cwd()
429 for folder in folders:
430 try:
431 os.chdir(folder)
432 symbol = Path.cwd().name
433 do = DatasetOptimizer(symbol)
434 if args.summary:
435 do.summary(args.summary)
436 elif args.old_setups:
437 do.test_old_paw_data()
438 elif args.best:
439 do.best()
440 else:
441 do.run()
442 finally:
443 os.chdir(home)