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

1import os 

2from typing import Any, Union 

3 

4import numpy as np 

5 

6from ase import Atoms 

7from ase.io.jsonio import read_json, write_json 

8from ase.parallel import parprint, world 

9 

10DFTCalculator = Any 

11 

12 

13def ensemble(energy: float, 

14 contributions: np.ndarray, 

15 xc: str, 

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

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

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

19 return ensemble.get_ensemble_energies() 

20 

21 

22class BEEFEnsemble: 

23 """BEEF type ensemble error estimation.""" 

24 

25 def __init__(self, 

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

27 e: float = None, 

28 contribs: np.ndarray = None, 

29 xc: str = None, 

30 verbose: bool = True): 

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

32 if atoms is None: 

33 assert e is not None 

34 assert contribs is not None 

35 assert xc is not None 

36 else: 

37 if isinstance(atoms, Atoms): 

38 calc = atoms.calc 

39 self.atoms = atoms 

40 else: 

41 calc = atoms 

42 self.atoms = calc.atoms 

43 self.calc = calc 

44 xc = self.calc.get_xc_functional() 

45 self.e = e 

46 self.contribs = contribs 

47 self.xc = xc 

48 self.verbose = verbose 

49 self.done = False 

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

51 self.beef_type = 'beefvdw' 

52 elif self.xc == 'mBEEF': 

53 self.beef_type = 'mbeef' 

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

55 self.beef_type = 'mbeefvdw' 

56 else: 

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

58 

59 def get_ensemble_energies(self, 

60 size: int = 2000, 

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

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

63 self.seed = seed 

64 if self.verbose: 

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

66 

67 if self.contribs is None: 

68 self.contribs = self.calc.get_nonselfconsistent_energies( 

69 self.beef_type) 

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

71 if self.beef_type == 'beefvdw': 

72 assert len(self.contribs) == 32 

73 coefs = self.get_beefvdw_ensemble_coefs(size, seed) 

74 elif self.beef_type == 'mbeef': 

75 assert len(self.contribs) == 64 

76 coefs = self.get_mbeef_ensemble_coefs(size, seed) 

77 elif self.beef_type == 'mbeefvdw': 

78 assert len(self.contribs) == 28 

79 coefs = self.get_mbeefvdw_ensemble_coefs(size, seed) 

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

81 self.done = True 

82 

83 if self.verbose: 

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

85 

86 return self.e + self.de 

87 

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

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

90 from ase.dft.pars_beefvdw import uiOmega as omega 

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

92 

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

94 RandV = generator.randn(31, size) 

95 

96 for j in range(size): 

97 v = RandV[:, j] 

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

99 if j == 0: 

100 ensemble_coefs = coefs_i 

101 else: 

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

103 PBEc_ens = -ensemble_coefs[:, 30] 

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

105 

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

107 """Perturbation coefficients of the mBEEF ensemble""" 

108 from ase.dft.pars_mbeef import uiOmega as omega 

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

110 

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

112 mu, sigma = 0.0, 1.0 

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

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

115 rand)[:]).T 

116 

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

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

119 from ase.dft.pars_mbeefvdw import uiOmega as omega 

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

121 

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

123 mu, sigma = 0.0, 1.0 

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

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

126 

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

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

129 generator = np.random.RandomState(seed) 

130 return s, v.T, generator 

131 

132 def write(self, fname): 

133 """Write ensemble data file""" 

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

135 fname += '.bee' 

136 assert self.done 

137 if world.rank == 0: 

138 if os.path.isfile(fname): 

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

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

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

142 write_json(fd, obj) 

143 

144 

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

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

147 fname += '.bee' 

148 with open(fname) as fd: 

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

150 if all: 

151 return e, de, contribs, seed, xc 

152 else: 

153 return e + de 

154 

155 

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

157 import warnings 

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

159 return BEEFEnsemble(*args, **kwargs)