Coverage for gpaw/directmin/fdpw/pz_localization.py: 50%
229 statements
« prev ^ index » next coverage.py v7.7.1, created at 2025-07-20 00:19 +0000
1"""
2Optimization of orbitals
3among occupied states only
5"""
7from gpaw.directmin.tools import get_n_occ, get_indices, expm_ed
8from gpaw.directmin.sd_etdm import LBFGS_P
9from gpaw.directmin.ls_etdm import StrongWolfeConditions as SWC
10from ase.units import Hartree
11import numpy as np
12import time
15class PZLocalization:
17 def __init__(self, odd_pot, wfs, maxiter=50, g_tol=5.0e-4):
19 self.odd_pot = odd_pot
20 self.n_kps = wfs.kd.nibzkpts
21 self.g_tol = g_tol / Hartree
22 self.dtype = wfs.dtype
23 self.get_en_and_grad_iters = 0
24 self.method = 'LBFGS'
25 self.line_search_method = 'AwcSwc'
26 self.max_iter_line_search = 6
27 self.n_counter = maxiter
28 self.eg_count = 0
29 self.total_eg_count = 0
30 self.run_count = 0
31 self.U_k = {}
32 self.Unew_k = {}
33 self.e_total = 0.0
34 self.n_occ = {}
35 for kpt in wfs.kpt_u:
36 k = self.n_kps * kpt.s + kpt.q
37 self.n_occ[k] = get_n_occ(kpt)[0]
38 dim1 = self.n_occ[k]
39 dim2 = self.n_occ[k]
40 self.U_k[k] = np.eye(dim1, dtype=self.dtype)
41 self.Unew_k[k] = np.eye(dim2, dtype=self.dtype)
43 def get_energy_and_gradients(self, a_k, wfs):
44 """
45 Energy E = E[A]. Gradients G_ij[A] = dE/dA_ij
46 Returns E[A] and G[A] at psi = exp(A).T kpt.psi
47 :param a_k: A
48 :return:
49 """
51 g_k = {}
52 self.e_total = 0.0
54 self.kappa = 0.0
55 for kpt in wfs.kpt_u:
56 k = self.n_kps * kpt.s + kpt.q
57 dim2 = self.Unew_k[k].shape[0]
58 if dim2 == 0:
59 g_k[k] = np.zeros_like(a_k[k])
60 continue
61 wfs.timer.start('Unitary matrix')
62 u_mat, evecs, evals = expm_ed(a_k[k], evalevec=True)
63 wfs.timer.stop('Unitary matrix')
64 self.Unew_k[k] = u_mat.copy()
65 kpt.psit_nG[:dim2] = \
66 np.tensordot(u_mat.T, self.psit_knG[k][:dim2], axes=1)
67 # calc projectors
68 wfs.pt.integrate(kpt.psit_nG, kpt.P_ani, kpt.q)
70 del u_mat
72 wfs.timer.start('Energy and gradients')
73 if self.odd_pot.name == 'ER_SIC':
74 g_k[k], e_sic, kappa1 = \
75 self.odd_pot.get_energy_and_gradients_inner_loop(
76 wfs, kpt, a_k[k], evals, evecs)
77 else:
78 g_k[k], e_sic, kappa1 = \
79 self.odd_pot.get_energy_and_gradients_inner_loop(
80 wfs, kpt, a_k[k], evals, evecs, exstate=False)
81 wfs.timer.stop('Energy and gradients')
82 if kappa1 > self.kappa:
83 self.kappa = kappa1
84 self.e_total += e_sic
86 self.kappa = wfs.kd.comm.max_scalar(self.kappa)
87 self.e_total = wfs.kd.comm.sum_scalar(self.e_total)
88 self.eg_count += 1
89 self.total_eg_count += 1
91 return self.e_total, g_k
93 def evaluate_phi_and_der_phi(self, a_k, p_k, alpha, wfs, dens,
94 phi=None, g_k=None):
95 """
96 phi = f(x_k + alpha_k*p_k)
97 der_phi = grad f(x_k + alpha_k*p_k) cdot p_k
98 :return: phi, der_phi, grad f(x_k + alpha_k*p_k)
99 """
100 if phi is None or g_k is None:
101 x_k = {k: a_k[k] + alpha * p_k[k] for k in a_k.keys()}
102 phi, g_k = \
103 self.get_energy_and_gradients(x_k, wfs)
104 del x_k
105 else:
106 pass
108 der_phi = 0.0
109 for k in p_k.keys():
110 il1 = get_indices(p_k[k].shape[0])
111 der_phi += np.dot(g_k[k][il1].conj(), p_k[k][il1]).real
113 der_phi = wfs.kd.comm.sum_scalar(der_phi)
115 return phi, der_phi, g_k
117 def get_search_direction(self, a_k, g_k, wfs):
119 # structure of vector is
120 # (x_1_up, x_2_up,..,y_1_up, y_2_up,..,
121 # x_1_down, x_2_down,..,y_1_down, y_2_down,.. )
123 a = {}
124 g = {}
126 for k in a_k.keys():
127 il1 = get_indices(a_k[k].shape[0])
128 a[k] = a_k[k][il1]
129 g[k] = g_k[k][il1]
131 p = self.sd.update_data(wfs, a, g, mode='lcao')
132 del a, g
134 p_k = {}
135 for k in p.keys():
136 p_k[k] = np.zeros_like(a_k[k])
137 il1 = get_indices(a_k[k].shape[0])
138 p_k[k][il1] = p[k]
139 # make it skew-hermitian
140 ind_l = np.tril_indices(p_k[k].shape[0], -1)
141 p_k[k][(ind_l[1], ind_l[0])] = -p_k[k][ind_l].conj()
142 del p
144 return p_k
146 def run(self, e_ks, wfs, dens, log, outer_counter=0,
147 small_random=True, randvalue=0.01, seed=None):
149 log = log
150 self.run_count += 1
151 self.counter = 0
152 self.eg_count = 0
153 # initial things
154 self.c_knm = {}
155 self.psit_knG = {}
157 rng = np.random.default_rng(seed)
159 for kpt in wfs.kpt_u:
160 k = self.n_kps * kpt.s + kpt.q
161 dim1 = self.U_k[k].shape[0]
162 self.psit_knG[k] = np.tensordot(
163 self.U_k[k].T, kpt.psit_nG[:dim1], axes=1)
165 a_k = {}
166 for kpt in wfs.kpt_u:
167 k = self.n_kps * kpt.s + kpt.q
168 d = self.Unew_k[k].shape[0]
169 if self.run_count == 1 and self.dtype == complex \
170 and small_random:
171 a = randvalue * rng.random((d, d)) * 1.0j
172 a = a - a.T.conj()
173 wfs.gd.comm.broadcast(a, 0)
174 a_k[k] = a
175 else:
176 a_k[k] = np.zeros(shape=(d, d), dtype=self.dtype)
178 self.sd = LBFGS_P(memory=20)
179 self.ls = SWC(
180 self.evaluate_phi_and_der_phi,
181 searchdirtype=self.method, use_descent_and_awc=True,
182 max_iter=self.max_iter_line_search)
184 threelasten = []
185 # get initial energy and gradients
186 self.e_total, g_k = self.get_energy_and_gradients(a_k, wfs)
187 threelasten.append(self.e_total)
188 g_max = g_max_norm(g_k, wfs)
189 if g_max < self.g_tol:
190 for kpt in wfs.kpt_u:
191 k = self.n_kps * kpt.s + kpt.q
192 dim1 = self.U_k[k].shape[0]
193 kpt.psit_nG[:dim1] = np.tensordot(
194 self.U_k[k].conj(), self.psit_knG[k], axes=1)
195 # calc projectors
196 wfs.pt.integrate(kpt.psit_nG, kpt.P_ani, kpt.q)
198 dim2 = self.Unew_k[k].shape[0]
199 if dim1 == dim2:
200 self.U_k[k] = self.U_k[k] @ self.Unew_k[k]
201 else:
202 u_oo = self.Unew_k[k]
203 u_ov = np.zeros(shape=(dim2, dim1 - dim2),
204 dtype=self.dtype)
205 u_vv = np.eye(dim1 - dim2, dtype=self.dtype)
206 unew = np.vstack([np.hstack([u_oo, u_ov]),
207 np.hstack([u_ov.T, u_vv])])
208 self.U_k[k] = self.U_k[k] @ unew
210 del self.psit_knG
211 del self.c_knm
212 if outer_counter is None:
213 return self.e_total, self.counter
214 else:
215 return self.e_total, outer_counter
217 # stuff which are needed for minim.
218 phi_0 = self.e_total
219 phi_old = None
220 der_phi_old = None
221 phi_old_2 = None
222 der_phi_old_2 = None
224 outer_counter += 1
225 if log is not None:
226 log_f(log, self.counter, self.kappa, e_ks, self.e_total,
227 outer_counter, g_max)
229 alpha = 1.0
230 not_converged = True
231 while not_converged:
233 # calculate search direction fot current As and Gs
234 p_k = self.get_search_direction(a_k, g_k, wfs)
236 # calculate derivative along the search direction
237 phi_0, der_phi_0, g_k = \
238 self.evaluate_phi_and_der_phi(a_k, p_k,
239 0.0, wfs, dens,
240 phi=phi_0, g_k=g_k)
241 if self.counter > 1:
242 phi_old = phi_0
243 der_phi_old = der_phi_0
245 # choose optimal step length along the search direction
246 # also get energy and gradients for optimal step
247 alpha, phi_0, der_phi_0, g_k = \
248 self.ls.step_length_update(
249 a_k, p_k, wfs, dens,
250 phi_0=phi_0, der_phi_0=der_phi_0,
251 phi_old=phi_old_2, der_phi_old=der_phi_old_2,
252 alpha_max=3.0, alpha_old=alpha)
254 # broadcast data is gd.comm > 1
255 if wfs.gd.comm.size > 1:
256 alpha_phi_der_phi = np.array([alpha, phi_0,
257 der_phi_0])
258 wfs.gd.comm.broadcast(alpha_phi_der_phi, 0)
259 alpha = alpha_phi_der_phi[0]
260 phi_0 = alpha_phi_der_phi[1]
261 der_phi_0 = alpha_phi_der_phi[2]
262 for kpt in wfs.kpt_u:
263 k = self.n_kps * kpt.s + kpt.q
264 if self.n_occ[k] == 0:
265 continue
266 wfs.gd.comm.broadcast(g_k[k], 0)
268 phi_old_2 = phi_old
269 der_phi_old_2 = der_phi_old
271 if alpha > 1.0e-10:
272 # calculate new matrices at optimal step lenght
273 a_k = {k: a_k[k] + alpha * p_k[k] for k in a_k.keys()}
274 g_max = g_max_norm(g_k, wfs)
276 # output
277 self.counter += 1
278 if outer_counter is not None:
279 outer_counter += 1
280 if log is not None:
281 log_f(
282 log, self.counter, self.kappa, e_ks, phi_0,
283 outer_counter, g_max)
285 not_converged = \
286 g_max > self.g_tol and \
287 self.counter < self.n_counter
289 threelasten.append(phi_0)
290 if len(threelasten) > 2:
291 threelasten = threelasten[-3:]
292 if threelasten[0] < threelasten[1] < \
293 threelasten[2]:
294 if log is not None:
295 log('Could not converge, leave the loop',
296 flush=True)
297 break
298 else:
299 break
301 if log is not None:
302 log('INNER LOOP FINISHED.\n')
303 log('Total number of e/g calls:' + str(self.eg_count))
305 for kpt in wfs.kpt_u:
306 k = self.n_kps * kpt.s + kpt.q
307 dim1 = self.U_k[k].shape[0]
308 kpt.psit_nG[:dim1] = np.tensordot(
309 self.U_k[k].conj(), self.psit_knG[k][:dim1],
310 axes=1)
311 # calc projectors
312 wfs.pt.integrate(kpt.psit_nG, kpt.P_ani, kpt.q)
313 dim2 = self.Unew_k[k].shape[0]
314 if dim1 == dim2:
315 self.U_k[k] = self.U_k[k] @ self.Unew_k[k]
316 else:
317 u_oo = self.Unew_k[k]
318 u_ov = np.zeros(shape=(dim2, dim1 - dim2),
319 dtype=self.dtype)
320 u_vv = np.eye(dim1 - dim2, dtype=self.dtype)
321 unew = np.vstack([np.hstack([u_oo, u_ov]),
322 np.hstack([u_ov.T, u_vv])])
323 self.U_k[k] = self.U_k[k] @ unew
325 del self.psit_knG
326 del self.c_knm
327 if outer_counter is None:
328 return self.e_total, self.counter
329 else:
330 return self.e_total, outer_counter
333def log_f(log, niter, kappa, e_ks, e_sic, outer_counter=None, g_max=np.inf):
335 t = time.localtime()
337 if niter == 0:
338 header0 = '\nINNER LOOP:\n'
339 header = ' Kohn-Sham SIC' \
340 ' Total \n' \
341 ' time energy: energy:' \
342 ' energy: Error: G_max:'
343 log(header0 + header)
345 if outer_counter is not None:
346 niter = outer_counter
348 log('iter: %3d %02d:%02d:%02d ' %
349 (niter,
350 t[3], t[4], t[5]
351 ), end='')
352 log('%11.6f %11.6f %11.6f %11.1e %11.1e' %
353 (Hartree * e_ks,
354 Hartree * e_sic,
355 Hartree * (e_ks + e_sic),
356 kappa,
357 Hartree * g_max), end='')
358 log(flush=True)
361def g_max_norm(g_k, wfs):
362 # get maximum of gradients
363 n_kps = wfs.kd.nibzkpts
365 max_norm = []
366 for kpt in wfs.kpt_u:
367 k = n_kps * kpt.s + kpt.q
368 dim = g_k[k].shape[0]
369 if dim == 0:
370 max_norm.append(0.0)
371 else:
372 max_norm.append(np.max(np.absolute(g_k[k])))
373 max_norm = np.max(np.asarray(max_norm))
374 g_max = wfs.world.max_scalar(max_norm)
376 return g_max