Coverage for gpaw/directmin/etdm_lcao.py: 85%

1""" 

2A class for finding optimal orbitals corresponding to a stationary point of 

3the energy functional using direct optimization and exponential transformation 

4in LCAO mode. 

5 

6It can be used with Kohn-Sham functionals and can include Perdew-Zunger 

7self-interaction correction (PZ-SIC) in the calculations. 

8 

9Ground state as well as variational excited state calculations can be 

10performed. Ground state calculations involve minimization of the energy 

11(direct minimization), while excited state calculations involve convergence 

12on a saddle point (direct optimization). 

13 

14Implementation of exponential transformation direct minimization (ETDM) for 

15ground state calculations LCAO mode: 

16 

17 Comput. Phys. Commun. 267, 108047 (2021) :doi:10.1016/j.cpc.2021.108047 

18 arXiv:2101.12597 [physics.comp-ph] 

19 

20GPAW implementations of direct optimization (DO) for variational excited state 

21calculations LCAO mode: 

22 

23 DO with preconditioned quasi-Newton algorithms and maximum overlap method 

24 (DO-MOM) 

25 J. Chem. Theory Comput. 16, 6968 (2020) :doi:10.1021/acs.jctc.0c00597 

26 arXiv:2006.15922 [physics.chem-ph] 

27 

28 DO with generalized mode following (DO-GMF) 

29 J. Chem. Theory Comput. 19, 3634 (2023) :doi:10.1021/acs.jctc.3c00178 

30 arXiv:2302.05912 [physics.chem-ph] 

31""" 

32 

33 

34import numpy as np 

35import warnings 

36from ase.utils import basestring 

37from gpaw.directmin.tools import expm_ed, expm_ed_unit_inv, random_a, \ 

38 sort_orbitals_according_to_occ, sort_orbitals_according_to_energies 

39from gpaw.directmin.lcao.etdm_helper_lcao import ETDMHelperLCAO 

40from gpaw.directmin.locfunc.localize_orbitals import localize_orbitals 

41from scipy.linalg import expm 

42from gpaw.directmin import search_direction, line_search_algorithm 

43from gpaw.directmin.tools import get_n_occ 

44from gpaw.directmin.derivatives import get_approx_analytical_hessian 

45from gpaw import BadParallelization 

46from copy import deepcopy 

47 

48 

49class LCAOETDM: 

50 

51 def __init__(self, 

52 excited_state=False, 

53 searchdir_algo=None, 

54 linesearch_algo=None, 

55 partial_diagonalizer='Davidson', 

56 update_ref_orbs_counter=20, 

57 update_ref_orbs_canonical=False, 

58 update_precond_counter=1000, 

59 use_prec=True, 

60 matrix_exp='pade-approx', 

61 representation='sparse', 

62 functional='ks', 

63 need_init_orbs=None, 

64 orthonormalization='gramschmidt', 

65 randomizeorbitals=None, 

66 checkgraderror=False, 

67 need_localization=True, 

68 localizationtype=None, 

69 localizationseed=None, 

70 constraints=None, 

71 subspace_convergence=5e-4 

72 ): 

