Coverage for gpaw/directmin/derivatives.py: 92%
660 statements
« prev ^ index » next coverage.py v7.7.1, created at 2025-07-14 00:18 +0000
1import numpy as np
2from gpaw.directmin.tools import get_n_occ, get_indices, random_a, \
3 sort_orbitals_according_to_occ, sort_orbitals_according_to_energies
4from ase.units import Hartree
5from gpaw.mpi import world
6from gpaw.io.logger import GPAWLogger
7from gpaw.typing import RNG
8from copy import deepcopy
9from typing import Any, Dict, Union
12class Derivatives:
14 def __init__(self, etdm, wfs, c_ref=None, a=None,
15 update_c_ref=False, eps=1.0e-7,
16 random_amat: Union[RNG, bool] = False):
17 """
18 :param etdm:
19 :param wfs:
20 :param c_ref: reference orbitals C_ref
21 :param a: skew-Hermitian matrix A
22 :param update_c_ref: if True update reference orbitals
23 :param eps: finite difference displacement
24 :param random_amat:
25 if True, use random matrix A with the `numpy.random` global RNG
26 if an RNG, use random matrix A with said RNG
27 """
29 self.eps = eps
31 # initialize vectors of elements matrix A
32 if a is None:
33 if wfs.mode == 'lcao':
34 self.a = {u: np.zeros_like(v, dtype=wfs.dtype)
35 for u, v in etdm.a_vec_u.items()}
36 else:
37 self.a = {u: np.zeros_like(v, dtype=wfs.dtype)
38 for u, v in etdm.U_k.items()}
40 if random_amat:
41 extra_kwargs: Dict[str, Any] = {}
42 if random_amat not in (True, ): # Explicitly specified RNG
43 extra_kwargs.update(rng=random_amat)
44 for kpt in wfs.kpt_u:
45 u = etdm.kpointval(kpt)
46 if wfs.mode == 'lcao':
47 shape = etdm.a_vec_u[u].shape
48 else:
49 shape = etdm.U_k[u].shape
50 a = random_a(shape, wfs.dtype, **extra_kwargs)
51 wfs.gd.comm.broadcast(a, 0)
52 self.a[u] = a
54 # initialize orbitals
55 if c_ref is None and wfs.mode == 'lcao':
56 self.c_ref = etdm.dm_helper.reference_orbitals
57 else:
58 self.c_ref = c_ref
60 # update ref orbitals if needed
61 if update_c_ref:
62 if wfs.mode == 'lcao':
63 etdm.rotate_wavefunctions(wfs, self.a, self.c_ref)
64 etdm.dm_helper.set_reference_orbitals(wfs, etdm.n_dim)
65 self.c_ref = etdm.dm_helper.reference_orbitals
66 self.a = {u: np.zeros_like(v, dtype=wfs.dtype)
67 for u, v in etdm.a_vec_u.items()}
68 else:
69 etdm.rotate_wavefunctions(wfs, self.a)
70 for kpt in wfs.kpt_u:
71 k = etdm.kpointval(kpt)
72 etdm.psit_knG[k] = kpt.psit_nG.copy()
73 self.a = {u: np.zeros_like(v, dtype=wfs.dtype)
74 for u, v in etdm.U_k.items()}
76 def get_analytical_derivatives(self, etdm, ham, wfs, dens,
77 what2calc='gradient'):
78 """
79 Calculate analytical gradient or approximation to the Hessian
80 with respect to the elements of a skew-Hermitian matrix
82 :param etdm:
83 :param ham:
84 :param wfs:
85 :param dens:
86 :param what2calc: calculate gradient or Hessian
87 :return: analytical gradient or Hessian
88 """
90 assert what2calc in ['gradient', 'hessian']
92 if what2calc == 'gradient':
93 # calculate analytical gradient
94 if wfs.mode == 'lcao':
95 analytical_der = etdm.get_energy_and_gradients(
96 self.a, ham, wfs, dens, self.c_ref)[1]
97 else:
98 analytical_der = etdm.get_energy_and_gradients(
99 self.a, wfs, dens, ham)[1]
100 else:
101 # Calculate analytical approximation to hessian
102 if wfs.mode == 'lcao':
103 ind_up = etdm.ind_up
104 analytical_der = np.hstack([get_approx_analytical_hessian(
105 kpt, etdm.dtype,
106 ind_up[etdm.kpointval(kpt)]).copy() for kpt in wfs.kpt_u])
107 else:
108 analytical_der = np.hstack([get_approx_analytical_hessian(
109 kpt, etdm.dtype).copy() for kpt in wfs.kpt_u])
110 analytical_der = construct_real_hessian(analytical_der)
111 analytical_der = np.diag(analytical_der)
113 return analytical_der
115 def get_numerical_derivatives(self, etdm, ham, wfs, dens,
116 what2calc='gradient'):
117 """
118 Calculate numerical gradient or Hessian with respect to
119 the elements of a skew-Hermitian matrix using central finite
120 differences
122 :param etdm:
123 :param ham:
124 :param wfs:
125 :param dens:
126 :param what2calc: calculate gradient or Hessian
127 :return: numerical gradient or Hessian
128 """
130 assert what2calc in ['gradient', 'hessian']
132 if wfs.mode == 'lcao':
133 numerical_der = self.get_numerical_derivatives_lcao(
134 etdm, ham, wfs, dens, what2calc=what2calc)
135 else:
136 if what2calc == 'gradient':
137 numerical_der = self.get_numerical_gradient_fdpw(
138 etdm, wfs, dens, ham, self.a)
139 else:
140 numerical_der = self.get_numerical_hessian_fdpw(
141 etdm, wfs, dens, ham, self.a)
143 return numerical_der
145 def get_numerical_derivatives_lcao(
146 self, etdm, ham, wfs, dens, what2calc='gradient'):
148 # total dimensionality if matrices are real
149 dim = sum([len(a) for a in self.a.values()])
150 steps = [1.0, 1.0j] if etdm.dtype == complex else [1.0]
151 use_energy_or_gradient = {'gradient': 0, 'hessian': 1}
153 matrix_exp = etdm.matrix_exp
154 if what2calc == 'gradient':
155 numerical_der = {u: np.zeros_like(v)
156 for u, v in self.a.items()}
157 else:
158 numerical_der = np.zeros(shape=(len(steps) * dim,
159 len(steps) * dim))
160 # have to use exact gradient when Hessian is calculated
161 etdm.matrix_exp = 'egdecomp'
163 row = 0
164 f = use_energy_or_gradient[what2calc]
165 for step in steps:
166 for kpt in wfs.kpt_u:
167 u = etdm.kpointval(kpt)
168 for i in range(len(self.a[u])):
169 a = self.a[u][i]
171 self.a[u][i] = a + step * self.eps
172 fplus = etdm.get_energy_and_gradients(
173 self.a, ham, wfs, dens, self.c_ref)[f]
175 self.a[u][i] = a - step * self.eps
176 fminus = etdm.get_energy_and_gradients(
177 self.a, ham, wfs, dens, self.c_ref)[f]
179 derf = apply_central_finite_difference_approx(
180 fplus, fminus, self.eps)
182 if what2calc == 'gradient':
183 numerical_der[u][i] += step * derf
184 else:
185 numerical_der[row] = construct_real_hessian(derf)
187 row += 1
188 self.a[u][i] = a
190 if what2calc == 'hessian':
191 etdm.matrix_exp = matrix_exp
193 return numerical_der
195 def get_numerical_gradient_fdpw(self, etdm, wfs, dens, ham, A_s):
196 dtype = etdm.dtype
197 h = [self.eps, -self.eps]
198 coef = [1.0, -1.0]
199 Gr_n_x = {}
200 Gr_n_y = {}
202 if dtype == complex:
203 for kpt in wfs.kpt_u:
204 k = etdm.n_kps * kpt.s + kpt.q
205 dim = A_s[k].shape[0]
206 iut = np.triu_indices(dim, 1)
207 dim_gr = iut[0].shape[0]
209 for z in range(2):
210 grad_num = np.zeros(shape=dim_gr,
211 dtype=etdm.dtype)
212 igr = 0
213 for i, j in zip(*iut):
214 A = A_s[k][i][j]
215 for l in range(2):
216 if z == 1:
217 if i == j:
218 A_s[k][i][j] = A + 1.0j * h[l]
219 else:
220 A_s[k][i][j] = A + 1.0j * h[l]
221 A_s[k][j][i] = -np.conjugate(
222 A + 1.0j * h[l])
223 else:
224 if i == j:
225 A_s[k][i][j] = A + 0.0j * h[l]
226 else:
227 A_s[k][i][j] = A + h[l]
228 A_s[k][j][i] = -np.conjugate(
229 A + h[l])
230 E = etdm.get_energy_and_gradients(
231 A_s, wfs, dens, ham)[0]
232 grad_num[igr] += E * coef[l]
233 grad_num[igr] *= 1.0 / (2.0 * self.eps)
234 if i == j:
235 A_s[k][i][j] = A
236 else:
237 A_s[k][i][j] = A
238 A_s[k][j][i] = -np.conjugate(A)
239 igr += 1
240 if z == 0:
241 Gr_n_x[k] = grad_num.copy()
242 else:
243 Gr_n_y[k] = grad_num.copy()
245 Gr_n = {k: (Gr_n_x[k] + 1.0j * Gr_n_y[k]) for k in
246 Gr_n_x.keys()}
247 else:
248 for kpt in wfs.kpt_u:
249 k = etdm.n_kps * kpt.s + kpt.q
250 dim = A_s[k].shape[0]
251 iut = np.triu_indices(dim, 1)
252 dim_gr = iut[0].shape[0]
253 grad_num = np.zeros(shape=dim_gr, dtype=etdm.dtype)
255 igr = 0
256 for i, j in zip(*iut):
257 A = A_s[k][i][j]
258 for l in range(2):
259 A_s[k][i][j] = A + h[l]
260 A_s[k][j][i] = -(A + h[l])
261 E = etdm.get_energy_and_gradients(
262 A_s, wfs, dens, ham)[0]
263 grad_num[igr] += E * coef[l]
264 grad_num[igr] *= 1.0 / (2.0 * self.eps)
265 A_s[k][i][j] = A
266 A_s[k][j][i] = -A
267 igr += 1
269 Gr_n_x[k] = grad_num.copy()
271 Gr_n = {k: (Gr_n_x[k]) for k in Gr_n_x.keys()}
273 return Gr_n
275 def get_numerical_hessian_fdpw(self, etdm, wfs, dens, ham, A_s):
276 h = [self.eps, -self.eps]
277 coef = [1.0, -1.0]
278 num_hes = {}
280 for kpt in wfs.kpt_u:
281 k = etdm.n_kps * kpt.s + kpt.q
282 dim = A_s[k].shape[0]
283 iut = np.tril_indices(dim, -1)
284 dim_gr = iut[0].shape[0]
285 hessian = np.zeros(shape=(dim_gr, dim_gr),
286 dtype=etdm.dtype)
287 ih = 0
288 for i, j in zip(*iut):
289 A = A_s[k][i][j]
290 for l in range(2):
291 A_s[k][i][j] = A + h[l]
292 A_s[k][j][i] = -(A + h[l])
293 g = etdm.get_energy_and_gradients(A_s, wfs, dens, ham)[1]
294 g = g[k][iut]
295 hessian[ih, :] += g * coef[l]
297 hessian[ih, :] *= 1.0 / (2.0 * self.eps)
299 A_s[k][i][j] = A
300 A_s[k][j][i] = -A
301 ih += 1
303 num_hes[k] = hessian.copy()
305 return num_hes
308class Davidson:
309 """
310 Finite difference generalized Davidson partial diagonalizer to obtain a
311 number of the eigenpairs with the smallest eigenvalues.
313 The following array indexation convention is used:
315 e: Target eigenpair
316 w: Krylov subspace
317 """
319 def __init__(self, etdm, logfile=None, fd_mode=None, m=None, h=None,
320 eps=None, cap_krylov=None, gmf=False,
321 accurate_first_pdiag=True, remember_sp_order=None,
322 sp_order=None, seed=None):
323 """
324 :param etdm: ETDM object for which the partial eigendecomposition
325 should be performed.
326 :param logfile: Name string of the Davidson log file. Use '-' for
327 stdout or None to discard.
328 :param fd_mode: Finite difference mode for partial Hessian evaluation.
329 Must be one of 'central' for central FD or 'forward'
330 for forward FD. Central FD uses two e/g evaluations per
331 Davidson step and target eigenpair with an error
332 scaling as O(h^2), forward FD uses one with O(h)
333 scaling.
334 :param m: Memory parameter indicating how large the Krylov space should
335 be able to become before resetting it with the Ritz vectors
336 of the previous step or terminating the calculation if
337 cap_krylov is True.
338 :param h: Displacement (in radians of orbital rotation) for finite
339 difference partial Hessian calculation.
340 :param eps: Convergence threshold for maximum component of the
341 residuals of the target eigenpairs.
342 :param cap_krylov: If True, terminate the calculation if the Krylov
343 space contains more than m vectors.
344 :param gmf: Toggle usage with generalized mode following instead of
345 stability analysis. The defaults and some actions will be
346 different.
347 :param accurate_first_pdiag: Approximate the target saddle point order
348 better by performing a more accurate first
349 partial diagonalization step at the
350 expense of more gradient evaluations
351 :param remember_sp_order: If True the number of target eigenpairs is
352 saved after converging the partial Hessian
353 eigendecomposition once and recovered for all
354 subsequent calculations. If False the number
355 of target eigenpairs is always gathered from
356 the diagonal Hessian approximation in ETDM.
357 :param sp_order: If given use this value for the number of target
358 eigenpairs instead of gathering it from the diagonal
359 Hessian approximation in ETDM.
360 :param seed: Seed for random perturbation of initial Krylov space.
361 """
363 self.name = 'Davidson'
364 self.gmf = gmf
365 self.etdm = etdm
366 self.fd_mode = fd_mode
367 self.remember_sp_order = remember_sp_order
368 self.sp_order = sp_order
369 self.log_sp_order_once = True
370 self.seed = seed
371 self.V_w = [] # Krylov subspace
372 self.C_we = [] # Preconditioner
373 self.M = [] # Matrix used to get preconditioner every step
374 self.W_w = [] # Hessian effect on Krylov subspace
375 # Rayleigh matrix, smaller representation of the
376 # diagonalization problem to solve:
377 self.H_ww = []
378 self.lambda_e = [] # Target eigenvalues
379 self.y_e = [] # Target eigenvectors in subspace representation
380 self.x_e = [] # Target eigenvectors
381 self.r_e = [] # Residuals of target eigenvectors
382 self.t_e = [] # Krylov space extension vectors
383 self.l = None # Number of target eigenpairs
384 self.h = h
385 self.m = m
386 self.converged_e = False
387 self.all_converged = False
388 self.error_e = []
389 self.n_iter = 0
390 self.eigenvalues = []
391 self.eigenvectors = []
392 self.reset = False
393 self.eps = eps
394 self.grad = None
395 self.cap_krylov = cap_krylov
396 self.dim_u = {}
397 self.dimtot = 0
398 self.nocc = {}
399 self.nbands = 0
400 self.c_ref = []
401 self.logfile = logfile
402 self.logger = GPAWLogger(world)
403 self.logger.fd = logfile
404 self.first_run = accurate_first_pdiag
405 self.lambda_w = [] # All eigenvalues
406 self.y_w = [] # All eigenvectors in subspace representation
407 self.x_w = [] # All eigenvectors
408 self.check_inputs()
410 def check_inputs(self):
411 defaults = self.set_defaults()
412 assert self.etdm.name == 'etdm-lcao', 'Check etdm.'
413 if self.logfile is not None:
414 assert isinstance(self.logfile, str), 'Check logfile.'
415 if self.fd_mode is None:
416 self.fd_mode = defaults['fd_mode']
417 else:
418 assert self.fd_mode in ['central', 'forward'], 'Check fd_mode.'
419 if self.m is None:
420 self.m = defaults['m']
421 else:
422 assert isinstance(self.m, int) or np.isinf(self.m), 'Check m.'
423 if self.h is None:
424 self.h = defaults['h']
425 else:
426 assert isinstance(self.h, float), 'Check h.'
427 if self.eps is None:
428 self.eps = defaults['eps']
429 else:
430 assert isinstance(self.eps, float), 'Check eps.'
431 if self.cap_krylov is None:
432 self.cap_krylov = defaults['cap_krylov']
433 else:
434 assert isinstance(self.cap_krylov, bool), 'Check cap_krylov.'
435 if self.remember_sp_order is None:
436 self.remember_sp_order = defaults['remember_sp_order']
437 else:
438 assert isinstance(self.remember_sp_order, bool), \
439 'Check remember_sp_order.'
440 if self.sp_order is not None:
441 assert isinstance(self.sp_order, int), 'Check sp_order.'
443 def set_defaults(self):
444 if self.gmf:
445 return {'fd_mode': 'forward',
446 'm': 300,
447 'h': 1e-3,
448 'eps': 1e-2,
449 'cap_krylov': True,
450 'remember_sp_order': True}
451 else:
452 return {'fd_mode': 'central',
453 'm': np.inf,
454 'h': 1e-3,
455 'eps': 1e-3,
456 'cap_krylov': False,
457 'remember_sp_order': False}
459 def todict(self):
460 return {'name': 'Davidson',
461 'logfile': self.logfile,
462 'fd_mode': self.fd_mode,
463 'm': self.m,
464 'h': self.h,
465 'eps': self.eps,
466 'cap_krylov': self.cap_krylov,
467 'gmf': self.gmf,
468 'remember_sp_order': self.remember_sp_order,
469 'sp_order': self.sp_order}
471 def introduce(self):
472 self.logger(
473 '|-------------------------------------------------------|')
474 self.logger(
475 '| Davidson partial diagonalizer |')
476 self.logger(
477 '|-------------------------------------------------------|\n',
478 flush=True)
480 def run(self, wfs, ham, dens, use_prev=False):
481 self.initialize(wfs, use_prev)
482 if not self.gmf:
483 sort_orbitals_according_to_occ(
484 wfs, self.etdm.constraints, update_mom=True)
485 self.n_iter = 0
486 self.c_ref = [deepcopy(wfs.kpt_u[x].C_nM)
487 for x in range(len(wfs.kpt_u))]
488 if self.fd_mode == 'forward' and self.grad is None:
489 self.obtain_grad_at_c_ref(wfs, ham, dens)
490 while not self.all_converged:
491 self.iterate(wfs, ham, dens)
492 if self.remember_sp_order:
493 if self.sp_order is None:
494 self.determine_sp_order_from_lambda()
495 self.logger(
496 'Saved target saddle point order as '
497 + str(self.sp_order) + ' for future partial '
498 'diagonalizations.', flush=True)
499 elif self.log_sp_order_once:
500 self.log_sp_order_once = False
501 self.logger(
502 'Using target saddle point order of '
503 + str(self.sp_order) + '.', flush=True)
504 if self.gmf:
505 self.obtain_x_in_krylov_subspace()
506 for k, kpt in enumerate(wfs.kpt_u):
507 kpt.C_nM = deepcopy(self.c_ref[k])
508 if not self.gmf:
509 sort_orbitals_according_to_energies(
510 ham, wfs, self.etdm.constraints)
511 self.first_run = False
513 def obtain_grad_at_c_ref(self, wfs, ham, dens):
514 a_vec_u = {}
515 n_dim = {}
516 for k, kpt in enumerate(wfs.kpt_u):
517 n_dim[k] = wfs.bd.nbands
518 a_vec_u[k] = np.zeros_like(self.etdm.a_vec_u[k])
519 self.grad = self.etdm.get_energy_and_gradients(
520 a_vec_u, ham, wfs, dens, self.c_ref)[1]
522 def determine_sp_order_from_lambda(self):
523 sp_order = 0
524 for i in range(len(self.lambda_w)):
525 if self.lambda_w[i] < 1e-8:
526 sp_order += 1
527 else:
528 break
529 self.sp_order = sp_order
530 if self.sp_order == 0:
531 self.sp_order = 1
533 def obtain_x_in_krylov_subspace(self):
534 self.x_w = []
535 for i in range(len(self.lambda_w)):
536 self.x_w.append(
537 self.V_w[:, :len(self.lambda_w)] @ self.y_w[i].T)
538 self.x_w = np.asarray(self.x_w).T
540 def initialize(self, wfs, use_prev=False):
541 """
542 This is separate from __init__ since the initial Krylov space is
543 obtained here every time a partial diagonalization is performed at
544 different electronic coordinates.
545 """
547 dimz = 2 if self.etdm.dtype == complex else 1
548 self.introduce()
549 self.reset = False
550 self.all_converged = False
551 self.l = 0
552 self.V_w = None
553 self.nbands = wfs.bd.nbands
554 appr_hess, appr_sp_order = self.estimate_spo_and_update_appr_hess(
555 wfs, use_prev=use_prev)
556 self.M = np.zeros(shape=self.dimtot * dimz)
557 for i in range(self.dimtot * dimz):
558 self.M[i] = np.real(appr_hess[i % self.dimtot])
559 if self.sp_order is not None:
560 self.l = self.sp_order
561 else:
562 self.l = appr_sp_order if self.gmf else appr_sp_order + 2
563 if self.l == 0:
564 self.l = 1
565 if self.l > self.dimtot * dimz:
566 self.l = self.dimtot * dimz
567 self.W_w = None
568 self.error_e = [np.inf for x in range(self.l)]
569 self.converged_e = [False for x in range(self.l)]
570 self.form_initial_krylov_subspace(
571 wfs, appr_hess, dimz, use_prev=use_prev)
572 text = 'Davidson will target the ' + str(self.l) + ' lowest eigenpairs'
573 if self.sp_order is None:
574 text += '.'
575 else:
576 text += ' as recovered from previous calculation.'
577 self.logger(text, flush=True)
579 def get_approximate_hessian_and_dim(self, wfs):
580 appr_hess = []
581 self.dimtot = 0
582 for k, kpt in enumerate(wfs.kpt_u):
583 hdia = get_approx_analytical_hessian(
584 kpt, self.etdm.dtype, ind_up=self.etdm.ind_up[k])
585 self.dim_u[k] = len(hdia)
586 self.dimtot += len(hdia)
587 appr_hess += list(hdia.copy())
588 self.nocc[k] = get_n_occ(kpt)[0]
589 return appr_hess
591 def estimate_spo_and_update_appr_hess(self, wfs, use_prev=False):
592 appr_sp_order = 0
593 appr_hess = self.get_approximate_hessian_and_dim(wfs)
594 if use_prev:
595 for i in range(len(self.lambda_w)):
596 if self.lambda_w[i] < -1e-4:
597 appr_sp_order += 1
598 if self.etdm.dtype == complex:
599 appr_hess[i] \
600 = self.lambda_w[i] + 1.0j * self.lambda_w[i]
601 else:
602 appr_hess[i] = self.lambda_w[i]
603 else:
604 for i in range(len(appr_hess)):
605 if np.real(appr_hess[i]) < -1e-4:
606 appr_sp_order += 1
607 return appr_hess, appr_sp_order
609 def form_initial_krylov_subspace(
610 self, wfs, appr_hess, dimz, use_prev=False):
611 rng = np.random.default_rng(self.seed)
612 wfs.timer.start('Initial Krylov subspace')
613 if use_prev:
614 self.randomize_krylov_subspace(rng, dimz)
615 else:
616 self.initialize_randomized_krylov_subspace(rng, dimz, appr_hess)
617 wfs.timer.start('Modified Gram-Schmidt')
618 self.V_w = mgs(self.V_w)
619 wfs.timer.stop('Modified Gram-Schmidt')
620 self.V_w = self.V_w.T
621 wfs.timer.stop('Initial Krylov subspace')
623 def randomize_krylov_subspace(self, rng, dimz, reps=1e-4):
624 self.V_w = deepcopy(self.x_w.T[: self.l])
625 for i in range(self.l):
626 for k in range(self.dimtot):
627 for l in range(dimz):
628 rand = np.zeros(shape=2)
629 if world.rank == 0:
630 rand[0] = rng.random()
631 rand[1] = 1 if rng.random() > 0.5 else -1
632 else:
633 rand[0] = 0.0
634 rand[1] = 0.0
635 world.broadcast(rand, 0)
636 self.V_w[i][l * self.dimtot + k] \
637 += rand[1] * reps * rand[0]
639 def initialize_randomized_krylov_subspace(
640 self, rng, dimz, appr_hess, reps=1e-4):
641 do_conj = False
643 # Just for F821
644 v = None
645 imin = None
647 self.V_w = []
648 for i in range(self.l):
649 if do_conj:
650 v[self.dimtot + imin] = -1.0
651 do_conj = False
652 else:
653 v = np.zeros(self.dimtot * dimz)
654 rdia = np.real(appr_hess).copy()
655 imin = int(np.where(rdia == min(rdia))[0][0])
656 rdia[imin] = np.inf
657 v[imin] = 1.0
658 if self.etdm.dtype == complex:
659 v[self.dimtot + imin] = 1.0
660 do_conj = True
661 for l in range(self.dimtot):
662 for m in range(dimz):
663 if l == imin:
664 continue
665 rand = np.zeros(shape=2)
666 if world.rank == 0:
667 rand[0] = rng.random()
668 rand[1] = 1 if rng.random() > 0.5 else -1
669 else:
670 rand[0] = 0.0
671 rand[1] = 0.0
672 world.broadcast(rand, 0)
673 v[m * self.dimtot + l] = rand[1] * reps * rand[0]
674 self.V_w.append(v / np.linalg.norm(v))
675 self.V_w = np.asarray(self.V_w)
677 def iterate(self, wfs, ham, dens):
678 self.t_e = []
679 self.evaluate_W(wfs, ham, dens)
680 wfs.timer.start('Rayleigh matrix formation')
681 self.H_ww = self.V_w.T @ self.W_w
682 wfs.timer.stop('Rayleigh matrix formation')
683 self.n_iter += 1
684 wfs.timer.start('Rayleigh matrix diagonalization')
685 eigv, eigvec = np.linalg.eigh(self.H_ww)
686 wfs.timer.stop('Rayleigh matrix diagonalization')
687 eigvec = eigvec.T
688 self.lambda_w = deepcopy(eigv)
689 self.y_w = deepcopy(eigvec)
690 self.lambda_e = eigv[: self.l]
691 self.y_e = eigvec[: self.l]
692 self.calculate_ritz_vectors(wfs)
693 self.calculate_residuals(wfs)
694 for i in range(self.l):
695 self.error_e[i] = np.abs(self.r_e[i]).max()
696 self.all_converged = True
697 for i in range(self.l):
698 self.converged_e[i] = self.error_e[i] < self.eps
699 self.all_converged = self.all_converged and self.converged_e[i]
700 if self.all_converged:
701 self.eigenvalues = deepcopy(self.lambda_e)
702 self.eigenvectors = deepcopy(self.x_e)
703 self.log()
704 return
705 self.calculate_preconditioner(wfs)
706 self.augment_krylov_subspace(wfs)
707 self.log()
709 def evaluate_W(self, wfs, ham, dens):
710 wfs.timer.start('FD Hessian vector product')
711 if self.W_w is None:
712 self.W_w = []
713 Vt = self.V_w.T
714 for i in range(len(Vt)):
715 self.W_w.append(self.get_fd_hessian(Vt[i], wfs, ham, dens))
716 self.reset = False
717 else:
718 added = len(self.V_w[0]) - len(self.W_w[0])
719 self.W_w = self.W_w.T.tolist()
720 Vt = self.V_w.T
721 for i in range(added):
722 self.W_w.append(self.get_fd_hessian(
723 Vt[-added + i], wfs, ham, dens))
724 self.W_w = np.asarray(self.W_w).T
725 wfs.timer.stop('FD Hessian vector product')
727 def calculate_ritz_vectors(self, wfs):
728 wfs.timer.start('Ritz vector calculation')
729 self.x_e = []
730 for i in range(self.l):
731 self.x_e.append(self.V_w @ self.y_e[i].T)
732 self.x_e = np.asarray(self.x_e)
733 wfs.timer.stop('Ritz vector calculation')
735 def calculate_residuals(self, wfs):
736 wfs.timer.start('Residual calculation')
737 self.r_e = []
738 for i in range(self.l):
739 self.r_e.append(
740 self.x_e[i] * self.lambda_e[i] - self.W_w @ self.y_e[i].T)
741 self.r_e = np.asarray(self.r_e)
742 wfs.timer.stop('Residual calculation')
744 def calculate_preconditioner(self, wfs):
745 n_dim = len(self.V_w)
746 wfs.timer.start('Preconditioner calculation')
747 self.C_we = np.zeros(shape=(self.l, n_dim))
748 for i in range(self.l):
749 self.C_we[i] = -np.abs(
750 np.repeat(self.lambda_e[i], n_dim) - self.M) ** -1
751 for l in range(len(self.C_we[i])):
752 if self.C_we[i][l] > -0.1 * Hartree:
753 self.C_we[i][l] = -0.1 * Hartree
754 wfs.timer.stop('Preconditioner calculation')
756 def augment_krylov_subspace(self, wfs):
757 wfs.timer.start('Krylov subspace augmentation')
758 self.get_new_krylov_subspace_directions(wfs)
759 self.V_w = np.asarray(self.V_w)
760 if self.cap_krylov:
761 if len(self.V_w) > self.m:
762 self.logger(
763 'Krylov space exceeded maximum size. Partial '
764 'diagonalization is not fully converged. Current size '
765 'is ' + str(len(self.V_w) - len(self.t_e)) + '. Size at '
766 'next step would be ' + str(len(self.V_w)) + '.',
767 flush=True)
768 self.all_converged = True
769 wfs.timer.start('Modified Gram-Schmidt')
770 self.V_w = mgs(self.V_w)
771 wfs.timer.stop('Modified Gram-Schmidt')
772 self.V_w = self.V_w.T
773 wfs.timer.stop('Krylov subspace augmentation')
775 def get_new_krylov_subspace_directions(self, wfs):
776 wfs.timer.start('New directions')
777 for i in range(self.l):
778 if not self.converged_e[i] or (self.gmf and self.first_run):
779 self.t_e.append(self.C_we[i] * self.r_e[i])
780 self.t_e = np.asarray(self.t_e)
781 if len(self.V_w[0]) <= self.m:
782 self.V_w = self.V_w.T.tolist()
783 for i in range(len(self.t_e)):
784 self.V_w.append(self.t_e[i])
785 elif not self.cap_krylov:
786 self.reset = True
787 self.V_w = deepcopy(self.x_e.tolist())
788 for i in range(len(self.t_e)):
789 self.V_w.append(self.t_e[i])
790 self.W_w = None
791 wfs.timer.stop('New directions')
793 def log(self):
794 self.logger('Dimensionality of Krylov space: '
795 + str(len(self.V_w[0]) - len(self.t_e)), flush=True)
796 if self.reset:
797 self.logger('Reset Krylov space', flush=True)
798 self.logger('\nEigenvalues:\n', flush=True)
799 text = ''
800 for i in range(self.l):
801 text += '%10d '
802 indices = text % tuple(range(1, self.l + 1))
803 self.logger(indices, flush=True)
804 text = ''
805 for i in range(self.l):
806 text += '%10.6f '
807 self.logger(text % tuple(self.lambda_e), flush=True)
808 self.logger('\nResidual maximum components:\n', flush=True)
809 self.logger(indices, flush=True)
810 text = ''
811 temp = list(np.round(deepcopy(self.error_e), 6))
812 for i in range(self.l):
813 text += '%10s '
814 temp[i] = str(temp[i])
815 if self.converged_e[i]:
816 temp[i] += 'c'
817 self.logger(text % tuple(temp), flush=True)
819 def get_fd_hessian(self, vin, wfs, ham, dens):
820 """
821 Get the dot product of the Hessian and a vector with a finite
822 difference approximation.
824 :param vin: The vector
825 :param wfs:
826 :param ham:
827 :param dens:
828 :return: Dot product vector of the Hessian and the vector.
829 """
831 v = self.h * vin
832 c_ref = deepcopy(self.c_ref)
833 gp = self.calculate_displaced_grad(wfs, ham, dens, c_ref, v)
834 hessi = []
835 if self.fd_mode == 'central':
836 gm = self.calculate_displaced_grad(wfs, ham, dens, c_ref, -1.0 * v)
837 for k in range(len(wfs.kpt_u)):
838 hessi += list((gp[k] - gm[k]) * 0.5 / self.h)
839 elif self.fd_mode == 'forward':
840 for k in range(len(wfs.kpt_u)):
841 hessi += list((gp[k] - self.grad[k]) / self.h)
842 for k, kpt in enumerate(wfs.kpt_u):
843 kpt.C_nM = c_ref[k] # LCAO-specific
844 dens.update(wfs)
845 if self.etdm.dtype == complex:
846 hessc = np.zeros(shape=(2 * self.dimtot))
847 hessc[: self.dimtot] = np.real(hessi)
848 hessc[self.dimtot:] = np.imag(hessi)
849 return hessc
850 else:
851 return np.asarray(hessi)
853 def calculate_displaced_grad(self, wfs, ham, dens, c_ref, v):
854 a_vec_u = {}
855 n_dim = {}
856 start = 0
857 end = 0
858 for k in range(len(wfs.kpt_u)):
859 a_vec_u[k] = np.zeros_like(self.etdm.a_vec_u[k])
860 n_dim[k] = wfs.bd.nbands
861 end += self.dim_u[k]
862 a_vec_u[k] += v[start: end]
863 if self.etdm.dtype == complex:
864 a_vec_u[k] += 1.0j * v[self.dimtot + start: self.dimtot + end]
865 start += self.dim_u[k]
866 return self.etdm.get_energy_and_gradients(
867 a_vec_u, ham, wfs, dens, c_ref)[1]
869 def break_instability(self, wfs, n_dim, c_ref, number,
870 initial_guess='displace', ham=None, dens=None):
871 """
872 Displaces orbital rotation coordinates in the direction of an
873 instability. Uses a fixed displacement or performs a line search.
875 :param wfs:
876 :param n_dim:
877 :param c_ref:
878 :param number: Instability index
879 :param initial_guess: How to displace. Can be one of the following:
880 displace: Use a fixed displacement; line_search: Performs a
881 backtracking line search.
882 :param ham:
883 :param dens:
884 """
886 assert self.converged_e, 'Davidson cannot break instabilities since' \
887 + ' the partial eigendecomposition has not been converged.'
888 assert len(self.lambda_e) >= number, 'Davidson cannot break' \
889 + ' instability no. ' + str(number) + ' since this eigenpair was' \
890 + 'not converged.'
891 assert self.lambda_e[number - 1] < 0.0, 'Eigenvector no. ' \
892 + str(number) + ' does not represent an instability.'
894 step = self.etdm.line_search.max_step
895 instability = step * self.x_e[number - 1]
896 if initial_guess == 'displace':
897 a_vec_u = self.displace(instability)
898 elif initial_guess == 'line_search':
899 assert ham is not None and dens is not None, 'Hamiltonian and' \
900 'density needed for' \
901 'line search.'
902 a_vec_u = self.do_line_search(
903 wfs, dens, ham, n_dim, c_ref, instability)
904 self.etdm.rotate_wavefunctions(wfs, a_vec_u, c_ref)
906 def displace(self, instability):
907 a_vec_u = deepcopy(self.etdm.a_vec_u)
908 start = 0
909 stop = 0
910 for k in a_vec_u.keys():
911 stop += self.dim_u[k]
912 a_vec_u[k] = instability[start: stop]
913 start += self.dim_u[k]
914 return a_vec_u
916 def do_line_search(self, wfs, dens, ham, n_dim, c_ref, instability):
917 a_vec_u = {}
918 p_vec_u = {}
919 start = 0
920 stop = 0
921 for k in a_vec_u.keys():
922 stop += self.dim_u[k]
923 p_vec_u[k] = instability[start: stop]
924 start += self.dim_u[k]
925 phi, g_vec_u = self.etdm.get_energy_and_gradients(
926 a_vec_u, ham, wfs, dens, c_ref)
927 der_phi = 0.0
928 for k in g_vec_u:
929 der_phi += g_vec_u[k].conj() @ p_vec_u[k]
930 der_phi = der_phi.real
931 der_phi = wfs.kd.comm.sum(der_phi)
932 alpha = self.etdm.line_search.step_length_update(
933 a_vec_u, p_vec_u, n_dim, ham, wfs, dens, c_ref,
934 phi_0=phi, der_phi_0=der_phi, phi_old=None,
935 der_phi_old=None, alpha_max=5.0, alpha_old=None,
936 kpdescr=wfs.kd)[0]
937 for k in a_vec_u.keys():
938 a_vec_u[k] = alpha * p_vec_u[k]
939 return a_vec_u
941 def estimate_sp_order(self, calc, method='appr-hess', target_more=1):
942 sort_orbitals_according_to_occ(
943 calc.wfs, self.etdm.constraints, update_mom=method == 'full-hess')
944 appr_hess, appr_sp_order = self.estimate_spo_and_update_appr_hess(
945 calc.wfs)
946 if calc.wfs.dtype == complex:
947 appr_sp_order *= 2
948 if method == 'full-hess':
949 constraints_copy = deepcopy(self.etdm.constraints)
950 self.etdm.constraints = [[] for _ in range(len(calc.wfs.kpt_u))]
951 self.sp_order = appr_sp_order + target_more
952 self.run(calc.wfs, calc.hamiltonian, calc.density)
953 appr_hess, appr_sp_order = self.estimate_spo_and_update_appr_hess(
954 calc.wfs, use_prev=True)
955 self.etdm.constraints = deepcopy(constraints_copy)
956 sort_orbitals_according_to_energies(
957 calc.hamiltonian, calc.wfs, self.etdm.constraints)
958 return appr_sp_order
961def mgs(vin):
962 """
963 Modified Gram-Schmidt orthonormalization
965 :param vin: Set of vectors.
966 :return: Orthonormal set of vectors.
967 """
969 v = deepcopy(vin)
970 q = np.zeros_like(v)
971 for i in range(len(v)):
972 q[i] = v[i] / np.linalg.norm(v[i])
973 for k in range(len(v)):
974 v[k] = v[k] - np.dot(np.dot(q[i].T, v[k]), q[i])
975 return q
978def construct_real_hessian(hess):
980 if hess.dtype == complex:
981 hess_real = np.hstack((np.real(hess), np.imag(hess)))
982 else:
983 hess_real = hess
985 return hess_real
988def apply_central_finite_difference_approx(fplus, fminus, eps):
990 if isinstance(fplus, dict) and isinstance(fminus, dict):
991 assert (len(fplus) == len(fminus))
992 derf = np.hstack([(fplus[k] - fminus[k]) * 0.5 / eps
993 for k in fplus.keys()])
994 elif isinstance(fplus, float) and isinstance(fminus, float):
995 derf = (fplus - fminus) * 0.5 / eps
996 else:
997 raise ValueError()
999 return derf
1002def get_approx_analytical_hessian(kpt, dtype, ind_up=None):
1003 """
1004 Calculate the following diagonal approximation to the Hessian:
1005 h_{lm, lm} = -2.0 * (eps_n[l] - eps_n[m]) * (f[l] - f[m])
1006 """
1008 f_n = kpt.f_n
1009 eps_n = kpt.eps_n
1010 if ind_up:
1011 il1 = list(ind_up)
1012 else:
1013 il1 = list(get_indices(eps_n.shape[0]))
1015 hess = np.zeros(len(il1[0]), dtype=dtype)
1016 x = 0
1017 for n, m in zip(*il1):
1018 df = f_n[n] - f_n[m]
1019 hess[x] = -2.0 * (eps_n[n] - eps_n[m]) * df
1020 if abs(hess[x]) < 1.0e-10:
1021 hess[x] = 0.0
1022 if dtype == complex:
1023 hess[x] += 1.0j * hess[x]
1024 x += 1
1026 return hess