1import numpy as np 

2import warnings 

3 

4from scipy.optimize import minimize 

5from ase.parallel import world 

6from ase.io.jsonio import write_json 

7from ase.optimize.optimize import Optimizer 

8from ase.optimize.gpmin.gp import GaussianProcess 

9from ase.optimize.gpmin.kernel import SquaredExponential 

10from ase.optimize.gpmin.prior import ConstantPrior 

11 

12 

13class GPMin(Optimizer, GaussianProcess): 

14 def __init__(self, atoms, restart=None, logfile='-', trajectory=None, 

15 prior=None, kernel=None, master=None, noise=None, weight=None, 

16 scale=None, force_consistent=None, batch_size=None, 

17 bounds=None, update_prior_strategy="maximum", 

18 update_hyperparams=False): 

19 """Optimize atomic positions using GPMin algorithm, which uses both 

20 potential energies and forces information to build a PES via Gaussian 

21 Process (GP) regression and then minimizes it. 

22 

23 Default behaviour: 

24 -------------------- 

25 The default values of the scale, noise, weight, batch_size and bounds 

26 parameters depend on the value of update_hyperparams. In order to get 

27 the default value of any of them, they should be set up to None. 

28 Default values are: 

29 

30 update_hyperparams = True 

31 scale : 0.3 

32 noise : 0.004 

33 weight: 2. 

34 bounds: 0.1 

35 batch_size: 1 

36 

37 update_hyperparams = False 

38 scale : 0.4 

39 noise : 0.005 

40 weight: 1. 

41 bounds: irrelevant 

42 batch_size: irrelevant 

43 

44 Parameters: 

45 ------------------ 

46 

47 atoms: Atoms object 

48 The Atoms object to relax. 

49 

50 restart: string 

51 JSON file used to store the training set. If set, file with 

52 such a name will be searched and the data in the file incorporated 

53 to the new training set, if the file exists. 

54 

55 logfile: file object or str 

56 If *logfile* is a string, a file with that name will be opened. 

57 Use '-' for stdout 

58 

59 trajectory: string 

60 File used to store trajectory of atomic movement. 

61 

62 master: boolean 

63 Defaults to None, which causes only rank 0 to save files. If 

64 set to True, this rank will save files. 

65 

66 force_consistent: boolean or None 

67 Use force-consistent energy calls (as opposed to the energy 

68 extrapolated to 0 K). By default (force_consistent=None) uses 

69 force-consistent energies if available in the calculator, but 

70 falls back to force_consistent=False if not. 

71 

72 prior: Prior object or None 

73 Prior for the GP regression of the PES surface 

74 See ase.optimize.gpmin.prior 

75 If *prior* is None, then it is set as the 

76 ConstantPrior with the constant being updated 

77 using the update_prior_strategy specified as a parameter 

78 

79 kernel: Kernel object or None 

80 Kernel for the GP regression of the PES surface 

81 See ase.optimize.gpmin.kernel 

82 If *kernel* is None the SquaredExponential kernel is used. 

83 Note: It needs to be a kernel with derivatives!!!!! 

84 

85 noise: float 

86 Regularization parameter for the Gaussian Process Regression. 

87 

88 weight: float 

89 Prefactor of the Squared Exponential kernel. 

90 If *update_hyperparams* is False, changing this parameter 

91 has no effect on the dynamics of the algorithm. 

92 

93 update_prior_strategy: string 

94 Strategy to update the constant from the ConstantPrior 

95 when more data is collected. It does only work when 

96 Prior = None 

97 

98 options: 

99 'maximum': update the prior to the maximum sampled energy 

100 'init' : fix the prior to the initial energy 

101 'average': use the average of sampled energies as prior 

102 

103 scale: float 

104 scale of the Squared Exponential Kernel 

105 

106 update_hyperparams: boolean 

107 Update the scale of the Squared exponential kernel 

108 every batch_size-th iteration by maximizing the 

109 marginal likelihood. 

110 

111 batch_size: int 

112 Number of new points in the sample before updating 

113 the hyperparameters. 

114 Only relevant if the optimizer is executed in update_hyperparams 

115 mode: (update_hyperparams = True) 

116 

117 bounds: float, 0<bounds<1 

118 Set bounds to the optimization of the hyperparameters. 

119 Let t be a hyperparameter. Then it is optimized under the 

120 constraint (1-bound)*t_0 <= t <= (1+bound)*t_0 

121 where t_0 is the value of the hyperparameter in the previous 

122 step. 

123 If bounds is False, no constraints are set in the optimization of 

124 the hyperparameters. 

125 

126 .. warning:: The memory of the optimizer scales as O(n²N²) where 

127 N is the number of atoms and n the number of steps. 

128 If the number of atoms is sufficiently high, this 

129 may cause a memory issue. 

130 This class prints a warning if the user tries to 

131 run GPMin with more than 100 atoms in the unit cell. 

132 """ 

133 # Warn the user if the number of atoms is very large 

134 if len(atoms) > 100: 

135 warning = ('Possible Memory Issue. There are more than ' 

136 '100 atoms in the unit cell. The memory ' 

137 'of the process will increase with the number ' 

138 'of steps, potentially causing a memory issue. ' 

139 'Consider using a different optimizer.') 

140 

141 warnings.warn(warning) 

142 

143 # Give it default hyperparameters 

144 if update_hyperparams: # Updated GPMin 

145 if scale is None: 