73 """Class for direct orbital optimization in LCAO mode. 

74 

75 Parameters 

76 ---------- 

77 excited_state: bool 

78 If False (default), use search direction and line search 

79 algorithms for ground state calculation ('l-bfgs-p' and 'swc-awc') 

80 if not specified by searchdir_algo linesearch_algo, and set 

81 need_init_orbs to True. Otherwise, use search direction and line 

82 search algorithms for excited state calculation ('l-sr1p' and 

83 'max-step'), and set need_init_orbs to False. 

84 searchdir_algo: str, dict or instance 

85 Search direction algorithm. Can be one of the algorithms available 

86 in sd_etdm.py: 

87 'sd': Steepest descent 

88 'fr-cg': Fletcher-Reeves conjugate gradient 

89 'l-bfgs': Limited-memory BFGS 

90 'l-bfgs-p': Limited-memory BFGS with preconditioner presented 

91 in :doi:`10.1016/j.cpc.2021.108047` (default when 

92 excited_state is False) 

93 'l-sr1p': limited-memory SR1 algorithm presented in 

94 :doi:`10.1021/acs.jctc.0c00597` (default when excited_state 

95 is True) 

96 The default memory for 'l-bfgs'/'l-bfgs-p' and 'l-sr1p' is 3 and 

97 20, respectively, and can be changed by supplying a dictionary: 

98 {'name': name, 'memory': memory}, where name should be 'l-bfgs', 

99 'l-bfgs-p' or 'l-sr1p' and memory should be an int. 

100 To use the generalized mode following (GMF) method for excited 

101 states, append '_gmf' to the search direction algorithm name. E.g. 

102 'l-bfgs-p_gmf'. 

103 linesearch_algo: str, dict or instance 

104 Line search algorithm. Can be one of the algorithms available 

105 in ls_etdm.py: 

106 'max-step': The quasi-Newton step is scaled if it exceeds a 

107 maximum step length (default when excited_state is True). 

108 The default maximum step length is 0.20, and can be changed 

109 by supplying a dictionary: 

110 {'name': 'max-step', 'max_step': max_step}, where max_step 

111 should be a float 

112 'swc-awc': Line search with Wolfe conditions (default when 

113 excited_state is False) 

114 partial_diagonalizer: 'str' or dict 

115 Algorithm for partial diagonalization of the electronic Hessian if 

116 GMF is used. Default is 'Davidson'. 

117 update_ref_orbs_counter: int 

118 When to update the coefficients of the reference orbitals. Default 

119 is 20 iterations. 

120 update_ref_orbs_canonical: bool 

121 If True, the coefficients of the reference orbitals are updated to 

122 canonical orbital coefficients, otherwise use the optimal orbital 

123 coefficients (default). 

124 use_prec: bool 

125 If True (default) use a preconditioner. The preconditioner is 

126 calculated as the inverse of a diagonal approximation of the 

127 Hessian (see :doi:`10.1021/j100322a012`) apart for 'l-bfgs-p', 

128 which uses the composite preconditioner presented in 

129 :doi:`10.1016/j.cpc.2021.108047`. 

130 update_precond_counter: int 

131 When to update the preconditioner. Default is 1000 iterations. 

132 representation: 'str' 

133 How to store the elements of the anti-Hermitian matrix for the 

134 matrix exponential. Can be one of 'sparse' (default), 'full', 

135 'u-invar'. The latter can be used only for unitary invariant 

136 functionals, such as Kohn-Sham functionals when the occupied 

137 orbitals have the same occupation number, but not for orbital 

138 density dependent functionals, such as when using PZ-SIC. 

139 matrix_exp: 'str' 

140 Algorithm for calculating the matrix exponential and the gradient 

141 with respect to the elements of the anti-Hermitian matrix for the 

142 exponential transformation. Can be one of 'pade-approx' (default), 

143 'egdecomp', 'egdecomp-u-invar' (the latter can be used only with 

144 'u-invar' representation). 

145 functional: str, dict or instance 

146 Type of functional. Can be one of: 

147 'ks': The functional as specified in the GPAW calculator is 

148 used (default) 

149 'pz-sic': Apply the Perdew-Zunger self-interaction correction 

150 on top of the functional as specified in the GPAW 

151 calculator. Dy default full SIC is applied. A scaling 

152 factor for SIC can be given by supplying a dictionary: 

153 functional={'name': 'pz-sic', 'scaling_factor': (a, a)}, 

154 where a is the scaling factor (float). 

155 need_init_orbs: bool 

156 If True (default when excited_state is False), obtain initial 

157 orbitals from eigendecomposition of the Hamiltonian matrix. If 

158 False (default when excited_state is True), use orbitals stored in 

159 wfs object as initial guess. 

160 orthonormalization: str 

161 Method to orthonormalize the orbitals. Can be one of 'gramschmidt' 

162 (Gram-Schmidt orthonormalization, default), 'loewdin' (Loewdin 

163 orthonormalization) or 'diag' (eigendecomposition of the 

164 Hamiltonian matrix). 

165 randomizeorbitals: numpy.random.Generator 

166 Optional RNG to add random noise to the initial guess orbitals, 

167 default is to not add noise to them 

168 checkgraderror: bool 

169 If True, can be used to check the error in the estimation of the 

170 gradient (only with representation 'full'). Default is False. 

171 need_localization: bool 

172 If True (default), localize initial guess orbitals. Requires a 

173 specification of localizationtype, otherwise it is set to False. 

174 Recommended for calculations with PZ-SIC. 

175 localizationtype: str 

176 Method for localizing the initial guess orbitals. Can be one of: 

177 'pz': Unitary optimization among occupied orbitals (subspace 

178 optimization) with PZ-SIC 

179 'pm': Pipek-Mezey localization (recommended for PZ-SIC) 

180 'fb' Foster-Boys localization 

181 Default is None, meaning that no localization is performed. 

182 localizationseed: int 

183 Seed for Pipek-Mezey or Foster-Boys localization. Default is 

184 None (no seed is used). 

185 constraints: list of lists of int 

186 List of constraints on the orbital rotation for each k-point. If 

187 a constraint is given as a pair of orbital indices, the rotation 

188 between these two orbitals is constrained. If a constraint is 

189 given as a single index, the orbital with this index is frozen. 

190 E.g.: calc.set(eigensolver=LCAOETDM(constraints=[[[i]], [[j]]])) 

191 will freeze orbital i in the first k-point and orbital j in the 

192 second k-point. 

193 subspace_convergence: float 

194 Tolerance on the norm of the gradient for convergence of the 

195 subspace optimization with PZ-SIC. 

196 """ 

