Coverage for ase / md / contour_exploration.py: 96.27%
161 statements
« prev ^ index » next coverage.py v7.13.3, created at 2026-02-04 10:20 +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 def rand_vect(self):
240 '''Returns a random (Natoms,3) vector'''
241 vect = self.rng.normal(size=(len(self._actual_atoms), 3))
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
383 if f is None:
384 f = atoms.get_forces()
386 # get the velocity vector and old kinetic energy for momentum rescaling
387 velocities = atoms.get_velocities()
388 KEold = atoms.get_kinetic_energy()
390 energy = atoms.get_potential_energy(force_consistent=True)
391 self.update_previous_energies(energy)
392 potentiostat_step_size = self.compute_potentiostat_step_size(f, energy)
394 self.N = normalize(f)
395 self.r = atoms.get_positions()
396 # remove velocity projection on forces
397 v_parallel = subtract_projection(velocities, self.N)
398 self.T = normalize(v_parallel)
400 if self.use_frenet_serret:
401 if self.Nold is not None and self.Told is not None:
402 dr = self._compute_update_with_fs(potentiostat_step_size)
403 else:
404 # we must have the old positions and vectors for an FS step
405 # if we don't, we can only do a small step
406 dr = self._compute_update_without_fs(
407 potentiostat_step_size,
408 scale=self.initialization_step_scale)
409 else: # of course we can run less accuratly without FS.
410 dr = self._compute_update_without_fs(potentiostat_step_size)
412 # now that dr is done, we check if there is translation
413 if self.remove_translation:
414 net_motion = dr.sum(axis=0) / len(atoms)
415 # print(net_motion)
416 dr = dr - net_motion
417 dr_unit = dr / np.linalg.norm(dr)
418 dr = dr_unit * self.step_size
420 # save old positions before update
421 self.Nold = self.N
422 self.rold = self.r
423 self.Told = self.T
425 # if we have constraints then this will do the first part of the
426 # RATTLE algorithm:
427 # If we can avoid using momenta, this will be simpler.
428 masses = atoms.get_masses()[:, np.newaxis]
429 atoms.set_positions(self.r + dr)
430 new_momenta = (atoms.get_positions() - self.r) * masses # / self.dt
432 # We need to store the momenta on the atoms before calculating
433 # the forces, as in a parallel Asap calculation atoms may
434 # migrate during force calculations, and the momenta need to
435 # migrate along with the atoms.
436 atoms.set_momenta(new_momenta, apply_constraint=False)
438 # Now we get the new forces!
439 f = atoms.get_forces(md=True)
441 # I don't really know if removing md=True from above will break
442 # compatibility with RATTLE, leaving it alone for now.
443 f_constrained = atoms.get_forces()
444 # but this projection needs the forces to be consistent with the
445 # constraints. We have to set the new velocities perpendicular so they
446 # get logged properly in the trajectory files.
447 vnew = subtract_projection(atoms.get_velocities(), f_constrained)
448 # using the md = True forces like this:
449 # vnew = subtract_projection(atoms.get_velocities(), f)
450 # will not work with constraints
451 atoms.set_velocities(vnew)
453 # rescaling momentum to maintain constant kinetic energy.
454 KEnew = atoms.get_kinetic_energy()
455 Ms = np.sqrt(KEold / KEnew) # Ms = Momentum_scale
456 atoms.set_momenta(Ms * atoms.get_momenta())
458 # Normally this would be the second part of RATTLE
459 # will be done here like this:
460 # atoms.set_momenta(atoms.get_momenta() + 0.5 * self.dt * f)
461 return f