1import numpy as np 

2from ase.optimize.optimize import Dynamics 

3 

4 

5def subtract_projection(a, b): 

6 '''returns new vector that removes vector a's projection vector b. Is 

7 also equivalent to the vector rejection.''' 

8 aout = a - np.vdot(a, b) / np.vdot(b, b) * b 

9 return aout 

10 

11 

12def normalize(a): 

13 '''Makes a unit vector out of a vector''' 

14 return a / np.linalg.norm(a) 

15 

16 

17class ContourExploration(Dynamics): 

18 

19 def __init__(self, atoms, 

20 maxstep=0.5, 

21 parallel_drift=0.1, 

22 energy_target=None, 

23 angle_limit=20, 

24 potentiostat_step_scale=None, 

25 remove_translation=False, 

26 use_frenet_serret=True, 

27 initialization_step_scale=1e-2, 

28 use_target_shift=True, target_shift_previous_steps=10, 

29 use_tangent_curvature=False, 

30 rng=np.random, 

31 force_consistent=None, 

32 trajectory=None, logfile=None, 

33 append_trajectory=False, loginterval=1): 

34 """Contour Exploration object. 

35 

36 Parameters: 

37 

38 atoms: Atoms object 

39 The Atoms object to operate on. Atomic velocities are required for 

40 the method. If the atoms object does not contain velocities, 

41 random ones will be applied. 

42 

43 maxstep: float 

44 Used to set the maximum distance an atom can move per 

45 iteration (default value is 0.5 Å). 

46 

47 parallel_drift: float 

48 The fraction of the update step that is parallel to the contour but 

49 in a random direction. Used to break symmetries. 

50 

51 energy_target: float 

52 The total system potential energy for that the potentiostat attepts 

53 to maintain. (defaults the initial potential energy) 

54 

55 angle_limit: float or None 

56 Limits the stepsize to a maximum change of direction angle using the 

57 curvature. Gives a scale-free means of tuning the stepsize on the 

58 fly. Typically less than 30 degrees gives reasonable results but 

59 lower angle limits result in higher potentiostatic accuracy. Units 

60 of degrees. (default 20°) 

61 

62 potentiostat_step_scale: float or None 

63 Scales the size of the potentiostat step. The potentiostat step is 

64 determined by linear extrapolation from the current potential energy 

65 to the target_energy with the current forces. A 

66 potentiostat_step_scale > 1.0 overcorrects and < 1.0 

67 undercorrects. By default, a simple heuristic is used to selected 

68 the valued based on the parallel_drift. (default None) 

69 

70 remove_translation: boolean 

71 When True, the net momentum is removed at each step. Improves 

72 potentiostatic accuracy slightly for bulk systems but should not be 

73 used with constraints. (default False) 

74 

75 use_frenet_serret: Bool 

76 Controls whether or not the Taylor expansion of the Frenet-Serret 

77 formulas for curved path extrapolation are used. Required for using 

78 angle_limit based step scalling. (default True) 

79 

80 initialization_step_scale: float 

81 Controls the scale of the initial step as a multiple of maxstep. 

82 (default 1e-2) 

83 

84 use_target_shift: boolean 

85 Enables shifting of the potentiostat target to compensate for 

86 systematic undercorrection or overcorrection by the potentiostat. 

87 Uses the average of the *target_shift_previous_steps* to prevent 

88 coupled occilations. (default True) 

89 

90 target_shift_previous_steps: int 

91 The number of pevious steps to average when using use_target_shift. 

92 (default 10) 

93 

94 use_tangent_curvature: boolean 

95 Use the velocity unit tangent rather than the contour normals from 

96 forces to compute the curvature. Usually not as accurate. 

97 (default False) 

98 

99 rng: a random number generator 

100 Lets users control the random number generator for the 

101 parallel_drift vector. (default numpy.random) 

102 

103 force_consistent: boolean 

104 (default None) 

105 

106 trajectory: Trajectory object or str (optional) 

107 Attach trajectory object. If *trajectory* is a string a 

108 Trajectory will be constructed. Default: None. 

109 

110 logfile: file object or str (optional) 

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

112 Use '-' for stdout. Default: None. 

113 

114 loginterval: int (optional) 

115 Only write a log line for every *loginterval* time steps. 

116 Default: 1 

117 

118 append_trajectory: boolean 

119 Defaults to False, which causes the trajectory file to be 

120 overwriten each time the dynamics is restarted from scratch. 

121 If True, the new structures are appended to the trajectory 

122 file instead. 

123 """ 

124 

125 if potentiostat_step_scale is None: 

