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