Coverage for eminus / band_minimizer.py: 98.94%
189 statements
« prev ^ index » next coverage.py v7.12.0, created at 2025-11-21 14:20 +0000
1# SPDX-FileCopyrightText: 2023 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
12from . import backend as xp
13from .dft import H, orth, orth_unocc
14from .energies import get_Eband
15from .logger import name
16from .minimizer import cg_method, cg_test, check_convergence, linmin_test
17from .utils import dotprod
20def scf_step_occ(scf, W):
21 """Perform one SCF step for an occupied band minimization calculation.
23 Args:
24 scf: SCF object.
25 W: Expansion coefficients of unconstrained wave functions in reciprocal space.
27 Returns:
28 Band energy.
29 """
30 atoms = scf.atoms
31 scf.Y = orth(atoms, W)
32 return get_Eband(scf, scf.Y, **scf._precomputed)
35def scf_step_unocc(scf, Z):
36 """Perform one SCF step for an unoccupied band minimization calculation.
38 Args:
39 scf: SCF object.
40 Z: Expansion coefficients of unconstrained wave functions in reciprocal space.
42 Returns:
43 Band energy.
44 """
45 atoms = scf.atoms
46 scf.D = orth_unocc(atoms, scf.Y, Z)
47 return get_Eband(scf, scf.D, **scf._precomputed)
50def get_grad_occ(scf, ik, spin, W, **kwargs):
51 """Calculate the occupied band energy gradient with respect to W.
53 Reference: Comput. Phys. Commun. 128, 1.
55 Args:
56 scf: SCF object.
57 ik: k-point index.
58 spin: Spin variable to track whether to do the calculation for spin up or down.
59 W: Expansion coefficients of unconstrained wave functions in reciprocal space.
61 Keyword Args:
62 **kwargs: See :func:`H`.
64 Returns:
65 Gradient.
66 """
67 atoms = scf.atoms
68 W = orth(atoms, W)
69 HW = H(scf, ik, spin, W, **kwargs)
70 WHW = W[ik][spin].conj().T @ HW
71 OW = atoms.O(W[ik][spin])
72 U = W[ik][spin].conj().T @ OW
73 invU = xp.linalg.inv(U)
74 U12 = xp.sqrtm(invU)
75 # grad E = (I - O(Y) Ydag) H(Y) U^-0.5
76 return atoms.kpts.wk[ik] * ((HW - OW @ WHW) @ U12)
79def get_grad_unocc(scf, ik, spin, Z, **kwargs):
80 """Calculate the unoccupied band energy gradient with respect to Z.
82 Reference: Comput. Phys. Commun. 128, 1.
84 Args:
85 scf: SCF object.
86 ik: k-point index.
87 spin: Spin variable to track whether to do the calculation for spin up or down.
88 Z: Expansion coefficients of unconstrained wave functions in reciprocal space.
90 Keyword Args:
91 **kwargs: See :func:`H`.
93 Returns:
94 Gradient.
95 """
96 atoms = scf.atoms
97 Y = scf.Y[ik][spin][:, scf.atoms.occ.f[ik][spin] > 0]
98 Ydag = Y.conj().T
99 # We need X12 later, so orthogonalize in-place and only the current state
100 rhoZ = Z[ik][spin] - Y @ Ydag @ atoms.O(Z[ik][spin])
101 X12 = xp.linalg.inv(xp.sqrtm(rhoZ.conj().T @ atoms.O(rhoZ)))
102 D = rhoZ @ X12
103 # Create the correct input shape for the Hamiltonian
104 D_tmp = [None] * len(Z)
105 D_tmp[ik] = xp.empty_like(Z[ik])
106 D_tmp[ik][spin] = D
107 HD = H(scf, ik, spin, D_tmp, **kwargs)
108 DHD = D.conj().T @ HD
109 I = xp.eye(Z[ik].shape[1])
110 # grad E = (I - O(Y) Ydag) (I - O(D) Ddag) H(D) X^-0.5
111 return atoms.kpts.wk[ik] * ((I - atoms.O(Y) @ Ydag) @ (HD - atoms.O(D) @ DHD) @ X12)
114@name("steepest descent minimization")
115def sd(
116 scf,
117 W,
118 Nit,
119 cost=scf_step_occ,
120 grad=get_grad_occ,
121 condition=check_convergence,
122 betat=3e-5,
123 **kwargs,
124):
125 """Steepest descent minimization algorithm for a fixed Hamiltonian.
127 Args:
128 scf: SCF object.
129 W: Expansion coefficients of unconstrained wave functions in reciprocal space.
130 Nit: Maximum number of SCF steps.
132 Keyword Args:
133 cost: Function that will run every SCF step.
134 grad: Function that calculates the respective gradient.
135 condition: Function to check and log the convergence condition.
136 betat: Step size.
137 **kwargs: Throwaway arguments.
139 Returns:
140 Band energies per SCF cycle and optimized expansion coefficients.
141 """
142 atoms = scf.atoms
143 costs = []
145 for _ in range(Nit):
146 c = cost(scf, W)
147 costs.append(c)
148 if condition(scf, "sd", costs):
149 break
150 for ik in range(atoms.kpts.Nk):
151 for spin in range(atoms.occ.Nspin):
152 g = grad(scf, ik, spin, W, **scf._precomputed)
153 W[ik][spin] = W[ik][spin] - betat * g
154 return costs, W
157@name("preconditioned line minimization")
158def pclm(
159 scf,
160 W,
161 Nit,
162 cost=scf_step_occ,
163 grad=get_grad_occ,
164 condition=check_convergence,
165 betat=3e-5,
166 precondition=True,
167 **kwargs,
168):
169 """Preconditioned line minimization algorithm for a fixed Hamiltonian.
171 Args:
172 scf: SCF object.
173 W: Expansion coefficients of unconstrained wave functions in reciprocal space.
174 Nit: Maximum number of SCF steps.
176 Keyword Args:
177 cost: Function that will run every SCF step.
178 grad: Function that calculates the respective gradient.
179 condition: Function to check and log the convergence condition.
180 betat: Step size.
181 precondition: Whether to use a preconditioner.
182 **kwargs: Throwaway arguments.
184 Returns:
185 Band energies per SCF cycle and optimized expansion coefficients.
186 """
187 atoms = scf.atoms
188 costs = []
190 linmin = xp.empty((atoms.kpts.Nk, atoms.occ.Nspin))
191 d = [xp.empty_like(Wk) for Wk in W]
193 if precondition:
194 method = "pclm"
195 else:
196 method = "lm"
198 for i in range(Nit):
199 c = cost(scf, W)
200 costs.append(c)
201 if condition(scf, method, costs, linmin):
202 break
203 for ik in range(atoms.kpts.Nk):
204 for spin in range(atoms.occ.Nspin):
205 g = grad(scf, ik, spin, W, **scf._precomputed)
206 if scf._log.level <= logging.DEBUG and i > 0:
207 linmin[ik][spin] = linmin_test(g, d[ik][spin])
208 if precondition:
209 d[ik][spin] = -atoms.K(g, ik)
210 else:
211 d[ik][spin] = -g
212 W[ik][spin] = W[ik][spin] + betat * d[ik][spin]
213 gt = grad(scf, ik, spin, W, **scf._precomputed)
214 beta = abs(betat * dotprod(g, d[ik][spin]) / dotprod(g - gt, d[ik][spin]))
215 W[ik][spin] = W[ik][spin] + beta * d[ik][spin]
216 return costs, W
219@name("line minimization")
220def lm(
221 scf,
222 W,
223 Nit,
224 cost=scf_step_occ,
225 grad=get_grad_occ,
226 condition=check_convergence,
227 betat=3e-5,
228 **kwargs,
229):
230 """Line minimization algorithm for a fixed Hamiltonian.
232 Args:
233 scf: SCF object.
234 W: Expansion coefficients of unconstrained wave functions in reciprocal space.
235 Nit: Maximum number of SCF steps.
237 Keyword Args:
238 cost: Function that will run every SCF step.
239 grad: Function that calculates the respective gradient.
240 condition: Function to check and log the convergence condition.
241 betat: Step size.
242 **kwargs: Throwaway arguments.
244 Returns:
245 Band energies per SCF cycle and optimized expansion coefficients.
246 """
247 return pclm(scf, W, Nit, cost, grad, condition, betat, precondition=False)
250@name("preconditioned conjugate-gradient minimization")
251def pccg(
252 scf,
253 W,
254 Nit,
255 cost=scf_step_occ,
256 grad=get_grad_occ,
257 condition=check_convergence,
258 betat=3e-5,
259 cgform=1,
260 precondition=True,
261):
262 """Preconditioned conjugate-gradient minimization algorithm for a fixed Hamiltonian.
264 Args:
265 scf: SCF object.
266 W: Expansion coefficients of unconstrained wave functions in reciprocal space.
267 Nit: Maximum number of SCF steps.
269 Keyword Args:
270 cost: Function that will run every SCF step.
271 grad: Function that calculates the respective gradient.
272 condition: Function to check and log the convergence condition.
273 betat: Step size.
274 cgform: Conjugate gradient form.
275 precondition: Whether to use a preconditioner.
277 Returns:
278 Band energies per SCF cycle and optimized expansion coefficients.
279 """
280 atoms = scf.atoms
281 costs = []
283 linmin = xp.empty((atoms.kpts.Nk, atoms.occ.Nspin))
284 cg = xp.empty((atoms.kpts.Nk, atoms.occ.Nspin))
285 norm_g = xp.empty((atoms.kpts.Nk, atoms.occ.Nspin))
286 d = [xp.empty_like(Wk) for Wk in W]
287 d_old = [xp.empty_like(Wk) for Wk in W]
288 g_old = [xp.empty_like(Wk) for Wk in W]
290 if precondition:
291 method = "pccg"
292 else:
293 method = "cg"
295 c = cost(scf, W)
296 costs.append(c)
297 condition(scf, method, costs)
298 # Do the first step without the linmin and cg tests, and without the cg_method
299 for ik in range(atoms.kpts.Nk):
300 for spin in range(atoms.occ.Nspin):
301 g = grad(scf, ik, spin, W, **scf._precomputed)
302 if precondition:
303 d[ik][spin] = -atoms.K(g, ik)
304 else:
305 d[ik][spin] = -g
306 W[ik][spin] = W[ik][spin] + betat * d[ik][spin]
307 gt = grad(scf, ik, spin, W, **scf._precomputed)
308 beta = abs(betat * dotprod(g, d[ik][spin]) / dotprod(g - gt, d[ik][spin]))
309 g_old[ik][spin], d_old[ik][spin] = g, d[ik][spin]
310 W[ik][spin] = W[ik][spin] + beta * d[ik][spin]
312 for _ in range(1, Nit):
313 c = cost(scf, W)
314 costs.append(c)
315 if condition(scf, method, costs, linmin, cg, norm_g):
316 break
317 for ik in range(atoms.kpts.Nk):
318 for spin in range(atoms.occ.Nspin):
319 g = grad(scf, ik, spin, W, **scf._precomputed)
320 # Calculate linmin and cg for each spin separately
321 if scf._log.level <= logging.DEBUG:
322 linmin[ik][spin] = linmin_test(g, d[ik][spin])
323 cg[ik][spin] = cg_test(atoms, ik, g, g_old[ik][spin], precondition)
324 beta, norm_g[ik][spin] = cg_method(
325 scf, ik, cgform, g, g_old[ik][spin], d_old[ik][spin], precondition
326 )
327 if precondition:
328 d[ik][spin] = -atoms.K(g, ik) + beta * d_old[ik][spin]
329 else:
330 d[ik][spin] = -g + beta * d_old[ik][spin]
331 W[ik][spin] = W[ik][spin] + betat * d[ik][spin]
332 gt = grad(scf, ik, spin, W, **scf._precomputed)
333 beta = abs(betat * dotprod(g, d[ik][spin]) / dotprod(g - gt, d[ik][spin]))
334 g_old[ik][spin], d_old[ik][spin] = g, d[ik][spin]
335 W[ik][spin] = W[ik][spin] + beta * d[ik][spin]
336 return costs, W
339@name("conjugate-gradient minimization")
340def cg(
341 scf,
342 W,
343 Nit,
344 cost=scf_step_occ,
345 grad=get_grad_occ,
346 condition=check_convergence,
347 betat=3e-5,
348 cgform=1,
349):
350 """Conjugate-gradient minimization algorithm for a fixed Hamiltonian.
352 Args:
353 scf: SCF object.
354 W: Expansion coefficients of unconstrained wave functions in reciprocal space.
355 Nit: Maximum number of SCF steps.
357 Keyword Args:
358 cost: Function that will run every SCF step.
359 grad: Function that calculates the respective gradient.
360 condition: Function to check and log the convergence condition.
361 betat: Step size.
362 cgform: Conjugate gradient form.
364 Returns:
365 Band energies per SCF cycle and optimized expansion coefficients.
366 """
367 return pccg(scf, W, Nit, cost, grad, condition, betat, cgform, precondition=False)
370@name("auto minimization")
371def auto(
372 scf,
373 W,
374 Nit,
375 cost=scf_step_occ,
376 grad=get_grad_occ,
377 condition=check_convergence,
378 betat=3e-5,
379 cgform=1,
380):
381 """Automatic precond. conjugate-gradient minimization algorithm for a fixed Hamiltonian.
383 This function chooses an sd step over the pccg step if the energy goes up.
385 Args:
386 scf: SCF object.
387 W: Expansion coefficients of unconstrained wave functions in reciprocal space.
388 Nit: Maximum number of SCF steps.
390 Keyword Args:
391 cost: Function that will run every SCF step.
392 grad: Function that calculates the respective gradient.
393 condition: Function to check and log the convergence condition.
394 betat: Step size.
395 cgform: Conjugate gradient form.
397 Returns:
398 Band energies per SCF cycle and optimized expansion coefficients.
399 """
400 atoms = scf.atoms
401 costs = []
403 linmin = xp.empty((atoms.kpts.Nk, atoms.occ.Nspin))
404 cg = xp.empty((atoms.kpts.Nk, atoms.occ.Nspin))
405 norm_g = xp.empty((atoms.kpts.Nk, atoms.occ.Nspin))
406 g = [xp.empty_like(Wk) for Wk in W]
407 d = [xp.empty_like(Wk) for Wk in W]
408 d_old = [xp.empty_like(Wk) for Wk in W]
409 g_old = [xp.empty_like(Wk) for Wk in W]
411 # Do the first step without the linmin and cg tests, and without the cg_method
412 for ik in range(atoms.kpts.Nk):
413 for spin in range(atoms.occ.Nspin):
414 g[ik][spin] = grad(scf, ik, spin, W, **scf._precomputed)
415 d[ik][spin] = -atoms.K(g[ik][spin], ik)
416 W[ik][spin] = W[ik][spin] + betat * d[ik][spin]
417 gt = grad(scf, ik, spin, W, **scf._precomputed)
418 beta = abs(
419 betat * dotprod(g[ik][spin], d[ik][spin]) / dotprod(g[ik][spin] - gt, d[ik][spin])
420 )
421 g_old[ik][spin], d_old[ik][spin] = g[ik][spin], d[ik][spin]
422 W[ik][spin] = W[ik][spin] + beta * d[ik][spin]
424 c = cost(scf, W)
425 costs.append(c)
426 if condition(scf, "pccg", costs):
427 return costs
429 for _ in range(1, Nit):
430 W_old = copy.deepcopy(W)
431 for ik in range(atoms.kpts.Nk):
432 for spin in range(atoms.occ.Nspin):
433 g[ik][spin] = grad(scf, ik, spin, W, **scf._precomputed)
434 # Calculate linmin and cg for each spin separately
435 if scf._log.level <= logging.DEBUG:
436 linmin[ik][spin] = linmin_test(g[ik][spin], d[ik][spin])
437 cg[ik][spin] = cg_test(atoms, ik, g[ik][spin], g_old[ik][spin])
438 beta, norm_g[ik][spin] = cg_method(
439 scf, ik, cgform, g[ik][spin], g_old[ik][spin], d_old[ik][spin]
440 )
441 d[ik][spin] = -atoms.K(g[ik][spin], ik) + beta * d_old[ik][spin]
442 W[ik][spin] = W[ik][spin] + betat * d[ik][spin]
443 gt = grad(scf, ik, spin, W, **scf._precomputed)
444 beta = abs(
445 betat
446 * dotprod(g[ik][spin], d[ik][spin])
447 / dotprod(g[ik][spin] - gt, d[ik][spin])
448 )
449 g_old[ik][spin], d_old[ik][spin] = g[ik][spin], d[ik][spin]
450 W[ik][spin] = W[ik][spin] + beta * d[ik][spin]
452 c = cost(scf, W)
453 # If the energy does not go down use the steepest descent step and recalculate the energy
454 if c > costs[-1]:
455 W = W_old
456 for ik in range(atoms.kpts.Nk):
457 for spin in range(atoms.occ.Nspin):
458 W[ik][spin] = W[ik][spin] - betat * g[ik][spin]
459 c = cost(scf, W)
460 costs.append(c)
461 # Do not print cg and linmin if we do the sd step
462 if condition(scf, "sd", costs, norm_g=norm_g):
463 break
464 else:
465 costs.append(c)
466 if condition(scf, "pccg", costs, linmin, cg, norm_g):
467 break
468 return costs, W
471#: Map minimizer names with their respective implementation.
472IMPLEMENTED = {
473 "sd": sd,
474 "lm": lm,
475 "pclm": pclm,
476 "cg": cg,
477 "pccg": pccg,
478 "auto": auto,
479}