197 

198 assert representation in ['sparse', 'u-invar', 'full'], 'Value Error' 

199 assert matrix_exp in ['egdecomp', 'egdecomp-u-invar', 'pade-approx'], \ 

200 'Value Error' 

201 if matrix_exp == 'egdecomp-u-invar': 

202 assert representation == 'u-invar', 'Use u-invar representation ' \ 

203 'with egdecomp-u-invar' 

204 assert orthonormalization in ['gramschmidt', 'loewdin', 'diag'], \ 

205 'Value Error' 

206 

207 self.sda = searchdir_algo 

208 self.lsa = linesearch_algo 

209 self.partial_diagonalizer = partial_diagonalizer 

210 self.localizationseed = localizationseed 

211 self.update_ref_orbs_counter = update_ref_orbs_counter 

212 self.update_ref_orbs_canonical = update_ref_orbs_canonical 

213 self.update_precond_counter = update_precond_counter 

214 self.use_prec = use_prec 

215 self.matrix_exp = matrix_exp 

216 self.localizationtype = localizationtype 

217 self.need_localization = need_localization 

218 self.need_init_orbs = need_init_orbs 

219 self.randomizeorbitals = randomizeorbitals 

220 self.representation = representation 

221 self.orthonormalization = orthonormalization 

222 self.constraints = constraints 

223 self.subspace_convergence = subspace_convergence 

224 self.released_subspace = False 

225 self.excited_state = excited_state 

226 self.name = 'etdm-lcao' 

227 self.iters = 0 

228 self.eg_count = 0 

229 self.subspace_iters = 0 

230 self.restart = False 

231 self.gmf = False 

232 self.check_inputs_and_init_search_algo() 

233 

234 self.checkgraderror = checkgraderror 

235 self._norm_commutator, self._norm_grad = 0., 0. 

236 self.error = 0 

237 self.e_sic = 0.0 

238 self.subspace_optimization = False 

239 

240 # these are things we cannot initialize now 

241 self.func_settings = functional 

242 self.dtype = None 

243 self.nkpts = None 

244 self.gd = None 

245 self.nbands = None 

246 

247 # values: vectors of the elements of matrices, keys: kpt number 

248 self.a_vec_u = {} # for the elements of the skew-Hermitian matrix A 

249 self.a_vec_oo_u = {} 

250 self.a_vec_ov_u = {} 

251 self.a_vec_all_u = {} 

252 self.g_vec_u = {} # for the elements of the gradient matrix G 

253 self.g_vec_oo_u = {} 

254 self.g_vec_ov_u = {} 

255 self.g_vec_all_u = {} 

256 self.evecs = {} # eigenvectors for i*a_vec_u 

257 self.evals = {} # eigenvalues for i*a_vec_u 

258 self.ind_up = {} 

259 self.ind_oo_up = {} 

260 self.ind_ov_up = {} 

261 self.ind_sparse_up = {} 

262 self.ind_all_up = {} 

263 self.n_dim = {} 

264 self.n_dim_oo = {} 

265 self.n_dim_all = {} 

266 self.alpha = 1.0 # step length 

267 self.phi_2i = [None, None] # energy at last two iterations 

268 self.der_phi_2i = [None, None] # energy gradient w.r.t. alpha 

269 self.hess = {} # approximate Hessian 

270 

271 # for mom 

272 self.initial_occupation_numbers = None 

273 

274 # in this attribute we store the object specific to each mode 

275 self.dm_helper = None 

276 

277 self.initialized = False 

278 

279 def check_inputs_and_init_search_algo(self): 

280 defaults = self.set_defaults() 

281 if self.sda is None: 

282 self.sda = defaults['searchdir_algo'] 

283 if self.lsa is None: 

284 self.lsa = defaults['linesearch_algo'] 

285 if self.need_init_orbs is None: 

286 self.need_init_orbs = defaults['need_init_orbs'] 

287 if self.localizationtype is None: 

288 self.need_localization = False 

289 

290 self.searchdir_algo = search_direction( 

291 self.sda, self, self.partial_diagonalizer) 

292 sd_name = self.searchdir_algo.name.split('_') 

293 if sd_name[0] == 'l-bfgs-p' and not self.use_prec: 

294 raise ValueError('Use l-bfgs-p with use_prec=True') 