146 scale = 0.3 

147 if noise is None: 

148 noise = 0.004 

149 if weight is None: 

150 weight = 2. 

151 if bounds is None: 

152 self.eps = 0.1 

153 elif bounds is False: 

154 self.eps = None 

155 else: 

156 self.eps = bounds 

157 if batch_size is None: 

158 self.nbatch = 1 

159 else: 

160 self.nbatch = batch_size 

161 else: # GPMin without updates 

162 if scale is None: 

163 scale = 0.4 

164 if noise is None: 

165 noise = 0.001 

166 if weight is None: 

167 weight = 1. 

168 if bounds is not None: 

169 warning = ('The parameter bounds is of no use ' 

170 'if update_hyperparams is False. ' 

171 'The value provided by the user ' 

172 'is being ignored.') 

173 warnings.warn(warning, UserWarning) 

174 if batch_size is not None: 

175 warning = ('The parameter batch_size is of no use ' 

176 'if update_hyperparams is False. ' 

177 'The value provided by the user ' 

178 'is being ignored.') 

179 warnings.warn(warning, UserWarning) 

180 

181 # Set the variables to something anyways 

182 self.eps = False 

183 self.nbatch = None 

184 

185 self.strategy = update_prior_strategy 

186 self.update_hp = update_hyperparams 

187 self.function_calls = 1 

188 self.force_calls = 0 

189 self.x_list = [] # Training set features 

190 self.y_list = [] # Training set targets 

191 

192 Optimizer.__init__(self, atoms, restart, logfile, 

193 trajectory, master, force_consistent) 

194 if prior is None: 

195 self.update_prior = True 

196 prior = ConstantPrior(constant=None) 

197 else: 

198 self.update_prior = False 

199 

200 if kernel is None: 

201 kernel = SquaredExponential() 

202 GaussianProcess.__init__(self, prior, kernel) 

203 self.set_hyperparams(np.array([weight, scale, noise])) 

204 

205 def acquisition(self, r): 

206 e = self.predict(r) 

207 return e[0], e[1:] 

208 

209 def update(self, r, e, f): 

210 """Update the PES 

211 

212 Update the training set, the prior and the hyperparameters. 

213 Finally, train the model 

214 """ 

215 # update the training set 

216 self.x_list.append(r) 

217 f = f.reshape(-1) 

218 y = np.append(np.array(e).reshape(-1), -f) 

219 self.y_list.append(y) 

220 

221 # Set/update the constant for the prior 

222 if self.update_prior: 

223 if self.strategy == 'average': 

224 av_e = np.mean(np.array(self.y_list)[:, 0]) 

225 self.prior.set_constant(av_e) 

226 elif self.strategy == 'maximum': 

227 max_e = np.max(np.array(self.y_list)[:, 0]) 

228 self.prior.set_constant(max_e) 

229 elif self.strategy == 'init': 

230 self.prior.set_constant(e) 

231 self.update_prior = False 

232 

233 # update hyperparams 

234 if (self.update_hp and self.function_calls % self.nbatch == 0 and 

235 self.function_calls != 0): 

236 self.fit_to_batch() 

237 

238 # build the model 

239 self.train(np.array(self.x_list), np.array(self.y_list)) 

240 

241 def relax_model(self, r0): 

242 result = minimize(self.acquisition, r0, method='L-BFGS-B', jac=True) 

243 if result.success: 

244 return result.x 

245 else: 

246 self.dump() 

247 raise RuntimeError("The minimization of the acquisition function " 

248 "has not converged") 

249 

250 def fit_to_batch(self): 

251 """Fit hyperparameters keeping the ratio noise/weight fixed""" 

252 ratio = self.noise/self.kernel.weight 

253 self.fit_hyperparameters(np.array(self.x_list), 

254 np.array(self.y_list), eps=self.eps) 

255 self.noise = ratio*self.kernel.weight 

256 

257 def step(self, f=None): 

258 atoms = self.atoms 

259 if f is None: 

260 f = atoms.get_forces() 

261 

262 fc = self.force_consistent 

263 r0 = atoms.get_positions().reshape(-1) 

264 e0 = atoms.get_potential_energy(force_consistent=fc) 

265 self.update(r0, e0, f) 

266 

267 r1 = self.relax_model(r0) 

268 self.atoms.set_positions(r1.reshape(-1, 3)) 

269 e1 = self.atoms.get_potential_energy(force_consistent=fc) 

270 f1 = self.atoms.get_forces() 

271 self.function_calls += 1 

272 self.force_calls += 1 

273 count = 0 

274 while e1 >= e0: 

275 self.update(r1, e1, f1) 

276 r1 = self.relax_model(r0) 

277 

278 self.atoms.set_positions(r1.reshape(-1, 3)) 

279 e1 = self.atoms.get_potential_energy(force_consistent=fc) 

280 f1 = self.atoms.get_forces() 

281 self.function_calls += 1 

282 self.force_calls += 1 

283 if self.converged(f1): 

284 break 

285 

286 count += 1 

287 if count == 30: 

288 raise RuntimeError("A descent model could not be built") 

289 self.dump() 

290 

291 def dump(self): 

292 """Save the training set""" 

293 if world.rank == 0 and self.restart is not None: 

294 with open(self.restart, 'wb') as fd: 

295 write_json(fd, (self.x_list, self.y_list)) 

296 

297 def read(self): 

298 self.x_list, self.y_list = self.load()