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 world, parprint 

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 def __init__(self, 

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

26 e: float = None, 

27 contribs: np.ndarray = None, 

28 xc: str = None, 

29 verbose: bool = True): 

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

31 if atoms is None: 

32 assert e is not None 

33 assert contribs is not None 

34 assert xc is not None 

35 else: 

36 if isinstance(atoms, Atoms): 

37 calc = atoms.calc 

38 self.atoms = atoms 

39 else: 

40 calc = atoms 

41 self.atoms = calc.atoms 

42 self.calc = calc 

43 xc = self.calc.get_xc_functional() 

44 self.e = e 

45 self.contribs = contribs 

46 self.xc = xc 

47 self.verbose = verbose 

48 self.done = False 

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

50 self.beef_type = 'beefvdw' 

51 elif self.xc == 'mBEEF': 

52 self.beef_type = 'mbeef' 

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

54 self.beef_type = 'mbeefvdw' 

55 else: 

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

57 

58 def get_ensemble_energies(self, 

59 size: int = 2000, 

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

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

62 self.seed = seed 

63 if self.verbose: 

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

65 

66 if self.contribs is None: 

67 self.contribs = self.calc.get_nonselfconsistent_energies( 

68 self.beef_type) 

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

70 if self.beef_type == 'beefvdw': 

71 assert len(self.contribs) == 32 

72 coefs = self.get_beefvdw_ensemble_coefs(size, seed) 

73 elif self.beef_type == 'mbeef': 

74 assert len(self.contribs) == 64 

75 coefs = self.get_mbeef_ensemble_coefs(size, seed) 

76 elif self.beef_type == 'mbeefvdw': 

77 assert len(self.contribs) == 28 

78 coefs = self.get_mbeefvdw_ensemble_coefs(size, seed) 

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

80 self.done = True 

81 

82 if self.verbose: 

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

84 

85 return self.e + self.de 

86 

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

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

89 from ase.dft.pars_beefvdw import uiOmega as omega 

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

91 

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

93 RandV = generator.randn(31, size) 

94 

95 for j in range(size): 

96 v = RandV[:, j] 

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

98 if j == 0: 

99 ensemble_coefs = coefs_i 

100 else: 

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

102 PBEc_ens = -ensemble_coefs[:, 30] 

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

104 

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

106 """Perturbation coefficients of the mBEEF ensemble""" 

107 from ase.dft.pars_mbeef import uiOmega as omega 

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

109 

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

111 mu, sigma = 0.0, 1.0 

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

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

114 rand)[:]).T 

115 

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

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

118 from ase.dft.pars_mbeefvdw import uiOmega as omega 

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

120 

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

122 mu, sigma = 0.0, 1.0 

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

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

125 

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

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

128 generator = np.random.RandomState(seed) 

129 return s, v.T, generator 

130 

131 def write(self, fname): 

132 """Write ensemble data file""" 

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

134 fname += '.bee' 

135 assert self.done 

136 if world.rank == 0: 

137 if os.path.isfile(fname): 

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

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

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

141 write_json(fd, obj) 

142 

143 

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

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

146 fname += '.bee' 

147 with open(fname, 'r') as fd: 

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

149 if all: 

150 return e, de, contribs, seed, xc 

151 else: 

152 return e + de 

153 

154 

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

156 import warnings 

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

158 return BEEFEnsemble(*args, **kwargs)