Coverage for /builds/ase/ase/ase/dft/bee.py: 90.91%

1# fmt: off 

2 

3import os 

4from typing import Any, Union 

5 

6import numpy as np 

7 

8from ase import Atoms 

9from ase.io.jsonio import read_json, write_json 

10from ase.parallel import parprint, world 

11 

12DFTCalculator = Any 

13 

14 

15def ensemble(energy: float, 

16 contributions: np.ndarray, 

17 xc: str, 

18 verbose: bool = False) -> np.ndarray: 

19 """Returns an array of ensemble total energies.""" 

20 ensemble = BEEFEnsemble(None, energy, contributions, xc, verbose) 

21 return ensemble.get_ensemble_energies() 

22 

23 

24class BEEFEnsemble: 

25 """BEEF type ensemble error estimation.""" 

26 

27 def __init__(self, 

28 atoms: Union[Atoms, DFTCalculator] = None, 

29 e: float = None, 

30 contribs: np.ndarray = None, 

31 xc: str = None, 

32 verbose: bool = True): 

33 if (atoms is not None or contribs is not None or xc is not None): 

34 if atoms is None: 

35 assert e is not None 

36 assert contribs is not None 

37 assert xc is not None 

38 else: 

39 if isinstance(atoms, Atoms): 

40 calc = atoms.calc 

41 self.atoms = atoms 

42 else: 

43 calc = atoms 

44 self.atoms = calc.atoms 

45 self.calc = calc 

46 xc = self.calc.get_xc_functional() 

47 self.e = e 

48 self.contribs = contribs 

49 self.xc = xc 

50 self.verbose = verbose 

51 self.done = False 

52 if self.xc in ['BEEF-vdW', 'BEEF', 'PBE']: 

53 self.beef_type = 'beefvdw' 

54 elif self.xc == 'mBEEF': 

55 self.beef_type = 'mbeef' 

56 elif self.xc == 'mBEEF-vdW': 

57 self.beef_type = 'mbeefvdw' 

58 else: 

59 raise NotImplementedError(f'No ensemble for xc = {self.xc}') 

60 

61 def get_ensemble_energies(self, 

62 size: int = 2000, 

63 seed: int = 0) -> np.ndarray: 

64 """Returns an array of ensemble total energies""" 

65 self.seed = seed 

66 if self.verbose: 

67 parprint(self.beef_type, 'ensemble started') 

68 

69 if self.contribs is None: 

70 self.contribs = self.calc.get_nonselfconsistent_energies( 

71 self.beef_type) 

72 self.e = self.calc.get_potential_energy(self.atoms) 

73 if self.beef_type == 'beefvdw': 

74 assert len(self.contribs) == 32 

75 coefs = self.get_beefvdw_ensemble_coefs(size, seed) 

76 elif self.beef_type == 'mbeef': 

77 assert len(self.contribs) == 64 

78 coefs = self.get_mbeef_ensemble_coefs(size, seed) 

79 elif self.beef_type == 'mbeefvdw': 

80 assert len(self.contribs) == 28 

81 coefs = self.get_mbeefvdw_ensemble_coefs(size, seed) 

82 self.de = np.dot(coefs, self.contribs) 

83 self.done = True 

84 

85 if self.verbose: 

86 parprint(self.beef_type, 'ensemble finished') 

87 

88 return self.e + self.de 

89 

90 def get_beefvdw_ensemble_coefs(self, size=2000, seed=0): 

91 """Perturbation coefficients of the BEEF-vdW ensemble""" 

92 from ase.dft.pars_beefvdw import uiOmega as omega 

93 assert np.shape(omega) == (31, 31) 

94 

95 W, V, generator = self.eigendecomposition(omega, seed) 

96 RandV = generator.randn(31, size) 

97 

98 for j in range(size): 

99 v = RandV[:, j] 

100 coefs_i = (np.dot(np.dot(V, np.diag(np.sqrt(W))), v)[:]) 

101 if j == 0: 

102 ensemble_coefs = coefs_i 

103 else: 

104 ensemble_coefs = np.vstack((ensemble_coefs, coefs_i)) 

105 PBEc_ens = -ensemble_coefs[:, 30] 

106 return (np.vstack((ensemble_coefs.T, PBEc_ens))).T 

107 

108 def get_mbeef_ensemble_coefs(self, size=2000, seed=0): 

109 """Perturbation coefficients of the mBEEF ensemble""" 

110 from ase.dft.pars_mbeef import uiOmega as omega 

111 assert np.shape(omega) == (64, 64) 

112 

113 W, V, generator = self.eigendecomposition(omega, seed) 

114 mu, sigma = 0.0, 1.0 

115 rand = np.array(generator.normal(mu, sigma, (len(W), size))) 

116 return (np.sqrt(2) * np.dot(np.dot(V, np.diag(np.sqrt(W))), 

117 rand)[:]).T 

118 

119 def get_mbeefvdw_ensemble_coefs(self, size=2000, seed=0): 

120 """Perturbation coefficients of the mBEEF-vdW ensemble""" 

121 from ase.dft.pars_mbeefvdw import uiOmega as omega 

122 assert np.shape(omega) == (28, 28) 

123 

124 W, V, generator = self.eigendecomposition(omega, seed) 

125 mu, sigma = 0.0, 1.0 

126 rand = np.array(generator.normal(mu, sigma, (len(W), size))) 

127 return (np.sqrt(2) * np.dot(np.dot(V, np.diag(np.sqrt(W))), rand)[:]).T 

128 

129 def eigendecomposition(self, omega, seed=0): 

130 _u, s, v = np.linalg.svd(omega) # unsafe: W, V = np.linalg.eig(omega) 

131 generator = np.random.RandomState(seed) 

132 return s, v.T, generator 

133 

134 def write(self, fname): 

135 """Write ensemble data file""" 

136 if not fname.endswith('.bee'): 

137 fname += '.bee' 

138 assert self.done 

139 if world.rank == 0: 

140 if os.path.isfile(fname): 

141 os.rename(fname, fname + '.old') 

142 obj = [self.e, self.de, self.contribs, self.seed, self.xc] 

143 with open(fname, 'w') as fd: 

144 write_json(fd, obj) 

145 

146 

147def readbee(fname: str, all: bool = False): 

148 if not fname.endswith('.bee'): 

149 fname += '.bee' 

150 with open(fname) as fd: 

151 e, de, contribs, seed, xc = read_json(fd, always_array=False) 

152 if all: 

153 return e, de, contribs, seed, xc 

154 else: 

155 return e + de 

156 

157 

158def BEEF_Ensemble(*args, **kwargs): 

159 import warnings 

160 warnings.warn('Please use BEEFEnsemble instead of BEEF_Ensemble.') 

161 return BEEFEnsemble(*args, **kwargs)