295 if len(sd_name) == 2: 

296 if sd_name[1] == 'gmf': 

297 self.searchdir_algo.name = sd_name[0] 

298 self.gmf = True 

299 self.g_vec_u_original = None 

300 self.pd = self.partial_diagonalizer 

301 

302 self.line_search = line_search_algorithm(self.lsa, 

303 self.evaluate_phi_and_der_phi, 

304 self.searchdir_algo) 

305 

306 def set_defaults(self): 

307 if self.excited_state: 

308 return {'searchdir_algo': 'l-sr1p', 

309 'linesearch_algo': 'max-step', 

310 'need_init_orbs': False} 

311 else: 

312 return {'searchdir_algo': 'l-bfgs-p', 

313 'linesearch_algo': 'swc-awc', 

314 'need_init_orbs': True} 

315 

316 def __repr__(self): 

317 

318 sda_name = self.searchdir_algo.name 

319 lsa_name = self.line_search.name 

320 if isinstance(self.func_settings, basestring): 

321 func_name = self.func_settings 

322 else: 

323 func_name = self.func_settings['name'] 

324 if self.gmf: 

325 if isinstance(self.pd, basestring): 

326 pd_name = self.pd 

327 else: 

328 pd_name = self.pd['name'] 

329 

330 add = '' 

331 pd_add = '' 

332 if self.gmf: 

333 add = ' with minimum mode following' 

334 pardi = {'Davidson': 'Finite difference generalized Davidson ' 

335 'algorithm'} 

336 pd_add = ' ' \ 

337 'Partial diagonalizer: {}\n'.format( 

338 pardi[pd_name]) 

339 

340 sds = {'sd': 'Steepest Descent', 

341 'fr-cg': 'Fletcher-Reeves conj. grad. method', 

342 'l-bfgs': 'L-BFGS algorithm', 

343 'l-bfgs-p': 'L-BFGS algorithm with preconditioning', 

344 'l-sr1p': 'Limited-memory SR1P algorithm'} 

345 

346 lss = {'max-step': 'step size equals one', 

347 'swc-awc': 'Inexact line search based on cubic interpolation,\n' 

348 ' strong and approximate Wolfe ' 

349 'conditions'} 

350 

351 repr_string = 'Direct minimisation' + add + ' using exponential ' \ 

352 'transformation.\n' 

353 repr_string += ' ' \ 

354 'Search ' \ 

355 'direction: {}\n'.format(sds[sda_name] + add) 

356 repr_string += pd_add 

357 repr_string += ' ' \ 

358 'Line ' \ 

359 'search: {}\n'.format(lss[lsa_name]) 

360 repr_string += ' ' \ 

361 'Preconditioning: {}\n'.format(self.use_prec) 

362 repr_string += ' ' \ 

363 'Orbital-density self-interaction ' \ 

364 'corrections: {}\n'.format(func_name) 

365 repr_string += ' ' \ 

366 'WARNING: do not use it for metals as ' \ 

367 'occupation numbers are\n' \ 

368 ' ' \ 

369 'not found variationally\n' 

370 

371 return repr_string 

372 

373 def initialize(self, gd, dtype, nbands, nkpts, nao, using_blacs, 

374 bd_comm_size, kpt_u): 

375 

376 assert nbands == nao, \ 

377 'Please, use: nbands=\'nao\'' 

378 if not bd_comm_size == 1: 

379 raise BadParallelization( 

380 'Band parallelization is not supported') 

381 if using_blacs: 

382 raise BadParallelization( 

383 'ScaLapack parallelization is not supported') 

384 

385 self.gd = gd 

386 self.dtype = dtype 

387 self.nbands = nbands 

388 self.nkpts = nkpts 

389 

390 for kpt in kpt_u: 

391 u = self.kpointval(kpt) 

392 # dimensionality of the problem. 

393 # this implementation rotates among all bands 

394 self.n_dim_all[u] = self.nbands 

395 

396 self.initialized = True 

397 

398 # need reference orbitals 

399 self.dm_helper = None 

400 

401 def initialize_dm_helper(self, wfs, ham, dens, log): 

402 

403 self.dm_helper = ETDMHelperLCAO( 

404 wfs, dens, ham, self.nkpts, self.func_settings, 

405 diagonalizer=None, 

406 orthonormalization=self.orthonormalization, 

407 need_init_orbs=self.need_init_orbs) 

408 self.need_init_orbs = self.dm_helper.need_init_orbs 

409 

410 # randomize orbitals? 

411 if self.randomizeorbitals is not None: 

412 for kpt in wfs.kpt_u: 

413 self.randomize_orbitals_kpt(wfs, kpt) 

414 self.randomizeorbitals = None 