126 # a heuristic guess since most systems will overshoot when there is 

127 # drift 

128 self.potentiostat_step_scale = 1.1 + 0.6 * parallel_drift 

129 else: 

130 self.potentiostat_step_scale = potentiostat_step_scale 

131 

132 self.rng = rng 

133 self.remove_translation = remove_translation 

134 self.use_frenet_serret = use_frenet_serret 

135 self.force_consistent = force_consistent 

136 self.use_tangent_curvature = use_tangent_curvature 

137 self.initialization_step_scale = initialization_step_scale 

138 self.maxstep = maxstep 

139 self.angle_limit = angle_limit 

140 self.parallel_drift = parallel_drift 

141 self.use_target_shift = use_target_shift 

142 

143 # These will be populated once self.step() is called, but could be set 

144 # after instantiating with ce = ContourExploration(...) like so: 

145 # ce.Nold = Nold 

146 # ce.r_old = atoms_old.get_positions() 

147 # ce.Told = Told 

148 # to resume a previous contour trajectory. 

149 

150 self.T = None 

151 self.Told = None 

152 self.N = None 

153 self.Nold = None 

154 self.r_old = None 

155 self.r = None 

156 

157 if energy_target is None: 

158 self.energy_target = atoms.get_potential_energy( 

159 force_consistent=self.force_consistent) 

160 else: 

161 self.energy_target = energy_target 

162 

163 # Initizing the previous steps at the target energy slows 

164 # target_shifting untill the system has had 

165 # 'target_shift_previous_steps' steps to equilibrate and should prevent 

166 # occilations. These need to be initialized before the initialize_old 

167 # step to prevent a crash 

168 self.previous_energies = np.full(target_shift_previous_steps, 

169 self.energy_target) 

170 

171 # these first two are purely for logging, 

172 # auto scaling will still occur 

173 # and curvature will still be found if use_frenet_serret == True 

174 self.step_size = 0.0 

175 self.curvature = 0 

176 

177 # loginterval exists for the MolecularDynamics class but not for 

178 # the more general Dynamics class 

179 Dynamics.__init__(self, atoms, 

180 logfile, trajectory, # loginterval, 

181 append_trajectory=append_trajectory, 

182 ) 

183 

184 # we need velocities or NaNs will be produced, 

185 # if none are provided we make random ones 

186 velocities = self.atoms.get_velocities() 

187 if np.linalg.norm(velocities) < 1e-6: 

188 # we have to pass dimension since atoms are not yet stored 

189 atoms.set_velocities(self.rand_vect()) 

190 

191 # Required stuff for Dynamics 

192 def todict(self): 

193 return {'type': 'contour-exploration', 

194 'dyn-type': self.__class__.__name__, 

195 'stepsize': self.step_size} 

196 

197 def run(self, steps=50): 

198 """ Call Dynamics.run and adjust max_steps """ 

199 self.max_steps = steps + self.nsteps 

200 return Dynamics.run(self) 

201 

202 def log(self): 

203 if self.logfile is not None: 

204 # name = self.__class__.__name__ 

205 if self.nsteps == 0: 

206 args = ( 

207 "Step", 

208 "Energy_Target", 

209 "Energy", 

210 "Curvature", 

211 "Step_Size", 

212 "Energy_Deviation_per_atom") 

213 msg = "# %4s %15s %15s %12s %12s %15s\n" % args 

214 self.logfile.write(msg) 

215 e = self.atoms.get_potential_energy( 

216 force_consistent=self.force_consistent) 

217 dev_per_atom = (e - self.energy_target) / len(self.atoms) 

218 args = ( 

219 self.nsteps, 

220 self.energy_target, 

221 e, 

222 self.curvature, 

223 self.step_size, 

224 dev_per_atom) 

225 msg = "%6d %15.6f %15.6f %12.6f %12.6f %24.9f\n" % args 

226 self.logfile.write(msg) 

227 

228 self.logfile.flush() 

229 

230 def rand_vect(self): 

231 '''Returns a random (Natoms,3) vector''' 

232 vect = self.rng.rand(len(self.atoms), 3) - 0.5 

233 return vect 

234 

235 def create_drift_unit_vector(self, N, T): 

236 '''Creates a random drift unit vector with no projection on N or T and 

237 with out a net translation so systems don't wander''' 

238 drift = self.rand_vect() 

239 drift = subtract_projection(drift, N) 

240 drift = subtract_projection(drift, T) 

241 # removes net translation, so systems don't wander 

242 drift = drift - drift.sum(axis=0) / len(self.atoms) 

243 D = normalize(drift) 

244 return D 

