Source code for

"""Module for calculations using the Maximum Overlap Method (MOM).


import numpy as np

from ase.units import Ha

from gpaw.occupations import FixedOccupationNumbers, ParallelLayout

[docs]def prepare_mom_calculation(calc, atoms, numbers, use_projections=True, update_numbers=True, use_fixed_occupations=False, width=0.0, niter_width_update=40, width_increment=0.0): """Helper function to prepare a calculator for a MOM calculation. calc: GPAW instance GPAW calculator object. atoms: ASE instance ASE atoms object. numbers: list (len=nspins) of lists (len=nbands) Occupation numbers (in the range from 0 to 1). Used for the initialization of the MOM reference orbitals. use_projections: bool If True (default), the occupied orbitals at iteration k are chosen as the orbitals ``|psi^(k)_m>`` with the biggest weights ``P_m`` evaluated as the projections onto the manifold of reference orbitals ``|psi_n>``: ``P_m = (Sum_n(|O_nm|^2))^0.5 (O_nm = <psi_n|psi^(k)_m>)`` see :doi:`10.1021/acs.jctc.7b00994`. If False, the weights are evaluated as: ``P_m = max_n(|O_nm|)``, see :doi:`10.1021/acs.jctc.0c00488`. update_numbers: bool If True (default), 'numbers' gets updated with the calculated occupation numbers, and when changing atomic positions the MOM reference orbitals will be initialized as the occupied orbitals found at convergence for the previous geometry. If False, when changing positions the MOM reference orbitals will be initialized as the orbitals of the previous geometry corresponding to the user-supplied 'numbers'. use_fixed_occupations: bool If False (default), the MOM algorithm is used. If True, fixed occupations will be used. width: float Width of Gaussian function in eV for smearing of holes and excited electrons. The holes and excited electrons are found with respect to the zero-width ground-state occupations. See :doi:`10.1021/acs.jctc.0c00597`. niter_width_update: int Number of iterations after which the width of the Gaussian smearing function is increased. width_increment: float How much to increase the width of the Gaussian smearing function. """ if calc.wfs is None: # We need the wfs object to initialize OccupationsMOM calc.initialize(atoms) occ_mom = OccupationsMOM(calc.wfs, numbers, use_projections, update_numbers, use_fixed_occupations, width, niter_width_update, width_increment) calc.set(occupations=occ_mom) calc.log(occ_mom)
class OccupationsMOM: """MOM occupation class. The occupation numbers are found using a maximum overlap criterion and then broadcasted in using the _calculate method of FixedOccupationNumbers. """ def __init__(self, wfs, numbers, use_projections=False, update_numbers=True, use_fixed_occupations=False, width=0.0, niter_width_update=10, width_increment=0.0): self.wfs = wfs self.numbers = np.array(numbers) self.use_projections = use_projections self.update_numbers = update_numbers self.use_fixed_occupations = use_fixed_occupations self.width = width / Ha self.niter_width_update = niter_width_update self.width_increment = width_increment / Ha parallel_layout = ParallelLayout(, self.wfs.kd.comm, self.occ = FixedOccupationNumbers(numbers, parallel_layout) self.extrapolate_factor = self.occ.extrapolate_factor = 'mom' self.iters = 0 self.initialized = False def todict(self): dct = {'name':, 'numbers': self.numbers, 'use_projections': self.use_projections, 'update_numbers': self.update_numbers, 'use_fixed_occupations': self.use_fixed_occupations} if self.width != 0.0: dct['width'] = self.width * Ha dct['niter_width_update'] = self.niter_width_update dct['width_increment'] = self.width_increment * Ha return dct def __str__(self): s = 'Excited-state calculation with Maximum Overlap Method\n' s += ' Gaussian smearing of holes and excited electrons: ' if self.width == 0.0: s += 'off\n' else: s += f'{self.width * Ha:.4f} eV\n' return s def calculate(self, nelectrons, eigenvalues, weights, fermi_levels_guess): if not self.initialized: # If MOM reference orbitals are not initialized yet (e.g. when # the calculation is initialized from atomic densities), update # the occupation numbers according to the user-supplied 'numbers' self.occ.f_sn = self.numbers.copy() self.initialize_reference_orbitals() else: self.occ.f_sn = self.update_occupations() self.iters += 1 f_qn, fermi_levels, e_entropy = self.occ.calculate(nelectrons, eigenvalues, weights, fermi_levels_guess) return f_qn, fermi_levels, e_entropy def initialize_reference_orbitals(self): if self.wfs.kpt_u[0].f_n is None: # If the occupation numbers are not already available # (e.g. when the calculation is initialized from atomic # densities) we first need to take a step of eigensolver # and update the occupation numbers according to the # 'user-supplied' numbers before initializing the MOM # reference orbitals return self.iters = 0 if self.use_fixed_occupations: self.initialized = True return # Initialize MOM reference orbitals for each equally # occupied subspace separately self.f_sn_unique = self.find_unique_occupation_numbers() if self.wfs.mode == 'lcao': self.c_ref = {} for kpt in self.wfs.kpt_u: self.c_ref[kpt.s] = {} for f_n_unique in self.f_sn_unique[kpt.s]: occupied = self.f_sn_unique[kpt.s][f_n_unique] self.c_ref[kpt.s][f_n_unique] = kpt.C_nM[occupied].copy() else: = {} self.p_an = {} for kpt in self.wfs.kpt_u:[kpt.s] = {} self.p_an[kpt.s] = {} for f_n_unique in self.f_sn_unique[kpt.s]: occupied = self.f_sn_unique[kpt.s][f_n_unique] # Pseudo wave functions[kpt.s][f_n_unique] = kpt.psit_nG[occupied].copy() # Atomic contributions times projector overlaps self.p_an[kpt.s][f_n_unique] = \ {a:[a].dO_ii, P_ni[occupied].T) for a, P_ni in kpt.P_ani.items()} self.initialized = True def update_occupations(self): if self.width != 0.0: if self.iters == 0: self.width_update_counter = 0 if self.iters % self.niter_width_update == 0: self.gauss_width = self.width + self.width_update_counter \ * self.width_increment self.width_update_counter += 1 if not self.use_fixed_occupations: f_sn = np.zeros_like(self.numbers) for kpt in self.wfs.kpt_u: # Compute weights with respect to each equally occupied # subspace of the reference orbitals and occupy orbitals # with biggest weights for f_n_unique in self.f_sn_unique[kpt.s]: occupied = self.f_sn_unique[kpt.s][f_n_unique] n_occ = len(f_sn[kpt.s][occupied]) unoccupied = f_sn[kpt.s] == 0 P = np.zeros(len(f_sn[kpt.s])) P[unoccupied] = self.calculate_weights(kpt, f_n_unique, unoccupied) P_max = np.argpartition(P, -n_occ)[-n_occ:] f_sn[kpt.s][P_max] = f_n_unique if self.update_numbers: self.numbers[kpt.s] = f_sn[kpt.s].copy() else: f_sn = self.numbers.copy() for kpt in self.wfs.kpt_u: if self.width != 0.0: orbs, f_sn_gs = self.find_hole_and_excited_orbitals(f_sn, kpt) if orbs: for o in orbs: mask, gauss = self.gaussian_smearing(kpt, f_sn_gs, o, self.gauss_width) f_sn_gs[mask] += (o[1] * gauss) f_sn[kpt.s] = f_sn_gs.copy() return f_sn def calculate_weights(self, kpt, f_n_unique, unoccupied): if self.wfs.mode == 'lcao': O =[kpt.s][f_n_unique].conj(),, kpt.C_nM[unoccupied].T)) else: # Pseudo wave function overlaps O = self.wfs.integrate([kpt.s][f_n_unique][:], kpt.psit_nG[unoccupied][:], True) # Atomic contributions O_corr = np.zeros_like(O) for a, p_a in self.p_an[kpt.s][f_n_unique].items(): O_corr +=[a][unoccupied].conj(), p_a).T O_corr = np.ascontiguousarray(O_corr) # Sum pseudo wave and atomic contributions O += O_corr if self.use_projections: P = np.sum(abs(O) ** 2, axis=0) P = P ** 0.5 else: P = np.amax(abs(O), axis=0) return P def find_hole_and_excited_orbitals(self, f_sn, kpt): # Zero-width occupations for ground state ne = int(f_sn[kpt.s].sum()) f_sn_gs = np.zeros_like(f_sn[kpt.s]) f_sn_gs[:ne] = 1.0 f_sn_diff = f_sn[kpt.s] - f_sn_gs # Select hole and excited orbitals idxs = np.where(np.abs(f_sn_diff) > 1e-5)[0] w = f_sn_diff[np.abs(f_sn_diff) > 1e-5] orbs = list(zip(idxs, w)) return orbs, f_sn_gs def gaussian_smearing(self, kpt, f_sn_gs, o, gauss_width): if o[1] < 0: mask = (f_sn_gs > 1e-8) else: mask = (f_sn_gs < 1e-8) e = kpt.eps_n[mask] de2 = -(e - kpt.eps_n[o[0]]) ** 2 gauss = (1 / (gauss_width * np.sqrt(2 * np.pi)) * np.exp(de2 / (2 * gauss_width ** 2))) gauss /= sum(gauss) return mask, gauss def find_unique_occupation_numbers(self): f_sn_unique = {} for kpt in self.wfs.kpt_u: f_sn_unique[kpt.s] = {} f_n = self.numbers[kpt.s] for f_n_unique in np.unique(f_n): if f_n_unique >= 1.0e-10: f_sn_unique[kpt.s][f_n_unique] = f_n == f_n_unique return f_sn_unique