415 

416 if not wfs.coefficients_read_from_file: 

417 wfs.calculate_occupation_numbers(dens.fixed) 

418 

419 self.initial_sort_orbitals(wfs) 

420 

421 # initialize matrices 

422 self.set_variable_matrices(wfs.kpt_u) 

423 

424 # localize orbitals? 

425 self.localize(wfs, dens, ham, log) 

426 

427 # MOM 

428 self.initialize_mom_reference_orbitals(wfs, dens) 

429 

430 for kpt in wfs.kpt_u: 

431 f_unique = np.unique(kpt.f_n) 

432 if len(f_unique) > 2 and self.representation == 'u-invar': 

433 warnings.warn("Use representation == 'sparse' when " 

434 "there are unequally occupied orbitals " 

435 "as the functional is not unitary invariant") 

436 

437 # set reference orbitals 

438 self.dm_helper.set_reference_orbitals(wfs, self.n_dim) 

439 

440 def set_variable_matrices(self, kpt_u): 

441 

442 # Matrices are sparse and Skew-Hermitian. 

443 # They have this structure: 

444 # A_BigMatrix = 

445 # 

446 # ( A_1 A_2 ) 

447 # ( -A_2.T.conj() 0 ) 

448 # 

449 # where 0 is a zero-matrix of size of (M-N) * (M-N) 

450 # 

451 # A_1 i skew-hermitian matrix of N * N, 

452 # N-number of occupied states 

453 # A_2 is matrix of size of (M-N) * N, 

454 # M - number of basis functions 

455 # 

456 # if the energy functional is unitary invariant 

457 # then A_1 = 0 

458 # (see Hutter J et. al, J. Chem. Phys. 101, 3862 (1994)) 

459 # 

460 # We will keep A_1 as we would like to work with metals, 

461 # SIC, and molecules with different occupation numbers. 

462 # this corresponds to 'sparse' representation 

463 # 

464 # Thus, for the 'sparse' we need to store upper 

465 # triangular part of A_1, and matrix A_2, so in total 

466 # (M-N) * N + N * (N - 1)/2 = N * (M - (N + 1)/2) elements 

467 # 

468 # we will store these elements as a vector and 

469 # also will store indices of the A_BigMatrix 

470 # which correspond to these elements. 

471 # 

472 # 'u-invar' corresponds to the case when we want to 

473 # store only A_2, that is this representation is sparser 

474 

475 for kpt in kpt_u: 

476 n_occ = get_n_occ(kpt)[0] 

477 u = self.kpointval(kpt) 

478 self.n_dim[u] = deepcopy(self.n_dim_all[u]) 

479 self.n_dim_oo[u] = n_occ 

480 # M - one dimension of the A_BigMatrix 

481 M = self.n_dim_all[u] 

482 i1_oo, i2_oo = [], [] 

483 for i in range(n_occ): 

484 for j in range(i + 1, n_occ): 

485 i1_oo.append(i) 

486 i2_oo.append(j) 

487 self.ind_oo_up[u] = (np.asarray(i1_oo), np.asarray(i2_oo)) 

488 i1_ov, i2_ov = [], [] 

489 for i in range(n_occ): 

490 for j in range(n_occ, M): 

491 i1_ov.append(i) 

492 i2_ov.append(j) 

493 self.ind_ov_up[u] = (np.asarray(i1_ov), np.asarray(i2_ov)) 

494 i1_sparse, i2_sparse = [], [] 

495 for i in range(n_occ): 

496 for j in range(i + 1, M): 

497 i1_sparse.append(i) 

498 i2_sparse.append(j) 

499 self.ind_sparse_up[u] = (np.asarray(i1_sparse), 

500 np.asarray(i2_sparse)) 

501 if self.representation == 'u-invar': 

502 self.ind_all_up[u] = self.ind_ov_up[u] 

503 if self.representation == 'sparse': 

504 self.ind_all_up[u] = self.ind_sparse_up[u] 

505 elif self.representation == 'full' and self.dtype == complex: 

506 # Take indices of all upper triangular and diagonal 

507 # elements of A_BigMatrix 

508 self.ind_all_up[u] = np.triu_indices(self.n_dim[u]) 

509 

510 shape_of_oo = len(self.ind_oo_up[u][0]) 

511 shape_of_ov = len(self.ind_ov_up[u][0]) 

512 shape_of_all = len(self.ind_all_up[u][0]) 

513 

514 self.a_vec_oo_u[u] = np.zeros(shape=shape_of_oo, dtype=self.dtype) 

515 self.a_vec_ov_u[u] = np.zeros(shape=shape_of_ov, dtype=self.dtype) 

