Coverage for eminus/band_minimizer.py: 98.95%
190 statements
« prev ^ index » next coverage.py v7.6.4, created at 2024-11-08 12:59 +0000
1# SPDX-FileCopyrightText: 2021 The eminus developers
2# SPDX-License-Identifier: Apache-2.0
3"""Minimization algorithms for fixed Hamiltonians.
5Similar to :mod:`eminus.minimizer` but for a fixed Hamiltonian the implementation can be simplified
6and made more performant.
7"""
9import copy
10import logging
12import numpy as np
13from scipy.linalg import inv, sqrtm
15from .dft import H, orth, orth_unocc
16from .energies import get_Eband
17from .logger import name
18from .minimizer import cg_method, cg_test, check_convergence, linmin_test
19from .utils import dotprod
22def scf_step_occ(scf, W):
23 """Perform one SCF step for an occupied band minimization calculation.
25 Args:
26 scf: SCF object.
27 W: Expansion coefficients of unconstrained wave functions in reciprocal space.
29 Returns:
30 Band energy.
31 """
32 atoms = scf.atoms
33 scf.Y = orth(atoms, W)
34 return get_Eband(scf, scf.Y, **scf._precomputed)
37def scf_step_unocc(scf, Z):
38 """Perform one SCF step for an unoccupied band minimization calculation.
40 Args:
41 scf: SCF object.
42 Z: Expansion coefficients of unconstrained wave functions in reciprocal space.
44 Returns:
45 Band energy.
46 """
47 atoms = scf.atoms
48 scf.D = orth_unocc(atoms, scf.Y, Z)
49 return get_Eband(scf, scf.D, **scf._precomputed)
52def get_grad_occ(scf, ik, spin, W, **kwargs):
53 """Calculate the occupied band energy gradient with respect to W.
55 Reference: Comput. Phys. Commun. 128, 1.
57 Args:
58 scf: SCF object.
59 ik: k-point index.
60 spin: Spin variable to track whether to do the calculation for spin up or down.
61 W: Expansion coefficients of unconstrained wave functions in reciprocal space.
63 Keyword Args:
64 **kwargs: See :func:`H`.
66 Returns:
67 Gradient.
68 """
69 atoms = scf.atoms
70 W = orth(atoms, W)
71 HW = H(scf, ik, spin, W, **kwargs)
72 WHW = W[ik][spin].conj().T @ HW
73 OW = atoms.O(W[ik][spin])
74 U = W[ik][spin].conj().T @ OW
75 invU = inv(U)
76 U12 = sqrtm(invU)
77 # grad E = (I - O(Y) Ydag) H(Y) U^-0.5
78 return atoms.kpts.wk[ik] * ((HW - OW @ WHW) @ U12)
81def get_grad_unocc(scf, ik, spin, Z, **kwargs):
82 """Calculate the unoccupied band energy gradient with respect to Z.
84 Reference: Comput. Phys. Commun. 128, 1.
86 Args:
87 scf: SCF object.
88 ik: k-point index.
89 spin: Spin variable to track whether to do the calculation for spin up or down.
90 Z: Expansion coefficients of unconstrained wave functions in reciprocal space.
92 Keyword Args:
93 **kwargs: See :func:`H`.
95 Returns:
96 Gradient.
97 """
98 atoms = scf.atoms
99 Y = scf.Y[ik][spin][:, scf.atoms.occ.f[ik][spin] > 0]
100 Ydag = Y.conj().T
101 # We need X12 later, so orthogonalize in-place and only the current state
102 rhoZ = Z[ik][spin] - Y @ Ydag @ atoms.O(Z[ik][spin])
103 X12 = inv(sqrtm(rhoZ.conj().T @ atoms.O(rhoZ)))
104 D = rhoZ @ X12
105 # Create the correct input shape for the Hamiltonian
106 D_tmp = [None] * len(Z)
107 D_tmp[ik] = np.empty_like(Z[ik])
108 D_tmp[ik][spin] = D
109 HD = H(scf, ik, spin, D_tmp, **kwargs)
110 DHD = D.conj().T @ HD
111 I = np.eye(Z[ik].shape[1])
112 # grad E = (I - O(Y) Ydag) (I - O(D) Ddag) H(D) X^-0.5
113 return atoms.kpts.wk[ik] * ((I - atoms.O(Y) @ Ydag) @ (HD - atoms.O(D) @ DHD) @ X12)
116@name("steepest descent minimization")
117def sd(
118 scf,
119 W,
120 Nit,
121 cost=scf_step_occ,
122 grad=get_grad_occ,
123 condition=check_convergence,
124 betat=3e-5,
125 **kwargs,
126):
127 """Steepest descent minimization algorithm for a fixed Hamiltonian.
129 Args:
130 scf: SCF object.
131 W: Expansion coefficients of unconstrained wave functions in reciprocal space.
132 Nit: Maximum number of SCF steps.
134 Keyword Args:
135 cost: Function that will run every SCF step.
136 grad: Function that calculates the respective gradient.
137 condition: Function to check and log the convergence condition.
138 betat: Step size.
139 **kwargs: Throwaway arguments.
141 Returns:
142 Band energies per SCF cycle and optimized expansion coefficients.
143 """
144 atoms = scf.atoms
145 costs = []
147 for _ in range(Nit):
148 c = cost(scf, W)
149 costs.append(c)
150 if condition(scf, "sd", costs):
151 break
152 for ik in range(atoms.kpts.Nk):
153 for spin in range(atoms.occ.Nspin):
154 g = grad(scf, ik, spin, W, **scf._precomputed)
155 W[ik][spin] = W[ik][spin] - betat * g
156 return costs, W
159@name("preconditioned line minimization")
160def pclm(
161 scf,
162 W,
163 Nit,
164 cost=scf_step_occ,
165 grad=get_grad_occ,
166 condition=check_convergence,
167 betat=3e-5,
168 precondition=True,
169 **kwargs,
170):
171 """Preconditioned line minimization algorithm for a fixed Hamiltonian.
173 Args:
174 scf: SCF object.
175 W: Expansion coefficients of unconstrained wave functions in reciprocal space.
176 Nit: Maximum number of SCF steps.
178 Keyword Args:
179 cost: Function that will run every SCF step.
180 grad: Function that calculates the respective gradient.
181 condition: Function to check and log the convergence condition.
182 betat: Step size.
183 precondition: Whether to use a preconditioner.
184 **kwargs: Throwaway arguments.
186 Returns:
187 Band energies per SCF cycle and optimized expansion coefficients.
188 """
189 atoms = scf.atoms
190 costs = []
192 linmin = np.empty((atoms.kpts.Nk, atoms.occ.Nspin))
193 d = [np.empty_like(Wk) for Wk in W]
195 if precondition:
196 method = "pclm"
197 else:
198 method = "lm"
200 for i in range(Nit):
201 c = cost(scf, W)
202 costs.append(c)
203 if condition(scf, method, costs, linmin):
204 break
205 for ik in range(atoms.kpts.Nk):
206 for spin in range(atoms.occ.Nspin):
207 g = grad(scf, ik, spin, W, **scf._precomputed)
208 if scf._log.level <= logging.DEBUG and i > 0:
209 linmin[ik][spin] = linmin_test(g, d[ik][spin])
210 if precondition:
211 d[ik][spin] = -atoms.K(g, ik)
212 else:
213 d[ik][spin] = -g
214 W[ik][spin] = W[ik][spin] + betat * d[ik][spin]
215 gt = grad(scf, ik, spin, W, **scf._precomputed)
216 beta = abs(betat * dotprod(g, d[ik][spin]) / dotprod(g - gt, d[ik][spin]))
217 W[ik][spin] = W[ik][spin] + beta * d[ik][spin]
218 return costs, W
221@name("line minimization")
222def lm(
223 scf,
224 W,
225 Nit,
226 cost=scf_step_occ,
227 grad=get_grad_occ,
228 condition=check_convergence,
229 betat=3e-5,
230 **kwargs,
231):
232 """Line minimization algorithm for a fixed Hamiltonian.
234 Args:
235 scf: SCF object.
236 W: Expansion coefficients of unconstrained wave functions in reciprocal space.
237 Nit: Maximum number of SCF steps.
239 Keyword Args:
240 cost: Function that will run every SCF step.
241 grad: Function that calculates the respective gradient.
242 condition: Function to check and log the convergence condition.
243 betat: Step size.
244 **kwargs: Throwaway arguments.
246 Returns:
247 Band energies per SCF cycle and optimized expansion coefficients.
248 """
249 return pclm(scf, W, Nit, cost, grad, condition, betat, precondition=False)
252@name("preconditioned conjugate-gradient minimization")
253def pccg(
254 scf,
255 W,
256 Nit,
257 cost=scf_step_occ,
258 grad=get_grad_occ,
259 condition=check_convergence,
260 betat=3e-5,
261 cgform=1,
262 precondition=True,
263):
264 """Preconditioned conjugate-gradient minimization algorithm for a fixed Hamiltonian.
266 Args:
267 scf: SCF object.
268 W: Expansion coefficients of unconstrained wave functions in reciprocal space.
269 Nit: Maximum number of SCF steps.
271 Keyword Args:
272 cost: Function that will run every SCF step.
273 grad: Function that calculates the respective gradient.
274 condition: Function to check and log the convergence condition.
275 betat: Step size.
276 cgform: Conjugate gradient form.
277 precondition: Whether to use a preconditioner.
279 Returns:
280 Band energies per SCF cycle and optimized expansion coefficients.
281 """
282 atoms = scf.atoms
283 costs = []
285 linmin = np.empty((atoms.kpts.Nk, atoms.occ.Nspin))
286 cg = np.empty((atoms.kpts.Nk, atoms.occ.Nspin))
287 norm_g = np.empty((atoms.kpts.Nk, atoms.occ.Nspin))
288 d = [np.empty_like(Wk) for Wk in W]
289 d_old = [np.empty_like(Wk) for Wk in W]
290 g_old = [np.empty_like(Wk) for Wk in W]
292 if precondition:
293 method = "pccg"
294 else:
295 method = "cg"
297 c = cost(scf, W)
298 costs.append(c)
299 condition(scf, method, costs)
300 # Do the first step without the linmin and cg tests, and without the cg_method
301 for ik in range(atoms.kpts.Nk):
302 for spin in range(atoms.occ.Nspin):
303 g = grad(scf, ik, spin, W, **scf._precomputed)
304 if precondition:
305 d[ik][spin] = -atoms.K(g, ik)
306 else:
307 d[ik][spin] = -g
308 W[ik][spin] = W[ik][spin] + betat * d[ik][spin]
309 gt = grad(scf, ik, spin, W, **scf._precomputed)
310 beta = abs(betat * dotprod(g, d[ik][spin]) / dotprod(g - gt, d[ik][spin]))
311 g_old[ik][spin], d_old[ik][spin] = g, d[ik][spin]
312 W[ik][spin] = W[ik][spin] + beta * d[ik][spin]
314 for _ in range(1, Nit):
315 c = cost(scf, W)
316 costs.append(c)
317 if condition(scf, method, costs, linmin, cg, norm_g):
318 break
319 for ik in range(atoms.kpts.Nk):
320 for spin in range(atoms.occ.Nspin):
321 g = grad(scf, ik, spin, W, **scf._precomputed)
322 # Calculate linmin and cg for each spin separately
323 if scf._log.level <= logging.DEBUG:
324 linmin[ik][spin] = linmin_test(g, d[ik][spin])
325 cg[ik][spin] = cg_test(atoms, ik, g, g_old[ik][spin], precondition)
326 beta, norm_g[ik][spin] = cg_method(
327 scf, ik, cgform, g, g_old[ik][spin], d_old[ik][spin], precondition
328 )
329 if precondition:
330 d[ik][spin] = -atoms.K(g, ik) + beta * d_old[ik][spin]
331 else:
332 d[ik][spin] = -g + beta * d_old[ik][spin]
333 W[ik][spin] = W[ik][spin] + betat * d[ik][spin]
334 gt = grad(scf, ik, spin, W, **scf._precomputed)
335 beta = abs(betat * dotprod(g, d[ik][spin]) / dotprod(g - gt, d[ik][spin]))
336 g_old[ik][spin], d_old[ik][spin] = g, d[ik][spin]
337 W[ik][spin] = W[ik][spin] + beta * d[ik][spin]
338 return costs, W
341@name("conjugate-gradient minimization")
342def cg(
343 scf,
344 W,
345 Nit,
346 cost=scf_step_occ,
347 grad=get_grad_occ,
348 condition=check_convergence,
349 betat=3e-5,
350 cgform=1,
351):
352 """Conjugate-gradient minimization algorithm for a fixed Hamiltonian.
354 Args:
355 scf: SCF object.
356 W: Expansion coefficients of unconstrained wave functions in reciprocal space.
357 Nit: Maximum number of SCF steps.
359 Keyword Args:
360 cost: Function that will run every SCF step.
361 grad: Function that calculates the respective gradient.
362 condition: Function to check and log the convergence condition.
363 betat: Step size.
364 cgform: Conjugate gradient form.
366 Returns:
367 Band energies per SCF cycle and optimized expansion coefficients.
368 """
369 return pccg(scf, W, Nit, cost, grad, condition, betat, cgform, precondition=False)
372@name("auto minimization")
373def auto(
374 scf,
375 W,
376 Nit,
377 cost=scf_step_occ,
378 grad=get_grad_occ,
379 condition=check_convergence,
380 betat=3e-5,
381 cgform=1,
382):
383 """Automatic precond. conjugate-gradient minimization algorithm for a fixed Hamiltonian.
385 This function chooses an sd step over the pccg step if the energy goes up.
387 Args:
388 scf: SCF object.
389 W: Expansion coefficients of unconstrained wave functions in reciprocal space.
390 Nit: Maximum number of SCF steps.
392 Keyword Args:
393 cost: Function that will run every SCF step.
394 grad: Function that calculates the respective gradient.
395 condition: Function to check and log the convergence condition.
396 betat: Step size.
397 cgform: Conjugate gradient form.
399 Returns:
400 Band energies per SCF cycle and optimized expansion coefficients.
401 """
402 atoms = scf.atoms
403 costs = []
405 linmin = np.empty((atoms.kpts.Nk, atoms.occ.Nspin))
406 cg = np.empty((atoms.kpts.Nk, atoms.occ.Nspin))
407 norm_g = np.empty((atoms.kpts.Nk, atoms.occ.Nspin))
408 g = [np.empty_like(Wk) for Wk in W]
409 d = [np.empty_like(Wk) for Wk in W]
410 d_old = [np.empty_like(Wk) for Wk in W]
411 g_old = [np.empty_like(Wk) for Wk in W]
413 # Do the first step without the linmin and cg tests, and without the cg_method
414 for ik in range(atoms.kpts.Nk):
415 for spin in range(atoms.occ.Nspin):
416 g[ik][spin] = grad(scf, ik, spin, W, **scf._precomputed)
417 d[ik][spin] = -atoms.K(g[ik][spin], ik)
418 W[ik][spin] = W[ik][spin] + betat * d[ik][spin]
419 gt = grad(scf, ik, spin, W, **scf._precomputed)
420 beta = abs(
421 betat * dotprod(g[ik][spin], d[ik][spin]) / dotprod(g[ik][spin] - gt, d[ik][spin])
422 )
423 g_old[ik][spin], d_old[ik][spin] = g[ik][spin], d[ik][spin]
424 W[ik][spin] = W[ik][spin] + beta * d[ik][spin]
426 c = cost(scf, W)
427 costs.append(c)
428 if condition(scf, "pccg", costs):
429 return costs
431 for _ in range(1, Nit):
432 W_old = copy.deepcopy(W)
433 for ik in range(atoms.kpts.Nk):
434 for spin in range(atoms.occ.Nspin):
435 g[ik][spin] = grad(scf, ik, spin, W, **scf._precomputed)
436 # Calculate linmin and cg for each spin separately
437 if scf._log.level <= logging.DEBUG:
438 linmin[ik][spin] = linmin_test(g[ik][spin], d[ik][spin])
439 cg[ik][spin] = cg_test(atoms, ik, g[ik][spin], g_old[ik][spin])
440 beta, norm_g[ik][spin] = cg_method(
441 scf, ik, cgform, g[ik][spin], g_old[ik][spin], d_old[ik][spin]
442 )
443 d[ik][spin] = -atoms.K(g[ik][spin], ik) + beta * d_old[ik][spin]
444 W[ik][spin] = W[ik][spin] + betat * d[ik][spin]
445 gt = grad(scf, ik, spin, W, **scf._precomputed)
446 beta = abs(
447 betat
448 * dotprod(g[ik][spin], d[ik][spin])
449 / dotprod(g[ik][spin] - gt, d[ik][spin])
450 )
451 g_old[ik][spin], d_old[ik][spin] = g[ik][spin], d[ik][spin]
452 W[ik][spin] = W[ik][spin] + beta * d[ik][spin]
454 c = cost(scf, W)
455 # If the energy does not go down use the steepest descent step and recalculate the energy
456 if c > costs[-1]:
457 W = W_old
458 for ik in range(atoms.kpts.Nk):
459 for spin in range(atoms.occ.Nspin):
460 W[ik][spin] = W[ik][spin] - betat * g[ik][spin]
461 c = cost(scf, W)
462 costs.append(c)
463 # Do not print cg and linmin if we do the sd step
464 if condition(scf, "sd", costs, norm_g=norm_g):
465 break
466 else:
467 costs.append(c)
468 if condition(scf, "pccg", costs, linmin, cg, norm_g):
469 break
470 return costs, W
473#: Map minimizer names with their respective implementation.
474IMPLEMENTED = {
475 "sd": sd,
476 "lm": lm,
477 "pclm": pclm,
478 "cg": cg,
479 "pccg": pccg,
480 "auto": auto,
481}