Coverage for gpaw/directmin/locfunc/localize_orbitals.py: 10%

1from gpaw.directmin.locfunc.etdm_localization_fdpw import FDPWETDMLocalize 

2from gpaw.directmin.locfunc.etdm_localization_lcao import LCAOETDMLocalize 

3from gpaw.directmin.fdpw.er_localization import ERLocalization as ERL 

4from gpaw.directmin.functional.fdpw import get_functional \ 

5 as get_functional_fdpw 

6from gpaw.pipekmezey.pipek_mezey_wannier import PipekMezey 

7from gpaw.pipekmezey.wannier_basic import WannierLocalization 

8import numpy as np 

9 

10 

11def localize_orbitals( 

12 wfs, dens, ham, log, localizationtype, tol=None, seed=None, 

13 func_settings=None): 

14 io = localizationtype 

15 

16 if io is None: 

17 return 

18 

19 locnames = io.lower().split('_') 

20 if tol is not None: 

21 tols = {name: tol for name in locnames} 

22 else: 

23 tols = {'er': 5.0e-5, 

24 'pm': 1.0e-10, 

25 'fb': 1.0e-10, 

26 'pz': 5.0e-4, 

27 'ks': 5.0e-4} 

28 

29 log("Initial localization: ...", flush=True) 

30 wfs.timer.start('Initial Localization') 

31 for name in locnames: 

32 tol = tols[name] 

33 if name == 'er': 

34 if wfs.mode == 'lcao': 

35 log('Edmiston-Ruedenberg localization is not supported ' 

36 'in LCAO', 

37 flush=True) 

38 continue 

39 log('Edmiston-Ruedenberg localization started', 

40 flush=True) 

41 dm = FDPWETDMLocalize( 

42 ERL(wfs, dens), wfs, 

43 maxiter=200, g_tol=tol, randval=0.1) 

44 dm.run(wfs, dens, seed=seed) 

45 log('Edmiston-Ruedenberg localization finished', 

46 flush=True) 

47 del dm 

48 elif name == 'pz': 

49 if wfs.mode != 'lcao': 

50 log('Perdew-Zunger localization started', flush=True) 

51 PZC = get_functional_fdpw(func_settings, wfs, dens, ham) 

52 dm = FDPWETDMLocalize( 

53 PZC, wfs, maxiter=200, g_tol=tol, randval=0.1) 

54 dm.run(wfs, dens, log=log) 

55 log('Perdew-Zunger localization finished', flush=True) 

56 else: 

57 dm = LCAOETDMLocalize( 

58 wfs.eigensolver, wfs, log, 

59 tol=wfs.eigensolver.subspace_convergence) 

60 dm.run(ham, dens) 

61 elif name == 'ks': 

62 log('ETDM minimization using occupied and virtual orbitals', 

63 flush=True) 

64 if wfs.mode == 'lcao': 

65 raise NotImplementedError 

66 else: 

67 KS = get_functional_fdpw('ks', wfs, dens, ham) 

68 dm = FDPWETDMLocalize(KS, wfs, maxiter=200, g_tol=tol, randval=0) 

69 dm.run(wfs, dens, ham=ham, log=log) 

70 log('ETDM minimization finished', flush=True) 

71 else: 

72 for k, kpt in enumerate(wfs.kpt_u): 

73 if sum(kpt.f_n) < 1.0e-3: 

74 continue 

75 if name == 'pm': 

76 if k == 0: 

77 log('Pipek-Mezey localization started', 

78 flush=True) 

79 lf_obj = PipekMezey( 

80 wfs=wfs, spin=kpt.s, dtype=wfs.dtype, seed=seed) 

81 lf_obj.localize(tolerance=tol) 

82 if k == 0: 

83 log('Pipek-Mezey localization finished', 

84 flush=True) 

85 U = np.ascontiguousarray( 

86 lf_obj.W_k[kpt.q].T) 

87 elif name == 'fb': 

88 if k == 0: 

89 log('Foster-Boys localization started', 

90 flush=True) 

91 lf_obj = WannierLocalization( 

92 wfs=wfs, spin=kpt.s, seed=seed) 

93 lf_obj.localize(tolerance=tol) 

94 if k == 0: 

95 log('Foster-Boys localization finished', 

96 flush=True) 

97 U = np.ascontiguousarray( 

98 lf_obj.U_kww[kpt.q].T) 

99 if wfs.dtype == float: 

100 U = U.real 

101 else: 

102 raise ValueError('Check localization type.') 

103 wfs.gd.comm.broadcast(U, 0) 

104 dim = U.shape[0] 

105 if wfs.mode == 'fd': 

106 kpt.psit_nG[:dim] = np.einsum( 

107 'ij,jkml->ikml', U, kpt.psit_nG[:dim]) 

108 elif wfs.mode == 'pw': 

109 kpt.psit_nG[:dim] = U @ kpt.psit_nG[:dim] 

110 else: 

111 kpt.C_nM[:dim] = U @ kpt.C_nM[:dim] 

112 

113 del lf_obj 

114 

115 wfs.timer.stop('Initial Localization') 

116 log("Done", flush=True)