516 self.a_vec_all_u[u] = np.zeros( 

517 shape=shape_of_all, dtype=self.dtype) 

518 self.g_vec_oo_u[u] = np.zeros(shape=shape_of_oo, dtype=self.dtype) 

519 self.g_vec_ov_u[u] = np.zeros(shape=shape_of_ov, dtype=self.dtype) 

520 self.g_vec_all_u[u] = np.zeros( 

521 shape=shape_of_all, dtype=self.dtype) 

522 self.evecs[u] = None 

523 self.evals[u] = None 

524 

525 # All constraints passed as a list of a single orbital index are 

526 # converted to all lists of orbital pairs involving that orbital 

527 # index 

528 if self.constraints: 

529 self.constraints[u] = convert_constraints( 

530 self.constraints[u], self.n_dim[u], 

531 len(kpt.f_n[kpt.f_n > 1e-10]), self.representation) 

532 

533 # If there are no degrees of freedom no need to optimize. 

534 # Indicated by setting n_dim to 0, as n_dim can never be 0 

535 # otherwise. 

536 for k in self.ind_all_up: 

537 if not self.ind_all_up[k][0].size \ 

538 or not self.ind_all_up[k][1].size: 

539 self.n_dim_all[k] = 0 # Skip full space optimization 

540 if not self.ind_oo_up[k][0].size \ 

541 or not self.ind_oo_up[k][1].size: 

542 self.n_dim_oo[k] = 0 # Skip PZ localization if requested 

543 

544 self.ind_up = deepcopy(self.ind_all_up) 

545 self.a_vec_u = deepcopy(self.a_vec_all_u) 

546 self.g_vec_u = deepcopy(self.g_vec_all_u) 

547 

548 # This conversion makes it so that constraint-related functions can 

549 # iterate through a list of no constraints rather than checking for 

550 # None every time 

551 if self.constraints is None: 

552 self.constraints = [[] for _ in range(len(kpt_u))] 

553 

554 self.iters = 1 

555 

556 def localize(self, wfs, dens, ham, log): 

557 if self.need_localization: 

558 localize_orbitals( 

559 wfs, dens, ham, log, self.localizationtype, 

560 seed=self.localizationseed) 

561 self.need_localization = False 

562 

563 def lock_subspace(self, subspace='oo'): 

564 self.subspace_optimization = True 

565 self.subspace_iters = 1 

566 if subspace == 'oo': 

567 self.n_dim = deepcopy(self.n_dim_oo) 

568 self.ind_up = deepcopy(self.ind_oo_up) 

569 self.a_vec_u = deepcopy(self.a_vec_oo_u) 

570 self.g_vec_u = deepcopy(self.g_vec_oo_u) 

571 elif subspace == 'ov': 

572 self.n_dim = deepcopy(self.n_dim_all) 

573 self.ind_up = deepcopy(self.ind_ov_up) 

574 self.a_vec_u = deepcopy(self.a_vec_ov_u) 

575 self.g_vec_u = deepcopy(self.g_vec_ov_u) 

576 self.alpha = 1.0 

577 self.phi_2i = [None, None] 

578 self.der_phi_2i = [None, None] 

579 

580 def release_subspace(self): 

581 self.subspace_optimization = False 

582 self.released_subspace = True 

583 self.n_dim = deepcopy(self.n_dim_all) 

584 self.ind_up = deepcopy(self.ind_all_up) 

585 self.a_vec_u = deepcopy(self.a_vec_all_u) 

586 self.g_vec_u = deepcopy(self.g_vec_all_u) 

587 self.alpha = 1.0 

588 self.phi_2i = [None, None] 

589 self.der_phi_2i = [None, None] 

590 

591 def iterate(self, ham, wfs, dens): 

592 """ 

593 One iteration of direct optimization 

594 for occupied orbitals 

595 

596 :param ham: 

597 :param wfs: 

598 :param dens: 

599 :return: 

600 """ 

601 with wfs.timer('Direct Minimisation step'): 

602 self.update_ref_orbitals(wfs, ham, dens) 

603 

604 a_vec_u = self.a_vec_u 

605 alpha = self.alpha 

606 phi_2i = self.phi_2i 

607 der_phi_2i = self.der_phi_2i 

608 c_ref = self.dm_helper.reference_orbitals 

609 

610 if self.iters == 1 or self.released_subspace or \ 

611 (self.subspace_optimization and self.subspace_iters == 1): 

612 phi_2i[0], g_vec_u = \ 

613 self.get_energy_and_gradients( 

614 a_vec_u, ham, wfs, dens, c_ref) 

615 else: 

616 g_vec_u = self.g_vec_u_original if self.gmf \ 

617 and not self.subspace_optimization else self.g_vec_u 