245 

246 def compute_step_contributions(self, potentiostat_step_size): 

247 '''Computes the orthogonal component sizes of the step so that the net 

248 step obeys the smaller of step_size or maxstep.''' 

249 if abs(potentiostat_step_size) < self.step_size: 

250 delta_s_perpendicular = potentiostat_step_size 

251 contour_step_size = np.sqrt( 

252 self.step_size**2 - potentiostat_step_size**2) 

253 delta_s_parallel = np.sqrt( 

254 1 - self.parallel_drift**2) * contour_step_size 

255 delta_s_drift = contour_step_size * self.parallel_drift 

256 

257 else: 

258 # in this case all priority goes to potentiostat terms 

259 delta_s_parallel = 0.0 

260 delta_s_drift = 0.0 

261 delta_s_perpendicular = np.sign( 

262 potentiostat_step_size) * self.step_size 

263 

264 return delta_s_perpendicular, delta_s_parallel, delta_s_drift 

265 

266 def _compute_update_without_fs(self, potentiostat_step_size, scale=1.0): 

267 '''Only uses the forces to compute an orthogonal update vector''' 

268 

269 # Without the use of curvature there is no way to estimate the 

270 # limiting step size 

271 self.step_size = self.maxstep * scale 

272 

273 delta_s_perpendicular, delta_s_parallel, delta_s_drift = \ 

274 self.compute_step_contributions( 

275 potentiostat_step_size) 

276 

277 dr_perpendicular = self.N * delta_s_perpendicular 

278 dr_parallel = delta_s_parallel * self.T 

279 

280 D = self.create_drift_unit_vector(self.N, self.T) 

281 dr_drift = D * delta_s_drift 

282 

283 dr = dr_parallel + dr_drift + dr_perpendicular 

284 dr = self.step_size * normalize(dr) 

285 return dr 

286 

287 def _compute_update_with_fs(self, potentiostat_step_size): 

288 '''Uses the Frenet–Serret formulas to perform curvature based 

289 extrapolation to compute the update vector''' 

290 # this should keep the dr clear of the constraints 

291 # by using the actual change, not a velocity vector 

292 delta_r = self.r - self.rold 

293 delta_s = np.linalg.norm(delta_r) 

294 # approximation of delta_s we use this incase an adaptive step_size 

295 # algo get used 

296 

297 delta_T = self.T - self.Told 

298 delta_N = self.N - self.Nold 

299 dTds = delta_T / delta_s 

300 dNds = delta_N / delta_s 

301 if self.use_tangent_curvature: 

302 curvature = np.linalg.norm(dTds) 

303 # on a perfect trajectory, the normal can be computed this way, 

304 # But the normal should always be tied to forces 

305 # N = dTds / curvature 

306 else: 

307 # normals are better since they are fixed to the reality of 

308 # forces. I see smaller forces and energy errors in bulk systems 

309 # using the normals for curvature 

310 curvature = np.linalg.norm(dNds) 

311 self.curvature = curvature 

312 

313 if self.angle_limit is not None: 

314 phi = np.pi / 180 * self.angle_limit 

315 self.step_size = np.sqrt(2 - 2 * np.cos(phi)) / curvature 

316 self.step_size = min(self.step_size, self.maxstep) 

317 

318 # now we can compute a safe step 

319 delta_s_perpendicular, delta_s_parallel, delta_s_drift = \ 

320 self.compute_step_contributions( 

321 potentiostat_step_size) 

322 

323 N_guess = self.N + dNds * delta_s_parallel 

324 T_guess = self.T + dTds * delta_s_parallel 

325 # the extrapolation is good at keeping N_guess and T_guess 

326 # orthogonal but not normalized: 

327 N_guess = normalize(N_guess) 

328 T_guess = normalize(T_guess) 

329 

330 dr_perpendicular = delta_s_perpendicular * (N_guess) 

331 

332 dr_parallel = delta_s_parallel * self.T * \ 

333 (1 - (delta_s_parallel * curvature)**2 / 6.0) \ 

334 + self.N * (curvature / 2.0) * delta_s_parallel**2 

335 

336 D = self.create_drift_unit_vector(N_guess, T_guess) 

337 dr_drift = D * delta_s_drift 

338 

339 # combine the components 

340 dr = dr_perpendicular + dr_parallel + dr_drift 

341 dr = self.step_size * normalize(dr) 

342 # because we guess our orthonormalization directions, 

343 # we renormalize to ensure a correct step size 

344 return dr 

345 

346 def update_previous_energies(self, energy): 

347 '''Updates the energy history in self.previous_energies to include the 

348 current energy.''' 

