Coverage for /builds/kinetik161/ase/ase/md/contour_exploration.py: 96.30%
162 statements
« prev ^ index » next coverage.py v7.2.7, created at 2023-12-10 11:04 +0000
1from typing import IO, Optional, Union
3import numpy as np
5from ase import Atoms
6from ase.optimize.optimize import Dynamics
9def subtract_projection(a, b):
10 '''returns new vector that removes vector a's projection vector b. Is
11 also equivalent to the vector rejection.'''
12 aout = a - np.vdot(a, b) / np.vdot(b, b) * b
13 return aout
16def normalize(a):
17 '''Makes a unit vector out of a vector'''
18 return a / np.linalg.norm(a)
21class ContourExploration(Dynamics):
23 def __init__(
24 self,
25 atoms: Atoms,
26 maxstep: float = 0.5,
27 parallel_drift: float = 0.1,
28 energy_target: Optional[float] = None,
29 angle_limit: Optional[float] = 20.0,
30 potentiostat_step_scale: Optional[float] = None,
31 remove_translation: bool = False,
32 use_frenet_serret: bool = True,
33 initialization_step_scale: float = 1e-2,
34 use_target_shift: bool = True,
35 target_shift_previous_steps: int = 10,
36 use_tangent_curvature: bool = False,
37 rng=np.random,
38 force_consistent: Optional[bool] = None,
39 trajectory: Optional[str] = None,
40 logfile: Optional[Union[IO, str]] = None,
41 append_trajectory: bool = False,
42 loginterval: int = 1,
43 ):
44 """Contour Exploration object.
46 Parameters:
48 atoms: Atoms object
49 The Atoms object to operate on. Atomic velocities are required for
50 the method. If the atoms object does not contain velocities,
51 random ones will be applied.
53 maxstep: float
54 Used to set the maximum distance an atom can move per
55 iteration (default value is 0.5 Å).
57 parallel_drift: float
58 The fraction of the update step that is parallel to the contour but
59 in a random direction. Used to break symmetries.
61 energy_target: float
62 The total system potential energy for that the potentiostat attepts
63 to maintain. (defaults the initial potential energy)
65 angle_limit: float or None
66 Limits the stepsize to a maximum change of direction angle using the
67 curvature. Gives a scale-free means of tuning the stepsize on the
68 fly. Typically less than 30 degrees gives reasonable results but
69 lower angle limits result in higher potentiostatic accuracy. Units
70 of degrees. (default 20°)
72 potentiostat_step_scale: float or None
73 Scales the size of the potentiostat step. The potentiostat step is
74 determined by linear extrapolation from the current potential energy
75 to the target_energy with the current forces. A
76 potentiostat_step_scale > 1.0 overcorrects and < 1.0
77 undercorrects. By default, a simple heuristic is used to selected
78 the valued based on the parallel_drift. (default None)
80 remove_translation: boolean
81 When True, the net momentum is removed at each step. Improves
82 potentiostatic accuracy slightly for bulk systems but should not be
83 used with constraints. (default False)
85 use_frenet_serret: Bool
86 Controls whether or not the Taylor expansion of the Frenet-Serret
87 formulas for curved path extrapolation are used. Required for using
88 angle_limit based step scalling. (default True)
90 initialization_step_scale: float
91 Controls the scale of the initial step as a multiple of maxstep.
92 (default 1e-2)
94 use_target_shift: boolean
95 Enables shifting of the potentiostat target to compensate for
96 systematic undercorrection or overcorrection by the potentiostat.
97 Uses the average of the *target_shift_previous_steps* to prevent
98 coupled occilations. (default True)
100 target_shift_previous_steps: int
101 The number of pevious steps to average when using use_target_shift.
102 (default 10)
104 use_tangent_curvature: boolean
105 Use the velocity unit tangent rather than the contour normals from
106 forces to compute the curvature. Usually not as accurate.
107 (default False)
109 rng: a random number generator
110 Lets users control the random number generator for the
111 parallel_drift vector. (default numpy.random)
113 force_consistent: boolean
114 (default None)
116 trajectory: Trajectory object or str (optional)
117 Attach trajectory object. If *trajectory* is a string a
118 Trajectory will be constructed. Default: None.
120 logfile: file object or str (optional)
121 If *logfile* is a string, a file with that name will be opened.
122 Use '-' for stdout. Default: None.
124 loginterval: int (optional)
125 Only write a log line for every *loginterval* time steps.
126 Default: 1
128 append_trajectory: boolean
129 Defaults to False, which causes the trajectory file to be
130 overwriten each time the dynamics is restarted from scratch.
131 If True, the new structures are appended to the trajectory
132 file instead.
133 """
135 if potentiostat_step_scale is None:
136 # a heuristic guess since most systems will overshoot when there is
137 # drift
138 self.potentiostat_step_scale = 1.1 + 0.6 * parallel_drift
139 else:
140 self.potentiostat_step_scale = potentiostat_step_scale
142 self.rng = rng
143 self.remove_translation = remove_translation
144 self.use_frenet_serret = use_frenet_serret
145 self.use_tangent_curvature = use_tangent_curvature
146 self.initialization_step_scale = initialization_step_scale
147 self.maxstep = maxstep
148 self.angle_limit = angle_limit
149 self.parallel_drift = parallel_drift
150 self.use_target_shift = use_target_shift
152 # These will be populated once self.step() is called, but could be set
153 # after instantiating with ce = ContourExploration(...) like so:
154 # ce.Nold = Nold
155 # ce.r_old = atoms_old.get_positions()
156 # ce.Told = Told
157 # to resume a previous contour trajectory.
159 self.T = None
160 self.Told = None
161 self.N = None
162 self.Nold = None
163 self.r_old = None
164 self.r = None
166 if energy_target is None:
167 self.energy_target = atoms.get_potential_energy(
168 force_consistent=True)
169 else:
170 self.energy_target = energy_target
172 # Initizing the previous steps at the target energy slows
173 # target_shifting untill the system has had
174 # 'target_shift_previous_steps' steps to equilibrate and should prevent
175 # occilations. These need to be initialized before the initialize_old
176 # step to prevent a crash
177 self.previous_energies = np.full(target_shift_previous_steps,
178 self.energy_target)
180 # these first two are purely for logging,
181 # auto scaling will still occur
182 # and curvature will still be found if use_frenet_serret == True
183 self.step_size = 0.0
184 self.curvature = 0
186 # loginterval exists for the MolecularDynamics class but not for
187 # the more general Dynamics class
188 Dynamics.__init__(self, atoms,
189 logfile, trajectory, # loginterval,
190 append_trajectory=append_trajectory,
191 )
193 self._actual_atoms = atoms
195 # we need velocities or NaNs will be produced,
196 # if none are provided we make random ones
197 velocities = self._actual_atoms.get_velocities()
198 if np.linalg.norm(velocities) < 1e-6:
199 # we have to pass dimension since atoms are not yet stored
200 atoms.set_velocities(self.rand_vect())
202 # Required stuff for Dynamics
203 def todict(self):
204 return {'type': 'contour-exploration',
205 'dyn-type': self.__class__.__name__,
206 'stepsize': self.step_size}
208 def run(self, steps=50):
209 """ Call Dynamics.run and adjust max_steps """
210 return Dynamics.run(self, steps=steps)
212 def log(self):
213 if self.logfile is not None:
214 # name = self.__class__.__name__
215 if self.nsteps == 0:
216 args = (
217 "Step",
218 "Energy_Target",
219 "Energy",
220 "Curvature",
221 "Step_Size",
222 "Energy_Deviation_per_atom")
223 msg = "# %4s %15s %15s %12s %12s %15s\n" % args
224 self.logfile.write(msg)
225 e = self._actual_atoms.get_potential_energy(force_consistent=True)
226 dev_per_atom = (e - self.energy_target) / len(self._actual_atoms)
227 args = (
228 self.nsteps,
229 self.energy_target,
230 e,
231 self.curvature,
232 self.step_size,
233 dev_per_atom)
234 msg = "%6d %15.6f %15.6f %12.6f %12.6f %24.9f\n" % args
235 self.logfile.write(msg)
237 self.logfile.flush()
239 def rand_vect(self):
240 '''Returns a random (Natoms,3) vector'''
241 vect = self.rng.random((len(self._actual_atoms), 3)) - 0.5
242 return vect
244 def create_drift_unit_vector(self, N, T):
245 '''Creates a random drift unit vector with no projection on N or T and
246 with out a net translation so systems don't wander'''
247 drift = self.rand_vect()
248 drift = subtract_projection(drift, N)
249 drift = subtract_projection(drift, T)
250 # removes net translation, so systems don't wander
251 drift = drift - drift.sum(axis=0) / len(self._actual_atoms)
252 D = normalize(drift)
253 return D
255 def compute_step_contributions(self, potentiostat_step_size):
256 '''Computes the orthogonal component sizes of the step so that the net
257 step obeys the smaller of step_size or maxstep.'''
258 if abs(potentiostat_step_size) < self.step_size:
259 delta_s_perpendicular = potentiostat_step_size
260 contour_step_size = np.sqrt(
261 self.step_size**2 - potentiostat_step_size**2)
262 delta_s_parallel = np.sqrt(
263 1 - self.parallel_drift**2) * contour_step_size
264 delta_s_drift = contour_step_size * self.parallel_drift
266 else:
267 # in this case all priority goes to potentiostat terms
268 delta_s_parallel = 0.0
269 delta_s_drift = 0.0
270 delta_s_perpendicular = np.sign(
271 potentiostat_step_size) * self.step_size
273 return delta_s_perpendicular, delta_s_parallel, delta_s_drift
275 def _compute_update_without_fs(self, potentiostat_step_size, scale=1.0):
276 '''Only uses the forces to compute an orthogonal update vector'''
278 # Without the use of curvature there is no way to estimate the
279 # limiting step size
280 self.step_size = self.maxstep * scale
282 delta_s_perpendicular, delta_s_parallel, delta_s_drift = \
283 self.compute_step_contributions(
284 potentiostat_step_size)
286 dr_perpendicular = self.N * delta_s_perpendicular
287 dr_parallel = delta_s_parallel * self.T
289 D = self.create_drift_unit_vector(self.N, self.T)
290 dr_drift = D * delta_s_drift
292 dr = dr_parallel + dr_drift + dr_perpendicular
293 dr = self.step_size * normalize(dr)
294 return dr
296 def _compute_update_with_fs(self, potentiostat_step_size):
297 '''Uses the Frenet–Serret formulas to perform curvature based
298 extrapolation to compute the update vector'''
299 # this should keep the dr clear of the constraints
300 # by using the actual change, not a velocity vector
301 delta_r = self.r - self.rold
302 delta_s = np.linalg.norm(delta_r)
303 # approximation of delta_s we use this incase an adaptive step_size
304 # algo get used
306 delta_T = self.T - self.Told
307 delta_N = self.N - self.Nold
308 dTds = delta_T / delta_s
309 dNds = delta_N / delta_s
310 if self.use_tangent_curvature:
311 curvature = np.linalg.norm(dTds)
312 # on a perfect trajectory, the normal can be computed this way,
313 # But the normal should always be tied to forces
314 # N = dTds / curvature
315 else:
316 # normals are better since they are fixed to the reality of
317 # forces. I see smaller forces and energy errors in bulk systems
318 # using the normals for curvature
319 curvature = np.linalg.norm(dNds)
320 self.curvature = curvature
322 if self.angle_limit is not None:
323 phi = np.pi / 180 * self.angle_limit
324 self.step_size = np.sqrt(2 - 2 * np.cos(phi)) / curvature
325 self.step_size = min(self.step_size, self.maxstep)
327 # now we can compute a safe step
328 delta_s_perpendicular, delta_s_parallel, delta_s_drift = \
329 self.compute_step_contributions(
330 potentiostat_step_size)
332 N_guess = self.N + dNds * delta_s_parallel
333 T_guess = self.T + dTds * delta_s_parallel
334 # the extrapolation is good at keeping N_guess and T_guess
335 # orthogonal but not normalized:
336 N_guess = normalize(N_guess)
337 T_guess = normalize(T_guess)
339 dr_perpendicular = delta_s_perpendicular * (N_guess)
341 dr_parallel = delta_s_parallel * self.T * \
342 (1 - (delta_s_parallel * curvature)**2 / 6.0) \
343 + self.N * (curvature / 2.0) * delta_s_parallel**2
345 D = self.create_drift_unit_vector(N_guess, T_guess)
346 dr_drift = D * delta_s_drift
348 # combine the components
349 dr = dr_perpendicular + dr_parallel + dr_drift
350 dr = self.step_size * normalize(dr)
351 # because we guess our orthonormalization directions,
352 # we renormalize to ensure a correct step size
353 return dr
355 def update_previous_energies(self, energy):
356 '''Updates the energy history in self.previous_energies to include the
357 current energy.'''
358 # np.roll shifts the values to keep nice sequential ordering.
359 self.previous_energies = np.roll(self.previous_energies, 1)
360 self.previous_energies[0] = energy
362 def compute_potentiostat_step_size(self, forces, energy):
363 '''Computes the potentiostat step size by linear extrapolation of the
364 potential energy using the forces. The step size can be positive or
365 negative depending on whether or not the energy is too high or too low.
366 '''
367 if self.use_target_shift:
368 target_shift = self.energy_target - np.mean(self.previous_energies)
369 else:
370 target_shift = 0.0
372 # deltaU is the potential error that will be corrected for
373 deltaU = energy - (self.energy_target + target_shift)
375 f_norm = np.linalg.norm(forces)
376 # can be positive or negative
377 potentiostat_step_size = (deltaU / f_norm) * \
378 self.potentiostat_step_scale
379 return potentiostat_step_size
381 def step(self, f=None):
382 atoms = self._actual_atoms
384 if f is None:
385 f = atoms.get_forces()
387 # get the velocity vector and old kinetic energy for momentum rescaling
388 velocities = atoms.get_velocities()
389 KEold = atoms.get_kinetic_energy()
391 energy = atoms.get_potential_energy(force_consistent=True)
392 self.update_previous_energies(energy)
393 potentiostat_step_size = self.compute_potentiostat_step_size(f, energy)
395 self.N = normalize(f)
396 self.r = atoms.get_positions()
397 # remove velocity projection on forces
398 v_parallel = subtract_projection(velocities, self.N)
399 self.T = normalize(v_parallel)
401 if self.use_frenet_serret:
402 if self.Nold is not None and self.Told is not None:
403 dr = self._compute_update_with_fs(potentiostat_step_size)
404 else:
405 # we must have the old positions and vectors for an FS step
406 # if we don't, we can only do a small step
407 dr = self._compute_update_without_fs(
408 potentiostat_step_size,
409 scale=self.initialization_step_scale)
410 else: # of course we can run less accuratly without FS.
411 dr = self._compute_update_without_fs(potentiostat_step_size)
413 # now that dr is done, we check if there is translation
414 if self.remove_translation:
415 net_motion = dr.sum(axis=0) / len(atoms)
416 # print(net_motion)
417 dr = dr - net_motion
418 dr_unit = dr / np.linalg.norm(dr)
419 dr = dr_unit * self.step_size
421 # save old positions before update
422 self.Nold = self.N
423 self.rold = self.r
424 self.Told = self.T
426 # if we have constraints then this will do the first part of the
427 # RATTLE algorithm:
428 # If we can avoid using momenta, this will be simpler.
429 masses = atoms.get_masses()[:, np.newaxis]
430 atoms.set_positions(self.r + dr)
431 new_momenta = (atoms.get_positions() - self.r) * masses # / self.dt
433 # We need to store the momenta on the atoms before calculating
434 # the forces, as in a parallel Asap calculation atoms may
435 # migrate during force calculations, and the momenta need to
436 # migrate along with the atoms.
437 atoms.set_momenta(new_momenta, apply_constraint=False)
439 # Now we get the new forces!
440 f = atoms.get_forces(md=True)
442 # I don't really know if removing md=True from above will break
443 # compatibility with RATTLE, leaving it alone for now.
444 f_constrained = atoms.get_forces()
445 # but this projection needs the forces to be consistent with the
446 # constraints. We have to set the new velocities perpendicular so they
447 # get logged properly in the trajectory files.
448 vnew = subtract_projection(atoms.get_velocities(), f_constrained)
449 # using the md = True forces like this:
450 # vnew = subtract_projection(atoms.get_velocities(), f)
451 # will not work with constraints
452 atoms.set_velocities(vnew)
454 # rescaling momentum to maintain constant kinetic energy.
455 KEnew = atoms.get_kinetic_energy()
456 Ms = np.sqrt(KEold / KEnew) # Ms = Momentum_scale
457 atoms.set_momenta(Ms * atoms.get_momenta())
459 # Normally this would be the second part of RATTLE
460 # will be done here like this:
461 # atoms.set_momenta(atoms.get_momenta() + 0.5 * self.dt * f)
462 return f