618 

619 make_pd = False 

620 if self.gmf and not self.subspace_optimization: 

621 with wfs.timer('Partial Hessian diagonalization'): 

622 self.searchdir_algo.update_eigenpairs( 

623 g_vec_u, wfs, ham, dens) 

624 # The diagonal Hessian approximation must be positive-definite 

625 make_pd = True 

626 

627 with wfs.timer('Preconditioning:'): 

628 precond = self.get_preconditioning( 

629 wfs, self.use_prec, make_pd=make_pd) 

630 

631 with wfs.timer('Get Search Direction'): 

632 # calculate search direction according to chosen 

633 # optimization algorithm (e.g. L-BFGS) 

634 p_vec_u = self.searchdir_algo.update_data( 

635 wfs, a_vec_u, g_vec_u, precond=precond, 

636 subspace=self.subspace_optimization) 

637 

638 # recalculate derivative with new search direction 

639 der_phi_2i[0] = 0.0 

640 for k in g_vec_u: 

641 der_phi_2i[0] += g_vec_u[k].conj() @ p_vec_u[k] 

642 der_phi_2i[0] = der_phi_2i[0].real 

643 der_phi_2i[0] = wfs.kd.comm.sum_scalar(der_phi_2i[0]) 

644 

645 alpha, phi_alpha, der_phi_alpha, g_vec_u = \ 

646 self.line_search.step_length_update( 

647 a_vec_u, p_vec_u, wfs, ham, dens, c_ref, phi_0=phi_2i[0], 

648 der_phi_0=der_phi_2i[0], phi_old=phi_2i[1], 

649 der_phi_old=der_phi_2i[1], alpha_max=5.0, alpha_old=alpha, 

650 kpdescr=wfs.kd) 

651 

652 if wfs.gd.comm.size > 1: 

653 with wfs.timer('Broadcast gradients'): 

654 alpha_phi_der_phi = np.array([alpha, phi_2i[0], 

655 der_phi_2i[0]]) 

656 wfs.gd.comm.broadcast(alpha_phi_der_phi, 0) 

657 alpha = alpha_phi_der_phi[0] 

658 phi_2i[0] = alpha_phi_der_phi[1] 

659 der_phi_2i[0] = alpha_phi_der_phi[2] 

660 for kpt in wfs.kpt_u: 

661 k = self.kpointval(kpt) 

662 wfs.gd.comm.broadcast(g_vec_u[k], 0) 

663 

664 # calculate new matrices for optimal step length 

665 for k in a_vec_u: 

666 a_vec_u[k] += alpha * p_vec_u[k] 

667 self.alpha = alpha 

668 self.g_vec_u = g_vec_u 

669 if self.subspace_optimization: 

670 self.subspace_iters += 1 

671 else: 

672 self.iters += 1 

673 

674 # and 'shift' phi, der_phi for the next iteration 

675 phi_2i[1], der_phi_2i[1] = phi_2i[0], der_phi_2i[0] 

676 phi_2i[0], der_phi_2i[0] = phi_alpha, der_phi_alpha, 

677 

678 if self.subspace_optimization: 

679 self.error = np.inf # Do not consider this converged! 

680 

681 def get_grad_norm(self): 

682 norm = 0.0 

683 for k in self.g_vec_u.keys(): 

684 norm += np.linalg.norm(self.g_vec_u[k]) 

685 return norm 

686 

687 def get_energy_and_gradients(self, a_vec_u, ham, wfs, dens, 

688 c_ref): 

689 

690 """ 

691 Energy E = E[C_ref exp(A)]. Gradients G_ij[C, A] = dE/dA_ij 

692 

693 :param wfs: 

694 :param ham: 

695 :param dens: 

696 :param a_vec_u: A 

697 :param c_ref: C_ref 

698 :return: 

699 """ 

700 

701 self.rotate_wavefunctions(wfs, a_vec_u, c_ref) 

702 

703 e_total = self.update_ks_energy(ham, wfs, dens) 

704 

705 with wfs.timer('Calculate gradients'): 

706 g_vec_u = {} 

707 self.error = 0.0 

708 self.e_sic = 0.0 # this is odd energy 

709 for kpt in wfs.kpt_u: 

710 k = self.kpointval(kpt) 

711 if self.n_dim[k] == 0: 

712 g_vec_u[k] = np.zeros_like(a_vec_u[k]) 

713 continue 

714 g_vec_u[k], error = self.dm_helper.calc_grad( 

715 wfs, ham, kpt, self.evecs[k], self.evals[k], 

716 self.matrix_exp, self.representation, self.ind_up[k], 

717 self.constraints[k]) 

718 