349 # np.roll shifts the values to keep nice sequential ordering. 

350 self.previous_energies = np.roll(self.previous_energies, 1) 

351 self.previous_energies[0] = energy 

352 

353 def compute_potentiostat_step_size(self, forces, energy): 

354 '''Computes the potentiostat step size by linear extrapolation of the 

355 potential energy using the forces. The step size can be positive or 

356 negative depending on whether or not the energy is too high or too low. 

357 ''' 

358 if self.use_target_shift: 

359 target_shift = self.energy_target - np.mean(self.previous_energies) 

360 else: 

361 target_shift = 0.0 

362 

363 # deltaU is the potential error that will be corrected for 

364 deltaU = energy - (self.energy_target + target_shift) 

365 

366 f_norm = np.linalg.norm(forces) 

367 # can be positive or negative 

368 potentiostat_step_size = (deltaU / f_norm) * \ 

369 self.potentiostat_step_scale 

370 return potentiostat_step_size 

371 

372 def step(self, f=None): 

373 atoms = self.atoms 

374 

375 if f is None: 

376 f = atoms.get_forces() 

377 

378 # get the velocity vector and old kinetic energy for momentum rescaling 

379 velocities = atoms.get_velocities() 

380 KEold = atoms.get_kinetic_energy() 

381 

382 energy = atoms.get_potential_energy( 

383 force_consistent=self.force_consistent) 

384 self.update_previous_energies(energy) 

385 potentiostat_step_size = self.compute_potentiostat_step_size(f, energy) 

386 

387 self.N = normalize(f) 

388 self.r = atoms.get_positions() 

389 # remove velocity projection on forces 

390 v_parallel = subtract_projection(velocities, self.N) 

391 self.T = normalize(v_parallel) 

392 

393 if self.use_frenet_serret: 

394 if self.Nold is not None and self.Told is not None: 

395 dr = self._compute_update_with_fs(potentiostat_step_size) 

396 else: 

397 # we must have the old positions and vectors for an FS step 

398 # if we don't, we can only do a small step 

399 dr = self._compute_update_without_fs( 

400 potentiostat_step_size, 

401 scale=self.initialization_step_scale) 

402 else: # of course we can run less accuratly without FS. 

403 dr = self._compute_update_without_fs(potentiostat_step_size) 

404 

405 # now that dr is done, we check if there is translation 

406 if self.remove_translation: 

407 net_motion = dr.sum(axis=0) / len(atoms) 

408 # print(net_motion) 

409 dr = dr - net_motion 

410 dr_unit = dr / np.linalg.norm(dr) 

411 dr = dr_unit * self.step_size 

412 

413 # save old positions before update 

414 self.Nold = self.N 

415 self.rold = self.r 

416 self.Told = self.T 

417 

418 # if we have constraints then this will do the first part of the 

419 # RATTLE algorithm: 

420 # If we can avoid using momenta, this will be simpler. 

421 masses = atoms.get_masses()[:, np.newaxis] 

422 atoms.set_positions(self.r + dr) 

423 new_momenta = (atoms.get_positions() - self.r) * masses # / self.dt 

424 

425 # We need to store the momenta on the atoms before calculating 

426 # the forces, as in a parallel Asap calculation atoms may 

427 # migrate during force calculations, and the momenta need to 

428 # migrate along with the atoms. 

429 atoms.set_momenta(new_momenta, apply_constraint=False) 

430 

431 # Now we get the new forces! 

432 f = atoms.get_forces(md=True) 

433 

434 # I don't really know if removing md=True from above will break 

435 # compatibility with RATTLE, leaving it alone for now. 

436 f_constrained = atoms.get_forces() 

437 # but this projection needs the forces to be consistent with the 

438 # constraints. We have to set the new velocities perpendicular so they 

439 # get logged properly in the trajectory files. 

440 vnew = subtract_projection(atoms.get_velocities(), f_constrained) 

441 # using the md = True forces like this: 

442 # vnew = subtract_projection(atoms.get_velocities(), f) 

443 # will not work with constraints 

444 atoms.set_velocities(vnew) 

445 

446 # rescaling momentum to maintain constant kinetic energy. 

447 KEnew = atoms.get_kinetic_energy() 

448 Ms = np.sqrt(KEold / KEnew) # Ms = Momentum_scale 

449 atoms.set_momenta(Ms * atoms.get_momenta()) 

450 

451 # Normally this would be the second part of RATTLE 

452 # will be done here like this: 

453 # atoms.set_momenta(atoms.get_momenta() + 0.5 * self.dt * f) 

454 return f