Coverage for gpaw/mom.py: 95%
184 statements
« prev ^ index » next coverage.py v7.7.1, created at 2025-07-14 00:18 +0000
1"""Module for calculations using the Maximum Overlap Method (MOM).
3See:
5 https://arxiv.org/abs/2102.06542,
6 :doi:`10.1021/acs.jctc.0c00597`.
7"""
9import numpy as np
11from ase.units import Ha
13from gpaw.occupations import FixedOccupationNumbers, ParallelLayout
14from scipy.optimize import linear_sum_assignment
17def prepare_mom_calculation(calc,
18 atoms,
19 numbers,
20 use_projections=True,
21 update_numbers=True,
22 use_fixed_occupations=False,
23 width=0.0,
24 niter_width_update=40,
25 width_increment=0.0):
26 """Helper function to prepare a calculator for a MOM calculation.
28 calc: GPAW instance
29 GPAW calculator object.
30 atoms: ASE instance
31 ASE atoms object.
32 numbers: list (len=nspins) of lists (len=nbands)
33 Occupation numbers (in the range from 0 to 1). Used
34 for the initialization of the MOM reference orbitals.
35 use_projections: bool
36 If True (default), the occupied orbitals at iteration k are
37 chosen as the orbitals ``|psi^(k)_m>`` with the biggest
38 weights ``P_m`` evaluated as the projections onto the manifold
39 of reference orbitals ``|psi_n>``: ``P_m = (Sum_n(|O_nm|^2))^0.5
40 (O_nm = <psi_n|psi^(k)_m>)`` see :doi:`10.1021/acs.jctc.7b00994`.
41 If False, try to match the orbitals directly based on
42 their overlaps.
43 update_numbers: bool
44 If True (default), 'numbers' gets updated with the calculated
45 occupation numbers, and when changing atomic positions
46 the MOM reference orbitals will be initialized as the
47 occupied orbitals found at convergence for the previous
48 geometry. If False, when changing positions the MOM
49 reference orbitals will be initialized as the orbitals
50 of the previous geometry corresponding to the user-supplied
51 'numbers'.
52 use_fixed_occupations: bool
53 If False (default), the MOM algorithm is used. If True,
54 fixed occupations will be used.
55 width: float
56 Width of Gaussian function in eV for smearing of holes
57 and excited electrons. The holes and excited electrons
58 are found with respect to the zero-width ground-state
59 occupations.
60 See :doi:`10.1021/acs.jctc.0c00597`.
61 niter_width_update: int
62 Number of iterations after which the width of the
63 Gaussian smearing function is increased.
64 width_increment: float
65 How much to increase the width of the Gaussian smearing
66 function.
67 """
69 new_gpaw = hasattr(calc, '_dft')
70 if not new_gpaw:
71 if calc.wfs is None:
72 # We need the wfs object to initialize OccupationsMOM
73 calc.initialize(atoms)
74 else:
75 if calc._dft is None:
76 calc.create_new_calculation(atoms)
78 occ_mom = OccupationsMOM(calc.wfs,
79 numbers,
80 use_projections,
81 update_numbers,
82 use_fixed_occupations,
83 width,
84 niter_width_update,
85 width_increment)
86 if not new_gpaw:
87 calc.set(occupations=occ_mom, _set_ok=True)
88 else:
89 calc.dft.scf_loop.occ_calc.occ = occ_mom
90 calc.dft.results = {}
92 calc.log(occ_mom)
94 return occ_mom
97class OccupationsMOM:
98 """MOM occupation class.
100 The occupation numbers are found using a maximum overlap
101 criterion and then broadcasted in occupations.py using the
102 _calculate method of FixedOccupationNumbers.
103 """
105 def __init__(self,
106 wfs,
107 numbers,
108 use_projections=False,
109 update_numbers=True,
110 use_fixed_occupations=False,
111 width=0.0,
112 niter_width_update=10,
113 width_increment=0.0):
114 self.wfs = wfs
115 self.numbers = np.array(numbers)
116 self.use_projections = use_projections
117 self.update_numbers = update_numbers
118 self.use_fixed_occupations = use_fixed_occupations
119 self.width = width / Ha
120 self.niter_width_update = niter_width_update
121 self.width_increment = width_increment / Ha
123 parallel_layout = ParallelLayout(self.wfs.bd,
124 self.wfs.kd.comm,
125 self.wfs.gd.comm)
126 self.occ = FixedOccupationNumbers(numbers, parallel_layout)
127 self.extrapolate_factor = self.occ.extrapolate_factor
129 self.name = 'mom'
130 self.iters = 0
131 self.initialized = False
133 def todict(self):
134 dct = {'name': self.name,
135 'numbers': self.numbers,
136 'use_projections': self.use_projections,
137 'update_numbers': self.update_numbers,
138 'use_fixed_occupations': self.use_fixed_occupations}
139 if self.width != 0.0:
140 dct['width'] = self.width * Ha
141 dct['niter_width_update'] = self.niter_width_update
142 dct['width_increment'] = self.width_increment * Ha
143 return dct
145 def __str__(self):
146 s = 'Excited-state calculation with Maximum Overlap Method\n'
147 s += ' Gaussian smearing of holes and excited electrons: '
148 if self.width == 0.0:
149 s += 'off\n'
150 else:
151 s += f'{self.width * Ha:.4f} eV\n'
152 return s
154 def calculate(self,
155 nelectrons,
156 eigenvalues,
157 weights,
158 fermi_levels_guess,
159 fix_fermi_level=False):
160 assert not fix_fermi_level
162 if not self.initialized:
163 # If MOM reference orbitals are not initialized yet (e.g. when
164 # the calculation is initialized from atomic densities), update
165 # the occupation numbers according to the user-supplied 'numbers'
166 self.occ.f_sn = self.numbers.copy()
167 self.initialize_reference_orbitals()
168 else:
169 self.occ.f_sn = self.update_occupations()
170 self.iters += 1
172 f_qn, fermi_levels, e_entropy = self.occ.calculate(nelectrons,
173 eigenvalues,
174 weights,
175 fermi_levels_guess)
176 return f_qn, fermi_levels, e_entropy
178 def initialize_reference_orbitals(self):
179 try:
180 f_n = self.wfs.kpt_u[0].f_n
181 except ValueError:
182 return
183 if f_n is None: # old gpaw
184 # If the occupation numbers are not already available
185 # (e.g. when the calculation is initialized from atomic
186 # densities) we first need to take a step of eigensolver
187 # and update the occupation numbers according to the
188 # 'user-supplied' numbers before initializing the MOM
189 # reference orbitals
190 return
192 self.iters = 0
194 if self.use_fixed_occupations:
195 self.initialized = True
196 return
198 # initial occupations numbers
199 self.f_sn0 = self.numbers.copy()
201 # Initialize MOM reference orbitals for each equally
202 # occupied subspace separately
203 self.subs_mask = self.find_unique_occupation_numbers()
204 if self.wfs.mode == 'lcao':
205 self.c_ref = {}
206 for kpt in self.wfs.kpt_u:
207 self.c_ref[kpt.s] = {}
209 if not self.use_projections:
210 self.c_ref[kpt.s]['all'] = kpt.C_nM[:].copy()
212 for fs_key in self.subs_mask[kpt.s]:
213 occupied = self.subs_mask[kpt.s][fs_key]
214 self.c_ref[kpt.s][fs_key] = kpt.C_nM[occupied].copy()
215 else:
216 self.wf = {}
217 self.p_an = {}
218 for kpt in self.wfs.kpt_u:
219 self.wf[kpt.s] = {}
220 self.p_an[kpt.s] = {}
222 if not self.use_projections:
223 # Pseudo wave functions
224 self.wf[kpt.s]['all'] = kpt.psit_nG[:].copy()
225 # Atomic contributions times projector overlaps
226 self.p_an[kpt.s]['all'] = \
227 {a: np.dot(self.wfs.setups[a].dO_ii, P_ni[:].T)
228 for a, P_ni in kpt.P_ani.items()}
230 for fs_key in self.subs_mask[kpt.s]:
231 occupied = self.subs_mask[kpt.s][fs_key]
232 # Pseudo wave functions
233 self.wf[kpt.s][fs_key] = kpt.psit_nG[occupied].copy()
234 # Atomic contributions times projector overlaps
235 self.p_an[kpt.s][fs_key] = \
236 {a: np.dot(self.wfs.setups[a].dO_ii, P_ni[occupied].T)
237 for a, P_ni in kpt.P_ani.items()}
239 self.initialized = True
241 def update_occupations(self):
242 if self.width != 0.0:
243 if self.iters == 0:
244 self.width_update_counter = 0
245 if self.iters % self.niter_width_update == 0:
246 self.gauss_width = self.width + self.width_update_counter \
247 * self.width_increment
248 self.width_update_counter += 1
250 if not self.use_fixed_occupations:
251 f_sn = np.zeros_like(self.numbers)
252 for kpt in self.wfs.kpt_u:
253 # Compute weights with respect to each equally occupied
254 # subspace of the reference orbitals and occupy orbitals
255 # with biggest weights
257 if not self.use_projections:
258 # match orbitals directly
259 # assign occupations based on match
261 # calculate overlap
262 O_nm = self.calculate_weights(kpt, 'all')
263 # find permutation maximizing overlap (absolute value)
264 _, col_ind = linear_sum_assignment(-np.abs(O_nm))
265 # map to the initial occupations
266 f_sn[kpt.s] = self.f_sn0[kpt.s][col_ind]
267 else:
268 # match occupation number subspaces to bands for
269 # maximizing weights (projections)
271 # number of subspaces
272 nsubs = len(self.subs_mask[kpt.s])
273 if nsubs == 1:
274 # we can just occupy the orbitals
275 # with largest projections
276 fs_key = list(self.subs_mask[kpt.s].keys())[0]
277 subs_mask = self.subs_mask[kpt.s][fs_key]
278 subs_size = np.sum(subs_mask)
279 P_m = self.calculate_weights(kpt, fs_key)
281 # we could use np.argpartition here
282 # but argsort is better readable
283 occ_mask = np.argsort(P_m)[::-1][:subs_size]
284 f_sn[kpt.s][occ_mask] = fs_key
285 else:
286 # for more than one subspace we need to figure out the
287 # assignment with the maximum projections
289 nband = len(self.numbers[kpt.s])
290 noccupied = np.sum(self.numbers[kpt.s] > 1.0e-10)
291 Ps_m = np.zeros((noccupied, nband))
293 fs = []
294 nn = 0
295 for ss, fs_key in enumerate(self.subs_mask[kpt.s]):
296 subs_mask = self.subs_mask[kpt.s][fs_key]
297 subs_size = np.sum(subs_mask)
298 # Ps_m ... projections of the subspace orbitals
299 # to the scf orbitals from the k-iteration
300 w = self.calculate_weights(kpt, fs_key)
301 Ps_m[nn: nn + subs_size, :] = w[None, :]
302 fs += subs_size * [fs_key]
303 nn += subs_size
305 # Optimize assigment of subspace occupations
306 # such that sum of the selected projections is maximal
307 row_ind, col_ind = linear_sum_assignment(-Ps_m)
309 # assign band index (=col_ind) with occupation number
310 assert (all(row_ind == np.arange(noccupied)))
311 f_sn[kpt.s][col_ind] = fs[:]
313 if self.update_numbers:
314 self.numbers[kpt.s] = f_sn[kpt.s].copy()
315 else:
316 f_sn = self.numbers.copy()
318 for kpt in self.wfs.kpt_u:
319 if self.width != 0.0:
320 orbs, f_sn_gs = self.find_hole_and_excited_orbitals(f_sn, kpt)
321 if orbs:
322 for o in orbs:
323 mask, gauss = self.gaussian_smearing(kpt,
324 f_sn_gs,
325 o,
326 self.gauss_width)
327 f_sn_gs[mask] += (o[1] * gauss)
328 f_sn[kpt.s] = f_sn_gs.copy()
330 return f_sn
332 def calculate_weights(self, kpt, fs_key):
333 if self.wfs.mode == 'lcao':
334 O = np.dot(self.c_ref[kpt.s][fs_key].conj(),
335 np.dot(kpt.S_MM, kpt.C_nM[:].T))
336 else:
337 # Pseudo wave function overlaps
338 O = self.wfs.integrate(self.wf[kpt.s][fs_key][:],
339 kpt.psit_nG[:][:], True)
341 # Atomic contributions
342 O_corr = np.zeros_like(O)
343 for a, p_a in self.p_an[kpt.s][fs_key].items():
344 O_corr += np.dot(kpt.P_ani[a][:].conj(), p_a).T
345 O_corr = np.ascontiguousarray(O_corr)
346 self.wfs.gd.comm.sum(O_corr)
348 # Sum pseudo wave and atomic contributions
349 O += O_corr
351 if fs_key == 'all':
352 # direct orbital matching
353 assert not self.use_projections
354 return O
355 else:
356 P = np.sum(abs(O)**2, axis=0)
357 P = P**0.5
359 return P
361 def find_hole_and_excited_orbitals(self, f_sn, kpt):
362 # Zero-width occupations for ground state
363 ne = int(f_sn[kpt.s].sum())
364 f_sn_gs = np.zeros_like(f_sn[kpt.s])
365 f_sn_gs[:ne] = 1.0
366 f_sn_diff = f_sn[kpt.s] - f_sn_gs
368 # Select hole and excited orbitals
369 idxs = np.where(np.abs(f_sn_diff) > 1e-5)[0]
370 w = f_sn_diff[np.abs(f_sn_diff) > 1e-5]
371 orbs = list(zip(idxs, w))
373 return orbs, f_sn_gs
375 def gaussian_smearing(self, kpt, f_sn_gs, o, gauss_width):
376 if o[1] < 0:
377 mask = (f_sn_gs > 1e-8)
378 else:
379 mask = (f_sn_gs < 1e-8)
381 e = kpt.eps_n[mask]
382 de2 = -(e - kpt.eps_n[o[0]]) ** 2
383 gauss = (1 / (gauss_width * np.sqrt(2 * np.pi)) *
384 np.exp(de2 / (2 * gauss_width ** 2)))
385 gauss /= sum(gauss)
387 return mask, gauss
389 def find_unique_occupation_numbers(self):
390 f_sn_unique = {}
391 for kpt in self.wfs.kpt_u:
392 f_sn_unique[kpt.s] = {}
393 f_n = self.numbers[kpt.s]
395 for fs_key in np.unique(f_n):
396 if fs_key >= 1.0e-10:
397 f_sn_unique[kpt.s][fs_key] = f_n == fs_key
399 return f_sn_unique