719 if hasattr(self.dm_helper.func, 'e_sic_by_orbitals'): 

720 self.e_sic \ 

721 += self.dm_helper.func.e_sic_by_orbitals[k].sum() 

722 

723 self.error += error 

724 self.error = wfs.kd.comm.sum_scalar(self.error) 

725 self.e_sic = wfs.kd.comm.sum_scalar(self.e_sic) 

726 

727 self.eg_count += 1 

728 

729 if self.representation == 'full' and self.checkgraderror: 

730 self._norm_commutator = 0.0 

731 for kpt in wfs.kpt_u: 

732 u = self.kpointval(kpt) 

733 a_mat = vec2skewmat(a_vec_u[u], self.n_dim[u], 

734 self.ind_up[u], wfs.dtype) 

735 g_mat = vec2skewmat(g_vec_u[u], self.n_dim[u], 

736 self.ind_up[u], wfs.dtype) 

737 

738 tmp = np.linalg.norm(g_mat @ a_mat - a_mat @ g_mat) 

739 if self._norm_commutator < tmp: 

740 self._norm_commutator = tmp 

741 

742 tmp = np.linalg.norm(g_mat) 

743 if self._norm_grad < tmp: 

744 self._norm_grad = tmp 

745 

746 return e_total + self.e_sic, g_vec_u 

747 

748 def update_ks_energy(self, ham, wfs, dens): 

749 """ 

750 Update Kohn-Sham energy 

751 It assumes the temperature is zero K. 

752 """ 

753 

754 dens.update(wfs) 

755 ham.update(dens, wfs, False) 

756 

757 return ham.get_energy(0.0, wfs, False) 

758 

759 def evaluate_phi_and_der_phi(self, a_vec_u, p_mat_u, alpha, 

760 wfs, ham, dens, c_ref, 

761 phi=None, g_vec_u=None): 

762 """ 

763 phi = f(x_k + alpha_k*p_k) 

764 der_phi = \\grad f(x_k + alpha_k*p_k) \\cdot p_k 

765 :return: phi, der_phi # floats 

766 """ 

767 if phi is None or g_vec_u is None: 

768 x_mat_u = {k: a_vec_u[k] + alpha * p_mat_u[k] for k in a_vec_u} 

769 phi, g_vec_u = self.get_energy_and_gradients( 

770 x_mat_u, ham, wfs, dens, c_ref) 

771 

772 # If GMF is used save the original gradient and invert the parallel 

773 # projection onto the eigenvectors with negative eigenvalues 

774 if self.gmf and not self.subspace_optimization: 

775 self.g_vec_u_original = deepcopy(g_vec_u) 

776 g_vec_u = self.searchdir_algo.invert_parallel_grad(g_vec_u) 

777 

778 der_phi = 0.0 

779 for k in p_mat_u: 

780 der_phi += g_vec_u[k].conj() @ p_mat_u[k] 

781 

782 der_phi = der_phi.real 

783 der_phi = wfs.kd.comm.sum_scalar(der_phi) 

784 

785 return phi, der_phi, g_vec_u 

786 

787 def update_ref_orbitals(self, wfs, ham, dens): 

788 """ 

789 Update reference orbitals 

790 

791 :param wfs: 

792 :param ham: 

793 :return: 

794 """ 

795 

796 if self.representation == 'full': 

797 badgrad = self._norm_commutator > self._norm_grad / 3. and \ 

798 self.checkgraderror 

799 else: 

800 badgrad = False 

801 counter = self.update_ref_orbs_counter 

802 if (self.iters % counter == 0 and self.iters > 1) or \ 

803 (self.restart and self.iters > 1) or badgrad: 

804 self.iters = 1 

805 if self.update_ref_orbs_canonical or self.restart: 

806 self.get_canonical_representation(ham, wfs, dens) 

807 else: 

808 self.set_ref_orbitals_and_a_vec(wfs) 

809 

810 # Erase memory of search direction algorithm 

811 self.searchdir_algo.reset() 

812 

813 def get_preconditioning(self, wfs, use_prec, make_pd=False): 

814 

815 if not use_prec: 

816 return None 

817 

818 if self.searchdir_algo.name == 'l-bfgs-p': 

819 beta0 = self.searchdir_algo.beta_0 

820 gamma = 0.25 

821 else: 

822 beta0 = 1.0 

823 gamma = 0.0 

824 

825 counter = self.update_precond_counter 

826 precond = {} 

827 for kpt in wfs.kpt_u: 

828 k = self.kpointval(kpt) 

829 w = kpt.weight / (3.0 - wfs.nspins) 

830 if self.iters % counter == 0 or self.iters == 1: 

831 self.hess[k] = get_approx_analytical_hessian(